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.

[jieyouxu]

Problem && Suggestion: I think we want at least test cases for:

• Positive test cases for what autodiff should be applied to.
• Negative test cases for when autodiff is misapplied to AST nodes and their diagnostics. In particular, closures, statements and expressions.
• #[autodiff] malformed attribute (where args?) error
• #[autodiff = ""] invalid attribute syntax
• #[autodiff()] where arg
• #[autodiff(df)] + fn df() <- what if I already have a df in value namespace?
• #[autodiff(df, Reverse)] + enum Foo { Reverse } + use Foo::Reverse; <- what if I already have a Reverse in type (enum variant decl) and value (enum variant constructor) namespace?
• #[autodiff(df)] <- is this minimally acceptable?
• #[autodiff(df, Reverse)] <- valid mode
• #[autodiff(df, Debug)] <- invalid mode
• #[autodiff(df, Forward, Reverse)] <- is this valid
• target fn has specified valid return type e.g. -> f64
• target fn has unspecified return type fn foo() {}
• target fn has specified but invalid return type e.g. -> Owo
• target fn has aliased f32/f64 return types (currently unsupported) -> F64Alias
• target fn has #[repr(transparent)] struct F64Trans { inner: f64 } return type (currently unsupported) -> F64Trans.

[ZuseZ4]

All of these cases now give error messages instead of ICEs.
I also made sure that we don't return after the first error, but continue to do something sufficiently okish that we can first expand all autodiff macros before we abort compilation.
We probably could include even better errors in the future, but I left those as fixme's for now.
Do they look good to you?

[ZuseZ4]

Well I'll go ahead and ping @ pnkfelix. This is the frontend for the new, experimental autodiff macro. The previous and the first comment have a summary of this work. Please let me know if you want any changes!

[jieyouxu]

(Actually I forgor to unassign)
r? jieyouxu

[ZuseZ4]

Ah, I wasn't aware that you also intended to approve it once it's ready, that makes things simpler.

[jieyouxu]

Thanks, I did another pass, and now most of the changes look ready for merge as a prototype. I tried to keep to the more significant feedback, further improvements can be done in the future.

[jieyouxu]

Question: what's the significance of offset being -1?

[ZuseZ4]

That's the bane of my existence. :D Enzyme uses a type layout representation which is suboptimal, it should move over to type trie's, but another student and I both gave up on our refactor PRs and the Enzyme core Lead dev also isn't actively working this. Luckily there are some AD users which seem to have the resources to eventually rewrite the whole typetree infrastructure, so I consider it just an implementation detail.

In less dramatic, -1 in Enzyme speech means everywhere.
that is {0:-1: Float} means at index 0 you have a ptr, if you dereference it it will be floats everywhere. Thus * f32.
If you have {-1:int} it means int's everywhere, e.g. [i32; N].
{0:-1:-1 float} then means one pointer at offset 0, if you dereference it there will be only pointers, if you dereference these new pointers they will point to array of floats.
Generally, it allows byte-specific descriptions.

This design has no way of handling recursive datastructures, it skips things that have more than 5 indirections and there are some hacks to make it handle gaps in layouts, as well as other issues with it. If Enzyme is slow at compile time than this is usually the culprit. Also it should be extended at some point to make use of const/mut knowledge, right now that get's lost once we have more than one indirection. To be fair I find it pretty cool that Enzyme is already so extremely fast while leaving some information like here still unused.

The middle-end PR will include tests for typetrees, since that's where we construct them.
I'll add this also to the docs, so I can link to them next time.

[jieyouxu]

Problem: we need to remember to flip these cfg(not(llvm_enzyme))? I still need a bit of help with understanding what these cfg(llvm_enzyme) do, and what are they supposed to be cfg() or cfg(not()) when this is intended to be fully ready to merge.

[jieyouxu]

Question: instead of guarding individual items, can this entire autodiff module be guarded?

[jieyouxu]

Discussion: the libs changes I'm not too sure about, is this needed for #[autodiff] for name resolution purposes? But this is not provided through the prelude right?

[ZuseZ4]

Since I can't answer directly:

Question: instead of guarding individual items, can this entire autodiff module be guarded?
It currently seems easier to have a dummy implementation which doesn't generate a placeholder function for df and give a nice error message.

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