From 1f6a9d427e42827094d3ab5e0fbff8095f7e0a31 Mon Sep 17 00:00:00 2001 From: midagedev Date: Wed, 30 Sep 2026 19:16:19 +0900 Subject: [PATCH] fix(mir-lower): give an aggregate niche payload's leaves typed slots An enum payload that is a struct or tuple with a pointer niche carrier and non-pointer leaves (`Option<(f32, f32, &mut T)>`, `Option<{&mut T, u64, u64}>`) gave only the pointer leaf a typed slot. The other leaves fell into `[N x i8]` filler bytes, so construct spilled the payload through a slot and the payload read loaded it back at another width. SROA cannot rebuild mixed-type slices, and the slot reached PTX as a local depot. Every leaf of such a payload now claims a typed slot at its rustc offset, reusing the carrier's slot where the types match, with construct and payload reads using insertvalue/extractvalue; size, alignment, field offsets and the niche encoding are unchanged, and a payload the rule cannot prove keeps the old lowering. Measured with one device-only probe source against each tree's backend: parent e027ac32 emits a 16 B `.local` depot in the `smooth_triplets` entry plus 16 B in the outlined `Cells::next_cell`; this commit emits none. `lease_borrow` and `lease_raw` emit none on either commit. Across the 241 examples under rustc-codegen-cuda/examples, only enum_niche_payload and set_discriminant_niche change PTX, and both lose their 16 B depot; vectorization and repr_u32_enum_stride are unchanged per entry. The #1332 slot-map test now expects the 16 bytes under the carrier as one i128 slot instead of [2 x i64] filler; alignment stays 16. Adds the enum_niche_payload example and the leaf-slot unit tests. Signed-off-by: midagedev --- .../src/convert/ops/aggregate/enums.rs | 665 +++++++++++++-- .../src/convert/types/enum_layout.rs | 540 ++++++++++++- .../examples/enum_niche_payload/Cargo.lock | 765 ++++++++++++++++++ .../examples/enum_niche_payload/Cargo.toml | 14 + .../examples/enum_niche_payload/README.md | 36 + .../examples/enum_niche_payload/src/main.rs | 283 +++++++ 6 files changed, 2197 insertions(+), 106 deletions(-) create mode 100644 cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/Cargo.lock create mode 100644 cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/Cargo.toml create mode 100644 cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/README.md create mode 100644 cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/src/main.rs diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/enums.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/enums.rs index 94aced0846..d91913981a 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/enums.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/enums.rs @@ -109,6 +109,17 @@ pub(crate) fn convert_construct_enum( // Insert every payload field that owns a struct slot; remember the // slotless ones for the memory pass below. + let storage_field_tys: Vec = { + let storage_type = llvm_struct_ty.deref(ctx); + let storage_struct = storage_type + .downcast_ref::() + .ok_or_else(|| { + pliron::input_error_noloc!( + "MirConstructEnumOp physical storage is not an LLVM struct" + ) + })?; + storage_struct.fields().collect() + }; let mut deferred: Vec<(usize, Value)> = Vec::new(); for (i, operand) in operands.into_iter().enumerate() { let flat = field_base + i; @@ -120,25 +131,37 @@ pub(crate) fn convert_construct_enum( slot_map.field_slots.len() ); }; + // A decomposed aggregate field: one insertvalue per leaf, extracted + // from the operand and coerced to its slot's storage type. No spill. + if let Some(leaves) = slot_map.field_leaf_slots.get(flat).and_then(|l| l.as_ref()) { + for (path, leaf_slot) in leaves { + let extract_op = llvm::ExtractValueOp::new(ctx, operand, path.clone())?; + rewriter.insert_operation(ctx, extract_op.get_operation()); + let leaf = extract_op.get_operation().deref(ctx).get_result(0); + let storage_slot = *leaf_slot as usize; + let Some(storage_ty) = storage_field_tys.get(storage_slot).copied() else { + return pliron::input_err_noloc!( + "MirConstructEnumOp physical field slot {} is out of range", + leaf_slot + ); + }; + let stored_leaf = coerce_enum_payload_storage(ctx, rewriter, leaf, storage_ty)?; + let insert_op = + llvm::InsertValueOp::new(ctx, current_struct, stored_leaf, vec![*leaf_slot]); + rewriter.insert_operation(ctx, insert_op.get_operation()); + current_struct = insert_op.get_operation().deref(ctx).get_result(0); + last_op = insert_op.get_operation(); + } + continue; + } match slot { Some(slot) => { - let storage_ty = { - let storage_type = llvm_struct_ty.deref(ctx); - let storage_struct = storage_type - .downcast_ref::() - .ok_or_else(|| { - pliron::input_error_noloc!( - "MirConstructEnumOp physical storage is not an LLVM struct" - ) - })?; - let storage_slot = *slot as usize; - if storage_slot >= storage_struct.num_fields() { - return pliron::input_err_noloc!( - "MirConstructEnumOp physical field slot {} is out of range", - slot - ); - } - storage_struct.field_type(storage_slot) + let storage_slot = *slot as usize; + let Some(storage_ty) = storage_field_tys.get(storage_slot).copied() else { + return pliron::input_err_noloc!( + "MirConstructEnumOp physical field slot {} is out of range", + slot + ); }; let stored_operand = coerce_enum_payload_storage(ctx, rewriter, operand, storage_ty)?; @@ -444,12 +467,26 @@ pub(crate) fn convert_get_discriminant( Ok(()) } +/// The LLVM type at an `extractvalue`/`insertvalue` index `path` into `ty`. +fn llvm_type_at_path(ctx: &Context, mut ty: TypeHandle, path: &[u32]) -> Option { + for &index in path { + ty = ty + .deref(ctx) + .downcast_ref::()? + .field_type(index as usize); + } + Some(ty) +} + /// Convert `mir.enum_payload` (reading a variant's field, e.g. the `x` /// in `E::A(x) => x`) to a payload-field read. /// -/// Three cases, decided by the [`EnumSlotMap`]: +/// Four cases, decided by the [`EnumSlotMap`]: /// /// - The field owns a struct slot: a plain `llvm.extractvalue`. +/// - The field was decomposed into leaf slots: rebuild the aggregate in +/// SSA, one `llvm.extractvalue` per leaf slot and one `llvm.insertvalue` +/// per leaf path, coercing each leaf back to its semantic type. No spill. /// - The field has no slot (its bytes are shared with a different-typed /// field of another variant): go through memory. Copy the enum to a /// stack slot, point at the field's byte position, and load it with @@ -506,6 +543,33 @@ pub(crate) fn convert_enum_payload( ); }; + // A decomposed aggregate field: rebuild the whole value leaf by leaf + // from the enum's SSA slots, with no round-trip through memory. + if let Some(leaves) = slot_map.field_leaf_slots.get(flat).and_then(|l| l.as_ref()) { + let semantic_ty = slot_map.field_llvm_types[flat]; + let undef_op = llvm::UndefOp::new(ctx, semantic_ty); + rewriter.insert_operation(ctx, undef_op.get_operation()); + let mut current = undef_op.get_operation().deref(ctx).get_result(0); + for (path, leaf_slot) in leaves { + let extract_op = llvm::ExtractValueOp::new(ctx, enum_val, vec![*leaf_slot])?; + rewriter.insert_operation(ctx, extract_op.get_operation()); + let stored_leaf = extract_op.get_operation().deref(ctx).get_result(0); + let Some(semantic_leaf_ty) = llvm_type_at_path(ctx, semantic_ty, path) else { + return pliron::input_err_noloc!( + "MirEnumPayloadOp leaf path {:?} does not resolve inside the field type", + path + ); + }; + let semantic_leaf = + coerce_enum_payload_storage(ctx, rewriter, stored_leaf, semantic_leaf_ty)?; + let insert_op = llvm::InsertValueOp::new(ctx, current, semantic_leaf, path.clone()); + rewriter.insert_operation(ctx, insert_op.get_operation()); + current = insert_op.get_operation().deref(ctx).get_result(0); + } + rewriter.replace_operation_with_values(ctx, op, vec![current]); + return Ok(()); + } + match slot { Some(slot) => { let extract_op = llvm::ExtractValueOp::new(ctx, enum_val, vec![slot])?; @@ -631,6 +695,19 @@ mod tests { ); } + /// `Option<(usize, &T)>` with the niche on the tuple's pointer: the + /// aggregate payload is decomposed into leaf slots (the `usize` leaf at + /// byte 0, the pointer leaf reusing the carrier at byte 8), so construct + /// and payload read rebuild the value entirely in SSA. + /// + /// This test used to pin the spill design ("the payload moves as one + /// unit through memory, never field by field"): construction and + /// extraction each spilled the enum and moved the whole tuple through + /// the slot. The leaf-slot rule (enum-typed-payload round) replaced that + /// with typed leaf traffic on purpose: byte-wise spills of mixed-type + /// filler were what SROA could not promote, leaving local-memory + /// traffic in `opt -O2` output. Both leaves still round-trip — now as + /// `extractvalue`/`insertvalue` pairs, with zero memory ops. #[test] fn nested_pointer_niche_tuple_construct_extract_and_discriminant_lower() { let mut ctx = make_ctx(); @@ -664,7 +741,6 @@ mod tests { let slot_map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); assert_eq!(slot_map.carrier_slot, Some(1)); assert_eq!(slot_map.field_slots, vec![None]); - let lowered_tuple = convert_type(&mut ctx, tuple_ty).unwrap(); let (module, block) = build_kernel(&mut ctx, vec![index, pointer], vec![]); let index_value = block.deref(&ctx).get_argument(0); @@ -729,79 +805,504 @@ mod tests { 1, "reading the pointer niche should inspect the carrier exactly once" ); + assert_eq!( + count_ops::(&ctx, &body), + 0, + "leaf slots keep construction and extraction out of memory entirely" + ); assert_eq!( count_ops::(&ctx, &body), - 3, - "construction and extraction should each spill the enum, plus one tuple payload store" + 0, + "the decomposed payload must not spill through a stack slot" ); assert_eq!( count_ops::(&ctx, &body), - 2, - "construction should reload the enum and extraction should load the tuple payload" + 0, + "the decomposed payload must be rebuilt from SSA slots, not loads" ); - // What this test is about is that the payload moves as one unit, never - // field by field. Assert that property directly instead of counting - // whole-aggregate accesses. - // - // Counting was the fragile form. The enum's physical storage here is - // `{i64, ptr}` -- the *same interned type* as the lowered payload tuple, - // since the niche carrier claims the pointer at byte 8 and the 8 bytes - // below it become one `i64` filler. So a count of `lowered_tuple`-typed - // accesses cannot separate the payload write from the enum spill, and - // the expected numbers move whenever that coincidence appears or - // disappears -- which has nothing to do with the property under test. - // - // A lowering that decomposed the payload is recognisable by what it - // emits instead: traffic in the tuple's *field* types. Look for that, - // and the assertion holds whatever the enum storage type happens to be. - let field_tys: Vec = lowered_tuple - .deref(&ctx) - .downcast_ref::() - .expect("the lowered tuple is an LLVM struct") - .fields() - .collect(); - - let store_tys: Vec = find_all::(&ctx, &body) - .iter() - .map(|store| store.get_operand_value(&ctx).get_type(&ctx)) - .collect(); - let load_tys: Vec = find_all::(&ctx, &body) + // Both leaves must actually move: the tuple's own construction, + // the enum construction, and the payload rebuild each insert both + // leaves, and the discriminant read extracts the carrier once. + // (The enum's physical storage here is `{i64, ptr}` -- the same + // interned type as the lowered payload tuple, since the carrier + // claims the pointer at byte 8 -- so the counts, not the aggregate + // types, carry the assertion.) + let mut inserted: Vec> = + insert_indices(&ctx, &find_all::(&ctx, &body)); + inserted.sort(); + assert_eq!( + inserted, + vec![vec![0], vec![0], vec![0], vec![1], vec![1], vec![1]], + "the tuple build, the enum build, and the payload rebuild each insert both leaves" + ); + let mut extracted: Vec> = find_all::(&ctx, &body) .iter() - .map(|load| { - load.get_operation() - .deref(&ctx) - .get_result(0) - .get_type(&ctx) - }) + .map(|extract| extract.indices(&ctx)) .collect(); - - let describe = |tys: &[TypeHandle]| -> Vec { - tys.iter() - .filter(|ty| field_tys.contains(ty)) - .map(|ty| ty.deref(&ctx).disp(&ctx).to_string()) - .collect() - }; - assert!( - describe(&store_tys).is_empty(), - "the payload must be stored whole, but these field-typed stores appear: {:?}", - describe(&store_tys) - ); - assert!( - describe(&load_tys).is_empty(), - "the payload must be read whole, but these field-typed loads appear: {:?}", - describe(&load_tys) + extracted.sort(); + assert_eq!( + extracted, + vec![vec![0], vec![0], vec![1], vec![1], vec![1]], + "each leaf is read from the operand tuple, from the enum slots, plus the carrier" ); + } - // And it must actually be moved: without this the checks above would - // also pass a lowering that emitted no payload traffic at all. - assert!( - store_tys.contains(&lowered_tuple), - "at least one store must move the complete {{i64, ptr}} payload" - ); - assert!( - load_tys.contains(&lowered_tuple), - "at least one load must read the complete {{i64, ptr}} payload" - ); + /// The `Option` shape ({&mut T, usize, usize}, niche on the + /// pointer): one fixture shared by the SSA and niche tests below. + fn rope_pair_enum(ctx: &mut Context) -> (TypeHandle, TypeHandle, TypeHandle) { + let logical: TypeHandle = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); + let word: TypeHandle = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); + let pointee: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); + let pointer: TypeHandle = MirPtrType::get_generic(ctx, pointee, false).into(); + let payload: TypeHandle = MirStructType::get_with_full_layout( + ctx, + "RopePair".into(), + vec!["rope".into(), "base".into(), "cache".into()], + vec![pointer, word, word], + vec![0, 1, 2], + vec![0, 8, 16], + 24, + 8, + ) + .into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + ctx, + "MaybeRopePair".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![24]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 24, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + (enum_ty, payload, pointer) + } + + /// The `Option<(f32, f32, &mut f32)>` shape: floats below the pointer + /// that carries the niche. + fn float_pair_enum(ctx: &mut Context) -> (TypeHandle, TypeHandle, TypeHandle, TypeHandle) { + use pliron::builtin::types::FP32Type; + let logical: TypeHandle = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); + let float: TypeHandle = FP32Type::get(ctx).into(); + let pointee: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); + let pointer: TypeHandle = MirPtrType::get_generic(ctx, pointee, false).into(); + let payload: TypeHandle = MirTupleType::get_with_layout( + ctx, + vec![float, float, pointer], + vec![0, 1, 2], + vec![0, 4, 8], + 16, + 8, + ) + .into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + ctx, + "MaybeFloatPair".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![16]), + ], + EnumEncoding { + tag_offset: 8, + total_size: 16, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + (enum_ty, payload, pointer, float) + } + + /// Construct-then-payload-read of a decomposed aggregate payload stays + /// in SSA: no `alloca`, no `store`, no `load`. The spill the old lowering + /// needed here (byte stores into filler, one wide load back) is exactly + /// what SROA cannot promote, and what survived `opt -O2` as local-memory + /// traffic in the rms_norm/rope kernels. The tagged no-niche pair is the + /// regression guard: it was already whole-slot SSA before the leaf rule. + #[test] + fn aggregate_payload_construct_and_read_stay_in_ssa() { + // Option: {&mut T, usize, usize} built from its fields. + { + let mut ctx = make_ctx(); + let (enum_ty, payload, pointer) = rope_pair_enum(&mut ctx); + let word: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let (module, block) = build_kernel(&mut ctx, vec![pointer, word, word], vec![]); + let rope = block.deref(&ctx).get_argument(0); + let base = block.deref(&ctx).get_argument(1); + let cache = block.deref(&ctx).get_argument(2); + let struct_op = Operation::new( + &mut ctx, + mir::MirConstructStructOp::get_concrete_op_info(), + vec![payload], + vec![rope, base, cache], + vec![], + 0, + ); + struct_op.insert_at_back(block, &ctx); + let struct_value = struct_op.deref(&ctx).get_result(0); + let construct = Operation::new( + &mut ctx, + mir::MirConstructEnumOp::get_concrete_op_info(), + vec![enum_ty], + vec![struct_value], + vec![], + 0, + ); + mir::MirConstructEnumOp::new(construct) + .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(1)); + construct.insert_at_back(block, &ctx); + let enum_value = construct.deref(&ctx).get_result(0); + let payload_read = Operation::new( + &mut ctx, + mir::MirEnumPayloadOp::get_concrete_op_info(), + vec![payload], + vec![enum_value], + vec![], + 0, + ); + mir::MirEnumPayloadOp::new(payload_read) + .set_attr_payload_variant_index(&ctx, VariantIndexAttr(1)); + mir::MirEnumPayloadOp::new(payload_read) + .set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); + payload_read.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + assert_eq!( + count_ops::(&ctx, &body), + 0, + "the RopePair payload must not spill through a stack slot" + ); + assert_eq!(count_ops::(&ctx, &body), 0); + assert_eq!(count_ops::(&ctx, &body), 0); + } + // Option<(f32, f32, &mut f32)>: tuple built from its elements. + { + let mut ctx = make_ctx(); + let (enum_ty, payload, pointer, float) = float_pair_enum(&mut ctx); + let (module, block) = build_kernel(&mut ctx, vec![float, float, pointer], vec![]); + let first = block.deref(&ctx).get_argument(0); + let second = block.deref(&ctx).get_argument(1); + let cell = block.deref(&ctx).get_argument(2); + let tuple_op = Operation::new( + &mut ctx, + mir::MirConstructTupleOp::get_concrete_op_info(), + vec![payload], + vec![first, second, cell], + vec![], + 0, + ); + tuple_op.insert_at_back(block, &ctx); + let tuple_value = tuple_op.deref(&ctx).get_result(0); + let construct = Operation::new( + &mut ctx, + mir::MirConstructEnumOp::get_concrete_op_info(), + vec![enum_ty], + vec![tuple_value], + vec![], + 0, + ); + mir::MirConstructEnumOp::new(construct) + .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(1)); + construct.insert_at_back(block, &ctx); + let enum_value = construct.deref(&ctx).get_result(0); + let payload_read = Operation::new( + &mut ctx, + mir::MirEnumPayloadOp::get_concrete_op_info(), + vec![payload], + vec![enum_value], + vec![], + 0, + ); + mir::MirEnumPayloadOp::new(payload_read) + .set_attr_payload_variant_index(&ctx, VariantIndexAttr(1)); + mir::MirEnumPayloadOp::new(payload_read) + .set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); + payload_read.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + assert_eq!( + count_ops::(&ctx, &body), + 0, + "the float-pair payload must not spill through a stack slot" + ); + assert_eq!(count_ops::(&ctx, &body), 0); + assert_eq!(count_ops::(&ctx, &body), 0); + } + // Option<(u32, u32)>: tagged, no niche, whole-slot payload — the + // regression guard that passes with and without the leaf rule. + { + let mut ctx = make_ctx(); + let logical: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let word: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let payload: TypeHandle = MirTupleType::get_with_layout( + &mut ctx, + vec![word, word], + vec![0, 1], + vec![0, 4], + 8, + 4, + ) + .into(); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "MaybeWords".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![payload], vec![4], vec![8]), + ], + 0, + 12, + 4, + ) + .into(); + let (module, block) = build_kernel(&mut ctx, vec![word, word], vec![]); + let first = block.deref(&ctx).get_argument(0); + let second = block.deref(&ctx).get_argument(1); + let tuple_op = Operation::new( + &mut ctx, + mir::MirConstructTupleOp::get_concrete_op_info(), + vec![payload], + vec![first, second], + vec![], + 0, + ); + tuple_op.insert_at_back(block, &ctx); + let tuple_value = tuple_op.deref(&ctx).get_result(0); + let construct = Operation::new( + &mut ctx, + mir::MirConstructEnumOp::get_concrete_op_info(), + vec![enum_ty], + vec![tuple_value], + vec![], + 0, + ); + mir::MirConstructEnumOp::new(construct) + .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(1)); + construct.insert_at_back(block, &ctx); + let enum_value = construct.deref(&ctx).get_result(0); + let payload_read = Operation::new( + &mut ctx, + mir::MirEnumPayloadOp::get_concrete_op_info(), + vec![payload], + vec![enum_value], + vec![], + 0, + ); + mir::MirEnumPayloadOp::new(payload_read) + .set_attr_payload_variant_index(&ctx, VariantIndexAttr(1)); + mir::MirEnumPayloadOp::new(payload_read) + .set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); + payload_read.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + assert_eq!(count_ops::(&ctx, &body), 0); + assert_eq!(count_ops::(&ctx, &body), 0); + assert_eq!(count_ops::(&ctx, &body), 0); + } + } + + /// The leaf rule must not disturb the niche encoding. `None` still + /// writes a null pointer into the carrier slot and the discriminant read + /// still inspects exactly that slot; `Some` still writes the payload's + /// pointer there — now as the leaf insert instead of a memory spill. + #[test] + fn pointer_niche_leaf_slots_keep_null_none_and_payload_pointer_some() { + let check = |carrier_slot: u32, pointer_path: Vec, is_struct_payload: bool| { + // None: the carrier write is IntToPtr(constant 0), and the + // discriminant reads the carrier slot once. + { + let mut ctx = make_ctx(); + let (enum_ty, payload) = if is_struct_payload { + let (enum_ty, payload, _) = rope_pair_enum(&mut ctx); + (enum_ty, payload) + } else { + let (enum_ty, payload, _, _) = float_pair_enum(&mut ctx); + (enum_ty, payload) + }; + let _ = &payload; + let logical: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let (module, block) = build_kernel(&mut ctx, vec![], vec![logical]); + let construct = Operation::new( + &mut ctx, + mir::MirConstructEnumOp::get_concrete_op_info(), + vec![enum_ty], + vec![], + vec![], + 0, + ); + mir::MirConstructEnumOp::new(construct) + .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(0)); + construct.insert_at_back(block, &ctx); + let enum_value = construct.deref(&ctx).get_result(0); + let get = Operation::new( + &mut ctx, + mir::MirGetDiscriminantOp::get_concrete_op_info(), + vec![logical], + vec![enum_value], + vec![], + 0, + ); + get.insert_at_back(block, &ctx); + let discriminant = get.deref(&ctx).get_result(0); + append_mir_return(&mut ctx, block, vec![discriminant]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + let carrier_inserts: Vec<_> = find_all::(&ctx, &body) + .into_iter() + .filter(|insert| insert.indices(&ctx) == vec![carrier_slot]) + .collect(); + assert_eq!( + carrier_inserts.len(), + 1, + "constructing None writes the carrier slot exactly once" + ); + let value = carrier_inserts[0] + .get_operation() + .deref(&ctx) + .get_operand(1); + let int_to_ptr = Operation::get_op::( + value.defining_op().expect("carrier value must be defined"), + &ctx, + ) + .expect("None's carrier write is a null pointer, not a payload"); + let bits = int_to_ptr.get_operation().deref(&ctx).get_operand(0); + let bits_constant = Operation::get_op::( + bits.defining_op().expect("null bits must be defined"), + &ctx, + ) + .expect("None's niche bits must be a constant"); + let attr = bits_constant.get_value(&ctx); + let integer = (&*attr as &dyn pliron::attribute::Attribute) + .downcast_ref::() + .expect("niche bits must be an integer"); + assert_eq!( + integer.value().to_u64(), + 0, + "None must encode as the null pointer niche" + ); + let ptr_to_ints = find_all::(&ctx, &body); + assert_eq!( + ptr_to_ints.len(), + 1, + "the discriminant reads the carrier once" + ); + let carrier_read = ptr_to_ints[0].get_operation().deref(&ctx).get_operand(0); + let extract = Operation::get_op::( + carrier_read + .defining_op() + .expect("carrier read must be defined"), + &ctx, + ) + .expect("the discriminant must read the carrier slot"); + assert_eq!(extract.indices(&ctx), vec![carrier_slot]); + } + // Some: the carrier slot receives the payload's pointer leaf. + { + let mut ctx = make_ctx(); + let (enum_ty, payload, operands, pointer) = if is_struct_payload { + let (enum_ty, payload, pointer) = rope_pair_enum(&mut ctx); + let word: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + (enum_ty, payload, vec![pointer, word, word], pointer) + } else { + let (enum_ty, payload, pointer, float) = float_pair_enum(&mut ctx); + (enum_ty, payload, vec![float, float, pointer], pointer) + }; + let _ = pointer; + let (module, block) = build_kernel(&mut ctx, operands.clone(), vec![]); + let values: Vec<_> = (0..operands.len()) + .map(|index| block.deref(&ctx).get_argument(index)) + .collect(); + let build_info = if is_struct_payload { + mir::MirConstructStructOp::get_concrete_op_info() + } else { + mir::MirConstructTupleOp::get_concrete_op_info() + }; + let payload_op = + Operation::new(&mut ctx, build_info, vec![payload], values, vec![], 0); + payload_op.insert_at_back(block, &ctx); + let payload_value = payload_op.deref(&ctx).get_result(0); + let construct = Operation::new( + &mut ctx, + mir::MirConstructEnumOp::get_concrete_op_info(), + vec![enum_ty], + vec![payload_value], + vec![], + 0, + ); + mir::MirConstructEnumOp::new(construct) + .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(1)); + construct.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + // Inserts at the carrier slot come from two builds here: the + // payload's own construction (same index inside the payload + // struct) and the enum's. Exactly one of them must receive + // the payload's pointer leaf as an SSA extract — never a + // constant or a load. + let carrier_inserts: Vec<_> = find_all::(&ctx, &body) + .into_iter() + .filter(|insert| insert.indices(&ctx) == vec![carrier_slot]) + .collect(); + assert!( + !carrier_inserts.is_empty(), + "constructing Some must write the carrier slot" + ); + let mut leaf_into_carrier = 0; + for insert in &carrier_inserts { + let value = insert.get_operation().deref(&ctx).get_operand(1); + let Some(def) = value.defining_op() else { + continue; + }; + let Some(extract) = Operation::get_op::(def, &ctx) else { + continue; + }; + assert_eq!( + extract.indices(&ctx), + pointer_path, + "the carrier receives the payload's pointer leaf, not a constant or a load" + ); + leaf_into_carrier += 1; + } + assert_eq!( + leaf_into_carrier, 1, + "exactly one carrier write carries the payload pointer leaf" + ); + } + }; + // RopePair: carrier at slot 0, pointer leaf at payload path [0]. + check(0, vec![0], true); + // Float pair: carrier at slot 2, pointer leaf at payload path [2]. + check(2, vec![2], false); } /// SetDiscriminant must use the slot map instead of assuming that the tag 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 3d6737580b..a64278e00f 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 @@ -19,6 +19,7 @@ use super::layout::{alignment_anchor_type, make_enum_filler_type, mir_type_conta use super::pointer_storage::{ PointerOverlapRejection, analyze_pointer_overlap, llvm_type_contains_pointer, }; +use super::struct_layout::{StructLayoutInfo, build_struct_slot_map}; use super::{ convert_type, is_zero_sized_type, llvm_type_contains_i1, llvm_type_is_byte_faithful, llvm_type_size_align, natural_struct_layout, @@ -27,6 +28,24 @@ use crate::convert::enum_payload_storage::{ MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES, enum_payload_storage_type, }; +/// A payload leaf's `insertvalue`/`extractvalue` path into the field's +/// lowered value, paired with the enum struct slot backing it. +pub(crate) type LeafSlot = (Vec, u32); + +/// The claim-index form used while [`build_enum_slot_map`] is still laying +/// claims down; converted to [`LeafSlot`] once slots exist. +type LeafClaim = (Vec, usize); + +/// One byte-range claim on the enum's storage: (offset, size, type). +type EnumClaim = (u64, u64, TypeHandle); + +/// A successful leaf decomposition: each leaf's claim placement, plus the +/// fresh claims the caller must add. +struct DecomposedPayload { + placement: Vec, + fresh_claims: Vec, +} + /// The LLVM struct for an enum, plus a map saying where the tag and each /// payload field ended up. /// @@ -43,16 +62,180 @@ pub(crate) struct EnumSlotMap { pub carrier_llvm_ty: Option, /// Which struct slot holds each payload field, in the flattened /// order of `MirEnumType::all_field_types`. `None` means the field - /// has no slot of its own: it is zero-sized, or its bytes are shared - /// with a different-typed field of another variant. Such fields are - /// read and written through memory at `field_offsets` instead. + /// has no slot of its own: it is zero-sized, its bytes are shared + /// with a different-typed field of another variant, or it was + /// decomposed leaf by leaf (check [`Self::field_leaf_slots`] first). + /// Slotless fields without leaves are read and written through + /// memory at `field_offsets` instead. pub field_slots: Vec>, + /// Per-leaf slots for decomposed payload fields, parallel to + /// [`Self::field_slots`]. `Some` exactly when the field is an + /// aggregate whose every non-ZST leaf owns a typed slot (a leaf may + /// reuse the niche carrier's slot or another field's slot): one + /// `(path into the field's lowered LLVM value, enum struct slot)` + /// pair per leaf, in walk order. Construction and payload reads then + /// rebuild the aggregate in SSA, one `insertvalue`/`extractvalue` + /// per leaf, with no spill through memory. + pub field_leaf_slots: Vec>>, /// Byte position of each payload field inside the enum. pub field_offsets: Vec, /// Converted LLVM type of each payload field. pub field_llvm_types: Vec, } +/// Collect the leaves of an aggregate payload's MIR type: every real field +/// at its rustc byte offset, recursing through structs and tuples. +/// +/// The walk is over the MIR type, not the lowered LLVM struct, on purpose: +/// a lowered struct's `[N x i8]` padding slots are indistinguishable from a +/// real `[u8; N]` field, and padding must stay filler while real bytes get +/// slots. Paths are the `extractvalue`/`insertvalue` index paths into the +/// payload's lowered value, taken from [`build_struct_slot_map`] so they +/// account for stripped ZSTs and inserted padding slots. Zero-sized fields +/// carry no bytes and are skipped. Returns false (caller fails closed) on a +/// layout the walk cannot trust: an unsizable leaf, a divergent (packed) +/// struct layout, or an aggregate without computable slot offsets. +fn collect_payload_leaves( + ctx: &mut Context, + mir_ty: TypeHandle, + base_offset: u64, + path: &[u32], + out: &mut Vec<(Vec, u64, TypeHandle)>, +) -> bool { + let layout = { + let ty_ref = mir_ty.deref(ctx); + if let Some(struct_ty) = ty_ref.downcast_ref::() { + Some(StructLayoutInfo::of_struct(struct_ty)) + } else { + ty_ref + .downcast_ref::() + .map(StructLayoutInfo::of_tuple) + } + }; + let Some(layout) = layout else { + // A leaf: scalars, pointers, arrays, unions, nested enums. + let Ok(converted) = convert_type(ctx, mir_ty) else { + return false; + }; + out.push((path.to_vec(), base_offset, converted)); + return true; + }; + let Ok(slot_map) = build_struct_slot_map(ctx, &layout) else { + return false; + }; + if slot_map.layout_diverges { + // Packed payloads sit at offsets a naturally laid-out claim cannot + // name; keep the whole-value memory round-trip for them. + return false; + } + let Some(slot_offsets) = slot_map.natural_slot_offsets else { + return false; + }; + for (decl_index, field_ty) in layout.field_types.iter().enumerate() { + let Some(Some(slot)) = slot_map.decl_to_llvm.get(decl_index) else { + continue; // zero-sized field: no bytes, no slot + }; + let Some(&field_offset) = slot_offsets.get(*slot as usize) else { + return false; + }; + let Some(field_offset) = base_offset.checked_add(field_offset) else { + return false; + }; + let mut child = path.to_vec(); + child.push(*slot); + if !collect_payload_leaves(ctx, *field_ty, field_offset, &child, out) { + return false; + } + } + true +} + +/// Try to place one colliding struct/tuple payload field leaf by leaf. +/// +/// Every leaf either reuses a claim that already holds the same type at the +/// same bytes (the niche carrier backing the payload's pointer leaf, or a +/// same-typed leaf another variant placed) or lands on bytes no claim covers +/// and gets its own typed claim at its rustc byte offset. Bytes no leaf +/// claims — the payload's own padding — stay filler. Returns the leaf +/// placements (path, claim index — fresh claims are numbered +/// `claims.len() + i`) plus the fresh claims to add. +/// +/// Returns `None` — the field keeps today's whole-value memory round-trip — +/// whenever the rule cannot prove a clean placement: a leaf only partially +/// covered by a claim, covered by a different type, misaligned, or +/// unsizable; a payload that is not a struct/tuple; or one whose lowered +/// storage carries `i1`, which stays on the canonical byte path +/// ([`super::llvm_byte_faithful_twin`]). +/// +/// This is what keeps construction and payload reads in SSA for the +/// `Option<{ptr, i64, i64}>` and `Option<(f32, f32, ptr)>` shapes: without +/// it, the non-pointer leaves fall into `[N x i8]` filler bytes, the value +/// round-trips through a spill slot (byte stores, one wide load), and SROA +/// cannot reassemble the wide value from the byte slices — the spill slot +/// survives `opt -O2` as local-memory traffic. +fn try_decompose_aggregate_payload( + ctx: &mut Context, + mir_field_ty: TypeHandle, + storage_ty: TypeHandle, + field_offset: u64, + claims: &[(u64, u64, TypeHandle)], +) -> Option { + { + // Only struct/tuple payloads decompose: their leaf count is + // proportional to source text. Scalars keep the shared-slot and + // memory paths; arrays keep the whole-value paths (and their own + // bounded-rewrite gate in `enum_payload_storage`). + let ty_ref = mir_field_ty.deref(ctx); + if !ty_ref.is::() && !ty_ref.is::() { + return None; + } + } + if llvm_type_contains_i1(ctx, storage_ty) { + // Bool-bearing payloads stay slotless so the canonical-byte + // zero-extension at the storage boundary keeps owning their bytes. + return None; + } + let mut leaves = Vec::new(); + if !collect_payload_leaves(ctx, mir_field_ty, field_offset, &[], &mut leaves) { + return None; + } + let mut placement = Vec::with_capacity(leaves.len()); + let mut fresh_claims = Vec::new(); + for (path, offset, leaf_ty) in leaves { + let leaf_storage = enum_payload_storage_type(ctx, leaf_ty).ok()?; + let (leaf_size, leaf_align) = llvm_type_size_align(ctx, leaf_storage)?; + if leaf_size == 0 { + continue; // zero-sized leaf: no bytes to place + } + if !offset.is_multiple_of(leaf_align.max(1)) { + // The leaf sits unaligned inside the payload (packed struct): + // not this rule's shape. + return None; + } + if let Some(claim_index) = claims + .iter() + .position(|&(o, s, t)| o == offset && s == leaf_size && t == leaf_storage) + { + placement.push((path, claim_index)); + continue; + } + let overlaps_claim = claims + .iter() + .any(|&(o, s, _)| offset < o.saturating_add(s) && o < offset.saturating_add(leaf_size)); + if overlaps_claim { + // A partial or different-typed overlap with another variant's + // bytes: fail closed, keep the memory round-trip. + return None; + } + placement.push((path, claims.len() + fresh_claims.len())); + fresh_claims.push((offset, leaf_size, leaf_storage)); + } + Some(DecomposedPayload { + placement, + fresh_claims, + }) +} + /// Build the LLVM struct for an enum, placing everything at the byte /// positions rustc chose. /// @@ -78,6 +261,10 @@ pub(crate) struct EnumSlotMap { /// - shared slot: another variant already placed the SAME type at the /// SAME position; both map to that slot. (If B were /// B(u32), A and B would simply share slot 1.) +/// - leaf slots: a colliding AGGREGATE whose every leaf either matches an +/// existing claim (the niche carrier backing its pointer +/// leaf) or lands on unclaimed bytes decomposes into one +/// typed slot per leaf; construct/payload then stay in SSA. /// - no slot: the bytes are taken by a different type (B's f32 vs /// A's u32 here). The field is still at byte 4, just not /// nameable as a struct field; reads and writes go @@ -235,6 +422,10 @@ pub(crate) fn build_enum_slot_map( }; let mut claim_of_field: Vec> = vec![None; field_llvm_types.len()]; + // Leaf placements for decomposed fields, in claim indices; converted to + // struct slots after Phase 2 lays the claims down. + let mut leaf_placement_of_field: Vec>> = + vec![None; field_llvm_types.len()]; let mut field_is_inhabited = Vec::with_capacity(field_llvm_types.len()); for (variant, count) in variant_field_counts.iter().enumerate() { field_is_inhabited.extend(std::iter::repeat_n( @@ -397,6 +588,23 @@ pub(crate) fn build_enum_slot_map( .filter(|&&(o, s, _)| offset < o + s && o < offset + size) .collect::>(); if !colliding_claims.is_empty() { + // A colliding aggregate may still be fully representable leaf by + // leaf: the pointer leaf reuses the carrier it coincides with and + // every integer/float leaf gets its own typed slot. That keeps + // construct and payload reads in SSA instead of a byte-wise spill + // SROA cannot promote. Anything the leaf walk cannot place + // exactly falls through to the whole-value paths below unchanged. + if let Some(decomposed) = try_decompose_aggregate_payload( + ctx, + all_field_types[flat], + storage_ty, + offset, + &claims, + ) { + claims.extend(decomposed.fresh_claims); + leaf_placement_of_field[flat] = Some(decomposed.placement); + continue; + } let has_pointer_overlap = llvm_type_contains_pointer(ctx, storage_ty) || colliding_claims .iter() @@ -525,6 +733,17 @@ pub(crate) fn build_enum_slot_map( .into_iter() .map(|c| c.map(|ci| slot_of_claim[ci])) .collect(); + let field_leaf_slots = leaf_placement_of_field + .into_iter() + .map(|placement| { + placement.map(|leaves| { + leaves + .into_iter() + .map(|(path, claim)| (path, slot_of_claim[claim])) + .collect() + }) + }) + .collect(); Ok(EnumSlotMap { llvm_struct_ty: llvm_types::StructType::get_unnamed( ctx, @@ -534,6 +753,7 @@ pub(crate) fn build_enum_slot_map( carrier_slot: carrier_claim.map(|claim| slot_of_claim[claim]), carrier_llvm_ty: carrier_ty, field_slots, + field_leaf_slots, field_offsets: all_field_offsets, field_llvm_types, }) @@ -1384,6 +1604,284 @@ mod tests { ); } + /// The `Option` shape: `{&mut T, usize, usize}` with the niche + /// on the pointer. Built once; the layout, slot, and lowering tests all + /// exercise it. + fn rope_pair_enum(ctx: &mut Context) -> TypeHandle { + let logical = mir_uint(ctx, 64); + let word = mir_uint(ctx, 64); + let pointee = mir_uint(ctx, 32); + let pointer: TypeHandle = MirPtrType::get_generic(ctx, pointee, false).into(); + let payload: TypeHandle = MirStructType::get_with_full_layout( + ctx, + "RopePair".into(), + vec!["rope".into(), "base".into(), "cache".into()], + vec![pointer, word, word], + vec![0, 1, 2], + vec![0, 8, 16], + 24, + 8, + ) + .into(); + MirEnumType::get_with_encoding( + ctx, + "MaybeRopePair".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![24]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 24, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into() + } + + /// The `Option<(f32, f32, &mut f32)>` shape: the two floats below the + /// pointer that carries the niche (the rms_norm `next_cell` return). + fn float_pair_enum(ctx: &mut Context) -> TypeHandle { + let logical = mir_uint(ctx, 64); + let float: TypeHandle = FP32Type::get(ctx).into(); + let pointee = mir_uint(ctx, 32); + let pointer: TypeHandle = MirPtrType::get_generic(ctx, pointee, false).into(); + let payload: TypeHandle = MirTupleType::get_with_layout( + ctx, + vec![float, float, pointer], + vec![0, 1, 2], + vec![0, 4, 8], + 16, + 8, + ) + .into(); + MirEnumType::get_with_encoding( + ctx, + "MaybeFloatPair".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![16]), + ], + EnumEncoding { + tag_offset: 8, + total_size: 16, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into() + } + + #[test] + fn aggregate_pointer_niche_payload_types_every_leaf() { + let mut ctx = make_ctx(); + let enum_ty = rope_pair_enum(&mut ctx); + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + let pointee = mir_uint(&mut ctx, 32); + let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); + let lowered_pointer = convert_type(&mut ctx, pointer).unwrap(); + let word = llvm_int(&mut ctx, 64); + assert_eq!( + struct_fields(&ctx, map.llvm_struct_ty), + vec![lowered_pointer, word, word], + "every leaf owns a typed slot: {{ptr, i64, i64}}, not {{ptr, [16 x i8]}}" + ); + assert_eq!( + llvm_type_size_align(&ctx, map.llvm_struct_ty), + Some((24, 8)), + "same size and alignment as rustc's layout and the old carrier+filler form" + ); + assert_eq!(map.carrier_slot, Some(0)); + } + + #[test] + fn float_pair_pointer_niche_payload_types_every_leaf() { + let mut ctx = make_ctx(); + let enum_ty = float_pair_enum(&mut ctx); + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + let float: TypeHandle = FP32Type::get(&ctx).into(); + let pointee = mir_uint(&mut ctx, 32); + let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); + let lowered_pointer = convert_type(&mut ctx, pointer).unwrap(); + assert_eq!( + struct_fields(&ctx, map.llvm_struct_ty), + vec![float, float, lowered_pointer], + "the float pair owns typed slots below the carrier: {{f32, f32, ptr}}, not {{i64, ptr}}" + ); + assert_eq!( + llvm_type_size_align(&ctx, map.llvm_struct_ty), + Some((16, 8)), + "same size and alignment as rustc's layout and the old filler form" + ); + assert_eq!(map.carrier_slot, Some(2)); + } + + #[test] + fn enum_slot_map_records_leaf_slots_for_decomposed_payloads() { + let mut ctx = make_ctx(); + let enum_ty = rope_pair_enum(&mut ctx); + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!(map.field_slots, vec![None]); + assert_eq!( + map.field_leaf_slots, + vec![Some(vec![(vec![0], 0), (vec![1], 1), (vec![2], 2),])], + "the pointer leaf reuses the carrier slot; both usize leaves get their own" + ); + + let enum_ty = float_pair_enum(&mut ctx); + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!(map.field_slots, vec![None]); + assert_eq!( + map.field_leaf_slots, + vec![Some(vec![(vec![0], 0), (vec![1], 1), (vec![2], 2),])], + "both floats get their own slots; the pointer leaf reuses the carrier" + ); + } + + #[test] + fn tagged_pair_payload_keeps_whole_struct_slot() { + // `Option<(u32, u32)>` has no niche, so the tuple never collides with + // anything: it keeps one whole-struct slot, before and after the + // leaf rule. A regression guard, not a leaf-rule case. + let mut ctx = make_ctx(); + let logical = mir_uint(&mut ctx, 32); + let word = mir_uint(&mut ctx, 32); + let payload: TypeHandle = + MirTupleType::get_with_layout(&mut ctx, vec![word, word], vec![0, 1], vec![0, 4], 8, 4) + .into(); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "MaybeWords".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![payload], vec![4], vec![8]), + ], + 0, + 12, + 4, + ) + .into(); + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!(map.field_slots, vec![Some(1)]); + assert_eq!(map.field_leaf_slots, vec![None]); + let tag = llvm_int(&mut ctx, 32); + let lowered_payload = convert_type(&mut ctx, payload).unwrap(); + assert_eq!( + struct_fields(&ctx, map.llvm_struct_ty), + vec![tag, lowered_payload] + ); + assert_eq!( + llvm_type_size_align(&ctx, map.llvm_struct_ty), + Some((12, 4)) + ); + } + + #[test] + fn overlapping_data_variants_keep_byte_carrier() { + // `A(u64)` / `B(f32, f32)`: both of B's fields overlap A's u64 with + // a different type. The leaf rule only types leaves no other variant + // uses differently, so B stays on the byte-carrier memory path. + { + let mut ctx = make_ctx(); + let tag = mir_uint(&mut ctx, 32); + let bits = mir_uint(&mut ctx, 64); + let float: TypeHandle = FP32Type::get(&ctx).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "BitsOrFloats".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::new_with_layout("Bits".into(), vec![bits], vec![8], vec![8]), + EnumVariant::new_with_layout( + "Floats".into(), + vec![float, float], + vec![8, 12], + vec![4, 4], + ), + ], + 0, + 16, + 8, + ) + .into(); + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!( + map.field_slots, + vec![Some(2), None, None], + "A's u64 owns its slot; B's floats share its bytes with another type" + ); + assert_eq!(map.field_leaf_slots, vec![None, None, None]); + let filler = llvm_int(&mut ctx, 32); + let word = llvm_int(&mut ctx, 64); + let tag_ty = llvm_int(&mut ctx, 32); + assert_eq!( + struct_fields(&ctx, map.llvm_struct_ty), + vec![tag_ty, filler, word], + "{{tag, filler, u64}}: B's bytes stay byte-carrier, not typed f32 slots" + ); + assert_eq!( + llvm_type_size_align(&ctx, map.llvm_struct_ty), + Some((16, 8)) + ); + } + // The aggregate form `A(u64)` / `B((f32, f32))` fails closed the + // same way: the tuple's leaves overlap A's differently typed bytes. + { + let mut ctx = make_ctx(); + let tag = mir_uint(&mut ctx, 32); + let bits = mir_uint(&mut ctx, 64); + let float: TypeHandle = FP32Type::get(&ctx).into(); + let payload: TypeHandle = MirTupleType::get_with_layout( + &mut ctx, + vec![float, float], + vec![0, 1], + vec![0, 4], + 8, + 4, + ) + .into(); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "BitsOrFloatPair".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::new_with_layout("Bits".into(), vec![bits], vec![8], vec![8]), + EnumVariant::new_with_layout("Floats".into(), vec![payload], vec![8], vec![8]), + ], + 0, + 16, + 8, + ) + .into(); + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!( + map.field_slots, + vec![Some(2), None], + "the tuple overlaps the u64 claim with different leaf types: no leaf slots" + ); + assert_eq!(map.field_leaf_slots, vec![None, None]); + } + } + #[test] fn enum_slot_map_three_field_tuple_follows_recorded_offsets() { // rustc lays out `(u32, f32, &T)` with the pointer first in memory: @@ -2059,36 +2557,30 @@ mod tests { 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"); - + // The leaf rule types the u128 leaf itself, so the 16 bytes below the + // carrier are one i128 slot instead of two i64 filler leaves. An i128 + // slot brings alignment 16 on its own, so the zero-length alignment + // anchor the filler form needed is gone. + assert_eq!(fields.len(), 3); assert_eq!( - filler.size(), - 2, - "16-byte enum filler should use two 64-bit leaves" - ); - - assert_eq!( - filler - .elem_type() + fields[0] .deref(&ctx) .downcast_ref::() .map(IntegerType::width), - Some(64), - "16-byte enum filler elements should be i64" + Some(128), + "the u128 payload leaf should own a typed i128 slot" + ); + assert_eq!(map.carrier_slot, Some(1)); + assert_eq!( + map.field_leaf_slots, + vec![Some(vec![(vec![0], 0), (vec![1], 1)])], + "the u128 leaf gets slot 0 and the pointer leaf reuses the carrier" ); 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" + "the leaf-typed payload must preserve rustc's enum size and alignment" ); } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/Cargo.lock new file mode 100644 index 0000000000..d64ea61988 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/Cargo.lock @@ -0,0 +1,765 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "bindgen" +version = "0.69.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "lazy_static", + "lazycell", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash", + "shlex", + "syn 2.0.119", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + +[[package]] +name = "cfg-if" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "clang-sys" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "157a8ba7b480713b56f4c09fd13fc3e0a22a5dfab8097ba61cbc5feef950788a" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ba106e844045ddc02c20768ac4b0b7cb45f2cc2ecbd833ac9f5dacb183c2055" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-bindings" +version = "0.4.0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-core" +version = "0.4.0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.4.0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn 2.0.119", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[package]] +name = "enum_niche_payload" +version = "0.1.0" +dependencies = [ + "cuda-core", + "cuda-device", + "cuda-host", + "ptx-parse", +] + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.191" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +dependencies = [ + "libloading", + "thiserror", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/Cargo.toml new file mode 100644 index 0000000000..fb6cca4f14 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "enum_niche_payload" +version = "0.1.0" +edition = "2024" +authors = ["NVIDIA Corporation"] +license = "Apache-2.0" + +[workspace] + +[dependencies] +cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-device = { path = "../../../cuda-device" } +cuda-host = { path = "../../../cuda-host" } +ptx-parse = { path = "../../../ptx-parse" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/README.md new file mode 100644 index 0000000000..be27bc4d66 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/README.md @@ -0,0 +1,36 @@ +# enum_niche_payload + +Regression test for an enum whose niche carrier shares bytes with an aggregate +payload: `Option<(f32, f32, &mut f32)>`, the rms_norm `next_cell` shape. + +Only the pointer leaf used to get a typed slot. The two floats fell into +`[N x i8]` filler, construction spilled them byte by byte, and the payload +read loaded them back as one wide value. SROA cannot rebuild mixed-type +slices, so the spill slot survived `opt -O2` as a local depot. + +The fix gives every leaf of such a payload a typed slot at its rustc byte +offset, reusing the carrier's slot where the types match, and builds and +reads the value with `insertvalue`/`extractvalue`. Size, alignment, field +offsets, and the niche encoding are unchanged; a payload the rule cannot +prove keeps the old lowering. + +## Depot + +Measured with the same device-only kernel source against each tree's backend: + +| entry | e027ac32 (parent) | with this fix | +|---|---|---| +| `smooth_triplets` | 16 B in the entry, 16 B in the outlined `Cells::next_cell` | none | + +`lease_borrow` and `lease_raw` are layout-neutral value checks, not part of +the depot criterion: neither emits a `.local` on the parent. + +## Run + +``` +cargo oxide run enum_niche_payload +cargo oxide build enum_niche_payload --verbose +crates/rustc-codegen-cuda/examples/enum_niche_payload/target/release/enum_niche_payload --verify-ptx +``` + +`--verify-ptx` fails on the parent: `smooth_triplets` carries the 16 B depot. diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/src/main.rs new file mode 100644 index 0000000000..0d428e0ba5 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/src/main.rs @@ -0,0 +1,283 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Reproducer for enum payload leaves that round-trip through local memory. +//! +//! An enum whose niche carrier shares its bytes with an aggregate payload +//! used to leave the payload's non-pointer leaves in `[N x i8]` filler: +//! construction spilled the value byte by byte, payload reads loaded it +//! back as one wide value, and SROA cannot reassemble mixed-type slices — +//! so the spill slot survived `opt -O2` as a `.local` depot. The leaf-slot +//! rule gives every payload leaf its own typed slot, and construction and +//! payload reads stay in SSA. +//! +//! Three shapes, one kernel each: +//! +//! - `smooth_triplets`: `Option<(f32, f32, &mut f32)>` driven by `while let` +//! (the rms_norm `next_cell` shape). The depot reproducer: 16 B in the +//! entry plus 16 B in the outlined `Cells::next_cell` on the parent, +//! none with the fix. +//! - `lease_borrow`: `Option` where `Lease` holds `&mut` to a +//! kernel-local struct (the rope `own` shape). +//! - `lease_raw`: the same lease copying the raw pointer it needs and +//! holding only a `PhantomData` borrow. +//! +//! Both lease shapes are layout-neutral value checks: neither emits a +//! `.local` on the parent either, so they do not measure this fix. +//! +//! Usage: +//! cargo oxide run enum_niche_payload +//! cargo oxide build enum_niche_payload --verbose +//! .../enum_niche_payload/target/release/enum_niche_payload --verify-ptx + +use core::marker::PhantomData; +use cuda_core::simt::LaunchConfig; +use cuda_core::{CudaContext, DeviceBuffer}; +use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; + +const TRIPLETS: usize = 4; + +/// (a) A lending cursor over non-overlapping triples: each `next_cell` +/// returns the first two elements by value and the third by borrow, all in +/// one `Option<(f32, f32, &mut f32)>` whose niche is the borrow. +struct Cells<'a> { + data: &'a mut [f32], +} + +impl<'a> Cells<'a> { + fn next_cell(&mut self) -> Option<(f32, f32, &mut f32)> { + let rest = core::mem::take(&mut self.data); + let (head, tail) = rest.split_at_mut_checked(3)?; + let (prev, mid) = head.split_at_mut_checked(1)?; + let (cur, cell) = mid.split_at_mut_checked(1)?; + self.data = tail; + Some((prev[0], cur[0], &mut cell[0])) + } +} + +/// The kernel-local state leased out by `own`/`own_raw`. +struct Stepper { + step: usize, + stride: usize, +} + +/// (b) The pair borrows the stepper: the enum payload carries a pointer to +/// a kernel-local struct, so its address rides through every `None`/`Some` +/// merge and the struct's storage escapes into the depot. +struct Lease<'r> { + src: &'r mut Stepper, + base: usize, +} + +/// (c) The pair copies the raw pointer it needs and holds only a +/// `PhantomData` borrow: same bytes for the enum (`{ptr, usize}`), but no +/// reference crosses the merge, so nothing forces the payload through +/// memory. +struct LeaseRaw<'r> { + src: *mut Stepper, + base: usize, + _borrow: PhantomData<&'r mut Stepper>, +} + +impl Stepper { + fn own(&mut self, thread: u32) -> Option> { + if thread % 2 != 0 { + return None; + } + Some(Lease { + src: self, + base: thread as usize, + }) + } + + fn own_raw(&mut self, thread: u32) -> Option> { + if thread % 2 != 0 { + return None; + } + Some(LeaseRaw { + src: self as *mut Stepper, + base: thread as usize, + _borrow: PhantomData, + }) + } +} + +#[cuda_module] +mod kernels { + use super::*; + + /// (a) `Option<(f32, f32, &mut f32)>` driven by `while let`: the float + /// pair below the pointer niche used to live in filler bytes. + #[kernel] + pub fn smooth_triplets(mut values: DisjointSlice, n: usize) { + if thread::index_1d().get() != 0 { + return; + } + // SAFETY: thread 0 alone owns the complete buffer for this launch; + // `n` is the element count the host passed for the same buffer. + let data = unsafe { core::slice::from_raw_parts_mut(values.as_mut_ptr(), n) }; + let mut cells = Cells { data }; + let mut visits = 0usize; + while let Some((prev, cur, cell)) = cells.next_cell() { + *cell += prev * cur; + visits += 1; + } + // SAFETY: one writer (thread 0), slot `n` reserved for the count. + unsafe { values.as_mut_ptr().add(n).write(visits as f32) }; + } + + /// (b) `Option` with `Lease` holding `&mut` to a kernel-local + /// struct. A layout-neutral value check: it emits no `.local` on the + /// parent either, so it does not measure the depot fix. + #[kernel] + pub fn lease_borrow(mut out: DisjointSlice) { + if thread::index_1d().get() != 0 { + return; + } + let mut stepper = Stepper { step: 1, stride: 7 }; + let mut acc = 0usize; + for i in 0..4u32 { + if let Some(lease) = stepper.own(i) { + lease.src.step = lease + .src + .step + .wrapping_mul(lease.src.stride) + .wrapping_add(lease.base); + acc = acc.wrapping_add(lease.src.step); + } + } + // SAFETY: one writer (thread 0); two u32 slots were allocated. + unsafe { + out.as_mut_ptr().write(stepper.step as u32); + out.as_mut_ptr().add(1).write(acc as u32); + } + } + + /// (c) The same lease with the raw pointer copied and only a + /// `PhantomData` borrow: `{ptr, usize}` payload, all leaves typed. A + /// layout-neutral value check, like (b). + #[kernel] + pub fn lease_raw(mut out: DisjointSlice) { + if thread::index_1d().get() != 0 { + return; + } + let mut stepper = Stepper { step: 1, stride: 7 }; + let mut acc = 0usize; + for i in 0..4u32 { + if let Some(lease) = stepper.own_raw(i) { + // SAFETY: `src` still points at this invocation's live + // kernel-local `stepper`, and thread 0 is the only writer. + unsafe { + (*lease.src).step = (*lease.src) + .step + .wrapping_mul((*lease.src).stride) + .wrapping_add(lease.base); + acc = acc.wrapping_add((*lease.src).step); + } + } + } + // SAFETY: one writer (thread 0); two u32 slots were allocated. + unsafe { + out.as_mut_ptr().write(stepper.step as u32); + out.as_mut_ptr().add(1).write(acc as u32); + } + } +} + +fn main() -> Result<(), Box> { + if std::env::args().any(|arg| arg == "--verify-ptx") { + return verify_ptx(); + } + + println!("=== enum_niche_payload ==="); + let ctx = CudaContext::new(0)?; + let stream = ctx.default_stream(); + let module = kernels::load(&ctx)?; + let one_thread = LaunchConfig::for_num_elems(1); + + // (a) smooth_triplets: cell[i] += prev * cur over 4 disjoint triples, + // then the visit count in slot n (buffer holds n + 1 elements). + let n = TRIPLETS * 3; + let mut upload: Vec = (0..n).map(|i| i as f32).collect(); + upload.push(0.0); + let mut buf = DeviceBuffer::from_host(&stream, &upload)?; + // SAFETY: launch shape/resources match the kernel; `n` triples plus the + // count slot match the buffer length the host uploaded. + unsafe { module.smooth_triplets(&stream, one_thread, &mut buf, n) }?; + let got = buf.to_host_vec(&stream)?; + let mut want = upload.clone(); + let mut visits = 0usize; + for t in 0..TRIPLETS { + let i = t * 3; + want[i + 2] += want[i] * want[i + 1]; + visits += 1; + } + for (i, (g, w)) in got[..n].iter().zip(&want).enumerate() { + assert_eq!(g, w, "smooth_triplets: triplet output {i}"); + } + assert_eq!( + got[n], visits as f32, + "smooth_triplets: visit count must survive the while-let loop" + ); + + // (b) and (c): the borrowed lease and the raw-pointer lease must agree. + let mut expected_step = 1usize; + let mut expected_acc = 0usize; + for i in 0..4u32 { + if i % 2 == 0 { + expected_step = expected_step.wrapping_mul(7).wrapping_add(i as usize); + expected_acc = expected_acc.wrapping_add(expected_step); + } + } + for name in ["lease_borrow", "lease_raw"] { + let zeros = vec![0u32; 2]; + let mut out = DeviceBuffer::from_host(&stream, &zeros)?; + // SAFETY: launch shape/resources match the kernel; the buffer owns + // exactly the two u32 slots each kernel writes. + if name == "lease_borrow" { + unsafe { module.lease_borrow(&stream, one_thread, &mut out) }?; + } else { + unsafe { module.lease_raw(&stream, one_thread, &mut out) }?; + } + let got = out.to_host_vec(&stream)?; + assert_eq!( + got, + vec![expected_step as u32, expected_acc as u32], + "{name}: the lease must update the stepper exactly as the host model" + ); + } + + println!("PASS: enum_niche_payload (all three shapes compute correctly)"); + Ok(()) +} + +/// Machine-checked depot criterion: every entry carries no `.local` traffic +/// with the fix. The parent fails on `smooth_triplets`, which keeps the +/// 16 B depot. +fn verify_ptx() -> Result<(), Box> { + let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("enum_niche_payload.ptx"); + let ptx = std::fs::read_to_string(&path)?; + let document = ptx_parse::Document::parse(&ptx)?; + + for (name, expects_depot) in [ + ("smooth_triplets", false), + ("lease_borrow", false), + ("lease_raw", false), + ] { + let definition = document + .definitions_named(name) + .find(|definition| definition.callable().kind() == ptx_parse::CallableKind::Entry) + .ok_or_else(|| format!("missing or incomplete PTX entry `{name}`"))?; + let has_depot = definition.text().contains(".local"); + if has_depot != expects_depot { + return Err(format!( + "{name} local-depot expectation failed: expected depot={expects_depot}, found depot={has_depot}" + ) + .into()); + } + } + println!("SUCCESS: depot expectations hold (no `.local` in any entry)"); + Ok(()) +}