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

[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.

[ZuseZ4]

@jieyouxu Now that the test infra has landed, is there anything left to do except updating the tests?
I gave TC a link to your comment here (#129458 (comment)) and he seemed to be fine with it. Shall we ping someone else?

[ZuseZ4]

I added the autodiff gating for all test files and moved the failing one to UI, they are not executed now without AD.
Also, I changed it to not be build by default, but that also means CI won't build / test this code anymore, but that's fine for now. I also fixed the UI / stderr test.

I also just added my former collaborator as co-author to the PR, so from my side it's ready.

[ZuseZ4]

There isn't a field to answer, so I'll say it here, yep I made two modules out of it @jieyouxu : #129458 (comment)

[bors]

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

[jieyouxu]

Thanks for the updates, this looks to be in a much better shape than before. I've left a ~final round of reviews, and this prototype looks like it's suitable for merge after addressing those. The diagnostics for invalid attributes seem somewhat rough, but the polishing can be improved in the future and is not blocking. The test coverage regarding to all the branches in the macro logic isn't very thorough, but that's fine for initial upstreaming of the prototype for now (but definitely will need to be improved if this ever wants to be stabilized).

[jieyouxu]

Remark: "failed to resolve" is a strange error here?

[ZuseZ4]

Ah, I missed answering that. Generally agreed, but as long as it errors I'm fine as it's one of the things that will require TypeInfo to catch all invalid usage combinations. Then again I think that if it wouldn't error Enzyme would handle it correctly, so if this bug magically disappeared then I can probably remove the error message and move it to the pretty tests.

[ZuseZ4]

@jieyouxu the last force-push only updated docs (for the -1 typetree info and to include an example of how expansion will look like), so I will assume that this should pass CI again. As before, let me know if I missed something, but I think I should have everything covered this time. Thanks a lot for taking the time to review this. The middle end should be fairly small from here on, and the rustc_cg_llvm changes are fairly big, but are also pretty straightforward due to including bindings and only forwarding rust configs to Enzyme configs (even though reviewers currently don't seem to be keen on it), so hopefully they will be a bit easier to handle.
I'll also write a small follow-up PR to change how we let users decide whether a function should return the original return value too, or only the gradient. It simplifies both the implementation and usability a bit (since it matches other autodiff tools like JAX). But that will again touch a lot of files (since as you mentioned in your proposal parsing rn is quite split up), so I don't want to have it in this PR. For most of the other smaller FIXMEs which don't require Enzyme core changes I hope to be able to pass it to interested potential contributors to get them onboarded, in which case I could even do the first round of reviews. Quite a few people already asked how they can help this project, I obviously just didn't want someone new to the project to have to touch 23 files like in this PR to get started.

[jieyouxu]

👍 That plan sounds fair to me.

[jieyouxu]

This looks mostly good to me. I just noticed one thing, please add a feature gate ui test to make sure we don't ICE or do funny things.

#![crate_type = "lib"]
#[autodiff]
//~^ ERROR cannot find attribute `autodiff` in this scope
//~| ERROR: `#[autodiff]` is experimental 
fn foo() {}

e.g.

rust/tests/ui/feature-gates/feature-gate-doc_notable_trait.rs

Line 1 in 6f4ae0f

#[doc(notable_trait)] //~ ERROR: `#[doc(notable_trait)]` is experimental