Enable autodiff with thin-lto #153687
Description
Let's assume a Rust code where a function foo calls bar, and bar calls baz.
autodiff works based on the chain rule from calculus, so in order to differentiate foo, we also need to differentiate bar and therefore baz. Those, however, might be in a different cgu or even crate. If Enzyme fails to find their LLVM-IR, compilation must fail. Importantly, Enzyme will never modify foo, bar, or baz. It creates full copies of those functions and then processes those copies, which should help. Enzyme knows to differentiate foo, because we generate declarations with magic names, called __enzyme_autodiff_<hash> or __enzyme_fwddiff_<hash>. When we call those declaration we pass a ptr to foo as the first argument. Enzyme will then differentiate foo, and replace those __enzyme_autodiff calls with calls to the newly generated, differentiated functions.
In theory, that would all be satisfied by embed-bitcode=true, but the easier solution so far was to just enforcing lto="fat", and run autodiff after merging everything into a single llvm module. In https://github.com/ZuseZ4/autodiff_thin_lto I created a test example across main.rs, lib.rs, and another test dependency on github. With thin-lto that currently runs into our fat-lto assertion. If we remove the assertion we run into either of two issues:
- With no further changes, compilation fails because we don't run enzyme and the declarations never get implemented:
rust-lld: error: undefined symbol: __enzyme_autodiff_RNvCs8v4yQ6vrPnh_3foo3bar
>>> referenced by main.rs:8 (src/main.rs:8)
>>> /home/manuel/prog/foo/target/debug/deps/foo-c10f7e02831e5066.57y0q35f2jngw94roh5l20tpp.1n5lzsr.rcgu.o:(foo::bar)- If we make sure that we run autodiff on some more
llvm_optimizecalls (coming from thin-lto), we now run into the issue that Enzyme can't find the IR of the function in mysimple_depgithub dependency when trying to differentiate a function in our library.
error: src/lib.rs:12:5: in function preprocess__RNvCscExB1zF26bR_3foo3__f1 double (double): Enzyme: No forward mode derivative found for _RNvCs8oqLYcHYtwO_10simple_dep1f
at context: %2 = call double @_RNvCs8oqLYcHYtwO_10simple_dep1f(double %0, double %0) #5, !dbg !19
error: could not compile `foo` (bin "foo")A year ago oli and I had a pair-programming session where we tried to enable thin-lto support based on this PR, but got stuck by some esoteric rustc bug a few layers into it and gave up. https://github.com/rust-lang/rust/pull/137570/changes
Since then, we moved autodiff to be an intrinsic, so now we always generate the __enzyme_autodiff calls, last year I had to make some changes in the PR to even codegen them for thin-lto. That's now automatically done. Generally, after #153379 we barely have any code left in the middle-end at all (just run rg autodiff compiler to check). So that pr isn't super relevant anymore and it's hopefully easier to achieve. Bjorn also cleaned up the lto handling, so I wanted to give it another try.
On the C++ side, there's an Enzyme plugin for lld, so it can run during link time. I guess we could use that, but I also feel like we should be able to do better and solve it within rustc? I tried to play around with compiler/rustc_codegen_llvm/src/back/write.rs and compiler/rustc_codegen_llvm/src/back/lto.rs, but kept falling back to either of the two errors above. Any recommendations on where to best add autodiff so we can access the bitcode of upstream libs? Potentially related: #153648
For runtime performance reasons, we currently also split the llvm opt pass pipelines into two runs and embed autodiff in between, but I think I can figure that out myself later how to copy that for thin-lto once I get to work in the first place. I'm also less worried about perf in the non-fat-lto case.