about summary refs log tree commit diff
path: root/compiler/rustc_codegen_llvm/src
diff options
context:
space:
mode:
Diffstat (limited to 'compiler/rustc_codegen_llvm/src')
-rw-r--r--compiler/rustc_codegen_llvm/src/abi.rs1
-rw-r--r--compiler/rustc_codegen_llvm/src/back/lto.rs2
-rw-r--r--compiler/rustc_codegen_llvm/src/builder.rs15
-rw-r--r--compiler/rustc_codegen_llvm/src/builder/autodiff.rs6
-rw-r--r--compiler/rustc_codegen_llvm/src/consts.rs4
-rw-r--r--compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs2
-rw-r--r--compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/native.rs4
-rw-r--r--compiler/rustc_codegen_llvm/src/intrinsic.rs8
-rw-r--r--compiler/rustc_codegen_llvm/src/lib.rs14
-rw-r--r--compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs206
-rw-r--r--compiler/rustc_codegen_llvm/src/typetree.rs122
-rw-r--r--compiler/rustc_codegen_llvm/src/va_arg.rs13
12 files changed, 380 insertions, 17 deletions
diff --git a/compiler/rustc_codegen_llvm/src/abi.rs b/compiler/rustc_codegen_llvm/src/abi.rs
index 861227f7c2a..1703cab942b 100644
--- a/compiler/rustc_codegen_llvm/src/abi.rs
+++ b/compiler/rustc_codegen_llvm/src/abi.rs
@@ -246,6 +246,7 @@ impl<'ll, 'tcx> ArgAbiExt<'ll, 'tcx> for ArgAbi<'tcx, Ty<'tcx>> {
                     scratch_align,
                     bx.const_usize(copy_bytes),
                     MemFlags::empty(),
+                    None,
                 );
                 bx.lifetime_end(llscratch, scratch_size);
             }
diff --git a/compiler/rustc_codegen_llvm/src/back/lto.rs b/compiler/rustc_codegen_llvm/src/back/lto.rs
index 78107d95e5a..5ac3a87c158 100644
--- a/compiler/rustc_codegen_llvm/src/back/lto.rs
+++ b/compiler/rustc_codegen_llvm/src/back/lto.rs
@@ -563,6 +563,8 @@ fn enable_autodiff_settings(ad: &[config::AutoDiff]) {
             config::AutoDiff::Enable => {}
             // We handle this below
             config::AutoDiff::NoPostopt => {}
+            // Disables TypeTree generation
+            config::AutoDiff::NoTT => {}
         }
     }
     // This helps with handling enums for now.
diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs
index a4dc4eb532f..5271d0b4bb8 100644
--- a/compiler/rustc_codegen_llvm/src/builder.rs
+++ b/compiler/rustc_codegen_llvm/src/builder.rs
@@ -2,6 +2,7 @@ use std::borrow::{Borrow, Cow};
 use std::ops::Deref;
 use std::{iter, ptr};
 
+use rustc_ast::expand::typetree::FncTree;
 pub(crate) mod autodiff;
 pub(crate) mod gpu_offload;
 
@@ -1107,11 +1108,12 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         src_align: Align,
         size: &'ll Value,
         flags: MemFlags,
+        tt: Option<FncTree>,
     ) {
         assert!(!flags.contains(MemFlags::NONTEMPORAL), "non-temporal memcpy not supported");
         let size = self.intcast(size, self.type_isize(), false);
         let is_volatile = flags.contains(MemFlags::VOLATILE);
-        unsafe {
+        let memcpy = unsafe {
             llvm::LLVMRustBuildMemCpy(
                 self.llbuilder,
                 dst,
@@ -1120,7 +1122,16 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
                 src_align.bytes() as c_uint,
                 size,
                 is_volatile,
-            );
+            )
+        };
+
+        // TypeTree metadata for memcpy is especially important: when Enzyme encounters
+        // a memcpy during autodiff, it needs to know the structure of the data being
+        // copied to properly track derivatives. For example, copying an array of floats
+        // 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);
         }
     }
 
diff --git a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs
index b66e3dfdeec..c3485f56391 100644
--- a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs
+++ b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs
@@ -1,6 +1,7 @@
 use std::ptr;
 
 use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, DiffActivity, DiffMode};
