autodiff: Handle slice-tailed DSTs in type trees

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -45,18 +45,13 @@ pub(crate) fn adjust_activity_to_abi<'tcx>(
     let mut del_activities = 0;
     for (i, ty) in sig.inputs().iter().enumerate() {
         if let Some(inner_ty) = ty.builtin_deref(true) {
-            if inner_ty.is_slice() {
+            let tail_ty = tcx.struct_tail_for_codegen(inner_ty, typing_env);
+            if let ty::Slice(element_ty) = tail_ty.kind() {
                 // Now we need to figure out the size of each slice element in memory to allow
                 // safety checks and usability improvements in the backend.
-                let sty = match inner_ty.builtin_index() {
-                    Some(sty) => sty,
-                    None => {
-                        panic!("slice element type unknown");
-                    }
-                };
                 let pci = PseudoCanonicalInput {
                     typing_env: TypingEnv::fully_monomorphized(),
-                    value: sty,
+                    value: *element_ty,
                 };
 
                 let layout = tcx.layout_of(pci);

compiler/rustc_middle/src/ty/typetree.rs

@@ -62,20 +62,28 @@ fn handle_indirection<'a>(
     let Some(inner_ty) = ty.builtin_deref(true) else {
         bug!("incorrect autodiff typetree handling for type: {}", ty);
     };
-    // slices are represented as `&'{erased} mut [f32]`
-    // This reads as a reference to a slice of f32.
-    // So we'd end up with ptr->RustSlice->f32 without this extra handling
-    if inner_ty.is_slice() {
-        if let ty::Slice(element_ty) = inner_ty.kind() {
-            let element_tree =
-                typetree_from_ty_impl_inner(tcx, *element_ty, depth + 1, visited, false);
-            return TypeTree(vec![Type {
-                offset: -1,
-                size: tcx.data_layout.pointer_size().bytes_usize(),
-                kind: Kind::RustSlice,
-                child: element_tree,
-            }]);
-        }
+    // A pointer to a slice-tailed DST is a fat pointer `{data, len}`. `RustSlice` describes both
+    // LLVM arguments, while its child describes the memory reached through `data`.
+    let typing_env = ty::TypingEnv::fully_monomorphized();
+    if let ty::Slice(element_ty) = tcx.struct_tail_for_codegen(inner_ty, typing_env).kind() {
+        // `layout.size` here is the sized prefix of `inner_ty`, not the slice element size.
+        // Direct slices, transparent wrappers (`OsStr`), and ZST-prefixed DSTs have no byte
+        // offset to preserve. Nonzero prefixes (e.g. `Header<[f32]>`) keep field offsets.
+        // ZST elements still take this path and yield an empty child TypeTree (size 0).
+        let child = if tcx
+            .layout_of(typing_env.as_query_input(inner_ty))
+            .is_ok_and(|layout| layout.size.bytes() == 0)
+        {
+            typetree_from_ty_impl_inner(tcx, *element_ty, depth + 1, visited, false)
+        } else {
+            typetree_from_ty_impl_inner(tcx, inner_ty, depth + 1, visited, true)
+        };
+        return TypeTree(vec![Type {
+            offset: -1,
+            size: tcx.data_layout.pointer_size().bytes_usize(),
+            kind: Kind::RustSlice,
+            child,
+        }]);
     }
 
     let child = typetree_from_ty_impl_inner(tcx, inner_ty, depth + 1, visited, true);
@@ -105,56 +113,60 @@ fn typetree_from_ty_impl_inner<'tcx>(
     }
     visited.push(ty);
 
-    match ty.kind() {
-        // See handle_indirection for an explanation on why we don't handle it here.
-        ty::Slice(..) => bug!("incorrect autodiff typetree handling for slice: {}", ty),
+    let tree = match ty.kind() {
+        // Direct slices are handled by `handle_indirection`. This arm describes a slice tail while
+        // recursing through a prefixed DST, so its caller can add the field offset.
+        ty::Slice(element_ty) => {
+            typetree_from_ty_impl_inner(tcx, *element_ty, depth + 1, visited, false)
+        }
         ty::Ref(..) | ty::RawPtr(..) => handle_indirection(ty, tcx, depth, visited),
         ty::Adt(def, _) if def.is_box() => handle_indirection(ty, tcx, depth, visited),
         ty::Array(element_ty, len_const) => {
             let len = len_const.try_to_target_usize(tcx).unwrap_or(0);
             if len == 0 {
-                return TypeTree::new();
-            }
-            let element_tree =
-                typetree_from_ty_impl_inner(tcx, *element_ty, depth + 1, visited, false);
-            let mut types = Vec::new();
-            for elem_type in &element_tree.0 {
-                types.push(Type::from_ty(-1, elem_type));
-            }
+                TypeTree::new()
+            } else {
+                let element_tree =
+                    typetree_from_ty_impl_inner(tcx, *element_ty, depth + 1, visited, false);
+                let mut types = Vec::new();
+                for elem_type in &element_tree.0 {
+                    types.push(Type::from_ty(-1, elem_type));
+                }
 
-            TypeTree(types)
+                TypeTree(types)
+            }
         }
         ty::Tuple(tuple_types) => {
             if tuple_types.is_empty() {
-                return TypeTree::new();
-            }
+                TypeTree::new()
+            } else {
+                let mut types = Vec::new();
+                let mut current_offset = 0;
 
-            let mut types = Vec::new();
-            let mut current_offset = 0;
+                for tuple_ty in tuple_types.iter() {
+                    let element_tree =
+                        typetree_from_ty_impl_inner(tcx, tuple_ty, depth + 1, visited, false);
 
-            for tuple_ty in tuple_types.iter() {
-                let element_tree =
-                    typetree_from_ty_impl_inner(tcx, tuple_ty, depth + 1, visited, false);
+                    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);
 
-                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 {
+                        let offset = if elem_type.offset == -1 {
+                            current_offset as isize
+                        } else {
+                            current_offset as isize + elem_type.offset
+                        };
+                        types.push(Type::from_ty(offset, elem_type));
+                    }
 
-                for elem_type in &element_tree.0 {
-                    let offset = if elem_type.offset == -1 {
-                        current_offset as isize
-                    } else {
-                        current_offset as isize + elem_type.offset
-                    };
-                    types.push(Type::from_ty(offset, elem_type));
+                    current_offset += element_layout;
                 }
 
-                current_offset += element_layout;
+                TypeTree(types)
             }
-
-            TypeTree(types)
         }
         ty::Adt(adt_def, args) if adt_def.is_struct() => {
             let struct_layout =
@@ -207,5 +219,9 @@ fn typetree_from_ty_impl_inner<'tcx>(
             TypeTree(vec![Type { offset, size, kind: enzyme_ty, child: TypeTree::new() }])
         }
         _ => TypeTree::new(),
-    }
+    };
+
+    let popped = visited.pop();
+    debug_assert_eq!(popped, Some(ty));
+    tree
 }

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

@@ -0,0 +1,19 @@
+//@ needs-enzyme
+//@ ignore-cross-compile
+
+use run_make_support::{llvm_filecheck, rfs, rustc};
+
+fn main() {
+    rustc()
+        .input("test.rs")
+        .arg("-Zautodiff=Enable,NoPostopt")
+        .opt_level("0")
+        .arg("-Clto=fat")
+        .emit("llvm-ir")
+        .run();
+
+    let ir = rfs::read("test.ll");
+    llvm_filecheck().patterns("slice-dst.check").check_prefix("OSSTR").stdin_buf(&ir).run();
+    llvm_filecheck().patterns("slice-dst.check").check_prefix("HEADER").stdin_buf(&ir).run();
+    llvm_filecheck().patterns("slice-dst.check").check_prefix("ZST").stdin_buf(&ir).run();
+}

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

@@ -0,0 +1,14 @@
+; Preserve element metadata for the `OsStr` regression from #160327.
+OSSTR-LABEL: define void @split_once(
+OSSTR-NOT: define
+OSSTR: call void @llvm.memcpy{{.*}}"enzyme_type"="{[0]:Pointer, [0,0]:Pointer, [0,0,-1]:Integer, [0,16]:Pointer, [0,16,-1]:Integer}"
+
+; Preserve both data layout and length metadata for a prefixed slice-tail DST.
+HEADER-LABEL: define{{.*}}@header_sum(
+HEADER-SAME: ptr{{.*}}"enzyme_type"="{[-1]:Pointer, [-1,0]:Float@float, [-1,4]:Float@float}"
+HEADER-SAME: i64 "enzyme_type"="{[0]:Integer}"
+
+; ZST elements produce no child metadata under the slice data pointer.
+ZST-LABEL: define{{.*}}@zst_slice_len(
+ZST-SAME: ptr{{.*}}"enzyme_type"="{[-1]:Pointer}"
+ZST-SAME: i64 "enzyme_type"="{[0]:Integer}"

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

@@ -0,0 +1,51 @@
+#![crate_type = "lib"]
+#![feature(autodiff)]
+
+use std::autodiff::autodiff_reverse;
+use std::ffi::OsStr;
+
+// Reduced from `clap_lex::OsStrExt::split_once`.
+#[no_mangle]
+#[inline(never)]
+pub fn split_once<'s>(arg: &'s OsStr, needle: &str) -> Option<(&'s OsStr, &'s OsStr)> {
+    let bytes = arg.as_encoded_bytes();
+    let index = bytes.windows(needle.len()).position(|window| window == needle.as_bytes())?;
+    let (first, second) = bytes.split_at(index + needle.len());
+    unsafe {
+        Some((
+            OsStr::from_encoded_bytes_unchecked(first),
+            OsStr::from_encoded_bytes_unchecked(second),
+        ))
+    }
+}
+
+#[repr(C)]
+pub struct Header<T: ?Sized> {
+    tag: f32,
+    data: T,
+}
+
+#[autodiff_reverse(d_header_sum, Duplicated, Active)]
+#[no_mangle]
+#[inline(never)]
+pub fn header_sum(value: &Header<[f32]>) -> f32 {
+    value.tag + value.data.iter().sum::<f32>()
+}
+
+#[no_mangle]
+pub fn exercise_header_sum(value: &Header<[f32]>, derivative: &mut Header<[f32]>) -> f32 {
+    d_header_sum(value, derivative, 1.0)
+}
+
+// ZST slice elements yield an empty child TypeTree; element size 0 is expected.
+#[autodiff_reverse(d_zst_slice_len, Duplicated, Active)]
+#[no_mangle]
+#[inline(never)]
+pub fn zst_slice_len(slice: &[()]) -> f32 {
+    slice.len() as f32
+}
+
+#[no_mangle]
+pub fn exercise_zst_slice_len(slice: &[()], derivative: &mut [()]) -> f32 {
+    d_zst_slice_len(slice, derivative, 1.0)
+}