add autodiff inline

[Shourya742]

closes: #138920

r? @ZuseZ4

try-job: dist-aarch64-linux

[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]

@bors delegate=@Shoruya742 with print comment adressed

[bors]

✌️ @Shoruya742, you can now approve this pull request!

If @ZuseZ4 told you to "r=me" after making some further change, please make that change, then do @bors r=@ZuseZ4

[Shourya742]

@bors r=ZuseZ4

[bors]

@Shourya742: 🔑 Insufficient privileges: Not in reviewers

[Shourya742]

@bors r=@ZuseZ4

[bors]

@Shourya742: 🔑 Insufficient privileges: Not in reviewers

[Shourya742]

@ZuseZ4 looks like I can't approve.

[ZuseZ4]

Strange,
@bors r+

[bors]

📌 Commit 48d05aa has been approved by ZuseZ4

It is now in the queue for this repository.

[jhpratt]

@bors r-

#140350 (comment)

[Shourya742]

Hey @ZuseZ4, it looks like there was a pointer mismatch. I updated all vanilla references to use *const c_str. Could we try running the job for dist-aarch64-linux?

[ZuseZ4]

Yes let's try that, but also when comparing it to other code it looks like existing functions (e.g. CreateAttrString) accept &str and use as_c_char_ptr from compiler/rustc_codegen_llvm/src/common.rs. They then forward both string and len to C++, see e.g. LLVMRustDIBuilderCreateFile in RustWrapper.cpp. There they use StringRef(Filename, FilenameLen) to be able to use Rust Strings in C++. That would allow you to get rid of the CString. Is there a reason why we can't copy it, or could you update the new wrappers to use it? That should also avoid the pointer type mismatch that we have.

@bors try

[bors]

⌛ Trying commit bbb9f6c with merge 93ac4be0c9a102e10028aa511237f84031510b8f...

[Shourya742]

Yes let's try that, but also when comparing it to other code it looks like existing functions (e.g. CreateAttrString) accept &str and use as_c_char_ptr from compiler/rustc_codegen_llvm/src/common.rs. They then forward both string and len to C++, see e.g. LLVMRustDIBuilderCreateFile in RustWrapper.cpp. There they use StringRef(Filename, FilenameLen) to be able to use Rust Strings in C++. That would allow you to get rid of the CString. Is there a reason why we can't copy it, or could you update the new wrappers to use it? That should also avoid the pointer type mismatch that we have.

@bors try

Let me check this — I was under the impression that with c_str, you don't really need to know the length explicitly, since it's terminated by a \0.

[ZuseZ4]

You're right, either solution should work. It's just the question of whether rustc users will need to allocate a CString which is zero terminated, or a Rust String which contains ptr + len. Working with pure Rust types would be nicer, and as I just noticed when checking the build failure, other LLVM wrappers also seem to make use of the ptr+len combination, so I think it would be good if we copy them. In that case you could also remove your CString and replace it with a normal string.

[bors]

☀️ Try build successful - checks-actions
Build commit: 93ac4be (93ac4be0c9a102e10028aa511237f84031510b8f)

[Shourya742]

@ZuseZ4, should we go ahead and push this in? I can follow up with another PR to replace CString with a Rust String and then use StringRef.