fix(shared): enforce typed shared-memory alignment

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<u64> {
-    let ty = ty.deref(ctx);
-    if let Some(array) = ty.downcast_ref::<MirArrayType>() {
-        return required_pointee_alignment(ctx, array.element_ty);
-    }
-    if let Some(tuple) = ty.downcast_ref::<MirTupleType>() {
-        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::<MirStructType>() {
-        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::<MirEnumType>() {
-        return (enumeration.abi_align() > 0).then(|| enumeration.abi_align());
-    }
-    if let Some(union) = ty.downcast_ref::<MirUnionType>() {
-        return (union.abi_align() > 0).then(|| union.abi_align());
-    }
-    if ty.is::<MirSliceType>() {
-        return Some(8);
-    }
-    if let Some(disjoint) = ty.downcast_ref::<MirDisjointSliceType>() {
-        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::<IntegerType>() {
-        let size = u64::from(integer.width()).div_ceil(8).max(1);
-        return Some(size.next_power_of_two().min(16));
-    }
-    if ty.is::<MirPtrType>() {
-        return Some(8);
-    }
-    if let Some(float) = type_cast::<dyn FloatTypeInterface>(&*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

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 <effective>
 /// ```
 ///
 /// # 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);

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<u64> {
+    let ty = ty.deref(ctx);
+    if let Some(array) = ty.downcast_ref::<MirArrayType>() {
+        return required_type_alignment(ctx, array.element_ty);
+    }
+    if let Some(tuple) = ty.downcast_ref::<MirTupleType>() {
+        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::<MirStructType>() {
+        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::<MirEnumType>() {
+        return (enumeration.abi_align() > 0).then(|| enumeration.abi_align());
+    }
+    if let Some(union) = ty.downcast_ref::<MirUnionType>() {
+        return (union.abi_align() > 0).then(|| union.abi_align());
+    }
+    if ty.is::<MirSliceType>() {
+        return Some(8);
+    }
+    if let Some(disjoint) = ty.downcast_ref::<MirDisjointSliceType>() {
+        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::<IntegerType>() {
+        let size = u64::from(integer.width()).div_ceil(8).max(1);
+        return Some(size.next_power_of_two().min(16));
+    }
+    if ty.is::<MirPtrType>() {
+        return Some(8);
+    }
+    if let Some(float) = type_cast::<dyn FloatTypeInterface>(&*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 {

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<DebugGlobalVariableInfo>,

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);