Prevent ICE at codegen_attrs.rs in autodiff validation by emitting nice error message

compiler/rustc_builtin_macros/messages.ftl

@@ -75,8 +75,10 @@ builtin_macros_autodiff_mode = unknown Mode: `{$mode}`. Use `Forward` or `Revers
 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}
+builtin_macros_autodiff_ret_activity = invalid return activity {$act} in {$mode} Mode
 builtin_macros_autodiff_ty_activity = {$act} can not be used for this type
 
+
 builtin_macros_autodiff_unknown_activity = did not recognize Activity: `{$act}`
 builtin_macros_bad_derive_target = `derive` may only be applied to `struct`s, `enum`s and `union`s
     .label = not applicable here

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -8,7 +8,8 @@ mod llvm_enzyme {
     use std::string::String;
 
     use rustc_ast::expand::autodiff_attrs::{
-        AutoDiffAttrs, DiffActivity, DiffMode, valid_input_activity, valid_ty_for_activity,
+        AutoDiffAttrs, DiffActivity, DiffMode, valid_input_activity, valid_ret_activity,
+        valid_ty_for_activity,
     };
     use rustc_ast::ptr::P;
     use rustc_ast::token::{Token, TokenKind};
@@ -632,10 +633,22 @@ mod llvm_enzyme {
                 errors = true;
             }
         }
+
+        if has_ret && !valid_ret_activity(x.mode, x.ret_activity) {
+            dcx.emit_err(errors::AutoDiffInvalidRetAct {
+                span,
+                mode: x.mode.to_string(),
+                act: x.ret_activity.to_string(),
+            });
+            // We don't set `errors = true` to avoid annoying type errors relative
+            // to the expanded macro type signature
+        }
+
         if errors {
             // This is not the right signature, but we can continue parsing.
             return (sig.clone(), new_inputs, idents, true);
         }
+
         let unsafe_activities = x
             .input_activity
             .iter()

compiler/rustc_builtin_macros/src/errors.rs

@@ -185,6 +185,15 @@ mod autodiff {
         pub(crate) act: String,
     }
 
+    #[derive(Diagnostic)]
+    #[diag(builtin_macros_autodiff_ret_activity)]
+    pub(crate) struct AutoDiffInvalidRetAct {
+        #[primary_span]
+        pub(crate) span: Span,
+        pub(crate) mode: String,
+        pub(crate) act: String,
+    }
+
     #[derive(Diagnostic)]
     #[diag(builtin_macros_autodiff_mode)]
     pub(crate) struct AutoDiffInvalidMode {

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -1,9 +1,7 @@
 use std::str::FromStr;
 
 use rustc_abi::ExternAbi;
-use rustc_ast::expand::autodiff_attrs::{
-    AutoDiffAttrs, DiffActivity, DiffMode, valid_input_activity, valid_ret_activity,
-};
+use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, DiffActivity, DiffMode};
 use rustc_ast::{MetaItem, MetaItemInner, attr};
 use rustc_attr_parsing::ReprAttr::ReprAlign;
 use rustc_attr_parsing::{AttributeKind, InlineAttr, InstructionSetAttr, OptimizeAttr};
@@ -930,15 +928,6 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> Option<AutoDiffAttrs> {
         }
     }
 
-    for &input in &arg_activities {
-        if !valid_input_activity(mode, input) {
-            span_bug!(attr.span(), "Invalid input activity {} for {} mode", input, mode);
-        }
-    }
-    if !valid_ret_activity(mode, ret_activity) {
-        span_bug!(attr.span(), "Invalid return activity {} for {} mode", ret_activity, mode);
-    }
-
     Some(AutoDiffAttrs { mode, ret_activity, input_activity: arg_activities })
 }
 

tests/ui/autodiff/autodiff_illegal.rs

@@ -12,41 +12,40 @@ use std::autodiff::autodiff;
 // We can't use Duplicated on scalars
 #[autodiff(df1, Reverse, Duplicated)]
 pub fn f1(x: f64) {
-//~^ ERROR     Duplicated can not be used for this type
+    //~^ ERROR     Duplicated can not be used for this type
     unimplemented!()
 }
 
 // Too many activities
 #[autodiff(df3, Reverse, Duplicated, Const)]
 pub fn f3(x: f64) {
-//~^^ ERROR     expected 1 activities, but found 2
+    //~^^ ERROR     expected 1 activities, but found 2
     unimplemented!()
 }
 
 // To few activities
 #[autodiff(df4, Reverse)]
 pub fn f4(x: f64) {
-//~^^ ERROR     expected 1 activities, but found 0
+    //~^^ ERROR     expected 1 activities, but found 0
     unimplemented!()
 }
 
 // We can't use Dual in Reverse mode
 #[autodiff(df5, Reverse, Dual)]
 pub fn f5(x: f64) {
-//~^^ ERROR     Dual can not be used in Reverse Mode
+    //~^^ ERROR     Dual can not be used in Reverse Mode
     unimplemented!()
 }
 
 // We can't use Duplicated in Forward mode
 #[autodiff(df6, Forward, Duplicated)]
 pub fn f6(x: f64) {
-//~^^ ERROR Duplicated can not be used in Forward Mode
-//~^^ ERROR Duplicated can not be used for this type
+    //~^^ ERROR Duplicated can not be used in Forward Mode
+    //~^^ ERROR Duplicated can not be used for this type
     unimplemented!()
 }
 
 fn dummy() {
-
     #[autodiff(df7, Forward, Dual)]
     let mut x = 5;
     //~^ ERROR autodiff must be applied to function
@@ -64,21 +63,21 @@ fn dummy() {
 // Malformed, where args?
 #[autodiff]
 pub fn f7(x: f64) {
-//~^ ERROR autodiff must be applied to function
+    //~^ ERROR autodiff must be applied to function
     unimplemented!()
 }
 
 // Malformed, where args?
 #[autodiff()]
 pub fn f8(x: f64) {
-//~^ ERROR autodiff requires at least a name and mode
+    //~^ ERROR autodiff requires at least a name and mode
     unimplemented!()
 }
 
 // Invalid attribute syntax
 #[autodiff = ""]
 pub fn f9(x: f64) {
-//~^ ERROR autodiff must be applied to function
+    //~^ ERROR autodiff must be applied to function
     unimplemented!()
 }
 
@@ -87,29 +86,33 @@ fn fn_exists() {}
 // We colide with an already existing function
 #[autodiff(fn_exists, Reverse, Active)]
 pub fn f10(x: f64) {
+<<<<<<< HEAD
+=======
 //~^^ ERROR the name `fn_exists` is defined multiple times [E0428]
+>>>>>>> 465e6968091de8a29af3d5315c8517b264c7df64
+    //~^^ ERROR the name `fn_exists` is defined multiple times [E0428]
     unimplemented!()
 }
 
 // Malformed, missing a mode
 #[autodiff(df11)]
 pub fn f11() {
-//~^ ERROR autodiff requires at least a name and mode
+    //~^ ERROR autodiff requires at least a name and mode
     unimplemented!()
 }
 
 // Invalid Mode
 #[autodiff(df12, Debug)]
 pub fn f12() {
-//~^^ ERROR unknown Mode: `Debug`. Use `Forward` or `Reverse`
+    //~^^ ERROR unknown Mode: `Debug`. Use `Forward` or `Reverse`
     unimplemented!()
 }
 
 // Invalid, please pick one Mode
 // or use two autodiff macros.
 #[autodiff(df13, Forward, Reverse)]
 pub fn f13() {
-//~^^ ERROR did not recognize Activity: `Reverse`
+    //~^^ ERROR did not recognize Activity: `Reverse`
     unimplemented!()
 }
 
@@ -130,7 +133,7 @@ type MyFloat = f32;
 // like THIR which has type information available.
 #[autodiff(df15, Reverse, Active, Active)]
 fn f15(x: MyFloat) -> f32 {
-//~^^ ERROR failed to resolve: use of undeclared type `MyFloat` [E0433]
+    //~^^ ERROR failed to resolve: use of undeclared type `MyFloat` [E0433]
     unimplemented!()
 }
 
@@ -141,7 +144,9 @@ fn f16(x: f32) -> MyFloat {
 }
 
 #[repr(transparent)]
-struct F64Trans { inner: f64 }
+struct F64Trans {
+    inner: f64,
+}
 
 // We would like to support `#[repr(transparent)]` f32/f64 wrapper in return type in the future
 #[autodiff(df17, Reverse, Active, Active)]
@@ -156,5 +161,24 @@ fn f18(x: F64Trans) -> f64 {
     unimplemented!()
 }
 
+// Invalid return activity
+#[autodiff(df19, Forward, Dual, Active)]
+fn f19(x: f32) -> f32 {
+    //~^^ ERROR invalid return activity Active in Forward Mode
+    unimplemented!()
+}
+
+#[autodiff(df20, Reverse, Active, Dual)]
+fn f20(x: f32) -> f32 {
+    //~^^ ERROR invalid return activity Dual in Reverse Mode
+    unimplemented!()
+}
+
+// Duplicated cannot be used as return activity
+#[autodiff(df21, Reverse, Active, Duplicated)]
+fn f21(x: f32) -> f32 {
+    //~^^ ERROR invalid return activity Duplicated in Reverse Mode
+    unimplemented!()
+}
 
 fn main() {}