Remove primal call and collect it in mono instead

Author: Sa4dUs

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -594,17 +594,11 @@ mod llvm_enzyme {
         span: Span,
         primal: Ident,
         idents: &[Ident],
-        errored: bool,
+        _errored: bool,
         generics: &Generics,
     ) -> P<ast::Block> {
-        let primal_call = gen_primal_call(ecx, span, primal, idents, generics);
-        let mut body = ecx.block(span, ThinVec::new());
-
-        // This uses primal args which won't be available if we errored before
-        if !errored {
-            body.stmts.push(ecx.stmt_semi(primal_call.clone()));
-        }
-
+        let _primal_call = gen_primal_call(ecx, span, primal, idents, generics);
+        let body = ecx.block(span, ThinVec::new());
         body
     }
 

compiler/rustc_monomorphize/src/collector.rs

@@ -205,6 +205,8 @@
 //! this is not implemented however: a mono item will be produced
 //! regardless of whether it is actually needed or not.
 
+mod autodiff;
+
 use std::cell::OnceCell;
 use std::path::PathBuf;
 
@@ -237,6 +239,8 @@ use rustc_span::source_map::{Spanned, dummy_spanned, respan};
 use rustc_span::{DUMMY_SP, Span};
 use tracing::{debug, instrument, trace};
 
+#[cfg(llvm_enzyme)]
+use crate::collector::autodiff::collect_enzyme_autodiff_source_fn;
 use crate::errors::{self, EncounteredErrorWhileInstantiating, NoOptimizedMir, RecursionLimit};
 
 #[derive(PartialEq)]
@@ -913,6 +917,9 @@ fn visit_instance_use<'tcx>(
         return;
     }
     if let Some(intrinsic) = tcx.intrinsic(instance.def_id()) {
+        #[cfg(llvm_enzyme)]
+        collect_enzyme_autodiff_source_fn(tcx, instance, intrinsic, output);
+
         if let Some(_requirement) = ValidityRequirement::from_intrinsic(intrinsic.name) {
             // The intrinsics assert_inhabited, assert_zero_valid, and assert_mem_uninitialized_valid will
             // be lowered in codegen to nothing or a call to panic_nounwind. So if we encounter any

compiler/rustc_monomorphize/src/collector/autodiff.rs

@@ -0,0 +1,38 @@
+use rustc_middle::bug;
+use rustc_middle::ty::{self, IntrinsicDef, TyCtxt};
+use tracing::debug;
+
+use crate::collector::{MonoItems, create_fn_mono_item};
+
+pub(crate) fn collect_enzyme_autodiff_source_fn<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    instance: ty::Instance<'tcx>,
+    intrinsic: IntrinsicDef,
+    output: &mut MonoItems<'tcx>,
+) {
+    if intrinsic.name != rustc_span::sym::enzyme_autodiff {
+        return;
+    };
+
+    debug!("enzyme_autodiff found");
+    let (primal, span) = match instance.args[0].kind() {
+        rustc_middle::infer::canonical::ir::GenericArgKind::Type(ty) => match ty.kind() {
+            ty::FnDef(def_id, substs) => {
+                let span = tcx.def_span(def_id);
+                let instance = ty::Instance::expect_resolve(
+                    tcx,
+                    ty::TypingEnv::non_body_analysis(tcx, def_id),
+                    *def_id,
+                    substs,
+                    span,
+                );
+
+                (instance, span)
+            }
+            _ => bug!("expected function"),
+        },
+        _ => bug!("expected type"),
+    };
+
+    output.push(create_fn_mono_item(tcx, primal, span));
+}