Move logic to a dedicated enzyme_autodiff intrinsic

Author: Sa4dUs

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -331,20 +331,23 @@ mod llvm_enzyme {
             .count() as u32;
         let (d_sig, new_args, idents, errored) = gen_enzyme_decl(ecx, &sig, &x, span);
 
-        // UNUSED
+        // TODO(Sa4dUs): Remove this and all the related logic
         let _d_body = gen_enzyme_body(
             ecx, &x, n_active, &sig, &d_sig, primal, &new_args, span, sig_span, idents, errored,
             &generics,
         );
 
+        let d_body =
+            call_enzyme_autodiff(ecx, primal, first_ident(&meta_item_vec[0]), span, &d_sig);
+
         // The first element of it is the name of the function to be generated
         let asdf = Box::new(ast::Fn {
             defaultness: ast::Defaultness::Final,
             sig: d_sig,
             ident: first_ident(&meta_item_vec[0]),
-            generics,
+            generics: generics.clone(),
             contract: None,
-            body: None, // This leads to an error when the ad function is inside a traits
+            body: Some(d_body),
             define_opaque: None,
         });
         let mut rustc_ad_attr =
@@ -431,18 +434,15 @@ mod llvm_enzyme {
             tokens: ts,
         });
 
-        let rustc_intrinsic_attr =
-            P(ast::NormalAttr::from_ident(Ident::with_dummy_span(sym::rustc_intrinsic)));
-        let new_id = ecx.sess.psess.attr_id_generator.mk_attr_id();
-        let intrinsic_attr = outer_normal_attr(&rustc_intrinsic_attr, new_id, span);
+        let vis_clone = vis.clone();
 
         let new_id = ecx.sess.psess.attr_id_generator.mk_attr_id();
         let d_attr = outer_normal_attr(&rustc_ad_attr, new_id, span);
         let d_annotatable = match &item {
             Annotatable::AssocItem(_, _) => {
                 let assoc_item: AssocItemKind = ast::AssocItemKind::Fn(asdf);
                 let d_fn = P(ast::AssocItem {
-                    attrs: thin_vec![d_attr, intrinsic_attr],
+                    attrs: thin_vec![d_attr],
                     id: ast::DUMMY_NODE_ID,
                     span,
                     vis,
@@ -452,15 +452,13 @@ mod llvm_enzyme {
                 Annotatable::AssocItem(d_fn, Impl { of_trait: false })
             }
             Annotatable::Item(_) => {
-                let mut d_fn =
-                    ecx.item(span, thin_vec![d_attr, intrinsic_attr], ItemKind::Fn(asdf));
+                let mut d_fn = ecx.item(span, thin_vec![d_attr], ItemKind::Fn(asdf));
                 d_fn.vis = vis;
 
                 Annotatable::Item(d_fn)
             }
             Annotatable::Stmt(_) => {
-                let mut d_fn =
-                    ecx.item(span, thin_vec![d_attr, intrinsic_attr], ItemKind::Fn(asdf));
+                let mut d_fn = ecx.item(span, thin_vec![d_attr], ItemKind::Fn(asdf));
                 d_fn.vis = vis;
 
                 Annotatable::Stmt(P(ast::Stmt {
@@ -474,7 +472,9 @@ mod llvm_enzyme {
             }
         };
 
-        return vec![orig_annotatable, d_annotatable];
+        let dummy_const_annotatable = gen_dummy_const(ecx, span, primal, sig, generics, vis_clone);
+
+        return vec![orig_annotatable, dummy_const_annotatable, d_annotatable];
     }
 
     // shadow arguments (the extra ones which were not in the original (primal) function), in reverse mode must be
@@ -495,6 +495,123 @@ mod llvm_enzyme {
         ty
     }
 
+    // Generate `enzyme_autodiff` intrinsic call
+    // ```
+    // std::intrinsics::enzyme_autodiff(source, diff, (args))
+    // ```
+    fn call_enzyme_autodiff(
+        ecx: &ExtCtxt<'_>,
+        primal: Ident,
+        diff: Ident,
+        span: Span,
+        d_sig: &FnSig,
+    ) -> P<ast::Block> {
+        let primal_path_expr = ecx.expr_path(ecx.path_ident(span, primal));
+        let diff_path_expr = ecx.expr_path(ecx.path_ident(span, diff));
+
+        let tuple_expr = ecx.expr_tuple(
+            span,
+            d_sig
+                .decl
+                .inputs
+                .iter()
+                .map(|arg| match arg.pat.kind {
+                    PatKind::Ident(_, ident, _) => ecx.expr_path(ecx.path_ident(span, ident)),
+                    _ => todo!(),
+                })
+                .collect::<ThinVec<_>>()
+                .into(),
+        );
+
+        let enzyme_path = ecx.path(
+            span,
+            vec![
+                Ident::from_str("std"),
+                Ident::from_str("intrinsics"),
+                Ident::from_str("enzyme_autodiff"),
+            ],
+        );
+        let call_expr = ecx.expr_call(
+            span,
+            ecx.expr_path(enzyme_path),
+            vec![primal_path_expr, diff_path_expr, tuple_expr].into(),
+        );
+
+        let block = ecx.block_expr(call_expr);
+
+        block
+    }
+
+    // Generate dummy const to prevent primal function
+    // from being optimized away before applying enzyme
+    // ```
+    // const _: () =
+    // {
+    //     #[used]
+    //     pub static DUMMY_PTR: fn_type = primal_fn;
+    // };
+    // ```
+    fn gen_dummy_const(
+        ecx: &ExtCtxt<'_>,
+        span: Span,
+        primal: Ident,
+        sig: FnSig,
+        generics: Generics,
+        vis: Visibility,
+    ) -> Annotatable {
+        // #[used]
+        let used_attr = P(ast::NormalAttr::from_ident(Ident::with_dummy_span(sym::used)));
+        let new_id = ecx.sess.psess.attr_id_generator.mk_attr_id();
+        let used_attr = outer_normal_attr(&used_attr, new_id, span);
+
+        // static DUMMY_PTR: <fn_type> = <primal_ident>
+        let static_ident = Ident::from_str_and_span("DUMMY_PTR", span);
+        let fn_ptr_ty = ast::TyKind::BareFn(Box::new(ast::BareFnTy {
+            safety: sig.header.safety,
+            ext: sig.header.ext,
+            generic_params: generics.params,
+            decl: sig.decl,
+            decl_span: sig.span,
+        }));
+        let static_ty = ecx.ty(span, fn_ptr_ty);
+
+        let static_expr = ecx.expr_path(ecx.path(span, vec![primal]));
+        let static_item_kind = ast::ItemKind::Static(Box::new(ast::StaticItem {
+            ident: static_ident,
+            ty: static_ty,
+            safety: ast::Safety::Default,
+            mutability: ast::Mutability::Not,
+            expr: Some(static_expr),
+            define_opaque: None,
+        }));
+
+        let static_item = ast::Item {
+            attrs: thin_vec![used_attr],
+            id: ast::DUMMY_NODE_ID,
+            span,
+            vis,
+            kind: static_item_kind,
+            tokens: None,
+        };
+
+        let block_expr = ecx.expr_block(Box::new(ast::Block {
+            stmts: thin_vec![ecx.stmt_item(span, P(static_item))],
+            id: ast::DUMMY_NODE_ID,
+            rules: ast::BlockCheckMode::Default,
+            span,
+            tokens: None,
+        }));
+
+        let const_item = ecx.item_const(
+            span,
+            Ident::from_str_and_span("_", span),
+            ecx.ty(span, ast::TyKind::Tup(thin_vec![])),
+            block_expr,
+        );
+
+        Annotatable::Item(const_item)
+    }
+
     // Will generate a body of the type:
     // ```
     // {

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -9,7 +9,7 @@ use rustc_middle::bug;
 use tracing::{debug, trace};
 
 use crate::back::write::llvm_err;
-use crate::builder::{Builder, OperandRef, PlaceRef, UNNAMED};
+use crate::builder::{Builder, PlaceRef, UNNAMED};
 use crate::context::SimpleCx;
 use crate::declare::declare_simple_fn;
 use crate::errors::{AutoDiffWithoutEnable, LlvmError};
@@ -199,7 +199,7 @@ pub(crate) fn generate_enzyme_call<'ll, 'tcx>(
     fn_to_diff: &'ll Value,
     outer_name: &str,
     ret_ty: &'ll Type,
-    fn_args: &[OperandRef<'tcx, &'ll Value>],
+    fn_args: &[&'ll Value],
     attrs: AutoDiffAttrs,
     dest: PlaceRef<'tcx, &'ll Value>,
 ) {
@@ -275,15 +275,13 @@ pub(crate) fn generate_enzyme_call<'ll, 'tcx>(
         args.push(cx.get_const_int(cx.type_i64(), attrs.width as u64));
     }
 
-    let outer_args: Vec<&llvm::Value> = fn_args.iter().map(|op| op.immediate()).collect();
-
     match_args_from_caller_to_enzyme(
         &cx,
         builder,
         attrs.width,
         &mut args,
         &attrs.input_activity,
-        &outer_args,
+        fn_args,
     );
 
     let call = builder.call(enzyme_ty, None, None, ad_fn, &args, None, None);

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -3,6 +3,7 @@ use std::cmp::Ordering;
 
 use rustc_abi::{Align, BackendRepr, ExternAbi, Float, HasDataLayout, Primitive, Size};
 use rustc_codegen_ssa::base::{compare_simd_types, wants_msvc_seh, wants_wasm_eh};
+use rustc_codegen_ssa::codegen_attrs::autodiff_attrs;
 use rustc_codegen_ssa::common::{IntPredicate, TypeKind};
 use rustc_codegen_ssa::errors::{ExpectedPointerMutability, InvalidMonomorphization};
 use rustc_codegen_ssa::mir::operand::{OperandRef, OperandValue};
@@ -198,48 +199,60 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                     &[ptr, args[1].immediate()],
                 )
             }
-            _ if tcx.has_attr(instance.def_id(), sym::rustc_autodiff) => {
-                // NOTE(Sa4dUs): This is a hacky way to get the autodiff items
-                // so we can focus on the lowering of the intrinsic call
-                let mut source_id = None;
-                let mut diff_attrs = None;
-                let items: Vec<_> = tcx.hir_body_owners().map(|i| i.to_def_id()).collect();
-
-                // Hacky way of getting primal-diff pair, only works for code with 1 autodiff call
-                for target_id in &items {
-                    let Some(target_attrs) = &tcx.codegen_fn_attrs(target_id).autodiff_item else {
-                        continue;
-                    };
+            sym::enzyme_autodiff => {
+                let val_arr: Vec<&'ll Value> = match args[2].val {
+                    crate::intrinsic::OperandValue::Ref(ref place_value) => {
+                        let mut ret_arr = vec![];
+                        let tuple_place = PlaceRef { val: *place_value, layout: args[2].layout };
 
-                    if target_attrs.is_source() {
-                        source_id = Some(*target_id);
-                    } else {
-                        diff_attrs = Some(target_attrs);
-                    }
-                }
+                        for i in 0..tuple_place.layout.layout.0.fields.count() {
+                            let field_place = tuple_place.project_field(self, i);
+                            let field_layout = tuple_place.layout.field(self, i);
+                            let llvm_ty = field_layout.llvm_type(self.cx);
 
-                if source_id.is_none() || diff_attrs.is_none() {
-                    bug!("could not find source_id={source_id:?} or diff_attrs={diff_attrs:?}");
-                }
+                            let field_val =
+                                self.load(llvm_ty, field_place.val.llval, field_place.val.align);
+
+                            ret_arr.push(field_val)
+                        }
 
-                let diff_attrs = diff_attrs.unwrap().clone();
+                        ret_arr
+                    }
+                    crate::intrinsic::OperandValue::Pair(v1, v2) => vec![v1, v2],
+                    OperandValue::Immediate(v) => vec![v],
+                    OperandValue::ZeroSized => bug!("unexpected `ZeroSized` arg"),
+                };
 
-                // Get source fn
-                let source_id = source_id.unwrap();
-                let fn_source = Instance::mono(tcx, source_id);
+                // Get source, diff, and attrs
+                let source_id = match fn_args.into_type_list(tcx)[0].kind() {
+                    ty::FnDef(def_id, _) => def_id,
+                    _ => bug!("invalid args"),
+                };
+                let fn_source = Instance::mono(tcx, *source_id);
                 let source_symbol =
                     symbol_name_for_instance_in_crate(tcx, fn_source.clone(), LOCAL_CRATE);
                 let fn_to_diff: Option<&'ll llvm::Value> = self.cx.get_function(&source_symbol);
                 let Some(fn_to_diff) = fn_to_diff else { bug!("could not find source function") };
 
+                let diff_id = match fn_args.into_type_list(tcx)[1].kind() {
+                    ty::FnDef(def_id, _) => def_id,
+                    _ => bug!("invalid args"),
+                };
+                let fn_diff = Instance::mono(tcx, *diff_id);
+                let diff_symbol =
+                    symbol_name_for_instance_in_crate(tcx, fn_diff.clone(), LOCAL_CRATE);
+
+                let diff_attrs = autodiff_attrs(tcx, *diff_id);
+                let Some(diff_attrs) = diff_attrs else { bug!("could not find autodiff attrs") };
+
                 // Build body
                 generate_enzyme_call(
                     self,
                     self.cx,
                     fn_to_diff,
-                    name.as_str(),
+                    &diff_symbol,
                     llret_ty,
-                    args,
+                    &val_arr,
                     diff_attrs.clone(),
                     result,
                 );

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -587,7 +587,7 @@ fn check_link_name_xor_ordinal(
 /// being differentiated). The other form is #[rustc_autodiff(Mode, ActivityList)] on top of the
 /// placeholder functions. We wrote the rustc_autodiff attributes ourself, so this should never
 /// panic, unless we introduced a bug when parsing the autodiff macro.
-fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> Option<AutoDiffAttrs> {
+pub fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> Option<AutoDiffAttrs> {
     let attrs = tcx.get_attrs(id, sym::rustc_autodiff);
 
     let attrs = attrs.filter(|attr| attr.has_name(sym::rustc_autodiff)).collect::<Vec<_>>();

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -135,6 +135,7 @@ fn intrinsic_operation_unsafety(tcx: TyCtxt<'_>, intrinsic_id: LocalDefId) -> hi
         | sym::round_ties_even_f32
         | sym::round_ties_even_f64
         | sym::round_ties_even_f128
+        | sym::enzyme_autodiff
         | sym::const_eval_select => hir::Safety::Safe,
         _ => hir::Safety::Unsafe,
     };
@@ -216,6 +217,7 @@ pub(crate) fn check_intrinsic_type(
 
             (n_tps, n_cts, inputs, output)
         }
+        sym::enzyme_autodiff => (4, 0, vec![param(0), param(1), param(2)], param(3)),
         sym::abort => (0, 0, vec![], tcx.types.never),
         sym::unreachable => (0, 0, vec![], tcx.types.never),
         sym::breakpoint => (0, 0, vec![], tcx.types.unit),

compiler/rustc_span/src/symbol.rs

@@ -915,6 +915,7 @@ symbols! {
         enumerate_method,
         env,
         env_CFG_RELEASE: env!("CFG_RELEASE"),
+        enzyme_autodiff,
         eprint_macro,
         eprintln_macro,
         eq,

library/core/src/intrinsics/mod.rs

@@ -3163,6 +3163,10 @@ pub const unsafe fn copysignf64(x: f64, y: f64) -> f64;
 #[rustc_intrinsic]
 pub const unsafe fn copysignf128(x: f128, y: f128) -> f128;
 
+#[rustc_nounwind]
+#[rustc_intrinsic]
+pub const fn enzyme_autodiff<F, G, T: crate::marker::Tuple, R>(f: F, df: G, args: T) -> R;
+
 /// Inform Miri that a given pointer definitely has a certain alignment.
 #[cfg(miri)]
 #[rustc_allow_const_fn_unstable(const_eval_select)]