Adding Autodiff support to rustc

[ZuseZ4]

Hello everyone. @bytesnake and I are currently working on integrating an llvm incubator project into rustc, which differentiates code in the calculus sense. We have used it successfully as a wrapper around rustc, but that comes with various limitations which we want to solve this way.

Requirements

We require the complete LLVM-IR, at least the one being used (possibly indirectly) by a function being differentiated. In general, that means that we have to merge all the compilation units and also rely on build-std, which worked pretty fine for us until now. Currently we also require debug-symbols, although we will change that later to rely on compiler internal layout details. For the best results we also require special / repeated optimizations, although these are optional to get started.

Implementation

We would like to make as little modifications as possible. In the rustc backend we therefore intent to:

1. Add enzyme as a submodule to the llvm project.
2. Add enzyme as an llvm subproject to x.py to allow building it, together with rust bindings
3. Add a trait file to rustc_codegen_ssa/src/traits/autodiff.rs
4. Update ModuleConfig (and equivalents earlier in the build chain) to include a boolean autodiff flag. If set to true we also check for build-std and embed-bitcode usage and panic if it isn't used.
5. Update rustc_codegen_llvm/src/lib.rs and rustc_codegen_ssa/src/back/write.rs by using cfg!() in order to:
   • During normal compilation process:
     • concatenate all modules (including std ones) using run-link
     • end with single module, run: optimizations, enzyme, optimizations2. Continue normal compilation. (Same case as if we compile a trivial crate without dependencies, so the rest of the compilation should work?)
   • During thin-lto:
     • complicated, skip for now (multiple modules)
   • During fat-lto:
     • We have one compilation artifact anyway, so that should be comparably easy?

All of the steps above do not affect users, so would probably only require a MCP. However, we do intend to expose this behaviour similar to the inline-asm group, by having a special macro adding a magic node. This way we can assure that autodiff will only be invoked using the specific macro, allowing us to catch invalid input as early as possible. This macro will also generate an unimplemented!() function which will later be filled when running Enzyme. You can see an example of such a macro here: autodiff/rosenbrock.rs at main · ZuseZ4/autodiff · GitHub. It demonstrates Reverse, Forward, and Forward-Vectorized (8-width) autodiff. We don't see the specific macro design as too relevant for now but would rather ask for some general feedback on our Implementation thoughts. Is there any fundamental flaw or can we just go ahead and come back for an RFC once we have finished moving our logic into rustc?