Skip to content
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

autodiff unnecessarily prevents inlining #138920

Open
ZuseZ4 opened this issue Mar 25, 2025 · 2 comments · May be fixed by #139308
Open

autodiff unnecessarily prevents inlining #138920

ZuseZ4 opened this issue Mar 25, 2025 · 2 comments · May be 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?

@Shourya742 Shourya742 linked a pull request Apr 3, 2025 that will close this issue
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