-
Notifications
You must be signed in to change notification settings - Fork 12.5k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Autodiff Upstreaming - enzyme frontend #129458
base: master
Are you sure you want to change the base?
Conversation
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
☔ The latest upstream changes (presumably #129563) made this pull request unmergeable. Please resolve the merge conflicts. |
let mut attr: ast::Attribute = ast::Attribute { | ||
kind: ast::AttrKind::Normal(rustc_ad_attr.clone()), | ||
//id: ast::DUMMY_TR_ID, | ||
id: ast::AttrId::from_u32(12341), // TODO: fix |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I would be happy for suggestions on what ID to use or create here.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Discussion (for other reviewers as well) [REV-ATTR-ID]: I feel like this needs to be a fresh AttrId
, like via ecx.sess.psess.attr_id_generator.mk_attr_id()
, otherwise we might get very kaboom behavior if the AttrId
s collide.
rust/compiler/rustc_ast/src/attr/mod.rs
Lines 48 to 53 in 0d63418
pub fn mk_attr_id(&self) -> AttrId { | |
let id = self.0.fetch_add(1, Ordering::Relaxed); | |
assert!(id != u32::MAX); | |
AttrId::from_u32(id) | |
} | |
} |
But I'm not too sure if this is right.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
compiler/rustc_ast/src/ast.rs
Outdated
@@ -2733,6 +2733,13 @@ impl FnRetTy { | |||
FnRetTy::Ty(ty) => ty.span, | |||
} | |||
} | |||
|
|||
pub fn has_ret(&self) -> bool { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Discussion [REV-1]: I find this somewhat misleading, because unspecified return type can still have non-unit ()
return types (e.g. closure type inference). Maybe has_specified_ret_ty
or something... What's the significance of distinguishing between specified and non-specified return types?
rust/compiler/rustc_ast/src/ast.rs
Lines 2719 to 2724 in 0d63418
pub enum FnRetTy { | |
/// Returns type is not specified. | |
/// | |
/// Functions default to `()` and closures default to inference. | |
/// Span points to where return type would be inserted. | |
Default(Span), |
I need to check what the call-site of this actually wants, seems a bit strange to want this.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is very much sketchy 👍
Ideally don't distinguish between -> ()
and no return, because in Rust these are precisely equal, but in any case, this should not be a helper on the AST that other users may accidentally use since it's 99.9% wrong.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't think closures are relevant enough to spend time on adding support for it, my macro rejects everything but functions.
I don't care about the difference between -> () and the implicit () return. It looks like it worked by coincidence before correctly, but I've added a check for explicit -> () return and also moved the function out of ast.rs
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Question: does it matter to autodiff then if the return type is a type-alias-to-unit or effectively-unit?
// e.g.
type Unit = ();
// or
#[repr(transparent)]
pub struct UnitButFancy { inner: () }
// or
pub struct Sandwich(());
Then the user may have written
pub fn foo() -> Unit /* or `UnitButFancy` or `Sandwich` */ {}
Would this be a problem?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't really know how the implementation would behave here.
The current mvp strongly relies on the user not doing "strange" things, I don't see this experiment as anywhere close of handling type aliases or malicious users correctly. Sicne there are still bugs left that users in good faith will run into, I'd just leave this as an unhandled case. This type doesn't have any floats in it so const would be the only reasonable annotation, but I don't think spending time on catching this would be a good investment.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
That seems fine, probably worth remarking this as a comment.
} | ||
} | ||
|
||
pub fn valid_ret_activity(mode: DiffMode, activity: DiffActivity) -> bool { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Question [REV-CTXT]: I can of course read the implementation below, but for someone in the future who's trying to change this code, what's the high-level description of what constitutes a "valid" ret activity? What's a DiffActivity
? Why is e.g. DiffMode::Forward
+ activity == DiffActivity::Dual
a valid ret activity but DiffMode::Forward
+ activity == DiffActivity::Const
not?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Suggestion: could we hoist definition of DiffActivity
and perhaps add some explanation for its purpose before its usage here? I think that would make it easier to follow.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Both of your examples are allowed, did you mean something else? I also have added some more docs and bundled all enum/struct definitions at the top, does that help?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Both of your examples are allowed, did you mean something else?
Err probably, substitute whatever I wrote with a combination that's not allowed (if there's any). In any case, could some invalid combination cause the enzyme backend to kaboom if we don't catch them here?
I also have added some more docs and bundled all enum/struct definitions at the top, does that help?
Yes that seems fine.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, every single invalid combination that we don't catch will make Enzyme go kaboom. If we're lucky we just run into an enzyme assertion with a more or less helpful error, in the worstcase we get a segfault.
tests/pretty/autodiff_reverse.rs
Outdated
|
||
// Test that reverse mode ad macros are expanded correctly. | ||
|
||
#[autodiff(df, Reverse, Duplicated, Const, Active)] |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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 adf
in value namespace?#[autodiff(df, Reverse)]
+enum Foo { Reverse }
+use Foo::Reverse;
<- what if I already have aReverse
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
.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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?
fb5fa71
to
9d8ec28
Compare
9d8ec28
to
a989f28
Compare
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
working my way through it. Already lead to some nice improvement, because I now return incorrect function signatures for df if I encounter a usage error, knowing that compilation will abort later anyway. By not erroring immediately we can catch more errors in one go.
|
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
The job Click to see the possible cause of the failure (guessed by this bot)
|
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 functiondf
(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. Iff
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:
to
I will add a few more tests once I figured out why rustc rebuilds every time I touch a test.
Tracking: