autodiff: tuple support in typetree

Author: KMJ-007

compiler/rustc_middle/src/ty/mod.rs

@@ -2334,5 +2334,41 @@ pub fn typetree_from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> TypeTree {
         }
     }
 
+    if let ty::Tuple(tuple_types) = ty.kind() {
+        if tuple_types.is_empty() {
+            return TypeTree::new();
+        }
+
+        let mut types = Vec::new();
+        let mut current_offset = 0;
+
+        for tuple_ty in tuple_types.iter() {
+            let element_tree = typetree_from_ty(tcx, tuple_ty);
+
+            let element_layout = tcx
+                .layout_of(ty::TypingEnv::fully_monomorphized().as_query_input(tuple_ty))
+                .ok()
+                .map(|layout| layout.size.bytes_usize())
+                .unwrap_or(0);
+
+            for elem_type in &element_tree.0 {
+                types.push(Type {
+                    offset: if elem_type.offset == -1 {
+                        current_offset as isize
+                    } else {
+                        current_offset as isize + elem_type.offset
+                    },
+                    size: elem_type.size,
+                    kind: elem_type.kind,
+                    child: elem_type.child.clone(),
+                });
+            }
+
+            current_offset += element_layout;
+        }
+
+        return TypeTree(types);
+    }
+
     TypeTree::new()
 }

tests/run-make/autodiff/type-trees/tuple-typetree/rmake.rs

@@ -0,0 +1,9 @@
+//@ needs-enzyme
+//@ ignore-cross-compile
+
+use run_make_support::{llvm_filecheck, rfs, rustc};
+
+fn main() {
+    rustc().input("test.rs").arg("-Zautodiff=Enable").emit("llvm-ir").run();
+    llvm_filecheck().patterns("tuple.check").stdin_buf(rfs::read("test.ll")).run();
+}

tests/run-make/autodiff/type-trees/tuple-typetree/test.rs

@@ -0,0 +1,15 @@
+#![feature(autodiff)]
+
+use std::autodiff::autodiff_reverse;
+
+#[autodiff_reverse(d_test, Duplicated, Active)]
+#[no_mangle]
+fn test_tuple(tuple: &(f64, f64, f64)) -> f64 {
+    tuple.0 + tuple.1 * 2.0 + tuple.2 * 3.0
+}
+
+fn main() {
+    let tuple = (1.0, 2.0, 3.0);
+    let mut d_tuple = (0.0, 0.0, 0.0);
+    let _result = d_test(&tuple, &mut d_tuple, 1.0);
+}

tests/run-make/autodiff/type-trees/tuple-typetree/tuple.check

@@ -0,0 +1,4 @@
+; Check that tuple TypeTree metadata is correctly generated
+; Should show Float@double at offsets 0, 8, 16 for (f64, f64, f64)
+
+CHECK: define{{.*}}"enzyme_type"="{[]:Float@double}"{{.*}}@test_tuple{{.*}}"enzyme_type"="{[]:Pointer}"