added typetree support for memcpy

KMJ-007 · KMJ-007 · commit a6a760ad36a1 · 2025-08-24T00:13:14.000Z
diff --git a/compiler/rustc_codegen_llvm/src/abi.rs b/compiler/rustc_codegen_llvm/src/abi.rs
@@ -244,6 +244,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/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;
 
@@ -1118,11 +1119,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,
@@ -1131,7 +1133,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/typetree.rs b/compiler/rustc_codegen_llvm/src/typetree.rs
@@ -1,7 +1,9 @@
+#[cfg(llvm_enzyme)]
 use std::ffi::{CString, c_char, c_uint};
 
 use rustc_ast::expand::typetree::{FncTree, TypeTree as RustTypeTree};
 
+#[cfg(llvm_enzyme)]
 use crate::attributes;
 use crate::llvm::{self, Value};
 
@@ -50,15 +52,6 @@ fn to_enzyme_typetree(
     enzyme_tt
 }
 
-#[cfg(not(llvm_enzyme))]
-fn to_enzyme_typetree(
-    _rust_typetree: RustTypeTree,
-    _data_layout: &str,
-    _llcx: &llvm::Context,
-) -> ! {
-    unimplemented!("TypeTree conversion not available without llvm_enzyme support")
-}
-
 // Attaches TypeTree information to LLVM function as enzyme_type attributes.
 #[cfg(llvm_enzyme)]
 pub(crate) fn add_tt<'ll>(
diff --git a/compiler/rustc_codegen_llvm/src/va_arg.rs b/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 {
diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs
@@ -1623,6 +1623,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);
diff --git a/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs b/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);
     }
 }
 
diff --git a/compiler/rustc_codegen_ssa/src/mir/statement.rs b/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 { .. }
diff --git a/compiler/rustc_codegen_ssa/src/traits/builder.rs b/compiler/rustc_codegen_ssa/src/traits/builder.rs
@@ -424,6 +424,7 @@ pub trait BuilderMethods<'a, 'tcx>:
         src_align: Align,
         size: Self::Value,
         flags: MemFlags,
+        tt: Option<rustc_ast::expand::typetree::FncTree>,
     );
     fn memmove(
         &mut self,
@@ -480,7 +481,7 @@ pub trait BuilderMethods<'a, 'tcx>:
             temp.val.store_with_flags(self, dst.with_type(layout), flags);
         } else if !layout.is_zst() {
             let bytes = self.const_usize(layout.size.bytes());
-            self.memcpy(dst.llval, dst.align, src.llval, src.align, bytes, flags);
+            self.memcpy(dst.llval, dst.align, src.llval, src.align, bytes, flags, None);
         }
     }
 
diff --git a/compiler/rustc_interface/src/tests.rs b/compiler/rustc_interface/src/tests.rs
@@ -764,7 +764,6 @@ fn test_unstable_options_tracking_hash() {
     tracked!(allow_features, Some(vec![String::from("lang_items")]));
     tracked!(always_encode_mir, true);
     tracked!(assume_incomplete_release, true);
-    tracked!(autodiff, vec![AutoDiff::Enable]);
     tracked!(autodiff, vec![AutoDiff::Enable, AutoDiff::NoTT]);
     tracked!(binary_dep_depinfo, true);
     tracked!(box_noalias, false);
diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs
@@ -2286,12 +2286,12 @@ pub fn typetree_from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> TypeTree {
         let child = typetree_from_ty(tcx, inner_ty);
         return TypeTree(vec![Type {
             offset: -1,
-            size: 8, // TODO(KMJ-007): Get actual pointer size from target
+            size: tcx.data_layout.pointer_size().bytes_usize(),
             kind: Kind::Pointer,
             child,
         }]);
     }
 
-    // TODO(KMJ-007): Handle arrays, slices, structs, and other complex types
+    // FIXME(KMJ-007): Handle arrays, slices, structs, and other complex types
     TypeTree::new()
 }
diff --git a/tests/codegen-llvm/autodiff/typetree.rs b/tests/codegen-llvm/autodiff/typetree.rs
@@ -30,4 +30,4 @@ fn main() {
     let output_ = d_simple(&x, &mut df_dx, 1.0);
     assert_eq!(output, output_);
     assert_eq!(2.0, df_dx);
-}
+}
diff --git a/tests/run-make/autodiff/type-trees/memcpy-typetree/memcpy-ir.check b/tests/run-make/autodiff/type-trees/memcpy-typetree/memcpy-ir.check
@@ -0,0 +1,8 @@
+; Check that enzyme_type attributes are present in the LLVM IR function definition
+; This verifies our TypeTree system correctly attaches metadata for Enzyme
+
+CHECK: define{{.*}}"enzyme_type"="{[]:Float@double}"{{.*}}@test_memcpy({{.*}}"enzyme_type"="{[]:Pointer}"
+
+; Check that llvm.memcpy exists (either call or declare)
+CHECK: {{(call|declare).*}}@llvm.memcpy
+
diff --git a/tests/run-make/autodiff/type-trees/memcpy-typetree/memcpy.check b/tests/run-make/autodiff/type-trees/memcpy-typetree/memcpy.check
@@ -0,0 +1,13 @@
+CHECK: force_memcpy
+
+CHECK: @llvm.memcpy.p0.p0.i64
+
+CHECK: test_memcpy - {[-1]:Float@double} |{[-1]:Pointer}:{}
+
+CHECK-DAG: ptr %{{[0-9]+}}: {[-1]:Pointer, [-1,0]:Float@double, [-1,8]:Float@double, [-1,16]:Float@double, [-1,24]:Float@double}
+
+CHECK-DAG: load double{{.*}}: {[-1]:Float@double}
+
+CHECK-DAG: fmul double{{.*}}: {[-1]:Float@double}
+
+CHECK-DAG: fadd double{{.*}}: {[-1]:Float@double}
diff --git a/tests/run-make/autodiff/type-trees/memcpy-typetree/memcpy.rs b/tests/run-make/autodiff/type-trees/memcpy-typetree/memcpy.rs
@@ -0,0 +1,36 @@
+#![feature(autodiff)]
+
+use std::autodiff::autodiff_reverse;
+use std::ptr;
+
+#[inline(never)]
+fn force_memcpy(src: *const f64, dst: *mut f64, count: usize) {
+    unsafe {
+        ptr::copy_nonoverlapping(src, dst, count);
+    }
+}
+
+#[autodiff_reverse(d_test_memcpy, Duplicated, Active)]
+#[no_mangle]
+fn test_memcpy(input: &[f64; 128]) -> f64 {
+    let mut local_data = [0.0f64; 128];
+
+    // Use a separate function to prevent inlining and optimization
+    force_memcpy(input.as_ptr(), local_data.as_mut_ptr(), 128);
+
+    // Sum only first few elements to keep the computation simple
+    local_data[0] * local_data[0]
+        + local_data[1] * local_data[1]
+        + local_data[2] * local_data[2]
+        + local_data[3] * local_data[3]
+}
+
+fn main() {
+    let input = [1.0; 128];
+    let mut d_input = [0.0; 128];
+    let result = test_memcpy(&input);
+    let result_d = d_test_memcpy(&input, &mut d_input, 1.0);
+
+    assert_eq!(result, result_d);
+    println!("Memcpy test passed: result = {}", result);
+}
diff --git a/tests/run-make/autodiff/type-trees/memcpy-typetree/rmake.rs b/tests/run-make/autodiff/type-trees/memcpy-typetree/rmake.rs
@@ -0,0 +1,39 @@
+//@ needs-enzyme
+//@ ignore-cross-compile
+
+use run_make_support::{llvm_filecheck, rfs, rustc};
+
+fn main() {
+    // First, compile to LLVM IR to check for enzyme_type attributes
+    let _ir_output = rustc()
+        .input("memcpy.rs")
+        .arg("-Zautodiff=Enable")
+        .arg("-Zautodiff=NoPostopt")
+        .opt_level("0")
+        .arg("--emit=llvm-ir")
+        .arg("-o")
+        .arg("main.ll")
+        .run();
+
+    // Then compile with TypeTree analysis output for the existing checks
+    let output = rustc()
+        .input("memcpy.rs")
+        .arg("-Zautodiff=Enable,PrintTAFn=test_memcpy")
+        .arg("-Zautodiff=NoPostopt")
+        .opt_level("3")
+        .arg("-Clto=fat")
+        .arg("-g")
+        .run();
+
+    let stdout = output.stdout_utf8();
+    let stderr = output.stderr_utf8();
+    let ir_content = rfs::read_to_string("main.ll");
+
+    rfs::write("memcpy.stdout", &stdout);
+    rfs::write("memcpy.stderr", &stderr);
+    rfs::write("main.ir", &ir_content);
+
+    llvm_filecheck().patterns("memcpy.check").stdin_buf(stdout).run();
+
+    llvm_filecheck().patterns("memcpy-ir.check").stdin_buf(ir_content).run();
+}
diff --git a/tests/run-make/autodiff/type-trees/nott-flag/rmake.rs b/tests/run-make/autodiff/type-trees/nott-flag/rmake.rs
@@ -23,14 +23,8 @@ fn main() {
         .run();
 
     // Verify NoTT version does NOT have enzyme_type attributes
-    llvm_filecheck()
-        .patterns("nott.check")
-        .stdin_buf(rfs::read("nott.ll"))
-        .run();
-    
+    llvm_filecheck().patterns("nott.check").stdin_buf(rfs::read("nott.ll")).run();
+
     // Verify TypeTree version DOES have enzyme_type attributes
-    llvm_filecheck()
-        .patterns("with_tt.check")
-        .stdin_buf(rfs::read("with_tt.ll"))
-        .run();
-}
+    llvm_filecheck().patterns("with_tt.check").stdin_buf(rfs::read("with_tt.ll")).run();
+}
diff --git a/tests/run-make/autodiff/type-trees/nott-flag/test.rs b/tests/run-make/autodiff/type-trees/nott-flag/test.rs
@@ -12,4 +12,4 @@ fn main() {
     let x = 2.0;
     let mut dx = 0.0;
     let _result = d_square(&x, &mut dx, 1.0);
-}
+}
diff --git a/tests/ui/autodiff/flag_nott.rs b/tests/ui/autodiff/flag_nott.rs
@@ -16,4 +16,4 @@ fn main() {
     let x = 2.0;
     let mut dx = 0.0;
     let result = d_square(&x, &mut dx, 1.0);
-}
+}

Original file line number	Diff line number	Diff line change
`@@ -244,6 +244,7 @@ impl<'ll, 'tcx> ArgAbiExt<'ll, 'tcx> for ArgAbi<'tcx, Ty<'tcx>> {`
`244`	`244`	`scratch_align,`
`245`	`245`	`bx.const_usize(copy_bytes),`
`246`	`246`	`MemFlags::empty(),`
	`247`	`+ None,`
`247`	`248`	`);`
`248`	`249`	`bx.lifetime_end(llscratch, scratch_size);`
`249`	`250`	`}`
Original file line number	Diff line number	Diff line change
`@@ -30,7 +30,7 @@ fn copy_intrinsic<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(`
`30`	`30`	`if allow_overlap {`
`31`	`31`	`bx.memmove(dst, align, src, align, size, flags);`
`32`	`32`	`} else {`
`33`		`- bx.memcpy(dst, align, src, align, size, flags);`
	`33`	`+ bx.memcpy(dst, align, src, align, size, flags, None);`
`34`	`34`	`}`
`35`	`35`	`}`
`36`	`36`
Original file line number	Diff line number	Diff line change
`@@ -90,7 +90,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {`
`90`	`90`	`let align = pointee_layout.align;`
`91`	`91`	`let dst = dst_val.immediate();`
`92`	`92`	`let src = src_val.immediate();`
`93`		`- bx.memcpy(dst, align, src, align, bytes, crate::MemFlags::empty());`
	`93`	`+ bx.memcpy(dst, align, src, align, bytes, crate::MemFlags::empty(), None);`
`94`	`94`	`}`
`95`	`95`	`mir::StatementKind::FakeRead(..)`
`96`	`96`	`\| mir::StatementKind::Retag { .. }`
Original file line number	Diff line number	Diff line change
`@@ -2286,12 +2286,12 @@ pub fn typetree_from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> TypeTree {`
`2286`	`2286`	`let child = typetree_from_ty(tcx, inner_ty);`
`2287`	`2287`	`return TypeTree(vec![Type {`
`2288`	`2288`	`offset: -1,`
`2289`		`- size: 8, // TODO(KMJ-007): Get actual pointer size from target`
	`2289`	`+ size: tcx.data_layout.pointer_size().bytes_usize(),`
`2290`	`2290`	`kind: Kind::Pointer,`
`2291`	`2291`	`child,`
`2292`	`2292`	`}]);`
`2293`	`2293`	`}`
`2294`	`2294`
`2295`		`- // TODO(KMJ-007): Handle arrays, slices, structs, and other complex types`
	`2295`	`+ // FIXME(KMJ-007): Handle arrays, slices, structs, and other complex types`
`2296`	`2296`	`TypeTree::new()`
`2297`	`2297`	`}`