memcpy compiling

Signed-off-by: Karan Janthe <[email protected]>

Author: KMJ-007

compiler/rustc_codegen_llvm/src/abi.rs

@@ -238,6 +238,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);
             }

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;
 
@@ -31,6 +32,7 @@ use tracing::{debug, instrument};
 
 use crate::abi::FnAbiLlvmExt;
 use crate::attributes;
+use crate::builder::autodiff::add_tt;
 use crate::common::Funclet;
 use crate::context::{CodegenCx, FullCx, GenericCx, SCx};
 use crate::llvm::{
@@ -1105,11 +1107,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,
@@ -1118,7 +1121,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 {
+            add_tt(self.cx().llmod, self.cx().llcx, memcpy, tt);
         }
     }
 

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -1,6 +1,8 @@
+use std::os::raw::{c_char, c_uint};
 use std::ptr;
 
 use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, AutoDiffItem, DiffActivity, DiffMode};
+use rustc_ast::expand::typetree::{FncTree, TypeTree as RustTypeTree};
 use rustc_codegen_ssa::ModuleCodegen;
 use rustc_codegen_ssa::common::TypeKind;
 use rustc_codegen_ssa::traits::BaseTypeCodegenMethods;
@@ -512,3 +514,128 @@ pub(crate) fn differentiate<'ll>(
 
     Ok(())
 }
+
+/// Converts a Rust TypeTree to Enzyme's internal TypeTree format
+///
+/// This function takes a Rust-side TypeTree (from rustc_ast::expand::typetree)
+/// and converts it to Enzyme's internal C++ TypeTree representation that
+/// Enzyme can understand during differentiation analysis.
+fn to_enzyme_typetree(
+    rust_typetree: RustTypeTree,
+    data_layout: &str,
+    llcx: &llvm::Context,
+) -> llvm::TypeTree {
+    // Start with an empty TypeTree
+    let mut enzyme_tt = llvm::TypeTree::new();
+
+    // Convert each Type in the Rust TypeTree to Enzyme format
+    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::Unknown => llvm::CConcreteType::DT_Unknown,
+        };
+
+        // Create a TypeTree for this specific type
+        let type_tt = llvm::TypeTree::from_type(concrete_type, llcx);
+
+        // Apply offset if specified
+        let type_tt = if rust_type.offset == -1 {
+            type_tt // -1 means everywhere/no specific offset
+        } else {
+            // Apply specific offset positioning
+            type_tt.shift(data_layout, rust_type.offset, rust_type.size as isize, 0)
+        };
+
+        // Merge this type into the main TypeTree
+        enzyme_tt = enzyme_tt.merge(type_tt);
+    }
+
+    enzyme_tt
+}
+
+/// Attaches TypeTree information to LLVM function as enzyme_type attributes.
+///
+/// This function converts Rust TypeTrees to Enzyme format and attaches them as
+/// LLVM string attributes. Enzyme reads these attributes during autodiff analysis
+/// to understand the memory layout and generate correct derivative code.
+///
+/// # Arguments
+/// * `llmod` - LLVM module containing the function
+/// * `llcx` - LLVM context for creating attributes  
+/// * `fn_def` - LLVM function to attach TypeTrees to
+/// * `tt` - Function TypeTree containing input and return type information
+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;
+
+    // Get LLVM data layout string for TypeTree conversion
+    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");
+
+    // Attribute name that Enzyme recognizes for TypeTree information
+    let attr_name = "enzyme_type";
+    let c_attr_name = std::ffi::CString::new(attr_name).unwrap();
+
+    // Attach TypeTree attributes to each input parameter
+    // Enzyme uses these to understand parameter memory layouts during differentiation
+    for (i, input) in inputs.iter().enumerate() {
+        unsafe {
+            // Convert Rust TypeTree to Enzyme's internal format
+            let enzyme_tt = to_enzyme_typetree(input.clone(), llvm_data_layout, llcx);
+
+            // Serialize TypeTree to string format that Enzyme can parse
+            let c_str = llvm::EnzymeTypeTreeToString(enzyme_tt.inner);
+            let c_str = std::ffi::CStr::from_ptr(c_str);
+
+            // Create LLVM string attribute with TypeTree information
+            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,
+            );
+
+            // Attach attribute to the specific function parameter
+            // Note: ArgumentPlace uses 0-based indexing, but LLVM uses 1-based for arguments
+            attributes::apply_to_llfn(fn_def, llvm::AttributePlace::Argument(i as u32), &[attr]);
+
+            // Free the C string to prevent memory leaks
+            llvm::EnzymeTypeTreeToStringFree(c_str.as_ptr());
+        }
+    }
+
+    // Attach TypeTree attribute to the return type
+    // Enzyme needs this to understand how to handle return value derivatives
+    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,
+        );
+
+        // Attach to function return type
+        attributes::apply_to_llfn(fn_def, llvm::AttributePlace::ReturnValue, &[ret_attr]);
+
+        // Free the C string
+        llvm::EnzymeTypeTreeToStringFree(c_str.as_ptr());
+    }
+}

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -2142,7 +2142,7 @@ unsafe extern "C" {
         SPFlags: DISPFlags,
         MaybeFn: Option<&'a Value>,
         TParam: &'a DIArray,
-        Decl: Option<&'a DIDescriptor>,
+        Decl: Option<&'a DIDescriptor>, 
     ) -> &'a DISubprogram;
 
     pub(crate) fn LLVMRustDIBuilderCreateMethod<'a>(
