Skip to content

autodiff unnecessarily prevents inlining #138920

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
ZuseZ4 opened this issue Mar 25, 2025 · 2 comments · Fixed by #139308
Closed

autodiff unnecessarily prevents inlining #138920

ZuseZ4 opened this issue Mar 25, 2025 · 2 comments · Fixed by #139308
Assignees
Labels
E-easy Call for participation: Easy difficulty. Experience needed to fix: Not much. Good first issue. F-autodiff `#![feature(autodiff)]`

Comments

@ZuseZ4
Copy link
Member

ZuseZ4 commented Mar 25, 2025

Using std::autodiff and therefore RUSTFLAGS=-Zautodiff=Enable in release mode currently alters the compilation pipeline in the following way:

  1. We run llvm's O3 opts (with slight modifications)
  2. We run the enzyme autodiff pass at the end of the O3 opt run above.
  3. We build and run a completely new PassManager and run llvm's O3 pipeline again, this time without modifications.

We have to make sure that functions being differentiated are not getting inlined during 1), since we still need them in step 2). Therefore they all receive noinline attributes.

The problem we now have is that in Step 3), we potentially want to inline the new function which was generated through enzyme/autodiff. However, all function which got differentiated still have the noinline attribute.
We kept track of which function got differentiated, so it should be comparably easy to go through the llvm-ir module, find all functions which got created during step 2), and remove their noonline attribute before we move on to step 3). This will improve performance for smaller functions.

Here is the second llvm_optimize call, we would want to remove noinline before this call (and after the first call 5 lines above).

write::llvm_optimize(cgcx, dcx, module, None, config, opt_level, opt_stage, stage)?;

This is a good issue to get started, so if you're interested please feel free to start a discussion here, I'm happy to help with the design or implementation.

Tracking:

@ZuseZ4 ZuseZ4 added E-easy Call for participation: Easy difficulty. Experience needed to fix: Not much. Good first issue. F-autodiff `#![feature(autodiff)]` labels Mar 25, 2025
@rustbot rustbot added the needs-triage This issue may need triage. Remove it if it has been sufficiently triaged. label Mar 25, 2025
@ZuseZ4 ZuseZ4 removed the needs-triage This issue may need triage. Remove it if it has been sufficiently triaged. label Mar 25, 2025
@Shourya742
Copy link
Contributor

@rustbot claim

@Shourya742
Copy link
Contributor

Hey @ZuseZ4, thanks for the detailed breakdown!
Do let me know if this is in correct direction or not. One possible way to handle this could be updating our FFI bindings to support iterating over functions, getting the module context, and setting attributes. Since attributes need a context, we might have to add LLVMGetModuleContext to fetch that first unless there's already another way to do it that we haven't considered? For function iteration, we'd probably need LLVMGetFirstFunction and LLVMGetNextFunction in our bindings, so we can loop through all functions in a module. Once we have a function, we'd use LLVMAttributeFunctionIndex (which, if I’ve got it right, is just 0) to apply attributes at the function level. To actually make everything alwaysinline, we'd likely use LLVMCreateStringAttribute to create the attribute in the right context and then attach it with LLVMAddAttributeAtIndex. In theory, that should force inlining, but I’m not totally sure. Does this approach make sense, or are we missing something?

jhpratt added a commit to jhpratt/rust that referenced this issue Apr 26, 2025
…-inline, r=ZuseZ4

add autodiff inline

closes: rust-lang#138920

r? `@ZuseZ4`
tgross35 added a commit to tgross35/rust that referenced this issue Apr 27, 2025
…-inline, r=ZuseZ4

add autodiff inline

closes: rust-lang#138920

r? ``@ZuseZ4``
bors added a commit to rust-lang-ci/rust that referenced this issue Apr 27, 2025
…nline, r=<try>

add autodiff inline

closes: rust-lang#138920

r? `@ZuseZ4`

try-job: dist-aarch64-linux
bors added a commit to rust-lang-ci/rust that referenced this issue Apr 28, 2025
…nline, r=<try>

add autodiff inline

closes: rust-lang#138920

r? `@ZuseZ4`

try-job: dist-aarch64-linux
ChrisDenton added a commit to ChrisDenton/rust that referenced this issue Apr 28, 2025
…-inline, r=ZuseZ4

add autodiff inline

closes: rust-lang#138920

r? `@ZuseZ4`

try-job: dist-aarch64-linux
ChrisDenton added a commit to ChrisDenton/rust that referenced this issue Apr 28, 2025
…-inline, r=ZuseZ4

add autodiff inline

closes: rust-lang#138920

r? ``@ZuseZ4``

try-job: dist-aarch64-linux
@bors bors closed this as completed in d4845e1 Apr 29, 2025
rust-timer added a commit to rust-lang-ci/rust that referenced this issue Apr 29, 2025
Rollup merge of rust-lang#139308 - Shourya742:2025-03-29-add-autodiff-inline, r=ZuseZ4

add autodiff inline

closes: rust-lang#138920

r? ```@ZuseZ4```

try-job: dist-aarch64-linux
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
E-easy Call for participation: Easy difficulty. Experience needed to fix: Not much. Good first issue. F-autodiff `#![feature(autodiff)]`
Projects
None yet
Development

Successfully merging a pull request may close this issue.

3 participants