Proposal
Enzyme is an LLVM incubator project that is able to vectorize and differentiate (in the calculus sense) a given function. Computing gradients efficiently is necessary for various algorithms such as backpropagation, Bayesian inference, uncertainty quantification, and probabilistic programming.
The reason why this work should be integrated into the rust compiler is that it has been shown both in theory and in practise (PyTorch 2.0, Jax/XLA, ...) that Automatic Differentiation (AD) of code is hard to do efficiently on unoptimized source code like Rust. Enzyme has shown a geometric mean speedup of 4.5x over mature pre-optimization AD tools. Enzyme has also been successfully used in combination with parallel paradigms like CUDA, ROCm, MPI and OpenMP. Enzyme also supports other languages, e.g., Julia, Fortran, compiled Python and C++. Handling LLVM-IR generated by such a variety of compilers does help to recognize and fix bugs early which otherwise would have been hard to discover using only Rust test cases.
The latest proof-of-concept rust-enzyme has demonstrated that our approach is mature enough and is able to solve all fundamental issues of earlier iterations. We, therefore, propose to conditionally enable Enzyme for nightly releases to allow more Rust developers to test the current state and provide valuable feedback.
We do intend to stabilize this work at some point because we believe that good AD and vectorization support is a fundamental building block of scientific computing and machine learning. Working around the lack of such a tool in our opinion has similar consequences as working around the availability of async, asm, or simd support for the corresponding groups. User written implementations would either lack performance and test coverage as explained above, need to re-implement various LLVM/Enzyme optimizations or transpile the user code from Rust into other languages that have AD support. As a consequence, we hope to support both existing projects and maintainers, as well as lowering the barrier for new contributors in this field.
We do intend to write a full RFC in more than a year from now, once we had sufficient users provide feedback on the Rust frontend of our work.