@@ -2673,4 +2673,144 @@ unsafe extern "C" {
 
     pub(crate) fn LLVMRustSetNoSanitizeAddress(Global: &Value);
     pub(crate) fn LLVMRustSetNoSanitizeHWAddress(Global: &Value);
+
+    // ========== ENZYME AUTODIFF FFI FUNCTIONS ==========
+
+    // Enzyme Type Tree Functions (minimal set for TypeTree support)
+    pub(crate) fn EnzymeNewTypeTree() -> CTypeTreeRef;
+    pub(crate) fn EnzymeFreeTypeTree(CTT: CTypeTreeRef);
+    pub(crate) fn EnzymeNewTypeTreeCT(arg1: CConcreteType, ctx: &Context) -> CTypeTreeRef;
+    pub(crate) fn EnzymeNewTypeTreeTR(arg1: CTypeTreeRef) -> CTypeTreeRef;
+    pub(crate) fn EnzymeMergeTypeTree(arg1: CTypeTreeRef, arg2: CTypeTreeRef);
+    pub(crate) fn EnzymeTypeTreeShiftIndiciesEq(
+        arg1: CTypeTreeRef,
+        data_layout: *const c_char,
+        offset: i64,
+        max_size: i64,
+        add_offset: u64,
+    );
+    pub(crate) fn EnzymeTypeTreeToString(arg1: CTypeTreeRef) -> *const c_char;
+    pub(crate) fn EnzymeTypeTreeToStringFree(arg1: *const c_char);
 }
+
+// ========== ENZYME TYPES AND ENUMS ==========
+
+// Type Tree Support for Autodiff
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+#[repr(u32)]
+#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
+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,
+}
+
+pub(crate) type CTypeTreeRef = *mut EnzymeTypeTree;
+
+#[repr(C)]
+#[derive(Debug, Copy, Clone)]
+pub(crate) struct EnzymeTypeTree {
+    _unused: [u8; 0],
+}
+
+
+
+
+
+
+
+// TypeTree wrapper for Rust-side type safety and memory management
+pub(crate) struct TypeTree {
+    pub(crate) inner: CTypeTreeRef,
+}
+
+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
+    }
+}
+
+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) }
+    }
+}
+
+

compiler/rustc_codegen_llvm/src/va_arg.rs

@@ -735,6 +735,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 {

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -1550,6 +1550,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     align,
                     bx.const_usize(copy_bytes),
                     MemFlags::empty(),
+                    None,
                 );
                 // ...and then load it with the ABI type.
                 llval = load_cast(bx, cast, llscratch, scratch_align);

compiler/rustc_codegen_ssa/src/mir/intrinsic.rs

@@ -30,7 +30,7 @@ fn copy_intrinsic<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     if allow_overlap {
         bx.memmove(dst, align, src, align, size, flags);
     } else {
-        bx.memcpy(dst, align, src, align, size, flags);
+        bx.memcpy(dst, align, src, align, size, flags, None);
     }
 }
 

compiler/rustc_codegen_ssa/src/mir/statement.rs

@@ -90,7 +90,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 let align = pointee_layout.align;
                 let dst = dst_val.immediate();
                 let src = src_val.immediate();
-                bx.memcpy(dst, align, src, align, bytes, crate::MemFlags::empty());
+                bx.memcpy(dst, align, src, align, bytes, crate::MemFlags::empty(), None);
             }
             mir::StatementKind::FakeRead(..)
             | mir::StatementKind::Retag { .. }