add test, cleanup

Author: ZuseZ4

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -43,7 +43,7 @@ use crate::errors::{
     AutoDiffWithoutEnable, AutoDiffWithoutLto, IntrinsicSignatureMismatch, IntrinsicWrongArch,
     OffloadWithoutEnable, OffloadWithoutFatLTO, UnknownIntrinsic,
 };
-use crate::intrinsic::ty::typetree::{fnc_typetrees, typetree_from_ty};
+use crate::intrinsic::ty::typetree::fnc_typetrees;
 use crate::llvm::{self, Type, Value};
 use crate::type_of::LayoutLlvmExt;
 use crate::va_arg::emit_va_arg;
@@ -1933,10 +1933,8 @@ fn get_args_from_tuple<'ll, 'tcx>(
                         let llvm_ty = field.layout.llvm_type(bx.cx);
                         let pair_val = bx.load(llvm_ty, field.val.llval, field.val.align);
 
-                        let extract_ty = field.layout.ty;
-                        let tt = typetree_from_ty(bx.tcx(), extract_ty);
-                        result.push(bx.extract_value(pair_val, 0, Some(tt.clone())));
-                        result.push(bx.extract_value(pair_val, 1, Some(tt)));
+                        result.push(bx.extract_value(pair_val, 0, None));
+                        result.push(bx.extract_value(pair_val, 1, None));
                         tuple_index += 1;
                     }
                     PassMode::Indirect { .. } => {

compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs

@@ -66,7 +66,7 @@ unsafe extern "C" {
         NameLen: libc::size_t,
     ) -> Option<&Value>;
 
-    pub(crate) fn LLVMRustIsIntrinsicCall(V: &Value) -> bool;
+    pub(crate) safe fn LLVMRustIsIntrinsicCall(V: &Value) -> bool;
     pub(crate) safe fn LLVMRustSupportsEnzymeMD(V: &Value) -> bool;
     pub(crate) fn LLVMRustSetEnzymeTypeMD(
         v: &Value,

compiler/rustc_codegen_llvm/src/typetree.rs

@@ -70,6 +70,12 @@ fn process_typetree_recursive(
     extra_ints
 }
 
+enum TTLocation {
+    Definition,
+    Callsite,
+    Intrinsic,
+}
+
 #[cfg_attr(not(feature = "llvm_enzyme"), allow(unused))]
 pub(crate) fn add_tt<'tcx, 'll>(
     llmod: &'ll llvm::Module,
@@ -104,11 +110,19 @@ pub(crate) fn add_tt<'tcx, 'll>(
     let attr_name = "enzyme_type";
     let c_attr_name = CString::new(attr_name).unwrap();
 
+    let tt_location: TTLocation = if llvm::LLVMRustIsIntrinsicCall(fn_def) {
+        TTLocation::Callsite
+    } else if llvm::LLVMRustSupportsEnzymeMD(fn_def) {
+        TTLocation::Intrinsic
+    } else {
+        TTLocation::Definition
+    };
+
     let mut offset = 0;
     for (i, input) in inputs.iter().enumerate() {
         let (enzyme_tt, extra_ints) = to_enzyme_typetree(&input, llvm_data_layout, llcx);
 
-        // This scope is just a visual reminder that we *must* drop the enzyme_wrapper before
+        // This scope is a simple solution since we *must* drop the enzyme_wrapper before
         // we drop any typetrees (mainly enzyme_tt and extra_ints). Drop calls can not accept
         // arguments like an enzyme_wrapper, so the typetree drop impl has to call get_instance
         // on the static enzyme instance, which is behind a Mutex. Therefore we'd deadlock if we
@@ -118,74 +132,59 @@ pub(crate) fn add_tt<'tcx, 'll>(
             let c_str = enzyme_wrapper.tree_to_cstr(enzyme_tt.inner);
 
             let attr = llvm::CreateAttrStringValueFromCStr(llcx, &c_attr_name, &c_str);
-            dbg!(&fn_def);
-            if llvm::LLVMRustSupportsEnzymeMD(fn_def) {
-                dbg!("extractvalue md");
-                let md = enzyme_wrapper.tree_to_md(enzyme_tt.inner, llcx);
-                unsafe {
-                    llvm::LLVMRustSetEnzymeTypeMD(fn_def, md.unwrap());
+            let arg_pos = llvm::AttributePlace::Argument(i as u32 + offset);
+            match tt_location {
+                TTLocation::Definition => {
+                    attributes::apply_to_llfn(fn_def, arg_pos, &[attr]);
+                }
+                TTLocation::Callsite => {
+                    attributes::apply_to_callsite(fn_def, arg_pos, &[attr]);
+                }
+                TTLocation::Intrinsic => {
+                    let md = enzyme_wrapper.tree_to_md(enzyme_tt.inner, llcx);
+                    unsafe {
+                        llvm::LLVMRustSetEnzymeTypeMD(fn_def, md.unwrap());
+                    }
                 }
-            } else if unsafe { llvm::LLVMRustIsIntrinsicCall(fn_def) } {
-                dbg!("callsite");
-                attributes::apply_to_callsite(
-                    fn_def,
-                    llvm::AttributePlace::Argument(i as u32 + offset),
-                    &[attr],
-                );
-            } else {
-                dbg!("llfn");
-                attributes::apply_to_llfn(
-                    fn_def,
-                    llvm::AttributePlace::Argument(i as u32 + offset),
-                    &[attr],
-                );
             }
             enzyme_wrapper.tree_to_string_free(c_str.as_ptr());
             for v in &extra_ints {
                 offset += 1;
                 let c_str = enzyme_wrapper.tree_to_cstr(v.inner);
                 let int_attr = llvm::CreateAttrStringValueFromCStr(llcx, &c_attr_name, &c_str);
-                dbg!(&fn_def);
-                if llvm::LLVMRustSupportsEnzymeMD(fn_def) {
-                    dbg!("extractvalue input(?)");
-                    let md = enzyme_wrapper.tree_to_md(enzyme_tt.inner, llcx);
-                    unsafe {
-                        llvm::LLVMRustSetEnzymeTypeMD(fn_def, md.unwrap());
+                let arg_pos = llvm::AttributePlace::Argument(i as u32 + offset);
+                match tt_location {
+                    TTLocation::Intrinsic => {
+                        let md = enzyme_wrapper.tree_to_md(enzyme_tt.inner, llcx);
+                        unsafe {
+                            llvm::LLVMRustSetEnzymeTypeMD(fn_def, md.unwrap());
+                        }
+                    }
+                    TTLocation::Definition => {
+                        attributes::apply_to_llfn(fn_def, arg_pos, &[int_attr]);
+                    }
+                    TTLocation::Callsite => {
+                        attributes::apply_to_callsite(fn_def, arg_pos, &[int_attr]);
                     }
-                } else if unsafe { llvm::LLVMRustIsIntrinsicCall(fn_def) } {
-                    dbg!("callsite input");
-                    attributes::apply_to_callsite(
-                        fn_def,
-                        llvm::AttributePlace::Argument(i as u32 + offset),
-                        &[int_attr],
-                    );
-                } else {
-                    dbg!("llfn");
-                    attributes::apply_to_llfn(
-                        fn_def,
-                        llvm::AttributePlace::Argument(i as u32 + offset),
-                        &[int_attr],
-                    );
                 }
                 enzyme_wrapper.tree_to_string_free(c_str.as_ptr());
             }
         }
     }
-    // We will only fail this if Rust types got lowered to LLVM in a way that we didn't predict.
-    // Error, so we can learn from our mistakes.
-    //if unsafe { !llvm::LLVMRustIsIntrinsicCall(fn_def) } {
-    //    dbg!("checking parameter count");
-    //    let expected = offset as usize + inputs.len();
-    //    let actual = llvm::count_params(fn_def) as usize;
-    //    if expected != actual {
-    //        tcx.dcx().warn(format!(
-    //            "autodiff type-tree failure. We expected {expected} LLVM argument(s), \
-    //             but the generated LLVM function has {actual} parameter(s)"
-    //        ));
-    //    }
-    //}
-
-    let (enzyme_tt, extra_ints) = to_enzyme_typetree(&ret_tt, llvm_data_layout, llcx);
+    // We will fail here if Rust types got lowered to LLVM in a way that we didn't predict.
+    // We Error, so we can learn from our mistakes.
+    if matches!(tt_location, TTLocation::Definition) {
+        let expected = offset as usize + inputs.len();
+        let actual = llvm::count_params(fn_def) as usize;
+        if expected != actual {
+            tcx.dcx().warn(format!(
+                "autodiff type-tree failure. We expected {expected} LLVM argument(s), \
+                 but the generated LLVM function has {actual} parameter(s)"
+            ));
+        }
+    }
+
+    let (enzyme_tt, _extra_ints) = to_enzyme_typetree(&ret_tt, llvm_data_layout, llcx);
     if ret_tt != RustTypeTree::new() {
         let enzyme_wrapper = EnzymeWrapper::get_instance();
         //if !extra_ints.is_empty() {
@@ -195,21 +194,20 @@ pub(crate) fn add_tt<'tcx, 'll>(
 
         let ret_attr = llvm::CreateAttrStringValueFromCStr(llcx, &c_attr_name, &c_str);
 
-        if llvm::LLVMRustSupportsEnzymeMD(fn_def) {
-            dbg!("extractvalue md");
-            let md = enzyme_wrapper.tree_to_md(enzyme_tt.inner, llcx);
-            unsafe {
-                llvm::LLVMRustSetEnzymeTypeMD(fn_def, md.unwrap());
+        let arg_pos = llvm::AttributePlace::ReturnValue;
+        match tt_location {
+            TTLocation::Definition => {
+                attributes::apply_to_llfn(fn_def, arg_pos, &[ret_attr]);
+            }
+            TTLocation::Callsite => {
+                attributes::apply_to_callsite(fn_def, arg_pos, &[ret_attr]);
+            }
+            TTLocation::Intrinsic => {
+                let md = enzyme_wrapper.tree_to_md(enzyme_tt.inner, llcx);
+                unsafe {
+                    llvm::LLVMRustSetEnzymeTypeMD(fn_def, md.unwrap());
+                }
             }
-        } else if unsafe { llvm::LLVMRustIsIntrinsicCall(fn_def) } {
-            dbg!("intrinsiccall");
-            attributes::apply_to_callsite(
-                fn_def,
-                llvm::AttributePlace::ReturnValue,
-                &[ret_attr],
-            );
-        } else {
-            attributes::apply_to_llfn(fn_def, llvm::AttributePlace::ReturnValue, &[ret_attr]);
         }
         enzyme_wrapper.tree_to_string_free(c_str.as_ptr());
     }

compiler/rustc_codegen_ssa/src/mir/operand.rs

@@ -22,9 +22,8 @@ use crate::common::IntPredicate;
 use crate::traits::*;
 
 
-use rustc_ast::expand::typetree::{TypeTree, FncTree};
+use rustc_ast::expand::typetree::TypeTree;
 use rustc_middle::ty::typetree::typetree_from_ty;
-//use rustc_middle::ty::typetree_from_ty;
 use crate::TyCtxt;
 use rustc_span::sym;
 

compiler/rustc_codegen_ssa/src/traits/builder.rs

@@ -2,7 +2,7 @@ use std::assert_matches;
 use std::ops::Deref;
 
 use rustc_abi::{Align, Scalar, Size, WrappingRange};
-use rustc_ast::expand::typetree::{TypeTree, FncTree};
+use rustc_ast::expand::typetree::{FncTree, TypeTree};
 use rustc_hir::attrs::AttributeKind;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrs;
 use rustc_middle::mir;
@@ -468,7 +468,6 @@ pub trait BuilderMethods<'a, 'tcx>:
         src_align: Align,
         size: Self::Value,
         flags: MemFlags,
-        //tt: Option<FncTree>,
     );
     fn memset(
         &mut self,
@@ -477,7 +476,6 @@ pub trait BuilderMethods<'a, 'tcx>:
         size: Self::Value,
         align: Align,
         flags: MemFlags,
-        //tt: Option<FncTree>,
     );
 
     // Produce a value from calling the `vscale` intrinsic (containing the `vscale` multiplier that
@@ -524,16 +522,7 @@ pub trait BuilderMethods<'a, 'tcx>:
             let tt = typetree_from_ty(self.tcx(), layout.ty);
             // We seem to pass all values to memcpy with one more indirection.
             let tt = tt.add_indirection();
-            dbg!(&tt);
-            use rustc_middle::ty::print::with_no_trimmed_paths;
-
-            with_no_trimmed_paths!({
-                eprintln!("memcpy ty = {:?}", layout.ty);
-            });
-            let fnc_tree = FncTree {
-                args: vec![tt.clone(), tt],
-                ret: TypeTree::new(),
-            };
+            let fnc_tree = FncTree { args: vec![tt.clone(), tt], ret: TypeTree::new() };
             let bytes = self.const_usize(layout.size.bytes());
             let bytes = if layout.peel_transparent_wrappers(self).ty.is_scalable_vector() {
                 let vscale = self.vscale(self.type_i64());

tests/run-make/autodiff/type-trees/iter/rmake.rs

@@ -0,0 +1,12 @@
+//@ needs-enzyme
+//@ ignore-cross-compile
+
+use run_make_support::{llvm_filecheck, rfs, rustc};
+
+fn main() {
+    rustc().input("window.rs").arg("-Zautodiff=Enable").arg("-Clto=fat").emit("llvm-ir").run();
+    //rustc().input("window.rs").arg("-Zautodiff=Enable,NoTT").arg("-Clto=fat").run_fail().assert_stderr_contains("Enzyme: Cannot deduce type of copy");
+    //rustc().input("window.rs").arg("-Zautodiff=Enable,NoTT").arg("-Clto=fat").arg("-O").run();
+
+    llvm_filecheck().patterns("window.check").stdin_buf(rfs::read("window.ll")).run();
+}

tests/run-make/autodiff/type-trees/iter/window.check

@@ -0,0 +1,24 @@
+; Check that array TypeTree metadata is correctly generated for the iter/window case
+; We check all three relevant cases. The metadata generated for extractvalue, 
+; the metadata for the memcpy call,
+; and the metadata for the function definition. 
+; At the time of writing, at least the first two where necessary to make compilation succeed in debug mode.
+
+CHECK-LABEL: ; <core::slice::iter::Windows<f64> as core::iter::traits::iterator::Iterator>::fold
+CHECK:   extractvalue { ptr, i64 } %[[IN:[0-9]+]], 0, !enzyme_type ![[DATA_TT:[0-9]+]]
+CHECK:   extractvalue { ptr, i64 } %[[IN]], 1, !enzyme_type ![[LEN_TT:[0-9]+]]
+
+CHECK-LABEL: ; window::main
+CHECK:  call void @llvm.memcpy.p0.p0.i64(ptr align 8 "enzyme_type"="{[0]:Pointer, [0,-1]:Float@double}" %vjp, ptr align 8 "enzyme_type"="{[0]:Pointer, [0,-1]:Float@double}" {{.*}}, i64 24, i1 false)
+
+
+CHECK: define {{.*}}void @f(ptr align 8 "enzyme_type"="{[-1]:Pointer, [-1,-1]:Float@double}" %x, ptr align 8 "enzyme_type"="{[-1]:Pointer, [-1,-1]:Float@double}" %0, ptr writeonly align 8 captures(none) "enzyme_type"="{[-1]:Pointer, [-1,-1]:Float@double}" %y)
+
+
+CHECK-DAG: ![[DATA_TT]] = !{!"Unknown", i32 -1, ![[PTR_TT:[0-9]+]]}
+CHECK-DAG: ![[PTR_TT]] = !{!"Pointer", i32 -1, ![[FLOAT_TT:[0-9]+]]}
+CHECK-DAG: ![[FLOAT_TT]] = !{!"Float@double"}
+
+CHECK-DAG: ![[LEN_TT]] = !{!"Unknown", i32 0, ![[INT_TT:[0-9]+]], i32 1, ![[INT_TT]], i32 2, ![[INT_TT]], i32 3, ![[INT_TT]], i32 4, ![[INT_TT]], i32 5, ![[INT_TT]], i32 6, ![[INT_TT]], i32 7, ![[INT_TT]]}
+CHECK-DAG: ![[INT_TT]] = !{!"Integer"}
+

tests/run-make/autodiff/type-trees/iter/window.rs

@@ -0,0 +1,55 @@
+#![feature(autodiff)]
+
+use std::autodiff::autodiff_reverse;
+
+// This tests verifies that Enzyme can differentiate the iterator and window version of the for
+// loops given below. Iterators (especially the windows use here) cause a lot of extra abstractions
+// and indirections. Without extra typetree hints, Enzyme failed to differentiate them in debug
+// mode.
+
+//@revisions: tt no_tt
+//@[tt] compile-flags: -Z autodiff=Enable
+//@[no_tt] compile-flags: -Z autodiff=Enable,NoTT
+//@[no_tt] build-fail
+
+#[unsafe(no_mangle)]
+#[inline(never)]
+#[autodiff_reverse(f_rev, 2, Duplicated, Const, Duplicated)]
+fn f(x: &[f64; 3], args: &[f64; 3], y: &mut [f64; 2]) {
+    y[0] = x.iter().map(|i| args[0] * i.powi(2)).sum();
+    y[1] = x
+        .windows(2)
+        .map(|w| (args[1] - w[0]).powi(2) + args[2] * (w[1] - w[0].powi(2)).powi(2))
+        .sum();
+    // The iterators above are equivalent to the two following for loops.
+    // for i in 0..3 {
+    //     y[0] += args[0] * x[i].powi(2);
+    // }
+    // for i in 0..2 {
+    //     y[1] += (args[1] - x[i]).powi(2) + args[2] * (x[i + 1] - x[i].powi(2)).powi(2);
+    // }
+}
+
+// Not generally recommended, but since we rewrite llvm-ir, it should be good enough.
+fn assert_abs_diff_eq<const N: usize>(x: &[f64; N], y: &[f64; N]) {
+    for i in 0..N {
+        assert_eq!(x[i], y[i]);
+    }
+}
+
+fn main() {
+    let x = [3.0, 5.0, 7.0];
+    let args = [2.0, 1.0, 100.0];
+
+    let mut vjp = ([0.0; 3], [0.0; 3]);
+    let mut y = [0.0; 2];
+    let mut dy = ([1.0, 0.0], [0.0, 1.0]);
+
+    f_rev(
+        &x, &mut vjp.0, &mut vjp.1, &args, &mut y, &mut dy.0, &mut dy.1,
+    );
+
+    assert_abs_diff_eq::<2>(&y, &[166.0, 34020.0]);
+    assert_abs_diff_eq::<3>(&vjp.0, &[12.0, 20.0, 28.0]);
+    assert_abs_diff_eq::<3>(&vjp.1, &[4804.0, 35208.0, -3600.0]);
+}