+use rustc_ast::expand::typetree::FncTree;
 use rustc_codegen_ssa::common::TypeKind;
 use rustc_codegen_ssa::traits::{BaseTypeCodegenMethods, BuilderMethods};
 use rustc_middle::ty::{Instance, PseudoCanonicalInput, TyCtxt, TypingEnv};
@@ -294,6 +295,7 @@ pub(crate) fn generate_enzyme_call<'ll, 'tcx>(
     fn_args: &[&'ll Value],
     attrs: AutoDiffAttrs,
     dest: PlaceRef<'tcx, &'ll Value>,
+    fnc_tree: FncTree,
 ) {
     // We have to pick the name depending on whether we want forward or reverse mode autodiff.
     let mut ad_name: String = match attrs.mode {
@@ -370,6 +372,10 @@ pub(crate) fn generate_enzyme_call<'ll, 'tcx>(
         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);
+    }
+
     let call = builder.call(enzyme_ty, None, None, ad_fn, &args, None, None);
 
     builder.store_to_place(call, dest.val);
diff --git a/compiler/rustc_codegen_llvm/src/consts.rs b/compiler/rustc_codegen_llvm/src/consts.rs
index a110ecbb75d..40375ef6510 100644
--- a/compiler/rustc_codegen_llvm/src/consts.rs
+++ b/compiler/rustc_codegen_llvm/src/consts.rs
@@ -564,7 +564,7 @@ impl<'ll> CodegenCx<'ll, '_> {
         let g = self.define_global(&sym, llty).unwrap_or_else(|| {
             bug!("symbol `{}` is already defined", sym);
         });
-        set_global_alignment(self, g, self.tcx.data_layout.i8_align.abi);
+        set_global_alignment(self, g, self.tcx.data_layout.i8_align);
         llvm::set_initializer(g, llval);
         llvm::set_linkage(g, llvm::Linkage::PrivateLinkage);
         llvm::set_section(g, c"__TEXT,__cstring,cstring_literals");
@@ -680,7 +680,7 @@ impl<'ll> CodegenCx<'ll, '_> {
         let methname_g = self.define_global(&methname_sym, methname_llty).unwrap_or_else(|| {
             bug!("symbol `{}` is already defined", methname_sym);
         });
-        set_global_alignment(self, methname_g, self.tcx.data_layout.i8_align.abi);
+        set_global_alignment(self, methname_g, self.tcx.data_layout.i8_align);
         llvm::set_initializer(methname_g, methname_llval);
         llvm::set_linkage(methname_g, llvm::Linkage::PrivateLinkage);
         llvm::set_section(
diff --git a/compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs b/compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs
index 4b74c04ed7a..1e4ace4ca92 100644
--- a/compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs
+++ b/compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs
@@ -1043,7 +1043,7 @@ fn create_member_type<'ll, 'tcx>(
             file_metadata,
             line_number,
             layout.size.bits(),
-            layout.align.abi.bits() as u32,
+            layout.align.bits() as u32,
             offset.bits(),
             flags,
             type_di_node,
diff --git a/compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/native.rs b/compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/native.rs
index 62d38d463ab..1ae6e6e5eec 100644
--- a/compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/native.rs
+++ b/compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/native.rs
@@ -289,7 +289,7 @@ fn build_enum_variant_part_di_node<'ll, 'tcx>(
                 file_metadata,
                 line_number,
                 enum_type_and_layout.size.bits(),
-                enum_type_and_layout.align.abi.bits() as u32,
+                enum_type_and_layout.align.bits() as u32,
                 DIFlags::FlagZero,
                 tag_member_di_node,
                 create_DIArray(DIB(cx), &[]),
@@ -449,7 +449,7 @@ fn build_enum_variant_member_di_node<'ll, 'tcx>(
             file_di_node,
             line_number,
             enum_type_and_layout.size.bits(),
-            enum_type_and_layout.align.abi.bits() as u32,
+            enum_type_and_layout.align.bits() as u32,
             Size::ZERO.bits(),
             discr,
             DIFlags::FlagZero,
diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs
index 50398a32142..467655b0bfc 100644
--- a/compiler/rustc_codegen_llvm/src/intrinsic.rs
+++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs
@@ -297,7 +297,7 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                 let align = if name == sym::unaligned_volatile_load {
                     1
                 } else {
-                    result.layout.align.abi.bytes() as u32
+                    result.layout.align.bytes() as u32
                 };
                 unsafe {
                     llvm::LLVMSetAlignment(load, align);
@@ -1047,7 +1047,7 @@ fn codegen_emcc_try<'ll, 'tcx>(
         // create an alloca and pass a pointer to that.
         let ptr_size = bx.tcx().data_layout.pointer_size();
         let ptr_align = bx.tcx().data_layout.pointer_align().abi;
-        let i8_align = bx.tcx().data_layout.i8_align.abi;
+        let i8_align = bx.tcx().data_layout.i8_align;
         // Required in order for there to be no padding between the fields.
         assert!(i8_align <= ptr_align);
         let catch_data = bx.alloca(2 * ptr_size, ptr_align);
@@ -1212,6 +1212,9 @@ fn codegen_autodiff<'ll, 'tcx>(
         &mut diff_attrs.input_activity,
     );
 
+    let fnc_tree =
+        rustc_middle::ty::fnc_typetrees(tcx, fn_source.ty(tcx, TypingEnv::fully_monomorphized()));
+
     // Build body
     generate_enzyme_call(
         bx,
@@ -1222,6 +1225,7 @@ fn codegen_autodiff<'ll, 'tcx>(
         &val_arr,
         diff_attrs.clone(),
         result,
+        fnc_tree,
     );
 }
 
diff --git a/compiler/rustc_codegen_llvm/src/lib.rs b/compiler/rustc_codegen_llvm/src/lib.rs
index 13bdb7cb1a2..2405a25c702 100644
--- a/compiler/rustc_codegen_llvm/src/lib.rs
+++ b/compiler/rustc_codegen_llvm/src/lib.rs
@@ -68,6 +68,7 @@ mod llvm_util;
 mod mono_item;
 mod type_;
 mod type_of;
+mod typetree;
 mod va_arg;
 mod value;
 
@@ -231,6 +232,10 @@ impl CodegenBackend for LlvmCodegenBackend {
         crate::DEFAULT_LOCALE_RESOURCE
     }
 
+    fn name(&self) -> &'static str {
+        "llvm"
+    }
+
     fn init(&self, sess: &Session) {
         llvm_util::init(sess); // Make sure llvm is inited
     }
@@ -349,7 +354,14 @@ impl CodegenBackend for LlvmCodegenBackend {
 
         // Run the linker on any artifacts that resulted from the LLVM run.
         // This should produce either a finished executable or library.
-        link_binary(sess, &LlvmArchiveBuilderBuilder, codegen_results, metadata, outputs);
+        link_binary(
+            sess,
+            &LlvmArchiveBuilderBuilder,
+            codegen_results,
+            metadata,
+            outputs,
+            self.name(),
+        );
     }
 }
 
diff --git a/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs
index 695435eb6da..e63043b2122 100644
--- a/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs
+++ b/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs
@@ -3,9 +3,36 @@
 use libc::{c_char, c_uint};
 
 use super::MetadataKindId;
-use super::ffi::{AttributeKind, BasicBlock, Metadata, Module, Type, Value};
+use super::ffi::{AttributeKind, BasicBlock, Context, Metadata, Module, Type, Value};
 use crate::llvm::{Bool, Builder};
 
+// TypeTree types
+pub(crate) type CTypeTreeRef = *mut EnzymeTypeTree;
+
+#[repr(C)]
+#[derive(Debug, Copy, Clone)]
+pub(crate) struct EnzymeTypeTree {
+    _unused: [u8; 0],
+}
+
+#[repr(u32)]
+#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
+#[allow(non_camel_case_types)]
+pub(crate) enum CConcreteType {
+    DT_Anything = 0,
+    DT_Integer = 1,
+    DT_Pointer = 2,
+    DT_Half = 3,
+    DT_Float = 4,
+    DT_Double = 5,
+    DT_Unknown = 6,
+    DT_FP128 = 9,
+}
+
+pub(crate) struct TypeTree {
+    pub(crate) inner: CTypeTreeRef,
+}
+
 #[link(name = "llvm-wrapper", kind = "static")]
 unsafe extern "C" {
     // Enzyme
@@ -68,10 +95,40 @@ pub(crate) mod Enzyme_AD {
 
     use libc::c_void;
 
+    use super::{CConcreteType, CTypeTreeRef, Context};
+
     unsafe extern "C" {
         pub(crate) fn EnzymeSetCLBool(arg1: *mut ::std::os::raw::c_void, arg2: u8);
         pub(crate) fn EnzymeSetCLString(arg1: *mut ::std::os::raw::c_void, arg2: *const c_char);
     }
+
+    // TypeTree functions
+    unsafe extern "C" {
+        pub(crate) fn EnzymeNewTypeTree() -> CTypeTreeRef;
+        pub(crate) fn EnzymeNewTypeTreeCT(arg1: CConcreteType, ctx: &Context) -> CTypeTreeRef;
+        pub(crate) fn EnzymeNewTypeTreeTR(arg1: CTypeTreeRef) -> CTypeTreeRef;
+        pub(crate) fn EnzymeFreeTypeTree(CTT: CTypeTreeRef);
+        pub(crate) fn EnzymeMergeTypeTree(arg1: CTypeTreeRef, arg2: CTypeTreeRef) -> bool;
+        pub(crate) fn EnzymeTypeTreeOnlyEq(arg1: CTypeTreeRef, pos: i64);
+        pub(crate) fn EnzymeTypeTreeData0Eq(arg1: CTypeTreeRef);
+        pub(crate) fn EnzymeTypeTreeShiftIndiciesEq(
+            arg1: CTypeTreeRef,
+            data_layout: *const c_char,
+            offset: i64,
+            max_size: i64,
+            add_offset: u64,
+        );
+        pub(crate) fn EnzymeTypeTreeInsertEq(
+            CTT: CTypeTreeRef,
+            indices: *const i64,
+            len: usize,
+            ct: CConcreteType,
+            ctx: &Context,
+        );
+        pub(crate) fn EnzymeTypeTreeToString(arg1: CTypeTreeRef) -> *const c_char;
+        pub(crate) fn EnzymeTypeTreeToStringFree(arg1: *const c_char);
+    }
+
     unsafe extern "C" {
         static mut EnzymePrintPerf: c_void;
         static mut EnzymePrintActivity: c_void;
@@ -141,6 +198,67 @@ pub(crate) use self::Fallback_AD::*;
 pub(crate) mod Fallback_AD {
     #![allow(unused_variables)]
 
+    use libc::c_char;
+
+    use super::{CConcreteType, CTypeTreeRef, Context};
+
+    // TypeTree function fallbacks
+    pub(crate) unsafe fn EnzymeNewTypeTree() -> CTypeTreeRef {
+        unimplemented!()
+    }
+
+    pub(crate) unsafe fn EnzymeNewTypeTreeCT(arg1: CConcreteType, ctx: &Context) -> CTypeTreeRef {
+        unimplemented!()
+    }
+
+    pub(crate) unsafe fn EnzymeNewTypeTreeTR(arg1: CTypeTreeRef) -> CTypeTreeRef {
+        unimplemented!()
+    }
+
+    pub(crate) unsafe fn EnzymeFreeTypeTree(CTT: CTypeTreeRef) {
+        unimplemented!()
+    }
+
+    pub(crate) unsafe fn EnzymeMergeTypeTree(arg1: CTypeTreeRef, arg2: CTypeTreeRef) -> bool {
+        unimplemented!()
+    }
+
+    pub(crate) unsafe fn EnzymeTypeTreeOnlyEq(arg1: CTypeTreeRef, pos: i64) {
+        unimplemented!()
+    }
+
+    pub(crate) unsafe fn EnzymeTypeTreeData0Eq(arg1: CTypeTreeRef) {
+        unimplemented!()
+    }
+
+    pub(crate) unsafe fn EnzymeTypeTreeShiftIndiciesEq(
+        arg1: CTypeTreeRef,
+        data_layout: *const c_char,
+        offset: i64,
+        max_size: i64,
+        add_offset: u64,
+    ) {
+        unimplemented!()
+    }
+
+    pub(crate) unsafe fn EnzymeTypeTreeInsertEq(
+        CTT: CTypeTreeRef,
+        indices: *const i64,
+        len: usize,
+        ct: CConcreteType,
+        ctx: &Context,
+    ) {
+        unimplemented!()
+    }
+
+    pub(crate) unsafe fn EnzymeTypeTreeToString(arg1: CTypeTreeRef) -> *const c_char {
+        unimplemented!()
+    }
+
+    pub(crate) unsafe fn EnzymeTypeTreeToStringFree(arg1: *const c_char) {
+        unimplemented!()
+    }
+
     pub(crate) fn set_inline(val: bool) {
         unimplemented!()
     }
@@ -169,3 +287,89 @@ pub(crate) mod Fallback_AD {
         unimplemented!()
     }
 }
+
+impl TypeTree {
+    pub(crate) fn new() -> TypeTree {
+        let inner = unsafe { EnzymeNewTypeTree() };
+        TypeTree { inner }
+    }
+
+    pub(crate) fn from_type(t: CConcreteType, ctx: &Context) -> TypeTree {
+        let inner = unsafe { EnzymeNewTypeTreeCT(t, ctx) };
+        TypeTree { inner }
+    }
+
+    pub(crate) fn merge(self, other: Self) -> Self {
+        unsafe {
+            EnzymeMergeTypeTree(self.inner, other.inner);
+        }
+        drop(other);
+        self
+    }
+
+    #[must_use]
+    pub(crate) fn shift(
+        self,
+        layout: &str,
+        offset: isize,
+        max_size: isize,
+        add_offset: usize,
+    ) -> Self {
+        let layout = std::ffi::CString::new(layout).unwrap();
+
+        unsafe {
+            EnzymeTypeTreeShiftIndiciesEq(
+                self.inner,
+                layout.as_ptr(),
+                offset as i64,
+                max_size as i64,
+                add_offset as u64,
+            );
+        }
+
+        self
+    }
+
+    pub(crate) fn insert(&mut self, indices: &[i64], ct: CConcreteType, ctx: &Context) {
+        unsafe {
+            EnzymeTypeTreeInsertEq(self.inner, indices.as_ptr(), indices.len(), ct, ctx);
+        }
+    }
+}
+
+impl Clone for TypeTree {
+    fn clone(&self) -> Self {
+        let inner = unsafe { EnzymeNewTypeTreeTR(self.inner) };
+        TypeTree { inner }
+    }
+}
+
+impl std::fmt::Display for TypeTree {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let ptr = unsafe { EnzymeTypeTreeToString(self.inner) };
+        let cstr = unsafe { std::ffi::CStr::from_ptr(ptr) };
+        match cstr.to_str() {
+            Ok(x) => write!(f, "{}", x)?,
+            Err(err) => write!(f, "could not parse: {}", err)?,
+        }
+
+        // delete C string pointer
+        unsafe {
+            EnzymeTypeTreeToStringFree(ptr);
+        }
+
+        Ok(())
+    }
+}
+
+impl std::fmt::Debug for TypeTree {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        <Self as std::fmt::Display>::fmt(self, f)
+    }
+}
+
+impl Drop for TypeTree {
+    fn drop(&mut self) {
+        unsafe { EnzymeFreeTypeTree(self.inner) }
+    }
+}
diff --git a/compiler/rustc_codegen_llvm/src/typetree.rs b/compiler/rustc_codegen_llvm/src/typetree.rs
new file mode 100644
index 00000000000..7e263503700
--- /dev/null
+++ b/compiler/rustc_codegen_llvm/src/typetree.rs
@@ -0,0 +1,122 @@
+use rustc_ast::expand::typetree::FncTree;
+#[cfg(feature = "llvm_enzyme")]
+use {
+    crate::attributes,
+    rustc_ast::expand::typetree::TypeTree as RustTypeTree,
+    std::ffi::{CString, c_char, c_uint},
+};
+
+use crate::llvm::{self, Value};
+
+#[cfg(feature = "llvm_enzyme")]
+fn to_enzyme_typetree(
+    rust_typetree: RustTypeTree,
+    _data_layout: &str,
+    llcx: &llvm::Context,
+) -> llvm::TypeTree {
+    let mut enzyme_tt = llvm::TypeTree::new();
+    process_typetree_recursive(&mut enzyme_tt, &rust_typetree, &[], llcx);
+    enzyme_tt
+}
+#[cfg(feature = "llvm_enzyme")]
+fn process_typetree_recursive(
+    enzyme_tt: &mut llvm::TypeTree,
+    rust_typetree: &RustTypeTree,
+    parent_indices: &[i64],
+    llcx: &llvm::Context,
+) {
+    for rust_type in &rust_typetree.0 {
+        let concrete_type = match rust_type.kind {
+            rustc_ast::expand::typetree::Kind::Anything => llvm::CConcreteType::DT_Anything,
+            rustc_ast::expand::typetree::Kind::Integer => llvm::CConcreteType::DT_Integer,
+            rustc_ast::expand::typetree::Kind::Pointer => llvm::CConcreteType::DT_Pointer,
+            rustc_ast::expand::typetree::Kind::Half => llvm::CConcreteType::DT_Half,
+            rustc_ast::expand::typetree::Kind::Float => llvm::CConcreteType::DT_Float,
+            rustc_ast::expand::typetree::Kind::Double => llvm::CConcreteType::DT_Double,
+            rustc_ast::expand::typetree::Kind::F128 => llvm::CConcreteType::DT_FP128,
+            rustc_ast::expand::typetree::Kind::Unknown => llvm::CConcreteType::DT_Unknown,
+        };
+
+        let mut indices = parent_indices.to_vec();
+        if !parent_indices.is_empty() {
+            indices.push(rust_type.offset as i64);
+        } else if rust_type.offset == -1 {
+            indices.push(-1);
+        } else {
+            indices.push(rust_type.offset as i64);
+        }
+
+        enzyme_tt.insert(&indices, concrete_type, llcx);
+
+        if rust_type.kind == rustc_ast::expand::typetree::Kind::Pointer
+            && !rust_type.child.0.is_empty()
+        {
+            process_typetree_recursive(enzyme_tt, &rust_type.child, &indices, llcx);
+        }
+    }
+}
+
+#[cfg(feature = "llvm_enzyme")]
+pub(crate) fn add_tt<'ll>(
+    llmod: &'ll llvm::Module,
+    llcx: &'ll llvm::Context,
+    fn_def: &'ll Value,
+    tt: FncTree,
+) {
+    let inputs = tt.args;
+    let ret_tt: RustTypeTree = tt.ret;
+
+    let llvm_data_layout: *const c_char = unsafe { llvm::LLVMGetDataLayoutStr(&*llmod) };
+    let llvm_data_layout =
+        std::str::from_utf8(unsafe { std::ffi::CStr::from_ptr(llvm_data_layout) }.to_bytes())
+            .expect("got a non-UTF8 data-layout from LLVM");
+
+    let attr_name = "enzyme_type";
+    let c_attr_name = CString::new(attr_name).unwrap();
+
+    for (i, input) in inputs.iter().enumerate() {
+        unsafe {
+            let enzyme_tt = to_enzyme_typetree(input.clone(), llvm_data_layout, llcx);
+            let c_str = llvm::EnzymeTypeTreeToString(enzyme_tt.inner);
+            let c_str = std::ffi::CStr::from_ptr(c_str);
+
+            let attr = llvm::LLVMCreateStringAttribute(
+                llcx,
+                c_attr_name.as_ptr(),
+                c_attr_name.as_bytes().len() as c_uint,
+                c_str.as_ptr(),
+                c_str.to_bytes().len() as c_uint,
+            );
+
+            attributes::apply_to_llfn(fn_def, llvm::AttributePlace::Argument(i as u32), &[attr]);
+            llvm::EnzymeTypeTreeToStringFree(c_str.as_ptr());
+        }
+    }
+
+    unsafe {
+        let enzyme_tt = to_enzyme_typetree(ret_tt, llvm_data_layout, llcx);
+        let c_str = llvm::EnzymeTypeTreeToString(enzyme_tt.inner);
+        let c_str = std::ffi::CStr::from_ptr(c_str);
+
+        let ret_attr = llvm::LLVMCreateStringAttribute(
+            llcx,
+            c_attr_name.as_ptr(),
+            c_attr_name.as_bytes().len() as c_uint,
+            c_str.as_ptr(),
+            c_str.to_bytes().len() as c_uint,
+        );
+
+        attributes::apply_to_llfn(fn_def, llvm::AttributePlace::ReturnValue, &[ret_attr]);
+        llvm::EnzymeTypeTreeToStringFree(c_str.as_ptr());
+    }
+}
+
+#[cfg(not(feature = "llvm_enzyme"))]
+pub(crate) fn add_tt<'ll>(
+    _llmod: &'ll llvm::Module,
+    _llcx: &'ll llvm::Context,
+    _fn_def: &'ll Value,
+    _tt: FncTree,
+) {
+    unimplemented!()
+}
diff --git a/compiler/rustc_codegen_llvm/src/va_arg.rs b/compiler/rustc_codegen_llvm/src/va_arg.rs
index ab08125217f..234366e491c 100644
--- a/compiler/rustc_codegen_llvm/src/va_arg.rs
+++ b/compiler/rustc_codegen_llvm/src/va_arg.rs
@@ -193,7 +193,7 @@ fn emit_aapcs_va_arg<'ll, 'tcx>(
     // the offset again.
 
     bx.switch_to_block(maybe_reg);
-    if gr_type && layout.align.abi.bytes() > 8 {
+    if gr_type && layout.align.bytes() > 8 {
         reg_off_v = bx.add(reg_off_v, bx.const_i32(15));
         reg_off_v = bx.and(reg_off_v, bx.const_i32(-16));
     }
@@ -291,7 +291,7 @@ fn emit_powerpc_va_arg<'ll, 'tcx>(
         bx.inbounds_ptradd(va_list_addr, bx.const_usize(1)) // fpr
     };
 
-    let mut num_regs = bx.load(bx.type_i8(), num_regs_addr, dl.i8_align.abi);
+    let mut num_regs = bx.load(bx.type_i8(), num_regs_addr, dl.i8_align);
 
     // "Align" the register count when the type is passed as `i64`.
     if is_i64 || (is_f64 && is_soft_float_abi) {
@@ -329,7 +329,7 @@ fn emit_powerpc_va_arg<'ll, 'tcx>(
         // Increase the used-register count.
         let reg_incr = if is_i64 || (is_f64 && is_soft_float_abi) { 2 } else { 1 };
         let new_num_regs = bx.add(num_regs, bx.cx.const_u8(reg_incr));
-        bx.store(new_num_regs, num_regs_addr, dl.i8_align.abi);
+        bx.store(new_num_regs, num_regs_addr, dl.i8_align);
 
         bx.br(end);
 
@@ -339,7 +339,7 @@ fn emit_powerpc_va_arg<'ll, 'tcx>(
     let mem_addr = {
         bx.switch_to_block(in_mem);
 
-        bx.store(bx.const_u8(max_regs), num_regs_addr, dl.i8_align.abi);
+        bx.store(bx.const_u8(max_regs), num_regs_addr, dl.i8_align);
 
         // Everything in the overflow area is rounded up to a size of at least 4.
         let overflow_area_align = Align::from_bytes(4).unwrap();
@@ -738,6 +738,7 @@ fn copy_to_temporary_if_more_aligned<'ll, 'tcx>(
             src_align,
             bx.const_u32(layout.layout.size().bytes() as u32),
             MemFlags::empty(),
+            None,
         );
         tmp
     } else {
@@ -760,7 +761,7 @@ fn x86_64_sysv64_va_arg_from_memory<'ll, 'tcx>(
     // byte boundary if alignment needed by type exceeds 8 byte boundary.
     // It isn't stated explicitly in the standard, but in practice we use
     // alignment greater than 16 where necessary.
-    if layout.layout.align.abi.bytes() > 8 {
+    if layout.layout.align.bytes() > 8 {
         unreachable!("all instances of VaArgSafe have an alignment <= 8");
     }
 
@@ -813,7 +814,7 @@ fn emit_xtensa_va_arg<'ll, 'tcx>(
     let va_ndx_offset = va_reg_offset + 4;
     let offset_ptr = bx.inbounds_ptradd(va_list_addr, bx.cx.const_usize(va_ndx_offset));
 
-    let offset = bx.load(bx.type_i32(), offset_ptr, bx.tcx().data_layout.i32_align.abi);
+    let offset = bx.load(bx.type_i32(), offset_ptr, bx.tcx().data_layout.i32_align);
     let offset = round_up_to_alignment(bx, offset, layout.align.abi);
 
     let slot_size = layout.size.align_to(Align::from_bytes(4).unwrap()).bytes() as i32;