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..100c3170e12a9 100644 --- a/compiler/rustc_middle/src/ty/typetree.rs +++ b/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 } 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..e0c8c87ca9e33 --- /dev/null +++ b/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(); +} 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..4b149c9ee090c --- /dev/null +++ b/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}" 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..e34a8c71f4cc5 --- /dev/null +++ b/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 { + 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) +} + +// 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) +}