Autodiff Upstreaming - rustc_codegen_llvm changes #130060
Conversation
|
r? @fee1-dead rustbot has assigned @fee1-dead. Use |
|
rustbot
added
S-waiting-on-review
|
This PR modifies If appropriate, please update Some changes occurred in cfg and check-cfg configuration cc @Urgau |
This comment has been minimized.
This comment has been minimized.
| @@ -176,6 +176,8 @@ pub(crate) fn default_configuration(sess: &Session) -> Cfg { | |||
| // NOTE: These insertions should be kept in sync with | |||
| // `CheckCfg::fill_well_known` below. | |||
|
|
|||
| ins_none!(sym::autodiff_fallback); | |||
There was a problem hiding this comment.
This shouldn't be insta stable, it should be at least gated behind nightly compiler.
| ins_none!(sym::autodiff_fallback); | |
| if sess.is_nightly_build() { | |
| ins_none!(sym::autodiff_fallback); | |
| } |
Please also follow all the steps regarding a new cfg as defined in the top of this file (as well as the tests files):
rust/compiler/rustc_session/src/config/cfg.rs
Lines 10 to 21 in e26b02a
|
r? compiler |
|
The job Click to see the possible cause of the failure (guessed by this bot) |
|
☔ The latest upstream changes (presumably #131237) made this pull request unmergeable. Please resolve the merge conflicts. |
|
r? compiler |
|
There's very little chance of this being merged in one PR with one commit of this size. You'll need to split this up into well-commented/motivated PRs that can be landed one at a time. I haven't spent much time looking at this PR, so I don't have any suggestions on how to split this up. I'd recommend finding someone on the compiler team who is interested in these changes and who you can work with to do the reviews. |
There was a problem hiding this comment.
I don't have time to review this, so just one drive-by note: It looks like a decent part of the extra FFI APIs in enzyme_ffi.rs are essentially duplicates of things that we already have bindings for under slightly different names and signatures. Like we already have LLVMRustAddFunctionAttributes/LLVMRustAddCallSiteAttributes and this introduces LLVMRustAddEnumAttributeAtIndex. It also looks like the code doesn't make use of the Builder abstraction and instead calls FFI APIs directly everywhere, which is probably also where the duplication comes from.
|
@nikic Do you mind if I drop the approach in this PR and add Enzyme back as an LLVM pass instead of using Enzyme as a library? Historically Julia (and Rust except for the very first experiments) used the lib approach, while Clang used the pass approach. The pass used to support fewer configurations and metadata information. We now added parser to read all metadata from the llvm-ir module and added missing config options, so Julia (Enzyme.jl) is now also moving to the pass since it simplifies the config a little. JuliaGPU/GPUCompiler.jl#636 (comment) To give a simplified example of the two approaches, the frontend rustc_builtin_macro fn f(some_args: types) {
// original user implementation, e.g. sin(x)
}
fn df(slightly_altered_args: similar_types) {
// macro generated placeholder
}I used to pass the llvm-ir function of (In reality, I would do the rewriting on llvm-ir, just using Rust for readability). If anyone has concerns it would be great to hear them now already, otherwise I'd try to get the new PR up soon. I've already done most of the work in the past since it also offers greater debugability due to having all relevant info in the llvm-ir module, instead of passing data directly in memory to a library call, which is hard to test and reproduce. |
Now that the autodiff/Enzyme backend is merged, this is an upstream PR for the
rustc_codegen_llvmchanges.It also includes small changes to three files under
compiler/rustc_ast, which overlap with my frontend PR (#129458).Here I only include minimal definitions of structs and enums to be able to build this backend code.
The same goes for minimal changes to
compiler/rustc_codegen_ssa, the majority of changes there will be in another PR, once either this or the frontend gets merged.We currently have 68 files left to merge, 19 in the frontend PR, 21 (+3 from the frontend) in this PR, and then ~30 in the middle-end.
This PR is large because it includes two of my three large files (~800 loc each). I could also first only upstream enzyme_ffi.rs, but I think people might want to see some use of these bindings in the same PR?
To already highlight the things which reviewers might want to discuss:
enzyme_ffi.rs: I do have a fallback module to make sure that we don't link rustc against Enzyme when we build rustc without autodiff support.add_panic_msg_to_globalwas a pain to write and I currently can't even use it. Enzyme writes gradients into shadow memory. Pass in one float scalar? We'll allocate and return an extra float telling you how this float affected the output. Pass in a slice of floats? We'll let you allocate the vector and pass in a mutable reference to a float slice, we'll then write the gradient into that slice. It should be at least as large as your original slice, so we check that and panic if not. Currently we panic silently, but I already generate a nicer panic message with this function. I just don't know how to print it to the user. yet. I discussed this with a few rustc devs and the best we could come up with (for now), was to look for mangled panic calls in the IR and pick one, which works surprisingly reliably. If someone knows a good way to clean this up and print the panic message I'm all in, otherwise I can remove the code that writes the nicer panic message and keep the silent panic, since it's enough for soundness. Especially since this PR is already a bit larger.SanitizeHWAddress: When differentiating C++, Enzyme can use TBAA to "understand" enums/unions, but for Rust we don't have this information. LLVM might to speculative loads which (without TBAA) confuse Enzyme, so we disable those with this attribute. This attribute is only set during the first opt run before Enzyme differentiates code. We then remove it again once we are done with autodiff and run the opt pipeline a second time. Since enums are everywhere in Rust, support for them is crucial, but if this looks too cursed I can remove these ~100 lines and keep them in my fork for now, we can then discuss them separately to make this PR simpler?Duplicated llvm-opt runs: Differentiating already optimized code (and being able to do additional optimizations on the fly, e.g. for GPU code) is the reason why Enzyme is so fast, so the compile time is acceptable for autodiff users: https://enzyme.mit.edu/talks/Publications/ (There are also algorithmic issues in Enzyme core which are more serious than running opt twice).
I assume that if we merge these minimal cg_ssa changes here already, I also need to fix the other backends (GCC and cliff) to have dummy implementations, correct?
I'm happy to split this PR up further if reviewers have recommendations on how to.
For the full implementation, see: #129175
Tracking: