added typetree support for memcpy

8 files changed, 103 insertions, 13 deletions

diff --git a/tests/codegen-llvm/autodiff/typetree.rs b/tests/codegen-llvm/autodiff/typetree.rs
index 3ad38d581b9..1cb0c2fb68b 100644
--- 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);
-}
\ No newline at end of file
+}
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
new file mode 100644
index 00000000000..3a59f06dda1
--- /dev/null
+++ 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
new file mode 100644
index 00000000000..ae70830297a
--- /dev/null
+++ 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}
\ No newline at end of file
diff --git a/tests/run-make/autodiff/type-trees/memcpy-typetree/memcpy.rs b/tests/run-make/autodiff/type-trees/memcpy-typetree/memcpy.rs
new file mode 100644
index 00000000000..3c1029190c8
--- /dev/null
+++ 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
new file mode 100644
index 00000000000..b4c650330fe
--- /dev/null
+++ 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
index bab863ca9ff..de540b990ca 100644
--- 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();
-}
\ No newline at end of file
+    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
index 5c634eea035..de3549c37c6 100644
--- 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);
-}
\ No newline at end of file
+}
diff --git a/tests/ui/autodiff/flag_nott.rs b/tests/ui/autodiff/flag_nott.rs
index 7a97d892cd8..faa9949fe81 100644
--- 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);
-}
\ No newline at end of file
+}