Autodiff Upstreaming - rustc_codegen_ssa, rustc_middle

[ZuseZ4]

This PR should not be merged until the rustc_codegen_llvm part is merged.
I will also alter it a little based on what get's shaved off from the cg_llvm PR,
and address some of the feedback I received in the other PR (including cleanups).

I am putting it already up to

1. Discuss with @jieyouxu if there is more work needed to add tests to this and
2. Pray that there is someone reviewing who can tell me why some of my autodiff invocations get lost.

Re 1: My test require fat-lto. I also modify the compilation pipeline. So if there are any other llvm-ir tests in the same compilation unit then I will likely break them. Luckily there are two groups who currently have the same fat-lto requirement for their GPU code which I have for my autodiff code and both groups have some plans to enable support for thin-lto. Once either that work pans out, I'll copy it over for this feature. I will also work on not changing the optimization pipeline for functions not differentiated, but that will require some thoughts and engineering, so I think it would be good to be able to run the autodiff tests isolated from the rest for now. Can you guide me here please?
For context, here are some of my tests in the samples folder: https://github.com/EnzymeAD/rustbook

Re 2: This is a pretty serious issue, since it effectively prevents publishing libraries making use of autodiff: EnzymeAD#173. For some reason my dummy code persists till the end, so the code which calls autodiff, deletes the dummy, and inserts the code to compute the derivative never gets executed. To me it looks like the rustc_autodiff attribute just get's dropped, but I don't know WHY? Any help would be super appreciated, as rustc queries look a bit voodoo to me.

Tracking:

• Tracking Issue for autodiff #124509

r? @jieyouxu

[ZuseZ4]

To expand on 2)
Assume you have the following code

#[autodiff(bar, Reverse, ...)]
fn foo(x: f32) -> f32 { x*x }

it will expand to

#[rustc_autodiff]
fn foo(x: f32) -> f32 {x*x}
#[rustc_autodiff(Reverse,...)]
fn bar(x: f32, scalar_factor: f32) -> (f32, f32) {
   // some_dummy_code()
}

Now I have some logic in this PR which picks up the rustc_autodiff attributes and passes them onto the backend, where for every single rustc_autodiff attribute with arguments we pick the function (thus bar here) and replace the dummy code with the right code to return the derivative. So bar would afterwards return (x*x, 2.0 * x). But as mentioned above, once you use autodiff in a library and call it in another module, the dummy code get's executed, as shown in the linked PR.
Any hints would be appreciated.

[jieyouxu]

I have some interim feedback for this draft PR

[jieyouxu]

Remark: this has no error messages if the autodiff options failed to parse, acceptable for unstable flag but still poor UX, unacceptable when it comes to stabilization time.

[Urgau]

It should be handled with a proper description in parse_autodiff above.

[ZuseZ4]

I assume there is some rustc function to generate suggestions when users make a typo?
I'll update it for now to print the unrecognized value.

[ZuseZ4]

I looked again at all the other options which accept more values (e.g. instrument_xray), and they have no error-handling either. Is there a preferred way to print errors here?

[oli-obk]

We can do this later imo.

[jieyouxu]

My test require fat-lto.

For codegen/assembly/ui tests you can make it build w/ fat LTO via

//@ compile-flags: -Clto=fat

I also modify the compilation pipeline. So if there are any other llvm-ir tests in the same compilation unit then I will likely break them.

I don't quite understand the implication of this. Is there some small example I can refer to?

I will also work on not changing the optimization pipeline for functions not differentiated, but that will require some thoughts and engineering, so I think it would be good to be able to run the autodiff tests isolated from the rest for now. Can you guide me here please?

Can you elaborate on what test conditions you need? What do you mean exactly by "isolated"? Can you not run autodiff but only if there is autodiff support, or do you mean like don't run by default even if there is autodiff support?

[ZuseZ4]

