use propper rustc error handler for type/act check

Author: ZuseZ4

compiler/rustc_ast/src/expand/autodiff_attrs.rs

@@ -3,6 +3,8 @@ use std::str::FromStr;
 use thin_vec::ThinVec;
 use std::fmt::{Display, Formatter};
 use crate::NestedMetaItem;
+use crate::ptr::P;
+use crate::{Ty, TyKind};
 
 #[allow(dead_code)]
 #[derive(Clone, Copy, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
@@ -40,7 +42,29 @@ pub fn valid_ret_activity(mode: DiffMode, activity: DiffActivity) -> bool {
         }
     }
 }
-
+fn is_ptr_or_ref(ty: &Ty) -> bool {
+    match ty.kind {
+        TyKind::Ptr(_) | TyKind::Ref(_, _) => true,
+        _ => false,
+    }
+}
+// TODO We should make this more robust to also
+// accept aliases of f32 and f64
+//fn is_float(ty: &Ty) -> bool {
+//    false
+//}
+pub fn valid_ty_for_activity(ty: &P<Ty>, activity: DiffActivity) -> bool {
+    if is_ptr_or_ref(ty) {
+        return activity == DiffActivity::Dual || activity == DiffActivity::DualOnly ||
+            activity == DiffActivity::Duplicated || activity == DiffActivity::DuplicatedOnly ||
+            activity == DiffActivity::Const;
+    }
+    true
+    //if is_scalar_ty(&ty) {
+    //    return activity == DiffActivity::Active || activity == DiffActivity::ActiveOnly ||
+    //        activity == DiffActivity::Const;
+    //}
+}
 pub fn valid_input_activity(mode: DiffMode, activity: DiffActivity) -> bool {
         return match mode {
             DiffMode::Inactive => false,

compiler/rustc_builtin_macros/messages.ftl

@@ -4,6 +4,7 @@ builtin_macros_alloc_must_statics = allocators must be statics
 builtin_macros_autodiff = autodiff must be applied to function
 builtin_macros_autodiff_not_build = this rustc version does not support autodiff
 builtin_macros_autodiff_mode_activity = {$act} can not be used in {$mode} Mode
+builtin_macros_autodiff_ty_activity = {$act} can not be used for this type
 
 builtin_macros_asm_clobber_abi = clobber_abi
 builtin_macros_asm_clobber_no_reg = asm with `clobber_abi` must specify explicit registers for outputs

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -3,7 +3,7 @@
 //use crate::util::check_autodiff;
 
 use crate::errors;
-use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, DiffActivity, DiffMode, valid_input_activity};
+use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, DiffActivity, DiffMode, valid_input_activity, valid_ty_for_activity};
 use rustc_ast::ptr::P;
 use rustc_ast::token::{Token, TokenKind};
 use rustc_ast::tokenstream::*;
@@ -175,25 +175,6 @@ fn assure_mut_ref(ty: &ast::Ty) -> ast::Ty {
     ty
 }
 
-// TODO We should make this more robust to also
-// accept aliases of f32 and f64
-#[cfg(llvm_enzyme)]
-fn is_float(ty: &ast::Ty) -> bool {
-    match ty.kind {
-        TyKind::Path(_, ref path) => {
-            let last = path.segments.last().unwrap();
-            last.ident.name == sym::f32 || last.ident.name == sym::f64
-        }
-        _ => false,
-    }
-}
-#[cfg(llvm_enzyme)]
-fn is_ptr_or_ref(ty: &ast::Ty) -> bool {
-    match ty.kind {
-        TyKind::Ptr(_) | TyKind::Ref(_, _) => true,
-        _ => false,
-    }
-}
 
 // The body of our generated functions will consist of two black_Box calls.
 // The first will call the primal function with the original arguments.
@@ -277,7 +258,7 @@ fn gen_primal_call(
 // zero-initialized by Enzyme). Active arguments are not handled yet.
 // Each argument of the primal function (and the return type if existing) must be annotated with an
 // activity.
-#[cfg(llvm_enzyme)]
+//#[cfg(llvm_enzyme)]
 fn gen_enzyme_decl(
     ecx: &ExtCtxt<'_>,
     sig: &ast::FnSig,
@@ -301,13 +282,17 @@ fn gen_enzyme_decl(
                 act: activity.to_string()
             });
         }
+        if !valid_ty_for_activity(&arg.ty, *activity) {
+            ecx.sess.dcx().emit_err(errors::AutoDiffInvalidTypeForActivity {
+                span: arg.ty.span,
+                act: activity.to_string()
+            });
+        }
         match activity {
             DiffActivity::Active => {
-                assert!(is_float(&arg.ty));
                 act_ret.push(arg.ty.clone());
             }
             DiffActivity::Duplicated => {
-                assert!(is_ptr_or_ref(&arg.ty));
                 let mut shadow_arg = arg.clone();
                 // We += into the shadow in reverse mode.
                 shadow_arg.ty = P(assure_mut_ref(&arg.ty));

compiler/rustc_builtin_macros/src/errors.rs

@@ -163,6 +163,13 @@ pub(crate) struct AllocMustStatics {
     #[primary_span]
     pub(crate) span: Span,
 }
+#[derive(Diagnostic)]
+#[diag(builtin_macros_autodiff_ty_activity)]
+pub(crate) struct AutoDiffInvalidTypeForActivity {
+    #[primary_span]
+    pub(crate) span: Span,
+    pub(crate) act: String,
+}
 
 #[derive(Diagnostic)]
 #[diag(builtin_macros_autodiff_mode_activity)]