From ac1e5f4e58041b9656bb4b129cb9cb665d8dcec4 Mon Sep 17 00:00:00 2001 From: Joao Roberto Date: Sun, 2 Aug 2026 13:19:20 -0300 Subject: [PATCH 1/2] Handle slice-tailed DSTs in autodiff type trees Use codegen tail discovery so wrappers such as OsStr keep slice element metadata without losing offsets for prefixed DSTs. Keep recursion tracking scoped to the current traversal path. --- .../src/builder/autodiff.rs | 11 +- compiler/rustc_middle/src/ty/typetree.rs | 112 ++++++++++-------- .../type-trees/slice-dst-typetree/rmake.rs | 18 +++ .../slice-dst-typetree/slice-dst.check | 9 ++ .../type-trees/slice-dst-typetree/test.rs | 38 ++++++ 5 files changed, 131 insertions(+), 57 deletions(-) create mode 100644 tests/run-make/autodiff/type-trees/slice-dst-typetree/rmake.rs create mode 100644 tests/run-make/autodiff/type-trees/slice-dst-typetree/slice-dst.check create mode 100644 tests/run-make/autodiff/type-trees/slice-dst-typetree/test.rs diff --git a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs index b0ac5c096efdf..8cefd8dc489a9 100644 --- a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs +++ b/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); diff --git a/compiler/rustc_middle/src/ty/typetree.rs b/compiler/rustc_middle/src/ty/typetree.rs index d4cda033a7e87..90fb0316fe39b 100644 --- a/compiler/rustc_middle/src/ty/typetree.rs +++ b/compiler/rustc_middle/src/ty/typetree.rs @@ -62,20 +62,26 @@ 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() { + let child = if tcx + .layout_of(typing_env.as_query_input(inner_ty)) + .is_ok_and(|layout| layout.size.bytes() == 0) + { + // Direct slices and transparent wrappers such as `OsStr` contain elements everywhere. + typetree_from_ty_impl_inner(tcx, *element_ty, depth + 1, visited, false) + } else { + // Preserve field offsets for a sized prefix before the slice tail. + 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 +111,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 +217,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 } diff --git a/tests/run-make/autodiff/type-trees/slice-dst-typetree/rmake.rs b/tests/run-make/autodiff/type-trees/slice-dst-typetree/rmake.rs new file mode 100644 index 0000000000000..c19202fa41fe5 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/slice-dst-typetree/rmake.rs @@ -0,0 +1,18 @@ +//@ 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(); +} diff --git a/tests/run-make/autodiff/type-trees/slice-dst-typetree/slice-dst.check b/tests/run-make/autodiff/type-trees/slice-dst-typetree/slice-dst.check new file mode 100644 index 0000000000000..6031b728213ae --- /dev/null +++ b/tests/run-make/autodiff/type-trees/slice-dst-typetree/slice-dst.check @@ -0,0 +1,9 @@ +; 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}" diff --git a/tests/run-make/autodiff/type-trees/slice-dst-typetree/test.rs b/tests/run-make/autodiff/type-trees/slice-dst-typetree/test.rs new file mode 100644 index 0000000000000..8cd4d7b57f270 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/slice-dst-typetree/test.rs @@ -0,0 +1,38 @@ +#![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 { + 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::() +} + +#[no_mangle] +pub fn exercise_header_sum(value: &Header<[f32]>, derivative: &mut Header<[f32]>) -> f32 { + d_header_sum(value, derivative, 1.0) +} From 18e0dd9aa8993a19e332fd080904a72270d18a0d Mon Sep 17 00:00:00 2001 From: Joao Roberto Date: Tue, 4 Aug 2026 15:31:15 -0300 Subject: [PATCH 2/2] Document zero-sized autodiff slice handling Clarify that slice-tail layout checks apply to the sized prefix rather than the slice element, and cover zero-sized slice elements in the type-tree run-make test. --- compiler/rustc_middle/src/ty/typetree.rs | 6 ++++-- .../autodiff/type-trees/slice-dst-typetree/rmake.rs | 1 + .../type-trees/slice-dst-typetree/slice-dst.check | 5 +++++ .../autodiff/type-trees/slice-dst-typetree/test.rs | 13 +++++++++++++ 4 files changed, 23 insertions(+), 2 deletions(-) diff --git a/compiler/rustc_middle/src/ty/typetree.rs b/compiler/rustc_middle/src/ty/typetree.rs index 90fb0316fe39b..100c3170e12a9 100644 --- a/compiler/rustc_middle/src/ty/typetree.rs +++ b/compiler/rustc_middle/src/ty/typetree.rs @@ -66,14 +66,16 @@ fn handle_indirection<'a>( // 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) { - // Direct slices and transparent wrappers such as `OsStr` contain elements everywhere. typetree_from_ty_impl_inner(tcx, *element_ty, depth + 1, visited, false) } else { - // Preserve field offsets for a sized prefix before the slice tail. typetree_from_ty_impl_inner(tcx, inner_ty, depth + 1, visited, true) }; return TypeTree(vec![Type { diff --git a/tests/run-make/autodiff/type-trees/slice-dst-typetree/rmake.rs b/tests/run-make/autodiff/type-trees/slice-dst-typetree/rmake.rs index c19202fa41fe5..e0c8c87ca9e33 100644 --- a/tests/run-make/autodiff/type-trees/slice-dst-typetree/rmake.rs +++ b/tests/run-make/autodiff/type-trees/slice-dst-typetree/rmake.rs @@ -15,4 +15,5 @@ fn main() { 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(); } diff --git a/tests/run-make/autodiff/type-trees/slice-dst-typetree/slice-dst.check b/tests/run-make/autodiff/type-trees/slice-dst-typetree/slice-dst.check index 6031b728213ae..4b149c9ee090c 100644 --- a/tests/run-make/autodiff/type-trees/slice-dst-typetree/slice-dst.check +++ b/tests/run-make/autodiff/type-trees/slice-dst-typetree/slice-dst.check @@ -7,3 +7,8 @@ OSSTR: call void @llvm.memcpy{{.*}}"enzyme_type"="{[0]:Pointer, [0,0]:Pointer, [ 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}" diff --git a/tests/run-make/autodiff/type-trees/slice-dst-typetree/test.rs b/tests/run-make/autodiff/type-trees/slice-dst-typetree/test.rs index 8cd4d7b57f270..e34a8c71f4cc5 100644 --- a/tests/run-make/autodiff/type-trees/slice-dst-typetree/test.rs +++ b/tests/run-make/autodiff/type-trees/slice-dst-typetree/test.rs @@ -36,3 +36,16 @@ pub fn header_sum(value: &Header<[f32]>) -> f32 { 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) +}