From 9966a0459117e5ae7317c863360050c84bdea08d Mon Sep 17 00:00:00 2001 From: Yunus Mutlu Date: Sat, 26 Sep 2026 23:05:39 +0300 Subject: [PATCH] fix(shared): enforce typed shared-memory alignment Signed-off-by: Yunus Mutlu --- cuda-oxide/crates/dialect-mir/src/ops/cast.rs | 60 +------- .../crates/dialect-mir/src/ops/memory.rs | 128 +++++++++++++++-- cuda-oxide/crates/dialect-mir/src/types.rs | 61 +++++++- cuda-oxide/crates/mir-lower/src/context.rs | 2 +- .../src/convert/ops/memory/shared.rs | 15 ++ .../src/convert/ops/memory/tests/shared.rs | 132 +++++++++++++++++- cuda-oxide/crates/mir-lower/src/lowering.rs | 38 ++++- 7 files changed, 361 insertions(+), 75 deletions(-) diff --git a/cuda-oxide/crates/dialect-mir/src/ops/cast.rs b/cuda-oxide/crates/dialect-mir/src/ops/cast.rs index fe796219d8..8e0cfbf1a1 100644 --- a/cuda-oxide/crates/dialect-mir/src/ops/cast.rs +++ b/cuda-oxide/crates/dialect-mir/src/ops/cast.rs @@ -11,7 +11,7 @@ use crate::attributes::{MirCastKindAttr, MirPointerKindAuthorityAttr}; use crate::types::{ MirArrayType, MirDisjointSliceType, MirEnumType, MirPointerCarrier, MirPointerKind, MirPtrType, MirSliceType, MirStructType, MirTupleType, MirUnionType, is_opaque_fn_pointer_type, - pointer_carriers_in_type, pointer_kinds_in_type, + pointer_carriers_in_type, pointer_kinds_in_type, required_type_alignment, }; use pliron::{ builtin::{ @@ -26,7 +26,7 @@ use pliron::{ op::Op, operation::Operation, result::Error, - r#type::{Typed, type_cast, type_impls}, + r#type::{Typed, type_impls}, value::Value, verify_err, }; @@ -248,62 +248,8 @@ fn is_promoted_empty_unique_ref_transition( } /// Required alignment of a promoted empty pointee. -/// -/// Even a zero-length reference must be aligned. Aggregate MIR types retain -/// rustc's exact ABI alignment; arrays inherit it from their element. Scalar -/// and pointer leaves use the NVPTX64 natural alignment modeled by lowering. -/// Unknown leaves fail closed instead of relying on a later lowering error. fn required_pointee_alignment(ctx: &Context, ty: pliron::r#type::TypeHandle) -> Option { - let ty = ty.deref(ctx); - if let Some(array) = ty.downcast_ref::() { - return required_pointee_alignment(ctx, array.element_ty); - } - if let Some(tuple) = ty.downcast_ref::() { - return if tuple.abi_align() > 0 { - Some(tuple.abi_align()) - } else if tuple.types.is_empty() && tuple.total_size == 0 { - Some(1) - } else { - None - }; - } - if let Some(structure) = ty.downcast_ref::() { - return if structure.abi_align > 0 { - Some(structure.abi_align) - } else if structure.field_types.is_empty() && structure.total_size == 0 { - Some(1) - } else { - None - }; - } - if let Some(enumeration) = ty.downcast_ref::() { - return (enumeration.abi_align() > 0).then(|| enumeration.abi_align()); - } - if let Some(union) = ty.downcast_ref::() { - return (union.abi_align() > 0).then(|| union.abi_align()); - } - if ty.is::() { - return Some(8); - } - if let Some(disjoint) = ty.downcast_ref::() { - let mut alignment = 8; - for &space_ty in &disjoint.space_tys { - alignment = alignment.max(required_pointee_alignment(ctx, space_ty)?); - } - return Some(alignment); - } - if let Some(integer) = ty.downcast_ref::() { - let size = u64::from(integer.width()).div_ceil(8).max(1); - return Some(size.next_power_of_two().min(16)); - } - if ty.is::() { - return Some(8); - } - if let Some(float) = type_cast::(&*ty) { - let size = u64::try_from(float.get_semantics().bits).ok()?.div_ceil(8); - return Some(size.next_power_of_two().min(16)); - } - None + required_type_alignment(ctx, ty) } /// Pair every pointer carrier in `target` with the carrier at the same diff --git a/cuda-oxide/crates/dialect-mir/src/ops/memory.rs b/cuda-oxide/crates/dialect-mir/src/ops/memory.rs index b59023dbf9..42af3e444e 100644 --- a/cuda-oxide/crates/dialect-mir/src/ops/memory.rs +++ b/cuda-oxide/crates/dialect-mir/src/ops/memory.rs @@ -1117,6 +1117,16 @@ impl Verify for MirSharedAllocOp { return verify_err!(op.loc(), "MirSharedAllocOp missing size attribute"); } + if let Some(alignment) = self.get_alignment_value(ctx) + && (alignment == 0 || !alignment.is_power_of_two()) + { + return verify_err!( + op.loc(), + "MirSharedAllocOp alignment must be a non-zero power of two, found {}", + alignment + ); + } + // Check result is a shared memory pointer let res = op.get_result(0); let res_ty = res.get_type(ctx); @@ -1341,6 +1351,86 @@ mod device_global_key_tests { } } +#[cfg(test)] +mod shared_alignment_verifier_tests { + use super::*; + use pliron::utils::apint::APInt; + use pliron::{ + builtin::{ + attributes::{IntegerAttr, TypeAttr}, + types::{IntegerType, Signedness}, + }, + common_traits::Verify, + operation::Operation, + }; + use std::num::NonZeroUsize; + + fn u64_attr(ctx: &Context, value: u64) -> IntegerAttr { + IntegerAttr::new( + IntegerType::get(ctx, 64, Signedness::Unsigned), + APInt::from_u64(value, NonZeroUsize::new(64).unwrap()), + ) + } + + fn shared_alloc(ctx: &mut Context, alignment: u64) -> MirSharedAllocOp { + let elem_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); + let result_ty = MirPtrType::get_shared(ctx, elem_ty, true); + let op = Operation::new( + ctx, + MirSharedAllocOp::get_concrete_op_info(), + vec![result_ty.into()], + vec![], + vec![], + 0, + ); + let alloc = MirSharedAllocOp::new(op); + alloc.set_attr_elem_type(ctx, TypeAttr::new(elem_ty)); + alloc.set_attr_size(ctx, u64_attr(ctx, 1)); + alloc.set_alignment_value(ctx, alignment); + alloc + } + + fn extern_shared(ctx: &mut Context, alignment: u64) -> MirExternSharedOp { + let elem_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); + let result_ty = MirPtrType::get_shared(ctx, elem_ty, true); + let op = Operation::new( + ctx, + MirExternSharedOp::get_concrete_op_info(), + vec![result_ty.into()], + vec![], + vec![], + 0, + ); + let extern_shared = MirExternSharedOp::new(op); + extern_shared.set_alignment_value(ctx, alignment); + extern_shared + } + + #[test] + fn shared_alloc_rejects_non_power_of_two_alignment() { + let mut ctx = Context::new(); + crate::register(&mut ctx); + + let invalid = shared_alloc(&mut ctx, 3); + assert!(invalid.verify(&ctx).is_err()); + + let valid = shared_alloc(&mut ctx, 8); + assert!(valid.verify(&ctx).is_ok()); + } + + #[test] + fn extern_shared_rejects_non_power_of_two_alignment() { + let mut ctx = Context::new(); + crate::register(&mut ctx); + + let invalid = extern_shared(&mut ctx, 3); + assert!(invalid.verify(&ctx).is_err()); + + let valid = extern_shared(&mut ctx, 8); + assert!(valid.verify(&ctx).is_ok()); + } +} + // ============================================================================ // MirExternSharedOp // ============================================================================ @@ -1354,18 +1444,21 @@ mod device_global_key_tests { /// All `MirExternSharedOp` instances in a kernel refer to the **same** underlying /// memory, with different byte offsets for partitioning. /// -/// This is lowered to an LLVM global with extern linkage and addrspace(3): +/// This is lowered to an LLVM global with extern linkage and addrspace(3). +/// The global alignment is the maximum effective requirement across the +/// function's dynamic shared-memory references and any propagated launch +/// contract: /// ```llvm -/// @__dynamic_smem = external addrspace(3) global [0 x i8], align 256 +/// @__dynamic_smem = external addrspace(3) global [0 x i8], align /// ``` /// /// # Attributes /// /// ```text -/// | Name | Type | Description | -/// |-----------------|-------------|------------------------------------------| -/// | `byte_offset` | IntegerAttr | Byte offset from start of dynamic smem | -/// | `mir_alignment` | IntegerAttr | Alignment hint (global always uses 256) | +/// | Name | Type | Description | +/// |--------------------|-------------|------------------------------------------------| +/// | `byte_offset` | IntegerAttr | Byte offset from start of dynamic smem | +/// | `extern_alignment` | IntegerAttr | Requested minimum alignment for this reference | /// ``` /// /// # Results @@ -1379,6 +1472,10 @@ mod device_global_key_tests { /// # Verification /// /// - Result must be a pointer type with shared address space (3). +/// - An explicit alignment must be a non-zero power of two. +/// +/// Lowering also raises the requested alignment to the pointee type's required +/// ABI alignment before computing the maximum alignment for the shared pool. #[pliron_op( name = "mir.extern_shared", format, @@ -1414,10 +1511,11 @@ impl MirExternSharedOp { self.set_attr_extern_byte_offset(ctx, offset_attr); } - /// Get alignment as u64 (returns 128 if not set). + /// Get the requested minimum alignment as u64. /// - /// Note: The actual global alignment is fixed at 256 bytes in mir-lower, - /// regardless of this attribute value. + /// Legacy or synthetic MIR without an explicit alignment attribute falls + /// back to 128 bytes. Importer-produced dynamic shared-memory operations + /// carry the source-level alignment explicitly. pub fn get_alignment_value(&self, ctx: &Context) -> u64 { self.get_attr_extern_alignment(ctx) .map(|attr| attr.value().to_u64()) @@ -1443,6 +1541,18 @@ impl Verify for MirExternSharedOp { fn verify(&self, ctx: &Context) -> Result<(), Error> { let op = &*self.get_operation().deref(ctx); + if let Some(alignment) = self + .get_attr_extern_alignment(ctx) + .map(|attr| attr.value().to_u64()) + && (alignment == 0 || !alignment.is_power_of_two()) + { + return verify_err!( + op.loc(), + "MirExternSharedOp alignment must be a non-zero power of two, found {}", + alignment + ); + } + // Check result is a shared memory pointer let res = op.get_result(0); let res_ty = res.get_type(ctx); diff --git a/cuda-oxide/crates/dialect-mir/src/types.rs b/cuda-oxide/crates/dialect-mir/src/types.rs index 4cdc437bf9..7135b89c0b 100644 --- a/cuda-oxide/crates/dialect-mir/src/types.rs +++ b/cuda-oxide/crates/dialect-mir/src/types.rs @@ -10,7 +10,7 @@ use pliron::context::Context; use pliron::derive::{format, pliron_type, type_interface_impl}; use pliron::location::Location; use pliron::result::Error; -use pliron::r#type::{Type, TypeHandle, TypedHandle}; +use pliron::r#type::{Type, TypeHandle, TypedHandle, type_cast}; use pliron::utils::apfloat::{self, GetSemantics, Semantics}; use pliron::{common_traits::Verify, verify_err}; @@ -2056,6 +2056,65 @@ impl Verify for MirEnumType { } } +/// Required ABI alignment of a MIR value type in bytes. +/// +/// Aggregate MIR types retain rustc's exact ABI alignment; arrays inherit it +/// from their element. Scalar and pointer leaves use the NVPTX64 natural +/// alignment modeled by lowering. Unknown leaves fail closed rather than +/// guessing an alignment that could make typed storage under-aligned. +pub fn required_type_alignment(ctx: &Context, ty: TypeHandle) -> Option { + let ty = ty.deref(ctx); + if let Some(array) = ty.downcast_ref::() { + return required_type_alignment(ctx, array.element_ty); + } + if let Some(tuple) = ty.downcast_ref::() { + return if tuple.abi_align() > 0 { + Some(tuple.abi_align()) + } else if tuple.types.is_empty() && tuple.total_size == 0 { + Some(1) + } else { + None + }; + } + if let Some(structure) = ty.downcast_ref::() { + return if structure.abi_align > 0 { + Some(structure.abi_align) + } else if structure.field_types.is_empty() && structure.total_size == 0 { + Some(1) + } else { + None + }; + } + if let Some(enumeration) = ty.downcast_ref::() { + return (enumeration.abi_align() > 0).then(|| enumeration.abi_align()); + } + if let Some(union) = ty.downcast_ref::() { + return (union.abi_align() > 0).then(|| union.abi_align()); + } + if ty.is::() { + return Some(8); + } + if let Some(disjoint) = ty.downcast_ref::() { + let mut alignment = 8; + for &space_ty in &disjoint.space_tys { + alignment = alignment.max(required_type_alignment(ctx, space_ty)?); + } + return Some(alignment); + } + if let Some(integer) = ty.downcast_ref::() { + let size = u64::from(integer.width()).div_ceil(8).max(1); + return Some(size.next_power_of_two().min(16)); + } + if ty.is::() { + return Some(8); + } + if let Some(float) = type_cast::(&*ty) { + let size = u64::try_from(float.get_semantics().bits).ok()?.div_ceil(8); + return Some(size.next_power_of_two().min(16)); + } + None +} + /// A pointer carrier embedded directly or recursively in a MIR value type. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct MirPointerCarrier { diff --git a/cuda-oxide/crates/mir-lower/src/context.rs b/cuda-oxide/crates/mir-lower/src/context.rs index f4fa8112ce..88c6f73388 100644 --- a/cuda-oxide/crates/mir-lower/src/context.rs +++ b/cuda-oxide/crates/mir-lower/src/context.rs @@ -94,7 +94,7 @@ pub struct SharedGlobalDeclaration { pub mir_elem_type: TypeHandle, /// Number of elements in the shared allocation. pub size: u64, - /// Explicit alignment, or zero when natural alignment is requested. + /// Effective physical storage alignment in bytes. pub alignment: u64, /// Source-level debug identity carried on the allocation, when present. pub debug_info: Option, diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/shared.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/shared.rs index fc413d25b5..fc66db6208 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/shared.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/shared.rs @@ -96,6 +96,21 @@ pub fn convert_shared_alloc_dc( // optional debug identity against the lowered element layout before it // can reach the declaration record or the emitted global. let llvm_elem_type = convert_type(ctx, mir_elem_type).map_err(anyhow_to_pliron)?; + + if alignment != 0 && !alignment.is_power_of_two() { + return Err(anyhow_to_pliron(anyhow::anyhow!( + "mir.shared_alloc alignment must be a power of two, found {}", + alignment + ))); + } + let required_alignment = dialect_mir::types::required_type_alignment(ctx, mir_elem_type) + .ok_or_else(|| { + anyhow_to_pliron(anyhow::anyhow!( + "cannot determine required alignment for mir.shared_alloc element type" + )) + })?; + let alignment = alignment.max(required_alignment); + let debug_info = debug_info .filter(|info| shared_debug_type_matches_physical(ctx, info, llvm_elem_type, size)); let debug_owner_function = debug_info.as_ref().and(debug_owner_function); diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/shared.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/shared.rs index 1e4702b076..aeb5981e0f 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/shared.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/shared.rs @@ -194,6 +194,62 @@ fn convert_shared_alloc_rejects_conflicting_key_declarations() { ); } +#[test] +fn shared_alloc_alignment_is_at_least_scalar_requirement() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + + append_shared_alloc_typed(&mut ctx, block, "under-aligned-i64", i64_ty, 1, None, 1); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let top = module_top_block(&ctx, module_ptr); + let global = top + .deref(&ctx) + .iter(&ctx) + .find_map(|op| Operation::get_op::(op, &ctx)) + .expect("expected shared global"); + assert_eq!( + global.get_alignment(&ctx), + Some(8), + "typed i64 shared storage must not be weakened below its ABI alignment" + ); +} + +#[test] +fn shared_alloc_alignment_preserves_mir_over_alignment() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let aligned_ty = over_aligned_tuple_ty(&mut ctx); + + append_shared_alloc_typed( + &mut ctx, + block, + "under-aligned-aggregate", + aligned_ty, + 1, + None, + 1, + ); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let top = module_top_block(&ctx, module_ptr); + let global = top + .deref(&ctx) + .iter(&ctx) + .find_map(|op| Operation::get_op::(op, &ctx)) + .expect("expected shared global"); + assert_eq!( + global.get_alignment(&ctx), + Some(32), + "MIR ABI over-alignment must survive shared-storage lowering" + ); +} + #[test] fn shared_alloc_repeated_debug_identity_is_attached_once() { let mut ctx = make_ctx(); @@ -649,11 +705,21 @@ fn append_extern_shared( byte_offset: u64, ) -> Ptr { let i8_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Signless).into(); - let shared_i8 = MirPtrType::get_shared(ctx, i8_ty, true); + append_extern_shared_typed(ctx, block, i8_ty, alignment, byte_offset) +} + +fn append_extern_shared_typed( + ctx: &mut Context, + block: Ptr, + pointee_type: TypeHandle, + alignment: u64, + byte_offset: u64, +) -> Ptr { + let shared_ptr = MirPtrType::get_shared(ctx, pointee_type, true); let op = Operation::new( ctx, mir::MirExternSharedOp::get_concrete_op_info(), - vec![shared_i8.into()], + vec![shared_ptr.into()], vec![], vec![], 0, @@ -665,6 +731,68 @@ fn append_extern_shared( op } +#[test] +fn extern_shared_alignment_is_at_least_pointee_requirement() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + + append_extern_shared_typed(&mut ctx, block, i64_ty, 1, 0); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let top = module_top_block(&ctx, module_ptr); + let global = top + .deref(&ctx) + .iter(&ctx) + .filter_map(|op| Operation::get_op::(op, &ctx)) + .find(|global| { + global + .get_symbol_name(&ctx) + .to_string() + .starts_with("__dynamic_smem_") + }) + .expect("expected dynamic shared-memory global"); + assert_eq!( + global.get_alignment(&ctx), + Some(8), + "typed i64 dynamic shared storage must not be weakened below its ABI alignment" + ); +} + +#[test] +fn extern_shared_alignment_uses_maximum_typed_requirement() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + let aligned_ty = over_aligned_tuple_ty(&mut ctx); + + append_extern_shared_typed(&mut ctx, block, i64_ty, 1, 0); + append_extern_shared_typed(&mut ctx, block, aligned_ty, 1, 0); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let top = module_top_block(&ctx, module_ptr); + let global = top + .deref(&ctx) + .iter(&ctx) + .filter_map(|op| Operation::get_op::(op, &ctx)) + .find(|global| { + global + .get_symbol_name(&ctx) + .to_string() + .starts_with("__dynamic_smem_") + }) + .expect("expected dynamic shared-memory global"); + assert_eq!( + global.get_alignment(&ctx), + Some(32), + "one dynamic pool must satisfy the strongest typed alignment requirement" + ); +} + /// A `mir.extern_shared` with a nonzero byte offset addresses into an extern /// `[0 x i8]` addrspace(3) global whose real size only exists at launch, so /// the lowering cannot promise LLVM's `inbounds` allocated-object contract diff --git a/cuda-oxide/crates/mir-lower/src/lowering.rs b/cuda-oxide/crates/mir-lower/src/lowering.rs index cd8310a489..80606ff769 100644 --- a/cuda-oxide/crates/mir-lower/src/lowering.rs +++ b/cuda-oxide/crates/mir-lower/src/lowering.rs @@ -49,7 +49,7 @@ use pliron::{ op::Op, operation::Operation, result::Result, - r#type::TypeHandle, + r#type::{TypeHandle, Typed}, value::Value, }; use reserved_oxide_symbols::{ @@ -314,7 +314,7 @@ pub fn convert_func( // Pre-scan MIR blocks for max dynamic shared memory alignment. // Must happen BEFORE inline_region empties the MIR region. let mir_blocks: Vec<_> = mir_region.deref(ctx).iter(ctx).collect(); - let body_max_align = compute_max_dynamic_smem_alignment(ctx, &mir_blocks); + let body_max_align = compute_max_dynamic_smem_alignment(ctx, &mir_blocks)?; let contract_min_align = dynamic_shared_alignment_attr(ctx, op); let max_align = match (body_max_align, contract_min_align) { (Some(body), Some(contract)) => Some(body.max(contract)), @@ -1274,7 +1274,7 @@ fn insert_op_sequentially( fn compute_max_dynamic_smem_alignment( ctx: &Context, mir_blocks: &[Ptr], -) -> Option { +) -> Result> { let mut max_alignment: Option = None; for mir_block in mir_blocks { @@ -1282,7 +1282,35 @@ fn compute_max_dynamic_smem_alignment( let op_id = Operation::get_opid(op, ctx); if op_id == dialect_mir::ops::MirExternSharedOp::get_opid_static() { let extern_shared = dialect_mir::ops::MirExternSharedOp::new(op); - let alignment = extern_shared.get_alignment_value(ctx); + let requested_alignment = extern_shared.get_alignment_value(ctx); + + if requested_alignment == 0 || !requested_alignment.is_power_of_two() { + return Err(anyhow_to_pliron(anyhow::anyhow!( + "mir.extern_shared alignment must be a non-zero power of two, found {}", + requested_alignment + ))); + } + + let result_ty = op.deref(ctx).get_result(0).get_type(ctx); + let result_ty_ref = result_ty.deref(ctx); + let ptr_ty = result_ty_ref + .downcast_ref::() + .ok_or_else(|| { + anyhow_to_pliron(anyhow::anyhow!( + "mir.extern_shared result is not a MIR pointer" + )) + })?; + + let required_alignment = dialect_mir::types::required_type_alignment( + ctx, + ptr_ty.pointee, + ) + .ok_or_else(|| { + anyhow_to_pliron(anyhow::anyhow!( + "cannot determine required alignment for mir.extern_shared pointee type" + )) + })?; + let alignment = requested_alignment.max(required_alignment); max_alignment = Some(match max_alignment { Some(current_max) => current_max.max(alignment), @@ -1292,7 +1320,7 @@ fn compute_max_dynamic_smem_alignment( } } - max_alignment + Ok(max_alignment) } // =====================================================================