Rollup merge of #140697 - Sa4dUs:split-autodiff, r=ZuseZ4

Split `autodiff` into `autodiff_forward` and `autodiff_reverse`
This PR splits `#[autodiff]` macro so `#[autodiff(df, Reverse, args)]` would become `#[autodiff_reverse(df, args)]` and `#[autodiff(df, Forward, args)]` would become `#[autodiff_forwad(df, args)]`.

Author: tgross35

compiler/rustc_builtin_macros/messages.ftl

@@ -56,7 +56,6 @@ builtin_macros_assert_requires_expression = macro requires an expression as an a
 
 builtin_macros_autodiff = autodiff must be applied to function
 builtin_macros_autodiff_missing_config = autodiff requires at least a name and mode
-builtin_macros_autodiff_mode = unknown Mode: `{$mode}`. Use `Forward` or `Reverse`
 builtin_macros_autodiff_mode_activity = {$act} can not be used in {$mode} Mode
 builtin_macros_autodiff_not_build = this rustc version does not support autodiff
 builtin_macros_autodiff_number_activities = expected {$expected} activities, but found {$found}

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -86,27 +86,23 @@ mod llvm_enzyme {
 ecx:&mutExtCtxt<'_>,
 meta_item:&ThinVec<MetaItemInner>,
 has_ret:bool,
+mode:DiffMode,
 ) -> AutoDiffAttrs{
 let dcx = ecx.sess.dcx();
-let mode = name(&meta_item[1]);
-letOk(mode) = DiffMode::from_str(&mode)else{
- dcx.emit_err(errors::AutoDiffInvalidMode{span: meta_item[1].span(), mode });
-returnAutoDiffAttrs::error();
-};
 
 // Now we check, whether the user wants autodiff in batch/vector mode, or scalar mode.
 // If he doesn't specify an integer (=width), we default to scalar mode, thus width=1.
-letmut first_activity = 2;
+letmut first_activity = 1;
 
-let width = iflet[_,_,x, ..] = &meta_item[..]
+let width = iflet[_, x, ..] = &meta_item[..]
  && letSome(x) = width(x)
 {
- first_activity = 3;
+ first_activity = 2;
 match x.try_into(){
 Ok(x) => x,
 Err(_) => {
  dcx.emit_err(errors::AutoDiffInvalidWidth{
-span: meta_item[2].span(),
+span: meta_item[1].span(),
 width: x,
 });
 returnAutoDiffAttrs::error();
@@ -165,6 +161,24 @@ mod llvm_enzyme {
  ts.push(TokenTree::Token(comma.clone(),Spacing::Alone));
 }
 
+pub(crate)fnexpand_forward(
+ecx:&mutExtCtxt<'_>,
+expand_span:Span,
+meta_item:&ast::MetaItem,
+item:Annotatable,
+) -> Vec<Annotatable>{
+expand_with_mode(ecx, expand_span, meta_item, item,DiffMode::Forward)
+}
+
+pub(crate)fnexpand_reverse(
+ecx:&mutExtCtxt<'_>,
+expand_span:Span,
+meta_item:&ast::MetaItem,
+item:Annotatable,
+) -> Vec<Annotatable>{
+expand_with_mode(ecx, expand_span, meta_item, item,DiffMode::Reverse)
+}
+
 /// We expand the autodiff macro to generate a new placeholder function which passes
 /// type-checking and can be called by users. The function body of the placeholder function will
 /// later be replaced on LLVM-IR level, so the design of the body is less important and for now
@@ -198,11 +212,12 @@ mod llvm_enzyme {
 /// ```
 /// FIXME(ZuseZ4): Once autodiff is enabled by default, make this a doc comment which is checked
 /// in CI.
-pub(crate)fnexpand(
+pub(crate)fnexpand_with_mode(
 ecx:&mutExtCtxt<'_>,
 expand_span:Span,
 meta_item:&ast::MetaItem,
 mutitem:Annotatable,
+mode:DiffMode,
 ) -> Vec<Annotatable>{
 ifcfg!(not(llvm_enzyme)){
  ecx.sess.dcx().emit_err(errors::AutoDiffSupportNotBuild{span: meta_item.span});
@@ -245,29 +260,41 @@ mod llvm_enzyme {
 // create TokenStream from vec elemtents:
 // meta_item doesn't have a .tokens field
 letmut ts:Vec<TokenTree> = vec![];
-if meta_item_vec.len() < 2{
-// At the bare minimum, we need a fnc name and a mode, even for a dummy function with no
-// input and output args.
+if meta_item_vec.len() < 1{
+// At the bare minimum, we need a fnc name.
  dcx.emit_err(errors::AutoDiffMissingConfig{span: item.span()});
 returnvec![item];
 }
 
-meta_item_inner_to_ts(&meta_item_vec[1],&mut ts);
+let mode_symbol = match mode {
+DiffMode::Forward => sym::Forward,
+DiffMode::Reverse => sym::Reverse,
+ _ => unreachable!("Unsupported mode: {:?}", mode),
+};
+
+// Insert mode token
+let mode_token = Token::new(TokenKind::Ident(mode_symbol,false.into()),Span::default());
+ ts.insert(0,TokenTree::Token(mode_token,Spacing::Joint));
+ ts.insert(
+1,
+TokenTree::Token(Token::new(TokenKind::Comma,Span::default()),Spacing::Alone),
+);
 
 // Now, if the user gave a width (vector aka batch-mode ad), then we copy it.
 // If it is not given, we default to 1 (scalar mode).
 let start_position;
 let kind:LitKind = LitKind::Integer;
 let symbol;
-if meta_item_vec.len() >= 3
- && letSome(width) = width(&meta_item_vec[2])
+if meta_item_vec.len() >= 2
+ && letSome(width) = width(&meta_item_vec[1])
 {
- start_position = 3;
+ start_position = 2;
  symbol = Symbol::intern(&width.to_string());
 }else{
- start_position = 2;
+ start_position = 1;
  symbol = sym::integer(1);
 }
+
 let l:Lit = Lit{ kind, symbol,suffix:None};
 let t = Token::new(TokenKind::Literal(l),Span::default());
 let comma = Token::new(TokenKind::Comma,Span::default());
@@ -289,7 +316,7 @@ mod llvm_enzyme {
  ts.pop();
 let ts:TokenStream = TokenStream::from_iter(ts);
 
-let x:AutoDiffAttrs = from_ast(ecx,&meta_item_vec, has_ret);
+let x:AutoDiffAttrs = from_ast(ecx,&meta_item_vec, has_ret, mode);
 if !x.is_active(){
 // We encountered an error, so we return the original item.
 // This allows us to potentially parse other attributes.
@@ -1017,4 +1044,4 @@ mod llvm_enzyme {
 }
 }
 
-pub(crate)use llvm_enzyme::expand;
+pub(crate)use llvm_enzyme::{expand_forward, expand_reverse};

compiler/rustc_builtin_macros/src/errors.rs

@@ -180,14 +180,6 @@ mod autodiff {
 pub(crate)act:String,
 }
 
-#[derive(Diagnostic)]
-#[diag(builtin_macros_autodiff_mode)]
-pub(crate)structAutoDiffInvalidMode{
-#[primary_span]
-pub(crate)span:Span,
-pub(crate)mode:String,
-}
-
 #[derive(Diagnostic)]
 #[diag(builtin_macros_autodiff_width)]
 pub(crate)structAutoDiffInvalidWidth{

compiler/rustc_builtin_macros/src/lib.rs

@@ -5,10 +5,10 @@
 #![allow(internal_features)]
 #![allow(rustc::diagnostic_outside_of_impl)]
 #![allow(rustc::untranslatable_diagnostic)]
+#![cfg_attr(not(bootstrap), feature(autodiff))]
 #![doc(html_root_url = "https://doc.rust-lang.org/nightly/nightly-rustc/")]
 #![doc(rust_logo)]
 #![feature(assert_matches)]
-#![feature(autodiff)]
 #![feature(box_patterns)]
 #![feature(decl_macro)]
 #![feature(if_let_guard)]
@@ -112,7 +112,8 @@ pub fn register_builtin_macros(resolver: &mut dyn ResolverExpand) {
 
  register_attr! {
  alloc_error_handler: alloc_error_handler::expand,
- autodiff: autodiff::expand,
+ autodiff_forward: autodiff::expand_forward,
+ autodiff_reverse: autodiff::expand_reverse,
  bench: test::expand_bench,
  cfg_accessible: cfg_accessible::Expander,
  cfg_eval: cfg_eval::expand,

compiler/rustc_passes/src/check_attr.rs

@@ -255,7 +255,7 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
 self.check_generic_attr(hir_id, attr, target,Target::Fn);
 self.check_proc_macro(hir_id, target,ProcMacroKind::Derive)
 }
-[sym::autodiff, ..] => {
+[sym::autodiff_forward, ..] | [sym::autodiff_reverse, ..] => {
 self.check_autodiff(hir_id, attr, span, target)
 }
 [sym::coroutine, ..] => {

compiler/rustc_span/src/symbol.rs

@@ -244,6 +244,7 @@ symbols! {
 FnMut,
 FnOnce,
 Formatter,
+Forward,
 From,
 FromIterator,
 FromResidual,
@@ -339,6 +340,7 @@ symbols! {
 Result,
 ResumeTy,
 Return,
+Reverse,
 Right,
 Rust,
 RustaceansAreAwesome,
@@ -522,7 +524,8 @@ symbols! {
  audit_that,
  augmented_assignments,
  auto_traits,
- autodiff,
+ autodiff_forward,
+ autodiff_reverse,
  automatically_derived,
  avx,
  avx10_target_feature,

library/core/src/lib.rs

@@ -225,10 +225,11 @@ pub mod assert_matches {
 
 // We don't export this through #[macro_export] for now, to avoid breakage.
 #[unstable(feature = "autodiff", issue = "124509")]
+#[cfg(not(bootstrap))]
 /// Unstable module containing the unstable `autodiff` macro.
 pubmod autodiff {
 #[unstable(feature = "autodiff", issue = "124509")]
-pubusecrate::macros::builtin::autodiff;
+pubusecrate::macros::builtin::{autodiff_forward, autodiff_reverse};
 }
 
 #[unstable(feature = "contracts", issue = "128044")]

library/core/src/macros/mod.rs

@@ -1519,20 +1519,41 @@ pub(crate) mod builtin {
 ($file:expr $(,)?) => {{/* compiler built-in */}};
 }
 
-/// Automatic Differentiation macro which allows generating a new function to compute
-/// the derivative of a given function. It may only be applied to a function.
-/// The expected usage syntax is
-/// `#[autodiff(NAME, MODE, INPUT_ACTIVITIES, OUTPUT_ACTIVITY)]`
-/// where:
-/// NAME is a string that represents a valid function name.
-/// MODE is any of Forward, Reverse, ForwardFirst, ReverseFirst.
-/// INPUT_ACTIVITIES consists of one valid activity for each input parameter.
-/// OUTPUT_ACTIVITY must not be set if we implicitly return nothing (or explicitly return
-/// `-> ()`). Otherwise it must be set to one of the allowed activities.
+/// This macro uses forward-mode automatic differentiation to generate a new function.
+/// It may only be applied to a function. The new function will compute the derivative
+/// of the function to which the macro was applied.
+///
+/// The expected usage syntax is:
+/// `#[autodiff_forward(NAME, INPUT_ACTIVITIES, OUTPUT_ACTIVITY)]`
+///
+/// - `NAME`: A string that represents a valid function name.
+/// - `INPUT_ACTIVITIES`: Specifies one valid activity for each input parameter.
+/// - `OUTPUT_ACTIVITY`: Must not be set if the function implicitly returns nothing
+/// (or explicitly returns `-> ()`). Otherwise, it must be set to one of the allowed activities.
+ #[unstable(feature = "autodiff", issue = "124509")]
+ #[allow_internal_unstable(rustc_attrs)]
+ #[rustc_builtin_macro]
+ #[cfg(not(bootstrap))]
+pub macro autodiff_forward($item:item){
+/* compiler built-in */
+}
+
+/// This macro uses reverse-mode automatic differentiation to generate a new function.
+/// It may only be applied to a function. The new function will compute the derivative
+/// of the function to which the macro was applied.
+///
+/// The expected usage syntax is:
+/// `#[autodiff_reverse(NAME, INPUT_ACTIVITIES, OUTPUT_ACTIVITY)]`
+///
+/// - `NAME`: A string that represents a valid function name.
+/// - `INPUT_ACTIVITIES`: Specifies one valid activity for each input parameter.
+/// - `OUTPUT_ACTIVITY`: Must not be set if the function implicitly returns nothing
+/// (or explicitly returns `-> ()`). Otherwise, it must be set to one of the allowed activities.
  #[unstable(feature = "autodiff", issue = "124509")]
  #[allow_internal_unstable(rustc_attrs)]
  #[rustc_builtin_macro]
-pub macro autodiff($item:item){
+ #[cfg(not(bootstrap))]
+pub macro autodiff_reverse($item:item){
 /* compiler built-in */
 }
 

library/std/src/lib.rs

@@ -276,12 +276,12 @@
 // tidy-alphabetical-start
 
 // stabilization was reverted after it hit beta
+#![cfg_attr(not(bootstrap), feature(autodiff))]
 #![feature(alloc_error_handler)]
 #![feature(allocator_internals)]
 #![feature(allow_internal_unsafe)]
 #![feature(allow_internal_unstable)]
 #![feature(asm_experimental_arch)]
-#![feature(autodiff)]
 #![feature(cfg_sanitizer_cfi)]
 #![feature(cfg_target_thread_local)]
 #![feature(cfi_encoding)]
@@ -636,12 +636,15 @@ pub mod simd {
 #[doc(inline)]
 pubusecrate::std_float::StdFloat;
 }
+
 #[unstable(feature = "autodiff", issue = "124509")]
+#[cfg(not(bootstrap))]
 /// This module provides support for automatic differentiation.
 pubmod autodiff {
 /// This macro handles automatic differentiation.
-pubuse core::autodiff::autodiff;
+pubuse core::autodiff::{autodiff_forward, autodiff_reverse};
 }
+
 #[stable(feature = "futures_api", since = "1.36.0")]
 pubmod task {
 //! Types and Traits for working with asynchronous tasks.

tests/codegen/autodiff/batched.rs

@@ -11,11 +11,11 @@
 
 #![feature(autodiff)]
 
-use std::autodiff::autodiff;
+use std::autodiff::autodiff_forward;
 
-#[autodiff(d_square3,Forward,Dual,DualOnly)]
-#[autodiff(d_square2,Forward,4,Dual,DualOnly)]
-#[autodiff(d_square1,Forward,4,Dual,Dual)]
+#[autodiff_forward(d_square3,Dual,DualOnly)]
+#[autodiff_forward(d_square2,4,Dual,DualOnly)]
+#[autodiff_forward(d_square1,4,Dual,Dual)]
 #[no_mangle]
 fnsquare(x:&f32) -> f32{
  x * x