Rollup merge of #138231 - Sa4dUs:autodiff-ice, r=ZuseZ4

Prevent ICE in autodiff validation by emitting user-friendly errors

This PR moves `valid_ret_activity` and `valid_input_activity` checks to the macro expansion phase in compiler/rustc_builtin_macros/src/autodiff.rs, replacing the following internal compiler error (ICE):
```
error: internal compiler error:
compiler/rustc_codegen_ssa/src/codegen_attrs.rs:935:13:
Invalid input activity Dual for Reverse mode
```
with a more user-friendly message.

The issue specifically affected the test file `tests/ui/autodiff/autodiff_illegal.rs`, impacting the functions `f5` and `f6`.

The ICE can be reproduced by following [Enzyme's Rustbook](https://enzymead.github.io/rustbook/installation.html) installation guide.

Additionally, this PR adds tests for invalid return activity in `autodiff_illegal.rs`, which previously triggered an unnoticed ICE before these fixes.

r? ``@oli-obk``

compiler/rustc_builtin_macros/messages.ftl

@@ -75,9 +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
     .label2 = not a `struct`, `enum` or `union`

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};
@@ -576,6 +577,8 @@ mod llvm_enzyme {
     //
     // Error handling: If the user provides an invalid configuration (incorrect numbers, types, or
     // both), we emit an error and return the original signature. This allows us to continue parsing.
+    // FIXME(Sa4dUs): make individual activities' span available so errors
+    // can point to only the activity instead of the entire attribute
     fn gen_enzyme_decl(
         ecx: &ExtCtxt<'_>,
         sig: &ast::FnSig,
@@ -623,10 +626,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::{
+use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, DiffActivity, DiffMode};
-    AutoDiffAttrs, DiffActivity, DiffMode, valid_input_activity, valid_ret_activity,
-};
 use rustc_ast::{MetaItem, MetaItemInner, attr};
 use rustc_attr_parsing::ReprAttr::ReprAlign;
 use rustc_attr_parsing::{AttributeKind, InlineAttr, InstructionSetAttr, OptimizeAttr};
@@ -874,15 +872,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 in Forward Mode
-//~^^ ERROR Duplicated can not be used for this type
+    //~^^ 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,21 +86,21 @@ fn fn_exists() {}
 // We colide with an already existing function
 #[autodiff(fn_exists, Reverse, Active)]
 pub fn f10(x: f64) {
-//~^^ ERROR the name `fn_exists` is defined multiple times [E0428]
+    //~^^ 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!()
 }
 
@@ -109,7 +108,7 @@ pub fn f12() {
 // 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 +129,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 +140,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 +157,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() {}

tests/ui/autodiff/autodiff_illegal.stderr

@@ -1,5 +1,5 @@
 error[E0658]: attributes on expressions are experimental
-  --> $DIR/autodiff_illegal.rs:54:5
+  --> $DIR/autodiff_illegal.rs:53:5
    |
 LL |     #[autodiff(df7, Forward, Dual)]
    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -53,25 +53,25 @@ LL | pub fn f6(x: f64) {
    |              ^^^
 
 error: autodiff must be applied to function
-  --> $DIR/autodiff_illegal.rs:51:5
+  --> $DIR/autodiff_illegal.rs:50:5
    |
 LL |     let mut x = 5;
    |     ^^^^^^^^^^^^^^
 
 error: autodiff must be applied to function
-  --> $DIR/autodiff_illegal.rs:55:5
+  --> $DIR/autodiff_illegal.rs:54:5
    |
 LL |     x = x + 3;
    |     ^
 
 error: autodiff must be applied to function
-  --> $DIR/autodiff_illegal.rs:60:5
+  --> $DIR/autodiff_illegal.rs:59:5
    |
 LL |     let add_one_v2 = |x: u32| -> u32 { x + 1 };
    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 error: autodiff must be applied to function
-  --> $DIR/autodiff_illegal.rs:66:1
+  --> $DIR/autodiff_illegal.rs:65:1
    |
 LL | / pub fn f7(x: f64) {
 LL | |
@@ -80,7 +80,7 @@ LL | | }
    | |_^
 
 error: autodiff requires at least a name and mode
-  --> $DIR/autodiff_illegal.rs:73:1
+  --> $DIR/autodiff_illegal.rs:72:1
    |
 LL | / pub fn f8(x: f64) {
 LL | |
@@ -89,7 +89,7 @@ LL | | }
    | |_^
 
 error: autodiff must be applied to function
-  --> $DIR/autodiff_illegal.rs:80:1
+  --> $DIR/autodiff_illegal.rs:79:1
    |
 LL | / pub fn f9(x: f64) {
 LL | |
@@ -98,7 +98,7 @@ LL | | }
    | |_^
 
 error[E0428]: the name `fn_exists` is defined multiple times
-  --> $DIR/autodiff_illegal.rs:88:1
+  --> $DIR/autodiff_illegal.rs:87:1
    |
 LL | fn fn_exists() {}
    | -------------- previous definition of the value `fn_exists` here
@@ -110,7 +110,7 @@ LL | #[autodiff(fn_exists, Reverse, Active)]
    = note: this error originates in the attribute macro `autodiff` (in Nightly builds, run with -Z macro-backtrace for more info)
 
 error: autodiff requires at least a name and mode
-  --> $DIR/autodiff_illegal.rs:96:1
+  --> $DIR/autodiff_illegal.rs:95:1
    |
 LL | / pub fn f11() {
 LL | |
@@ -119,19 +119,43 @@ LL | | }
    | |_^
 
 error: unknown Mode: `Debug`. Use `Forward` or `Reverse`
-  --> $DIR/autodiff_illegal.rs:102:18
+  --> $DIR/autodiff_illegal.rs:101:18
    |
 LL | #[autodiff(df12, Debug)]
    |                  ^^^^^
 
 error: did not recognize Activity: `Reverse`
-  --> $DIR/autodiff_illegal.rs:110:27
+  --> $DIR/autodiff_illegal.rs:109:27
    |
 LL | #[autodiff(df13, Forward, Reverse)]
    |                           ^^^^^^^
 
+error: invalid return activity Active in Forward Mode
+  --> $DIR/autodiff_illegal.rs:161:1
+   |
+LL | #[autodiff(df19, Forward, Dual, Active)]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   |
+   = note: this error originates in the attribute macro `autodiff` (in Nightly builds, run with -Z macro-backtrace for more info)
+
+error: invalid return activity Dual in Reverse Mode
+  --> $DIR/autodiff_illegal.rs:167:1
+   |
+LL | #[autodiff(df20, Reverse, Active, Dual)]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   |
+   = note: this error originates in the attribute macro `autodiff` (in Nightly builds, run with -Z macro-backtrace for more info)
+
+error: invalid return activity Duplicated in Reverse Mode
+  --> $DIR/autodiff_illegal.rs:174:1
+   |
+LL | #[autodiff(df21, Reverse, Active, Duplicated)]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   |
+   = note: this error originates in the attribute macro `autodiff` (in Nightly builds, run with -Z macro-backtrace for more info)
+
 error[E0433]: failed to resolve: use of undeclared type `MyFloat`
-  --> $DIR/autodiff_illegal.rs:131:1
+  --> $DIR/autodiff_illegal.rs:130:1
    |
 LL | #[autodiff(df15, Reverse, Active, Active)]
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ use of undeclared type `MyFloat`
@@ -139,14 +163,14 @@ LL | #[autodiff(df15, Reverse, Active, Active)]
    = note: this error originates in the attribute macro `autodiff` (in Nightly builds, run with -Z macro-backtrace for more info)
 
 error[E0433]: failed to resolve: use of undeclared type `F64Trans`
-  --> $DIR/autodiff_illegal.rs:153:1
+  --> $DIR/autodiff_illegal.rs:154:1
    |
 LL | #[autodiff(df18, Reverse, Active, Active)]
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ use of undeclared type `F64Trans`
    |
    = note: this error originates in the attribute macro `autodiff` (in Nightly builds, run with -Z macro-backtrace for more info)
 
-error: aborting due to 19 previous errors
+error: aborting due to 22 previous errors
 
 Some errors have detailed explanations: E0428, E0433, E0658.
 For more information about an error, try `rustc --explain E0428`.