https://github.com/rust-lang/rust/pull/130060/files#diff-a56b374664e290a55d70fa80e456b6280913830b382b73fb70c4483d3d4cf246
adjust's the llvm opt pipeline (if autodiff is enabled at build time and used).
We have a first opt run which skips the late llvm opts (which tend to increase code size), and runs opt a second time (now with the full pipeline) once autodiff is done. I will manage to not make it optimize unrelated code in the future, but for now it means other functions are now optimized 1.5 times by llvm. And in reality llvm opts don't really run to a fixpoint, so that is highly likely to change the IR.

[jieyouxu]

functions are now optimized 1.5 times by llvm. And in reality llvm opts don't really run to a fixpoint, so that is highly likely to change the IR.

This just means that for now, you'll have to gate autodiff-related tests with //@ needs-autodiff or somehow allow compiletest to determine that the llvm used is built with autodiff support and the test is exercising said autodiff support.

Note that this cannot break and should not modify code that does not use autodiff at all, which is indicated by codegen tests that do not use / opt-in to autodiff support.

[rustbot]

This PR modifies config.example.toml.

If appropriate, please update CONFIG_CHANGE_HISTORY in src/bootstrap/src/utils/change_tracker.rs.

Some changes occurred in coverage instrumentation.

cc @Zalathar

Some changes occurred in cfg and check-cfg configuration

cc @Urgau

These commits modify the Cargo.lock file. Unintentional changes to Cargo.lock can be introduced when switching branches and rebasing PRs.

If this was unintentional then you should revert the changes before this PR is merged.
Otherwise, you can ignore this comment.

[ZuseZ4]

I rebased now that the other autodiff PR got merged, fixed all conflicts, and got it to compile locally.
I will work through the existing feedback over the next days.

[bors]

☔ The latest upstream changes (presumably #136030) made this pull request unmergeable. Please resolve the merge conflicts.

[ZuseZ4]

@oli-obk Just to keep track, so far we have 3/4 things which should be fixed here. Potentially not all in this specific PR, but preferably before enabling it for default nightly builds.

1. (Most orthogonal) the macro getting lost when used in dependencies.

2. Performance: Have two opt runs, with AD being applied at the end of the first run. I need to look again at it, last time I wasn't 100% sure how to trigger these two runs after the refactoring.

3. Build setup: Don't force users to pass RUSTFLAGS="-Z llvm-plugins=/home/manuel/prog/rust-working/build/x86_64-unknown-linux-gnu/enzyme/build/Enzyme/LLVMEnzyme-19.so -C passes=enzyme. I'll probably ask on zulip/bootstrap for help with that.

4. (of course I'll also need to apply the other code quality recommendations above).

After thinking about it for a bit I'll probably just do 4) for now, and leave 1-3 for a follow-up PR, just for the sake of having a working version upstream. I'll need to talk to jieyouxu to see if we then add test's already here, or in the follow-up PR where I fix 3).

[ZuseZ4]

I'll need to add the fn-ptr error message and left a few questions.
I decided to leave the bootstrapping change for an extra PR, so the tests will probably also need to wait for that (since we use a -Z llvm-plugin flag which requires the absolute path to Enzyme, something which we probably don't want to add as a temporary workaround to the test infra).

[ZuseZ4]

Let's go!
@bors r=@oli-obk

[bors]

📌 Commit 48d4a6c has been approved by oli-obk

It is now in the queue for this repository.

[ZuseZ4]

Since it hasn't been added to a rollup yet, I quickly removed an unwanted print statement and especially adjusted a type check in my wrapper, which wasn't working correctly by checking the return type of a function, instead of a call. Wrappers are fun. But now even higher-order derivatives from my (not yet upstream) testset pass again!

[ZuseZ4]

@bors r=@oli-obk

[bors]

📌 Commit 1f30517 has been approved by oli-obk

It is now in the queue for this repository.

[lqd]

Hey @ZuseZ4, it looks like this PR had significant perf implications on a small number of benchmarks, as you can see in the post-merge perf results for the rollup where it landed. It could be related to LTO, from a quick look.

I don’t see a pre-merge perf run here, so I assume this is not expected. Could you take a look? Thanks!

Let’s also cc the reviewer @oli-obk.

[oli-obk]

Oh yea that was absolutely not expected. looking...

[oli-obk]

perf regression fixed by #136413