add autodiff inline

[Shourya742]

closes: #138920

r? @ZuseZ4

[ZuseZ4]

Thank you for looking into this!

There are two things to improve here, first we want to write that in safer Rust, and second this currently adds alwaysinline everywhere. However, we only want it for the functions we generated. I'd probably start with making it safer, that should be more straightforward, this PR is a good start: #135581

As a rule of thumb I often noticed during my autodiff work, that if you want to add something to an LLVM component, then rustc likely already has wrappers for it. If you want to get something out of an LLVM component, chances are you have to write those safe wrappers, since rustc usually just lowers things.

Here is an example I found after using rg (ripgrep) to find inline related matches in rustc_codegen_llvm:

rust/compiler/rustc_codegen_llvm/src/attributes.rs

Line 33 in 9e14530

fn inline_attr<'ll>(cx: &CodegenCx<'ll, '_>, inline: InlineAttr) -> Option<&'ll Attribute> {

SimpleCx and CodegenCx are safe wrappers around llvm modules to help you handle things safely. The CodegenCx is more powerful, but also needs more inputs to be constructed. You probably won't be able to generate a full CodegenCx, since you don't have a TypeContext (tcx) available. Instead search for the SimpleCx, you should be able to create one instead. In the next steps we can then use the SimpleCx, to safely replace the 4 LLVM functions which you are currently introducing.

(The way I usually find safe wrappers systematically is by first looking up the llvm function I need, then searching for existing uses (wrappers) in rustc, and then going up the wrapper chain till I find one that I can use.)

[bors]

☔ The latest upstream changes (presumably #139396) made this pull request unmergeable. Please resolve the merge conflicts.

[Shourya742]

@rustbot review

[ZuseZ4]

Thanks, that looks pretty good now. I'm a bit worried about regressing this, though, since just missing an inline will probably not be noticed immediately. So a tests/codegen/autodiff/inline.rs test would be good. Once it landed you should be able to use #139700 the NoPostopt flag from this PR, and check that the inline attribute is where we expect it to be. https://llvm.org/docs/CommandGuide/lit.html
You can try to copy from the other tests (just add the NoPostopt` flag in the first line), and check for the inline attribute on top of the function. Feel free to push a draft test, and I can help to clean it up.

[Shourya742]

Thanks, that looks pretty good now. I'm a bit worried about regressing this, though, since just missing an inline will probably not be noticed immediately. So a tests/codegen/autodiff/inline.rs test would be good. Once it landed you should be able to use #139700 the NoPostopt flag from this PR, and check that the inline attribute is where we expect it to be. https://llvm.org/docs/CommandGuide/lit.html You can try to copy from the other tests (just add the NoPostopt` flag in the first line), and check for the inline attribute on top of the function. Feel free to push a draft test, and I can help to clean it up.

Hello @ZuseZ4, I’ve locally rebased my PR onto #139700, and in the generated IR, I’m seeing the following:

; Function Attrs: alwaysinline noinline
declare double @__enzyme_autodiff_ZN6inline8d_square17h021c74e92c259cdeE(...) local_unnamed_addr #8

This looks correct to me. I used -Z autodiff=NoPostopt as you suggested.

[ZuseZ4]

oh, the test is helpful. I'm not sure what LLVM would do, but I think specifying noinline alwaysinline on the same function will be confusing either for LLVM or developers, probably both. I wouldn't want to rely on whatever behavior it has.

Can you (before adding alwaysinline) go through fn attrs, and assert that there is a noinline attr, with a comment that this check isn't strictly necessary, but just a guard to remind us if this code path ever changes? Then remove the no inline attr before adding always inpine. That should make it more robust.

[Shourya742]

Hi @ZuseZ4, I added LLVMRemoveEnumAttributeAtIndex to remove the noinline attribute from llfn, and also implemented a utility method has_attr based on LLVMRustHasAttributeAtIndex. However, I’m not seeing the noinline attribute being removed from the function. Not sure why that’s happening—any suggestions?

[Shourya742]

@rustbot review

[ZuseZ4]

Ok, I've experimented a bit.
First, in one location you use 0, where it should be !0 (or rather use the AttributePlace enum which has a to_cuint() function which lowers correctly and is nicer to use.
I've also realized that you were modifying the __enzyme declaration, not the actual function which got generated by enzyme. I'm not sure why the declaration is still around (IMHO enzyme should clean up after itself and remove unused declarations). I've updated the test to test that the actual function which we differentiated lost it's noinline attribute.

In general, your __enzyme string matching function was a bit risky, since we should always consider that users give functions arbitrary names. I've instead added some code earlier to the pipeline, which marks source and target function with an enzyme_marker attribute. This simplifies the backend, where we now can just iterate through all functions and skip the ones without an enzyme marker.

As to why your function didn't remove the attribute, I'm not fully sure, but I remember that I ran into the same bug a while ago. I ended up copying me old code here: https://github.com/EnzymeAD/rust/blob/322f2226c1f672c9b5e934b15d255ae0d66bd0e2/compiler/rustc_codegen_llvm/src/back/write.rs#L1195
It indeed ended up working, you can see that the noinline attribute disappears.

I pushed my code here. Sorry that it's quite messy, I cherry picked both my flags PR and your code (and messed up the cherry picking on top of that). Also, it's not using any pretty wrappers, but it hopefully shows how to remove the attribute.
master...EnzymeAD:rust:noinline-imprv

Either in this or a follow-up PR we should also be a bit more precise, since e.g. the diffe_ function which got generated by enzyme still has the noinline attribute. It's a newly generated function so there's no way to mark it with enzyme_marker before, but we know that our wrapper does nothing but call the diffe_ function. So whenever we see a function marked with enzyme_marker, we could iterate through the instructions in the body to get the diffe function call, and from there the diffe definition. Then we can also strip the noinline from there. But as mentioned, we don't have to do anything in one PR.

[ZuseZ4]

Ok, I think also remembered again why your approach didn't work.
If you look at the definition of the AttributeKind, it says that it will match the C++ RustAttributeKind.
It does not match the LLVM AttributeKind, for which there is another wrapper there on the C++ side in compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp (around line 282).
So if you pass the Rust attribute directly to the LLVM* functions, it will break. If you pass it to the LLVMRust functions and use a fromRust(attr), then it will be matched correctly. You can check it yourself by looking at
build/x86_64-unknown-linux-gnu/llvm/include/llvm/IR/Attributes.inc, where we have NoInline = 32,.
(So for fun you could try to remove the StackProtect attribute, which on the rust side has value 32. Maybe removing StackProtect will remove noinline (but don't do that in the final code). FFI is fun..
Also feel free to add a note to the AttributeKind enum in compiler/rustc_codegen_llvm/src/llvm/ffi.rs, making clear that it must not be passed to LLVM* functions, and only works for LLVMRust functions. (And I'd probably let that be verified by another reviewer)

[rustbot]

Some changes occurred in compiler/rustc_codegen_llvm/src/builder/autodiff.rs

cc @ZuseZ4

Some changes occurred in compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs

cc @ZuseZ4

[Shourya742]

Ok, I think also remembered again why your approach didn't work. If you look at the definition of the AttributeKind, it says that it will match the C++ RustAttributeKind. It does not match the LLVM AttributeKind, for which there is another wrapper there on the C++ side in compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp (around line 282). So if you pass the Rust attribute directly to the LLVM* functions, it will break. If you pass it to the LLVMRust functions and use a fromRust(attr), then it will be matched correctly. You can check it yourself by looking at build/x86_64-unknown-linux-gnu/llvm/include/llvm/IR/Attributes.inc, where we have NoInline = 32,. (So for fun you could try to remove the StackProtect attribute, which on the rust side has value 32. Maybe removing StackProtect will remove noinline (but don't do that in the final code). FFI is fun.. Also feel free to add a note to the AttributeKind enum in compiler/rustc_codegen_llvm/src/llvm/ffi.rs, making clear that it must not be passed to LLVM* functions, and only works for LLVMRust functions. (And I'd probably let that be verified by another reviewer)

That did the trick—fromRust was the missing piece. Now remove_from_llfn is successfully removing the attributes. Thanks for jogging the memory and pointing out the subtle mismatch—FFI really does keep you humble.

[Shourya742]

Either in this or a follow-up PR we should also be a bit more precise, since e.g. the diffe_ function which got generated by enzyme still has the noinline attribute. It's a newly generated function so there's no way to mark it with enzyme_marker before, but we know that our wrapper does nothing but call the diffe_ function. So whenever we see a function marked with enzyme_marker, we could iterate through the instructions in the body to get the diffe function call, and from there the diffe definition. Then we can also strip the noinline from there. But as mentioned, we don't have to do anything in one PR.

Thanks a lot for the thorough explanation!. The marker-based approach you added is a big improvement. Totally agree that relying on string matching for __enzyme is fragile, especially with user-defined names in play. Using enzyme_marker early in the pipeline to annotate both source and target functions makes things much more robust and simplifies things downstream—nice solution!. I'm happy to help clean things up and follow through with refining the approach—especially for that diffe_ function case you mentioned. That sounds like a good candidate for a follow-up PR where we can safely walk the body and strip noinline from any directly-called differentiated functions. Appreciate all the context on this!

[ZuseZ4]

Overall looks good, just left some nits. Can you clean up the history now, maybe bring it down to 2-3 commits (or just squash it all if you don't want to deal with it)?

[Shourya742]

Hey @ZuseZ4,
I’ve addressed the nits and restructured the commits. This is currently blocked by: #139700.
Everything else is good to go!