Naive impl of generate enzyme call

Sa4dUs · Sa4dUs · commit 23cb884ede2f · 2025-06-24T08:01:02.000Z
Now I use functions

squash
diff --git a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs
@@ -4,13 +4,13 @@ use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, AutoDiffItem, DiffActivit
 use rustc_codegen_ssa::ModuleCodegen;
 use rustc_codegen_ssa::back::write::ModuleConfig;
 use rustc_codegen_ssa::common::TypeKind;
-use rustc_codegen_ssa::traits::BaseTypeCodegenMethods;
+use rustc_codegen_ssa::traits::{BaseTypeCodegenMethods, BuilderMethods};
 use rustc_errors::FatalError;
 use rustc_middle::bug;
 use tracing::{debug, trace};
 
 use crate::back::write::llvm_err;
-use crate::builder::{SBuilder, UNNAMED};
+use crate::builder::{Builder, OperandRef, PlaceRef, UNNAMED};
 use crate::context::SimpleCx;
 use crate::declare::declare_simple_fn;
 use crate::errors::{AutoDiffWithoutEnable, LlvmError};
@@ -19,7 +19,7 @@ use crate::llvm::{Metadata, True};
 use crate::value::Value;
 use crate::{CodegenContext, LlvmCodegenBackend, ModuleLlvm, attributes, llvm};
 
