autodiff: slice support in typetree

Author: KMJ-007

compiler/rustc_middle/src/ty/mod.rs

@@ -2333,5 +2333,12 @@ pub fn typetree_from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> TypeTree {
         }
     }
 
+    if ty.is_slice() {
+        if let ty::Slice(element_ty) = ty.kind() {
+            let element_tree = typetree_from_ty(tcx, *element_ty);
+            return element_tree;
+        }
+    }
+
     TypeTree::new()
 }

tests/run-make/autodiff/type-trees/slice-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("slice.check").stdin_buf(rfs::read("test.ll")).run();
+}

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

@@ -0,0 +1,4 @@
+; Check that slice TypeTree metadata is correctly generated
+; Should show Float@double for slice elements
+
+CHECK: define{{.*}}"enzyme_type"="{[]:Float@double}"{{.*}}@test_slice{{.*}}"enzyme_type"="{[]:Pointer}"

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

@@ -0,0 +1,16 @@
+#![feature(autodiff)]
+
+use std::autodiff::autodiff_reverse;
+
+#[autodiff_reverse(d_test, Duplicated, Active)]
+#[no_mangle]
+fn test_slice(slice: &[f64]) -> f64 {
+    slice.iter().sum()
+}
+
+fn main() {
+    let arr = [1.0, 2.0, 3.0, 4.0, 5.0];
+    let slice = &arr[..];
+    let mut d_slice = [0.0; 5];
+    let _result = d_test(slice, &mut d_slice[..], 1.0);
+}