fixing typetree metadata issue

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

Author: KMJ-007

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -504,6 +504,9 @@ pub(crate) fn differentiate<'ll>(
                 },
             ));
         };
+        // Attach TypeTree metadata to the source function before calling Enzyme
+        let fnc_tree = FncTree { args: item.inputs.clone(), ret: item.output.clone() };
+        add_tt(cx.llmod, cx.llcx, fn_def, fnc_tree);
 
         generate_enzyme_call(&cx, fn_def, fn_target, item.attrs.clone());
     }

tests/run-make/autodiff/type-trees/type-analysis/memcpy/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 the differentiated function also has proper enzyme_type attributes
+
+CHECK: @diffetest_memcpy({{.*}}"enzyme_type"="{[]:Pointer}"{{.*}}"enzyme_type"="{[]:Pointer}"

tests/run-make/autodiff/type-trees/type-analysis/memcpy/memcpy.rs

@@ -12,12 +12,15 @@ fn test_memcpy(input: &[f64; 8]) -> f64 {
         ptr::copy_nonoverlapping(input.as_ptr(), local_data.as_mut_ptr(), 8);
     }
 
-    let mut result = 0.0;
-    for i in 0..8 {
-        result += local_data[i] * local_data[i];
-    }
-
-    result
+    // type tree does not support loops
+    local_data[0] * local_data[0]
+        + local_data[1] * local_data[1]
+        + local_data[2] * local_data[2]
+        + local_data[3] * local_data[3]
+        + local_data[4] * local_data[4]
+        + local_data[5] * local_data[5]
+        + local_data[6] * local_data[6]
+        + local_data[7] * local_data[7]
 }
 
 fn main() {

tests/run-make/autodiff/type-trees/type-analysis/memcpy/rmake.rs

@@ -1,12 +1,22 @@
 //@ needs-enzyme
 //@ ignore-cross-compile
 
-use std::fs;
-
 use run_make_support::{llvm_filecheck, rfs, rustc};
 
 fn main() {
-    // Compile the Rust file with the required flags, capturing both stdout and stderr
+    // 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("3")
+        .arg("-Clto=fat")
+        .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")
@@ -18,11 +28,13 @@ fn main() {
 
     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);
 
-    // Write the outputs to files
-    rfs::write("memcpy.stdout", stdout);
-    rfs::write("memcpy.stderr", stderr);
+    llvm_filecheck().patterns("memcpy.check").stdin_buf(stdout).run();
 
-    // Run FileCheck on the stdout using the check file
-    llvm_filecheck().patterns("memcpy.check").stdin_buf(rfs::read("memcpy.stdout")).run();
+    llvm_filecheck().patterns("memcpy-ir.check").stdin_buf(ir_content).run();
 }