Autodiff Upstreaming - enzyme frontend

[ZuseZ4]

This is an upstream PR for the autodiff rustc_builtin_macro that is part of the autodiff feature.

For the full implementation, see: #129175

Content:
It contains a new #[autodiff(<args>)] rustc_builtin_macro, as well as a #[rustc_autodiff] builtin attribute.
The autodiff macro is applied on function f and will expand to a second function df (name given by user).
It will add a dummy body to df to make sure it type-checks. The body will later be replaced by enzyme on llvm-ir level,
we therefore don't really care about the content. Most of the changes (700 from 1.2k) are in compiler/rustc_builtin_macros/src/autodiff.rs, which expand the macro. Nothing except expansion is implemented for now.
I have a fallback implementation for relevant functions in case that rustc should be build without autodiff support. The default for now will be off, although we want to flip it later (once everything landed) to on for nightly. For the sake of CI, I have flipped the defaults, I'll revert this before merging.

Dummy function Body:
The first line is an inline_asm nop to make inlining less likely (I have additional checks to prevent this in the middle end of rustc. If f gets inlined too early, we can't pass it to enzyme and thus can't differentiate it.
If df gets inlined too early, the call site will just compute this dummy code instead of the derivatives, a correctness issue. The following black_box lines make sure that none of the input arguments is getting optimized away before we replace the body.

Motivation:
The user facing autodiff macro can verify the user input. Then I write it as args to the rustc_attribute, so from here on I can know that these values should be sensible. A rustc_attribute also turned out to be quite nice to attach this information to the corresponding function and carry it till the backend.
This is also just an experiment, I expect to adjust the user facing autodiff macro based on user feedback, to improve usability.

As a simple example of what this will do, we can see this expansion:
From:

#[autodiff(df, Reverse, Duplicated, Const, Active)]
pub fn f1(x: &[f64], y: f64) -> f64 {
    unimplemented!()
}

to

#[rustc_autodiff]
#[inline(never)]
pub fn f1(x: &[f64], y: f64) -> f64 {
    ::core::panicking::panic("not implemented")
}
#[rustc_autodiff(Reverse, Duplicated, Const, Active,)]
#[inline(never)]
pub fn df(x: &[f64], dx: &mut [f64], y: f64, dret: f64) -> f64 {
    unsafe { asm!("NOP"); };
    ::core::hint::black_box(f1(x, y));
    ::core::hint::black_box((dx, dret));
    ::core::hint::black_box(f1(x, y))
}

I will add a few more tests once I figured out why rustc rebuilds every time I touch a test.

Tracking:

• Tracking Issue for autodiff #124509

try-job: dist-x86_64-msvc

[rustbot]

r? @pnkfelix

rustbot has assigned @pnkfelix.
They will have a look at your PR within the next two weeks and either review your PR or reassign to another reviewer.

Use r? to explicitly pick a reviewer

[bors]

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

[jieyouxu]

Thanks for the work on autodiff! I have some feedback and questions. As you may have gathered, I am not very knowledgeable about autodiff. If I ask some questions for more context / explanation, it might be good to encode them as comments in the impl itself for more context. So that if someone else (or even yourself) comes back later to try to change this impl, they are better equipped to figure out what this is doing.

EDIT: please ignore panic! -> bug! suggestions as that might not be available yet in macro expansion here.

[jieyouxu]

That doesn't seem great, since it fails in the post-optimization test suite. If this PR really did not change anything that is actually used, that might signify some miscompilation in PGO, or something sinister like that.

This is what I'm concerned about too, because AFAICT this PR really only touches the front-end things, and the code paths are gated behind unstable feature gates. Nothing is sticking out to me as something that can explode compile times.

[jieyouxu]

In theory, if this part is wrong, then we may generate a test that never terminates due to UB. I'm not sure exactly on which test (suite) the job gets stuck on. It seems unlikely for crashes tests because they don't enable the autodiff feature gates?

[jieyouxu]

That doesn't seem great, since it fails in the post-optimization test suite. If this PR really did not change anything that is actually used, that might signify some miscompilation in PGO, or something sinister like that.

Do you have any "insider" infra info on which test suite this job gets stuck on? And do you know if any of the CI jobs have enzyme available (esp. dist-x86_64-msvc)?

[Kobzol]

Hmm, I only have the same info as you, the log of the CI run. From that it's pretty clear where it is wrong - it's Testing stage0 compiletest suite=crashes mode=crashes, and specifically the test tests\crashes\23707.rs, which does not seem to terminate.

[jieyouxu]

That's the funny part. That crashes test looks like a trait system thing, which this PR doesn't really touch. Worst part is I can't repro the failure locally (it passes just fine).

[Kobzol]

Might be interesting to see if this also happens on Linux. You'd need to disable DIST_TRY_BUILD in jobs.yml and then do a normal try build.

[jieyouxu]

Doing that over in #131685.

[jieyouxu]

Didn't fail in try job. So what if we try again. Surely changing nothing will lead to a different outcome.

@bors r+

[bors]

💡 This pull request was already approved, no need to approve it again.

• There's another pull request that is currently being tested, blocking this pull request: Rollup of 9 pull requests #131711

[bors]

📌 Commit 7c37d2d has been approved by jieyouxu

It is now in the queue for this repository.

[bors]

⌛ Testing commit 7c37d2d with merge 785c830...

[bors]

☀️ Test successful - checks-actions
Approved by: jieyouxu
Pushing 785c830 to master...

[jieyouxu]

...? ok bors whatever u say

[rust-timer]

Finished benchmarking commit (785c830): comparison URL.

Overall result: ✅ improvements - no action needed

@rustbot label: -perf-regression

Instruction count

This is the most reliable metric that we have; it was used to determine the overall result at the top of this comment. However, even this metric can sometimes exhibit noise.

	mean	range	count
Regressions ❌ (primary)	-	-	0
Regressions ❌ (secondary)	-	-	0
Improvements ✅ (primary)	-	-	0
Improvements ✅ (secondary)	-0.5%	[-0.8%, -0.2%]	2
All ❌✅ (primary)	-	-	0

Max RSS (memory usage)

Results (secondary 1.3%)

This is a less reliable metric that may be of interest but was not used to determine the overall result at the top of this comment.

	mean	range	count
Regressions ❌ (primary)	-	-	0
Regressions ❌ (secondary)	4.0%	[2.2%, 6.6%]	5
Improvements ✅ (primary)	-	-	0
Improvements ✅ (secondary)	-3.1%	[-5.6%, -1.9%]	3
All ❌✅ (primary)	-	-	0

Cycles

This benchmark run did not return any relevant results for this metric.

Binary size

Results (primary 0.5%)

This is a less reliable metric that may be of interest but was not used to determine the overall result at the top of this comment.

	mean	range	count
Regressions ❌ (primary)	0.5%	[0.5%, 0.5%]	1
Regressions ❌ (secondary)	-	-	0
Improvements ✅ (primary)	-	-	0
Improvements ✅ (secondary)	-	-	0
All ❌✅ (primary)	0.5%	[0.5%, 0.5%]	1

Bootstrap: 781.98s -> 782.278s (0.04%)
Artifact size: 332.66 MiB -> 332.62 MiB (-0.01%)