From 55b5c3980a0c1d6c0c6b94ae0bc13401dce05635 Mon Sep 17 00:00:00 2001 From: Yunus Mutlu Date: Fri, 25 Sep 2026 15:21:44 +0300 Subject: [PATCH] fix(mir-lower): preserve address spaces through 16-byte enum fillers Signed-off-by: Yunus Mutlu --- .../src/convert/types/enum_layout.rs | 135 ++++++++++++++++++ .../mir-lower/src/convert/types/layout.rs | 39 +++-- .../examples/vectorization/src/main.rs | 119 +++++++++++++++ 3 files changed, 278 insertions(+), 15 deletions(-) diff --git a/cuda-oxide/crates/mir-lower/src/convert/types/enum_layout.rs b/cuda-oxide/crates/mir-lower/src/convert/types/enum_layout.rs index 25eb043d10..38446868e6 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/types/enum_layout.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/types/enum_layout.rs @@ -2011,4 +2011,139 @@ mod tests { "{error}" ); } + + #[test] + fn enum_slot_map_pointer_niche_with_u128_payload_preserves_align16() { + let mut ctx = make_ctx(); + + let logical = mir_uint(&mut ctx, 8); + let wide = mir_uint(&mut ctx, 128); + let u32_ty = mir_uint(&mut ctx, 32); + let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); + + let payload: TypeHandle = MirTupleType::get_with_layout( + &mut ctx, + vec![wide, pointer], + vec![0, 1], + vec![0, 16], + 32, + 16, + ) + .into(); + + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "MaybeWideAndPointer".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![32]), + ], + EnumEncoding { + tag_offset: 16, + total_size: 32, + abi_align: 16, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + + let fields = struct_fields(&ctx, map.llvm_struct_ty); + + // Required enum alignment is 16, while the 16-byte filler itself is + // deliberately represented as two i64 leaves with natural alignment 8. + // build_enum_slot_map supplies the zero-length alignment anchor separately. + assert_eq!(fields.len(), 4); + + let filler_ref = fields[1].deref(&ctx); + let filler = filler_ref + .downcast_ref::() + .expect("16-byte enum filler should lower to an LLVM array"); + + assert_eq!( + filler.size(), + 2, + "16-byte enum filler should use two 64-bit leaves" + ); + + assert_eq!( + filler + .elem_type() + .deref(&ctx) + .downcast_ref::() + .map(IntegerType::width), + Some(64), + "16-byte enum filler elements should be i64" + ); + + assert_eq!( + llvm_type_size_align(&ctx, map.llvm_struct_ty), + Some((32, 16)), + "typed 16-byte filler must preserve rustc's enum size and alignment" + ); + } + + #[test] + fn enum_filler_16_bytes_falls_back_when_i64_alignment_is_not_layout_neutral() { + let mut ctx = make_ctx(); + + let filler = make_enum_filler_type(&mut ctx, 0, 16, 4); + + let filler_ref = filler.deref(&ctx); + + let array = filler_ref + .downcast_ref::() + .expect("fallback enum filler should remain an LLVM byte array"); + + assert_eq!(array.size(), 16, "fallback must retain all 16 bytes"); + + assert_eq!( + array + .elem_type() + .deref(&ctx) + .downcast_ref::() + .map(IntegerType::width), + Some(8), + "alignment-sensitive fallback must remain [16 x i8]" + ); + + assert_eq!( + llvm_type_size_align(&ctx, filler), + Some((16, 1)), + "fallback byte array must remain byte-aligned" + ); + } + + #[test] + fn enum_filler_16_bytes_uses_two_i64s_at_align8() { + let mut ctx = make_ctx(); + + let filler = make_enum_filler_type(&mut ctx, 0, 16, 8); + + let filler_ref = filler.deref(&ctx); + let array = filler_ref + .downcast_ref::() + .expect("16-byte filler should use an LLVM array"); + + assert_eq!(array.size(), 2); + + assert_eq!( + array + .elem_type() + .deref(&ctx) + .downcast_ref::() + .map(IntegerType::width), + Some(64) + ); + + assert_eq!(llvm_type_size_align(&ctx, filler), Some((16, 8))); + } } diff --git a/cuda-oxide/crates/mir-lower/src/convert/types/layout.rs b/cuda-oxide/crates/mir-lower/src/convert/types/layout.rs index 54083e2424..6a4821257e 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/types/layout.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/types/layout.rs @@ -35,23 +35,26 @@ pub(super) fn make_padding_type(ctx: &mut Context, size: u64) -> TypeHandle { /// longer follow it, and the access lands in the generic window (`ld.v2.b64`) /// instead of `ld.global.v2.b64`. /// -/// One integer of the same width is the same bytes in one leaf. It is only -/// legal where it moves nothing: +/// A typed filler carries the same bytes with far fewer aggregate leaves when +/// doing so is layout-neutral: /// -/// - the width is 2, 4 or 8 bytes, so the integer is byte-faithful (a multiple -/// of 8 bits, hence no padding of its own). 16 is deliberately excluded: -/// `i128` is legal LLVM but lowers to register pairs on NVPTX, so widening -/// that far trades one win for another cost and wants its own measurement; -/// - `offset` is a multiple of the width, so LLVM inserts no gap ahead of the -/// field and every later field keeps its byte offset; -/// - the width does not exceed the enum's alignment, or the struct's natural -/// alignment would rise above rustc's. +/// - 2, 4, or 8 bytes use one integer of the same width when `offset` is a +/// multiple of that width and the width does not exceed the enum alignment; +/// - 16 bytes use `[2 x i64]` when the offset is 8-byte aligned and the enum +/// alignment is at least 8. Using `i128` here would impose 16-byte natural +/// alignment, so the two-i64 representation also covers 8-byte-aligned +/// Rust layouts without changing their ABI. /// -/// Anything else keeps the byte array. The filler is a *vehicle*, not a claim: -/// nothing reads it field-wise, because a payload that has no typed slot -/// round-trips through memory as a whole aggregate, and both ends of that -/// round trip are byte-exact. `build_enum_slot_map`'s own size and alignment -/// assertions — hard errors, not debug checks — backstop all three conditions. +/// The 16-byte case also avoids representing the filler as sixteen independent +/// byte leaves. Keeping it as two typed 64-bit leaves allows aggregate +/// scalarization to preserve pointer address-space information through enum +/// control flow. +/// +/// Anything that does not satisfy these layout constraints keeps the byte +/// array. The filler is a *vehicle*, not a claim: nothing reads it field-wise, +/// because a payload that has no typed slot round-trips through memory as a +/// whole aggregate. `build_enum_slot_map`'s size and alignment checks are hard +/// errors and backstop these choices. pub(super) fn make_enum_filler_type( ctx: &mut Context, offset: u64, @@ -62,6 +65,12 @@ pub(super) fn make_enum_filler_type( let width = u32::try_from(size * 8).expect("size is at most 8, so width is at most 64"); return IntegerType::get(ctx, width, Signedness::Signless).into(); } + + if size == 16 && offset.is_multiple_of(8) && abi_align >= 8 { + let i64_ty = IntegerType::get(ctx, 64, Signedness::Signless); + return llvm_types::ArrayType::get(ctx, i64_ty.into(), 2).into(); + } + make_padding_type(ctx, size) } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/src/main.rs index 6a5cf6115b..76d0ea85ea 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/src/main.rs @@ -155,6 +155,91 @@ mod view_kernels { empty.as_mut_ptr() } + /// Regression coverage for a 16-byte enum filler beside a pointer niche. + /// + /// The `[u64; 2]` payload occupies 16 bytes before the slice pointer. Enum + /// storage must represent that filler without obscuring the slice pointer's + /// global address space; otherwise the final vector load becomes generic + /// `ld.v4.*` instead of `ld.global.v4.*`. + #[kernel] + pub fn f32x4_view_copy_through_wide_option( + input: &[f32], + mut output: DisjointSlice, + lo: u64, + hi: u64, + ) { + let i = thread::index_1d().get(); + + let out_len = output.len(); + let out_flat = unsafe { core::slice::from_raw_parts_mut(output.as_mut_ptr(), out_len) }; + + let Some(input_quads) = vector::as_vectors::(input) else { + return; + }; + + let Some(out_quads) = vector::as_vectors_mut::(out_flat) else { + return; + }; + + let cookie = [lo, hi]; + + let wrapped = if lo == u64::MAX { + None + } else { + Some((cookie, input_quads)) + }; + + let Some((_cookie, quads)) = wrapped else { + return; + }; + + if i < quads.len() && i < out_quads.len() { + out_quads[i] = quads[i]; + } + } + + /// Regression coverage for a 16-byte enum filler beside a mutable pointer niche. + /// + /// This mirrors `f32x4_view_copy_through_wide_option`, but places the output + /// slice behind the wide enum payload. Losing its address space would turn the + /// vector store into generic `st.v4.*` instead of `st.global.v4.*`. + #[kernel] + pub fn f32x4_view_store_through_wide_option( + input: &[f32], + mut output: DisjointSlice, + lo: u64, + hi: u64, + ) { + let i = thread::index_1d().get(); + + let Some(input_quads) = vector::as_vectors::(input) else { + return; + }; + + let out_len = output.len(); + let out_flat = unsafe { core::slice::from_raw_parts_mut(output.as_mut_ptr(), out_len) }; + + let Some(out_quads) = vector::as_vectors_mut::(out_flat) else { + return; + }; + + let cookie = [lo, hi]; + + let wrapped = if lo == u64::MAX { + None + } else { + Some((cookie, out_quads)) + }; + + let Some((_cookie, out_quads)) = wrapped else { + return; + }; + + if i < input_quads.len() && i < out_quads.len() { + out_quads[i] = input_quads[i]; + } + } + /// Copy one `F32x4` quad per thread between flat `f32` buffers through /// checked [`vector::as_vectors`] views. The quad is moved whole, never /// decomposed into lanes, so both the load and the store fuse into @@ -326,6 +411,40 @@ fn main() { } } } + + // A 16-byte enum filler next to the slice pointer must not hide the + // pointer's global address space from NVPTX address-space inference. + let wide_option_body = kernel_body(&document, "f32x4_view_copy_through_wide_option"); + + match find_128bit_global(wide_option_body, "ld") { + Some(line) => { + println!("f32x4_view_copy_through_wide_option: {line}"); + } + None => { + errors += 1; + println!( + " !! f32x4_view_copy_through_wide_option: \ + 16-byte enum filler lost the 128-bit global load" + ); + } + } + + // The mutable slice variant must preserve the same address-space information + // for stores as the immutable slice variant does for loads. + let wide_store_body = kernel_body(&document, "f32x4_view_store_through_wide_option"); + + match find_128bit_global(wide_store_body, "st") { + Some(line) => { + println!("f32x4_view_store_through_wide_option: {line}"); + } + None => { + errors += 1; + println!( + " !! f32x4_view_store_through_wide_option: \ + 16-byte enum filler lost the 128-bit global store" + ); + } + } } else { println!( "PTX shape check skipped: the embedded bundle has no PTX text to \