Fix typetree generation for differentiated functions

compiler/rustc_ast/src/expand/typetree.rs

@@ -28,6 +28,9 @@ pub enum Kind {
     Anything,
     Integer,
     Pointer,
+    // We prefer to directly lower to things that our Enzyme backend supports.
+    // However, it's sometimes convenient to pass ptr+int as one type.
+    RustSlice,
     Half,
     Float,
     Double,
@@ -57,6 +60,9 @@ impl TypeTree {
         }
         Self(ints)
     }
+    pub fn add_indirection(self) -> Self {
+        Self(vec![Type { offset: 0, size: 1, kind: Kind::Pointer, child: self }])
+    }
 }
 
 #[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, StableHash)]
@@ -72,6 +78,11 @@ pub struct Type {
     pub kind: Kind,
     pub child: TypeTree,
 }
+impl Type {
+    pub fn from_ty(offset: isize, other: &Type) -> Self {
+        Self { offset, size: other.size, kind: other.kind, child: other.child.clone() }
+    }
+}
 
 impl Type {
     pub fn add_offset(self, add: isize) -> Self {

compiler/rustc_codegen_llvm/src/builder.rs

@@ -1180,7 +1180,7 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         // vs. copying a struct with mixed types requires different derivative handling.
         // The TypeTree tells Enzyme exactly what memory layout to expect.
         if let Some(tt) = tt {
-            crate::typetree::add_tt(self.cx().llmod, self.cx().llcx, memcpy, tt);
+            crate::typetree::add_tt(self, memcpy, tt);
         }
     }
 

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -143,7 +143,6 @@ pub(crate) fn adjust_activity_to_abi<'tcx>(
 // 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, 'tcx>(
-    cx: &SimpleCx<'ll>,
     builder: &mut Builder<'_, 'll, 'tcx>,
     width: u32,
     args: &mut Vec<&'ll Value>,
@@ -157,6 +156,7 @@ fn match_args_from_caller_to_enzyme<'ll, 'tcx>(
     // 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()?
+    let cx = &builder.scx;
     let mut outer_pos: usize = 0;
     let mut activity_pos = 0;
 
@@ -292,8 +292,7 @@ fn match_args_from_caller_to_enzyme<'ll, 'tcx>(
 // FIXME(ZuseZ4): `outer_fn` should include upstream safety checks to
 // cover some assumptions of enzyme/autodiff, which could lead to UB otherwise.
 pub(crate) fn generate_enzyme_call<'ll, 'tcx>(
-    builder: &mut Builder<'_, 'll, 'tcx>,
-    cx: &SimpleCx<'ll>,
+    bx: &mut Builder<'_, 'll, 'tcx>,
     fn_to_diff: &'ll Value,
     outer_name: &str,
     ret_ty: &'ll Type,
@@ -303,6 +302,7 @@ pub(crate) fn generate_enzyme_call<'ll, 'tcx>(
     dest_place: Option<PlaceValue<&'ll Value>>,
     fnc_tree: FncTree,
 ) -> IntrinsicResult<'tcx, &'ll Value> {
+    let cx: &SimpleCx<'ll> = &bx.scx;
     // We have to pick the name depending on whether we want forward or reverse mode autodiff.
     let mut ad_name: String = match attrs.mode {
         DiffMode::Forward => "__enzyme_fwddiff",
@@ -369,34 +369,27 @@ pub(crate) fn generate_enzyme_call<'ll, 'tcx>(
         args.push(cx.get_const_int(cx.type_i64(), attrs.width as u64));
     }
 
-    match_args_from_caller_to_enzyme(
-        &cx,
-        builder,
-        attrs.width,
-        &mut args,
-        &attrs.input_activity,
-        fn_args,
-    );
+    match_args_from_caller_to_enzyme(bx, attrs.width, &mut args, &attrs.input_activity, fn_args);
 
     if !fnc_tree.args.is_empty() || !fnc_tree.ret.0.is_empty() {
-        crate::typetree::add_tt(cx.llmod, cx.llcx, fn_to_diff, fnc_tree);
+        crate::typetree::add_tt(&bx, fn_to_diff, fnc_tree);
     }
 
-    let call = builder.call(enzyme_ty, None, None, ad_fn, &args, None, None);
+    let call = bx.call(enzyme_ty, None, None, ad_fn, &args, None, None);
 
-    let fn_ret_ty = builder.cx.val_ty(call);
-    if fn_ret_ty == builder.cx.type_void() || fn_ret_ty == builder.cx.type_struct(&[], false) {
+    let fn_ret_ty = bx.cx.val_ty(call);
+    if fn_ret_ty == bx.cx.type_void() || fn_ret_ty == bx.cx.type_struct(&[], false) {
         // If we return void or an empty struct, then our caller (due to how we generated it)
         // does not expect a return value. As such, we have no pointer (or place) into which
         // we could store our value, and would store into an undef, which would cause UB.
         // As such, we just ignore the return value in those cases.
         IntrinsicResult::Operand(OperandValue::ZeroSized)
     } else if let Some(dest_place) = dest_place {
-        builder.store_to_place(call, dest_place);
+        bx.store_to_place(call, dest_place);
         IntrinsicResult::WroteIntoPlace
     } else {
         IntrinsicResult::Operand(
-            OperandRef::from_immediate_or_packed_pair(builder, call, dest_layout).val,
+            OperandRef::from_immediate_or_packed_pair(bx, call, dest_layout).val,
         )
     }
 }

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -225,7 +225,7 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                 )
             }
             sym::autodiff => {
-                return codegen_autodiff(self, tcx, instance, args, result_layout, result_place);
+                return codegen_autodiff(self, instance, args, result_layout, result_place);
             }
             sym::offload => {
                 if tcx.sess.opts.unstable_opts.offload.is_empty() {
@@ -1743,12 +1743,12 @@ fn codegen_retag_inner<'ll, 'tcx>(
 
 fn codegen_autodiff<'ll, 'tcx>(
     bx: &mut Builder<'_, 'll, 'tcx>,
-    tcx: TyCtxt<'tcx>,
     instance: ty::Instance<'tcx>,
     args: &[OperandRef<'tcx, &'ll Value>],
     result_layout: ty::layout::TyAndLayout<'tcx>,
     result_place: Option<PlaceValue<&'ll Value>>,
 ) -> IntrinsicResult<'tcx, &'ll Value> {
+    let tcx = bx.tcx;
     if !tcx.sess.opts.unstable_opts.autodiff.contains(&rustc_session::config::AutoDiff::Enable) {
         let _ = tcx.dcx().emit_almost_fatal(AutoDiffWithoutEnable);
     }
@@ -1815,7 +1815,6 @@ fn codegen_autodiff<'ll, 'tcx>(
     // Build body
     generate_enzyme_call(
         bx,
-        bx.cx,
         fn_to_diff,
         &diff_symbol,
         llret_ty,

compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs

@@ -66,6 +66,7 @@ unsafe extern "C" {
         NameLen: libc::size_t,
     ) -> Option<&Value>;
 
+    pub(crate) safe fn LLVMRustIsCall(V: &Value) -> bool;
 }
 
 unsafe extern "C" {
@@ -292,6 +293,11 @@ pub(crate) mod Enzyme_AD {
             unsafe { (self.EnzymeTypeTreeToString)(tree) }
         }
 
+        pub(crate) fn tree_to_cstr(&self, tree: *mut EnzymeTypeTree) -> &std::ffi::CStr {
+            let c_str = self.tree_to_string(tree);
+            unsafe { std::ffi::CStr::from_ptr(c_str) }
+        }
+
         pub(crate) fn tree_to_string_free(&self, ch: *const c_char) {
             unsafe { (self.EnzymeTypeTreeToStringFree)(ch) }
         }

compiler/rustc_codegen_llvm/src/llvm/mod.rs

@@ -76,6 +76,21 @@ pub(crate) fn CreateAttrStringValue<'ll>(
         )
     }
 }
+pub(crate) fn CreateAttrStringValueFromCStr<'ll>(
+    llcx: &'ll Context,
+    attr: &std::ffi::CStr,
+    value: &std::ffi::CStr,
+) -> &'ll Attribute {
+    unsafe {
+        LLVMCreateStringAttribute(
+            llcx,
+            (*attr).as_ptr(),
+            (*attr).to_bytes().len() as c_uint,
+            (*value).as_ptr(),
+            (*value).to_bytes().len() as c_uint,
+        )
+    }
+}
 
 pub(crate) fn CreateAttrString<'ll>(llcx: &'ll Context, attr: &str) -> &'ll Attribute {
     unsafe {