Skip to content

autodiff unnecessarily prevents inlining #138920

@ZuseZ4

Description

@ZuseZ4

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:

Metadata

Metadata

Assignees

Labels

E-easyCall for participation: Easy difficulty. Experience needed to fix: Not much. Good first issue.F-autodiff`#![feature(autodiff)]`

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions