fix(shared): enforce typed shared-memory alignment

cuda-oxide/crates/cuda-device/src/shared.rs

@@ -80,7 +80,8 @@ use core::ops::{Index, IndexMut};
 ///
 /// - `T`: Element type (f32, f64, i32, etc.)
 /// - `N`: Array size (must be known at compile time)
-/// - `ALIGN`: Alignment in bytes (0 = natural alignment, 128 for TMA destinations)
+/// - `ALIGN`: Alignment in bytes. `0` means natural alignment; any other value
+///   must be a power of two (128 for TMA destinations).
 ///
 /// # Alignment
 ///
@@ -96,6 +97,18 @@ use core::ops::{Index, IndexMut};
 /// static mut TMA_TILE: SharedArray<f32, 256, 128> = SharedArray::UNINIT;
 /// ```
 ///
+/// The contract the compiler enforces:
+///
+/// - `ALIGN` is `0` or a non-zero power of two. Anything else is a compile
+///   error that names the use site.
+/// - Storage is never declared below `align_of::<T>()`. A request below it is
+///   raised silently, as `#[repr(align(N))]` raises a struct; a larger request
+///   is honored. So `SharedArray<u64, 2, 1>` is declared at 8 bytes, and a
+///   `#[repr(C, align(16))]` element at natural alignment is declared at 16.
+/// - Aggregate, nested-array and `f16` elements at natural alignment declare
+///   their natural alignment in PTX (`.align 4` for `[f32; N]`, `.align 2` for
+///   `f16`). Older releases printed a blanket `.align 8` for them.
+///
 /// # Soundness: `!Sync`
 ///
 /// `SharedArray` is intentionally `!Sync`. GPU shared memory is concurrently
@@ -305,7 +318,8 @@ pub fn __dynamic_shared_alignment<const ALIGN: usize>() {
 /// # Type Parameters
 ///
 /// - `T`: Element type for the returned pointer
-/// - `ALIGN`: Base alignment in bytes (default: 16, matching nvcc)
+/// - `ALIGN`: Base alignment in bytes (default: 16, matching nvcc). Must be a
+///   non-zero power of two.
 ///
 /// # Alignment
 ///
@@ -331,6 +345,18 @@ pub fn __dynamic_shared_alignment<const ALIGN: usize>() {
 /// - `128`: Required for TMA (Tensor Memory Accelerator) operations
 /// - `256`: Cache-line friendly for some architectures
 ///
+/// The contract the compiler enforces:
+///
+/// - `ALIGN` is a non-zero power of two. Anything else, including `0`, is a
+///   compile error that names the use site.
+/// - A kernel has one dynamic pool. It is declared at the strongest typed
+///   requirement over the kernel's `get`, `get_raw` and `offset` uses, each
+///   taken as the larger of its `ALIGN` and `align_of::<T>()`, combined with
+///   the launch contract's alignment. `DynamicSharedArray::<u64, 1>::get()`
+///   alone declares the pool at 8 bytes.
+/// - The declaration aligns the pool base. Keeping an `offset(byte_offset)`
+///   result aligned for `T` is still the caller's responsibility.
+///
 /// # Usage
 ///
 /// ```rust,ignore
@@ -484,8 +510,9 @@ impl<T, const ALIGN: usize> DynamicSharedArray<T, ALIGN> {
     /// Ensure `byte_offset` is properly aligned for type `T`. For example,
     /// `f32` requires 4-byte alignment, `f64` requires 8-byte alignment.
     ///
-    /// The base of dynamic shared memory is aligned to `ALIGN` bytes (default 16),
-    /// so to get an N-byte aligned pointer, ensure `byte_offset % N == 0`.
+    /// The base of dynamic shared memory is aligned to at least `ALIGN` bytes
+    /// (default 16) and never below `align_of::<T>()`, so to get an N-byte
+    /// aligned pointer, ensure `byte_offset % N == 0`.
     #[inline(never)]
     pub fn offset(byte_offset: usize) -> *mut T {
         // The byte_offset is used by the compiler to compute the pointer.

cuda-oxide/crates/dialect-mir/src/ops/memory.rs

@@ -1116,6 +1116,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);
@@ -1353,18 +1363,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                                    |
+/// |----------------------|-------------|------------------------------------------------|
+/// | `extern_byte_offset` | IntegerAttr | Byte offset from start of dynamic smem         |
+/// | `extern_alignment`   | IntegerAttr | Requested minimum alignment for this reference |
 /// ```
 ///
 /// # Results
@@ -1378,6 +1391,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,
@@ -1413,10 +1430,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())
@@ -1442,6 +1460,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);
@@ -1468,6 +1498,105 @@ impl Verify for MirExternSharedOp {
     }
 }
 
+#[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
+    }
+
+    /// The rejection must name the op and the offending value: that line is
+    /// what the user sees at the use site.
+    fn assert_rejected(result: Result<(), Error>, expected: &str) {
+        let error = result.expect_err("an invalid alignment must fail verification");
+        let message = error.to_string();
+        assert!(
+            message.contains(expected),
+            "expected `{expected}` in the diagnostic, got: {message}"
+        );
+    }
+
+    #[test]
+    fn shared_alloc_rejects_non_power_of_two_alignment() {
+        let mut ctx = Context::new();
+        crate::register(&mut ctx);
+
+        assert_rejected(
+            shared_alloc(&mut ctx, 3).verify(&ctx),
+            "MirSharedAllocOp alignment must be a non-zero power of two, found 3",
+        );
+        assert_rejected(
+            shared_alloc(&mut ctx, 0).verify(&ctx),
+            "MirSharedAllocOp alignment must be a non-zero power of two, found 0",
+        );
+        assert!(shared_alloc(&mut ctx, 8).verify(&ctx).is_ok());
+    }
+
+    #[test]
+    fn extern_shared_rejects_non_power_of_two_alignment() {
+        let mut ctx = Context::new();
+        crate::register(&mut ctx);
+
+        assert_rejected(
+            extern_shared(&mut ctx, 3).verify(&ctx),
+            "MirExternSharedOp alignment must be a non-zero power of two, found 3",
+        );
+        assert_rejected(
+            extern_shared(&mut ctx, 0).verify(&ctx),
+            "MirExternSharedOp alignment must be a non-zero power of two, found 0",
+        );
+        assert!(extern_shared(&mut ctx, 8).verify(&ctx).is_ok());
+    }
+}
+
 /// Register memory operations into the given context.
 pub fn register(ctx: &mut Context) {
     MirAllocaOp::register(ctx);

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

@@ -9,7 +9,7 @@ use super::common::{anyhow_to_pliron, counter_named_global};
 use crate::context::{
     SharedGlobalDeclaration, SharedGlobalKind, SharedGlobalRecord, SharedGlobalsMap,
 };
-use crate::convert::types::{convert_type, llvm_type_size_align};
+use crate::convert::types::{convert_type, llvm_type_size_align, required_storage_align};
 use crate::helpers;
 use llvm_export::ops as llvm;
 use llvm_export::ops::GlobalOpExt;
@@ -96,6 +96,22 @@ 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)?;
+
+    // Storage is never declared below the element's ABI alignment: a smaller
+    // request is raised, a larger one is honored. The verifier has already
+    // rejected a request that is not a power of two.
+    let required_alignment = required_storage_align(ctx, mir_elem_type, llvm_elem_type)
+        .ok_or_else(|| {
+            let elem_display = mir_elem_type.deref(ctx).disp(ctx).to_string();
+            pliron::create_error!(
+                op.deref(ctx).loc(),
+                pliron::result::ErrorKind::VerificationFailed,
+                "mir.shared_alloc element type `{elem_display}` has no known ABI alignment; \
+                 refusing to guess the storage alignment"
+            )
+        })?;
+    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);