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?