-fn get_params(fnc: &Value) -> Vec<&Value> {
+pub(crate) fn _get_params(fnc: &Value) -> Vec<&Value> {
     let param_num = llvm::LLVMCountParams(fnc) as usize;
     let mut fnc_args: Vec<&Value> = vec![];
     fnc_args.reserve(param_num);
@@ -30,7 +30,7 @@ fn get_params(fnc: &Value) -> Vec<&Value> {
     fnc_args
 }
 
-fn has_sret(fnc: &Value) -> bool {
+pub(crate) fn has_sret(fnc: &Value) -> bool {
     let num_args = llvm::LLVMCountParams(fnc) as usize;
     if num_args == 0 {
         false
@@ -49,9 +49,9 @@ fn has_sret(fnc: &Value) -> bool {
 // need to match those.
 // FIXME(ZuseZ4): This logic is a bit more complicated than it should be, can we simplify it
 // using iterators and peek()?
-fn match_args_from_caller_to_enzyme<'ll>(
+fn match_args_from_caller_to_enzyme<'ll, 'tcx>(
     cx: &SimpleCx<'ll>,
-    builder: &SBuilder<'ll, 'll>,
+    builder: &mut Builder<'_, 'll, 'tcx>,
     width: u32,
     args: &mut Vec<&'ll llvm::Value>,
     inputs: &[DiffActivity],
@@ -201,7 +201,7 @@ fn match_args_from_caller_to_enzyme<'ll>(
 // Beyond sret, this article describes our challenges nicely:
 // <https://yorickpeterse.com/articles/the-mess-that-is-handling-structure-arguments-and-returns-in-llvm/>
 // I.e. (i32, f32) will get merged into i64, but we don't handle that yet.
-fn compute_enzyme_fn_ty<'ll>(
+pub(crate) fn compute_enzyme_fn_ty<'ll>(
     cx: &SimpleCx<'ll>,
     attrs: &AutoDiffAttrs,
     fn_to_diff: &'ll Value,
@@ -289,11 +289,14 @@ fn compute_enzyme_fn_ty<'ll>(
 /// [^1]: <https://enzyme.mit.edu/getting_started/CallingConvention/>
 // FIXME(ZuseZ4): `outer_fn` should include upstream safety checks to
 // cover some assumptions of enzyme/autodiff, which could lead to UB otherwise.
-fn generate_enzyme_call<'ll>(
+pub(crate) fn generate_enzyme_call<'ll, 'tcx>(
+    builder: &mut Builder<'_, 'll, 'tcx>,
     cx: &SimpleCx<'ll>,
     fn_to_diff: &'ll Value,
     outer_fn: &'ll Value,
+    fn_args: &[OperandRef<'tcx, &'ll Value>],
     attrs: AutoDiffAttrs,
+    dest: PlaceRef<'tcx, &'ll Value>,
 ) {
     // We have to pick the name depending on whether we want forward or reverse mode autodiff.
     let mut ad_name: String = match attrs.mode {
@@ -366,14 +369,6 @@ fn generate_enzyme_call<'ll>(
         let enzyme_marker_attr = llvm::CreateAttrString(cx.llcx, "enzyme_marker");
         attributes::apply_to_llfn(outer_fn, Function, &[enzyme_marker_attr]);
 
-        // first, remove all calls from fnc
-        let entry = llvm::LLVMGetFirstBasicBlock(outer_fn);
-        let br = llvm::LLVMRustGetTerminator(entry);
-        llvm::LLVMRustEraseInstFromParent(br);
-
-        let last_inst = llvm::LLVMRustGetLastInstruction(entry).unwrap();
-        let mut builder = SBuilder::build(cx, entry);
-
         let num_args = llvm::LLVMCountParams(&fn_to_diff);
         let mut args = Vec::with_capacity(num_args as usize + 1);
         args.push(fn_to_diff);
@@ -389,40 +384,20 @@ fn generate_enzyme_call<'ll>(
         }
 
         let has_sret = has_sret(outer_fn);
-        let outer_args: Vec<&llvm::Value> = get_params(outer_fn);
+        let outer_args: Vec<&llvm::Value> = fn_args.iter().map(|op| op.immediate()).collect();
         match_args_from_caller_to_enzyme(
             &cx,
-            &builder,
+            builder,
             attrs.width,
             &mut args,
             &attrs.input_activity,
             &outer_args,
             has_sret,
         );
 
-        let call = builder.call(enzyme_ty, ad_fn, &args, None);
-
-        // This part is a bit iffy. LLVM requires that a call to an inlineable function has some
-        // metadata attached to it, but we just created this code oota. Given that the
-        // differentiated function already has partly confusing metadata, and given that this
-        // affects nothing but the auttodiff IR, we take a shortcut and just steal metadata from the
-        // dummy code which we inserted at a higher level.
-        // FIXME(ZuseZ4): Work with Enzyme core devs to clarify what debug metadata issues we have,
-        // and how to best improve it for enzyme core and rust-enzyme.
-        let md_ty = cx.get_md_kind_id("dbg");
-        if llvm::LLVMRustHasMetadata(last_inst, md_ty) {
-            let md = llvm::LLVMRustDIGetInstMetadata(last_inst)
-                .expect("failed to get instruction metadata");
-            let md_todiff = cx.get_metadata_value(md);
-            llvm::LLVMSetMetadata(call, md_ty, md_todiff);
-        } else {
-            // We don't panic, since depending on whether we are in debug or release mode, we might
-            // have no debug info to copy, which would then be ok.
-            trace!("no dbg info");
-        }
+        let call = builder.call(enzyme_ty, None, None, ad_fn, &args, None, None);
 
-        // Now that we copied the metadata, get rid of dummy code.
-        llvm::LLVMRustEraseInstUntilInclusive(entry, last_inst);
+        builder.store_to_place(call, dest.val);
 
         if cx.val_ty(call) == cx.type_void() || has_sret {
             if has_sret {
@@ -445,10 +420,10 @@ fn generate_enzyme_call<'ll>(
                 llvm::LLVMBuildStore(&builder.llbuilder, call, sret_ptr);
             }
             builder.ret_void();
-        } else {
-            builder.ret(call);
         }
 
+        builder.store_to_place(call, dest.val);
+
         // Let's crash in case that we messed something up above and generated invalid IR.
         llvm::LLVMRustVerifyFunction(
             outer_fn,
@@ -463,6 +438,7 @@ pub(crate) fn differentiate<'ll>(
     diff_items: Vec<AutoDiffItem>,
     _config: &ModuleConfig,
 ) -> Result<(), FatalError> {
+    // TODO(Sa4dUs): delete all this logic
     for item in &diff_items {
         trace!("{}", item);
     }
@@ -482,7 +458,7 @@ pub(crate) fn differentiate<'ll>(
     for item in diff_items.iter() {
         let name = item.source.clone();
         let fn_def: Option<&llvm::Value> = cx.get_function(&name);
-        let Some(fn_def) = fn_def else {
+        let Some(_fn_def) = fn_def else {
             return Err(llvm_err(
                 diag_handler.handle(),
                 LlvmError::PrepareAutoDiff {
@@ -494,7 +470,7 @@ pub(crate) fn differentiate<'ll>(
         };
         debug!(?item.target);
         let fn_target: Option<&llvm::Value> = cx.get_function(&item.target);
-        let Some(fn_target) = fn_target else {
+        let Some(_fn_target) = fn_target else {
             return Err(llvm_err(
                 diag_handler.handle(),
                 LlvmError::PrepareAutoDiff {
@@ -505,7 +481,7 @@ pub(crate) fn differentiate<'ll>(
             ));
         };
 
-        generate_enzyme_call(&cx, fn_def, fn_target, item.attrs.clone());
+        // generate_enzyme_call(&cx, fn_def, fn_target, item.attrs.clone());
     }
 
     // FIXME(ZuseZ4): support SanitizeHWAddress and prevent illegal/unsupported opts
diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs
@@ -9,17 +9,19 @@ use rustc_codegen_ssa::mir::operand::{OperandRef, OperandValue};
 use rustc_codegen_ssa::mir::place::{PlaceRef, PlaceValue};
 use rustc_codegen_ssa::traits::*;
 use rustc_hir as hir;
+use rustc_hir::def_id::LOCAL_CRATE;
 use rustc_middle::mir::BinOp;
 use rustc_middle::ty::layout::{FnAbiOf, HasTyCtxt, HasTypingEnv, LayoutOf};
-use rustc_middle::ty::{self, GenericArgsRef, Ty};
+use rustc_middle::ty::{self, GenericArgsRef, Instance, Ty};
 use rustc_middle::{bug, span_bug};
 use rustc_span::{Span, Symbol, sym};
-use rustc_symbol_mangling::mangle_internal_symbol;
+use rustc_symbol_mangling::{mangle_internal_symbol, symbol_name_for_instance_in_crate};
 use rustc_target::spec::{HasTargetSpec, PanicStrategy};
 use tracing::debug;
 
 use crate::abi::FnAbiLlvmExt;
 use crate::builder::Builder;
+use crate::builder::autodiff::generate_enzyme_call;
 use crate::context::CodegenCx;
 use crate::llvm::{self, Metadata};
 use crate::type_::Type;
@@ -200,20 +202,56 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
             _ if tcx.has_attr(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;
+                    };
 
-                // `diff_items` is empty even when autodiff is enabled, and if we're here,
-                // it's because some function was marked as intrinsic and had the `rustc_autodiff` attr
-                let diff_items = tcx.collect_and_partition_mono_items(()).autodiff_items;
+                    if target_attrs.is_source() {
+                        source_id = Some(*target_id);
+                    } else {
+                        diff_attrs = Some(target_attrs);
+                    }
+                }
 
-                // this shouldn't happen?
-                if diff_items.is_empty() {
-                    bug!("no autodiff items found for {def_id:?}");
+                if source_id.is_none() || diff_attrs.is_none() {
+                    bug!("could not find source_id={source_id:?} or diff_attrs={diff_attrs:?}");
                 }
 
-                // TODO(Sa4dUs): generate the enzyme call itself, based on the logic in `builder.rs`
+                let diff_attrs = diff_attrs.unwrap().clone();
+
+                // Get source fn
+                let source_id = source_id.unwrap();
+                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") };
+
+                // Declare target fn
+                let target_symbol =
+                    symbol_name_for_instance_in_crate(tcx, instance.clone(), LOCAL_CRATE);
+                let fn_abi = self.cx.fn_abi_of_instance(instance, ty::List::empty());
+                let outer_fn: &'ll Value =
+                    self.cx.declare_fn(&target_symbol, fn_abi, Some(instance));
+
+                // Build body
+                generate_enzyme_call(
+                    self,
+                    self.cx,
+                    fn_to_diff,
+                    outer_fn,
+                    args, // This argument was not in the original `generate_enzyme_call`, now it's included because `get_params` is not working anymore
+                    diff_attrs.clone(),
+                    result,
+                );
 
-                // Just gen the fallback body for now
-                return Err(ty::Instance::new_raw(def_id, instance.args));
+                return Ok(());
             }
             sym::is_val_statically_known => {
                 let intrinsic_type = args[0].layout.immediate_llvm_type(self.cx);
diff --git a/tests/codegen/autodiff/scalar.rs b/tests/codegen/autodiff/scalar.rs
@@ -2,6 +2,7 @@
 //@ no-prefer-dynamic
 //@ needs-enzyme
 #![feature(autodiff)]
+#![feature(intrinsics)]
 
 use std::autodiff::autodiff_reverse;
 
