autodiff: avoid need for -Zautodiff flag
Nobody has claimed this yet.
- Dominant language
- Rust
- Stars
- 119k
- Forks
- 16.2k
- PR merge metrics
- PR metrics pending
Description
Currently we there are two blockers for using std::autodiff - we require setting lto=fat and we require RUSTFLAGS=-Zautodiff=Enable.
There's an open issue to get rid of the lto requirement, this one is about the -Z flag.
The reason we need both lies in our compilation model. As we see in Nikita Popov's blog, there are two parts of LLVMs opt pipeline: https://www.npopov.com/2023/04/07/LLVM-middle-end-pipeline.html#summary
This is our compilation pipeline: (module simplification + module optimization)
For good performance, autodiff(Enzyme) must run in between those two (see Enzyme papers).
If a user called autodiff, Enzyme will generate new code. That code must also be optimized, our full optimization pipeline is therefore
(module simplification + Enzyme) - (module simplification + module optimization)
(brackets indicate a shared module pass manager)
autodiff will (and must, to be usable) compute derivatives across function, module, or crate boundaries. Therefore all crates must first pass the module simplification, but stop before the module optimization. At that stage (when compiling dependencies) we don't know yet based on the code itself if autodiff is used somewhere downstream. The -Zautodiff gives us the information, and it forces the new compilation behaviour accross the whole dependency tree. We then run fat-lto which combines everything into one LLVM module, and run Enzyme, plus module simplification+optimization.
Let's say we now want to always support autodiff, without the -Z flag.
It would be comparably easy to track if the Enzyme invocation actually had any work to do. If not, we could (in the future) skip the second module simplification and go directly to the module optimization step. This should get rid of the compile time overhead of the -Zautodiff flag, and we could just always support autodiff. There's a chance though that codegen will suffer a little bit, if we have two independent LLVM invocations for the module simplification and module optimization, instead of a combined one where some analysis results could be shared.
Iff that turns out to be the case, we could look into scheduling some of the valuable analysis passes from the module simplifcation phase ahead of our module optimization phase to recover that performance. So we would end up with either of
(module simplification+Enzyme) - (module optimization) // Enzyme was a no-op
(module simplification+Enzyme) - (module simplification+module optimization) // Enzyme had work to do.
This split is btw. also where we would (likely) end up if we were to port some LLVM optimizations into a future Rust LIR. I think the hope is to mostly have a smarter module simplification via LIR, I don't think people want to reproduce the module optimization part.
Another benefit is that we could simplify the codegen_llvm backend a little bit, since we currently support both compilation pipelines, the normal LLVM opt pipeline if the user does not set the flag, and the split, two stage pipeline if the user sets the autodiff flag.
The downside is that in the future all code would go through the more complex two stage opt pipeline, whereas right now autodiff/Enzyme has no impact at all on the compilation pipeline, if users don't use autodiff.
I'm not saying we necessarily should land these changes, but those would be the needed steps if we want to go that route.
Either way we first should land support for the non-fat-lto case, and then benchmark the perf impact of such change in different build configs.
So for now I'm mainly writing this out to keep track of it. cc @oli-obk @traviscross
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the issue's compilation-pipeline description and the linked LLVM middle-end pipeline article, then inspect the existing Enzyme and autodiff handling in rustc's LLVM codegen. Before implementation, establish the non-fat-LTO prerequisite and benchmark the proposed pipeline across different build configurations; the issue is done only when the design and performance impact support removing the flag.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100