fix

Author: 39ali

crates/rustc_codegen_spirv/src/abi.rs

@@ -948,25 +948,7 @@ fn trans_intrinsic_type<'tcx>(
             }
 
             // Generic arg 0: component type T
-            let component_type = match args.type_at(0).kind() {
-                TyKind::Float(FloatTy::F32) => SpirvType::Float(32).def(span, cx),
-                TyKind::Float(FloatTy::F64) => SpirvType::Float(64).def(span, cx),
-                TyKind::Int(IntTy::I8) => SpirvType::Integer(8, true).def(span, cx),
-                TyKind::Int(IntTy::I16) => SpirvType::Integer(16, true).def(span, cx),
-                TyKind::Int(IntTy::I32) => SpirvType::Integer(32, true).def(span, cx),
-                TyKind::Int(IntTy::I64) => SpirvType::Integer(64, true).def(span, cx),
-                TyKind::Uint(UintTy::U8) => SpirvType::Integer(8, false).def(span, cx),
-                TyKind::Uint(UintTy::U16) => SpirvType::Integer(16, false).def(span, cx),
-                TyKind::Uint(UintTy::U32) => SpirvType::Integer(32, false).def(span, cx),
-                TyKind::Uint(UintTy::U64) => SpirvType::Integer(64, false).def(span, cx),
-                _ => {
-                    return Err(cx.tcx.dcx().span_err(
-                        span,
-                        "unsupported component type for #[spirv(cooperative_matrix)]: \
-                         must be f32, f64, i8, i16, i32, i64, u8, u16, u32, or u64",
-                    ));
-                }
-            };
+            let component_type = cx.layout_of(args.type_at(0)).spirv_type(span, cx);
 
             // Const generic 1: USE (MatrixA=0, MatrixB=1, MatrixAccumulator=2)
             // Const generic 2: ROWS

crates/spirv-std/src/cooperative_matrix.rs

@@ -5,22 +5,47 @@
 //! -C target-feature=+CooperativeMatrixKHR,+ext:SPV_KHR_cooperative_matrix
 //! ```
 //!
+//! See the [SPV_KHR_cooperative_matrix specification] for full details.
+//!
+//! [SPV_KHR_cooperative_matrix specification]: https://github.khronos.org/SPIRV-Registry/extensions/KHR/SPV_KHR_cooperative_matrix.html
 #[cfg(target_arch = "spirv")]
 use core::arch::asm;
 use core::marker::PhantomData;
 use core::mem::MaybeUninit;
 
+/// Matrix role in a cooperative multiply-accumulate operation (`D = A × B + C`).
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+#[repr(u32)]
+pub enum MatrixUse {
+    /// Input operand A.
+    MatrixA = 0,
+    /// Input operand B.
+    MatrixB = 1,
+    /// Accumulator / result.
+    MatrixAccumulator = 2,
+}
+
+/// Memory layout for cooperative matrix load/store operations.
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+#[repr(u32)]
+pub enum MatrixLayout {
+    /// Rows are stored contiguously.
+    RowMajor = 0,
+    /// Columns are stored contiguously.
+    ColumnMajor = 1,
+}
+
 /// Matrix role: input operand A in D = A × B + C.
-pub const MATRIX_A: u32 = 0;
+pub const MATRIX_A: u32 = MatrixUse::MatrixA as u32;
 /// Matrix role: input operand B in D = A × B + C.
-pub const MATRIX_B: u32 = 1;
+pub const MATRIX_B: u32 = MatrixUse::MatrixB as u32;
 /// Matrix role: accumulator / result in D = A × B + C.
-pub const MATRIX_ACCUMULATOR: u32 = 2;
+pub const MATRIX_ACCUMULATOR: u32 = MatrixUse::MatrixAccumulator as u32;
 
 /// Memory layout: rows are stored contiguously.
-pub const ROW_MAJOR: u32 = 0;
+pub const ROW_MAJOR: MatrixLayout = MatrixLayout::RowMajor;
 /// Memory layout: columns are stored contiguously.
-pub const COLUMN_MAJOR: u32 = 1;
+pub const COLUMN_MAJOR: MatrixLayout = MatrixLayout::ColumnMajor;
 
 /// A cooperative matrix distributed across the subgroup.
 ///
@@ -29,7 +54,7 @@ pub const COLUMN_MAJOR: u32 = 1;
 ///
 /// # Type parameters
 /// - `T`: element type (`f32`, `f64`, `i32`, `u32`, `i8`, `u8`, etc.)
-/// - `USE`: matrix role — one of [`MATRIX_A`], [`MATRIX_B`], [`MATRIX_ACCUMULATOR`]
+/// - `USE`: matrix role — one of [`MatrixUse::MatrixA`], [`MatrixUse::MatrixB`], [`MatrixUse::MatrixAccumulator`] cast to `u32`
 /// - `ROWS`: number of rows
 /// - `COLS`: number of columns
 ///
@@ -46,22 +71,28 @@ pub struct CooperativeMatrix<T, const USE: u32, const ROWS: u32, const COLS: u32
 }
 
 impl<T, const USE: u32, const ROWS: u32, const COLS: u32> CooperativeMatrix<T, USE, ROWS, COLS> {
-    /// Load a cooperative matrix tile from `ptr` using `layout` and `stride`.
+    /// Load a cooperative matrix through a pointer.
     ///
-    /// - `ptr`: pointer to the first element of the tile in memory
-    /// - `layout`: [`ROW_MAJOR`] or [`COLUMN_MAJOR`]
-    /// - `stride`: distance (in elements) between the start of consecutive
-    ///   rows (row-major) or columns (column-major) in the full matrix
+    /// `slice` must point into an array. `layout` specifies whether the matrix
+    /// is stored in row-major ([`MatrixLayout::RowMajor`]) or column-major
+    /// ([`MatrixLayout::ColumnMajor`]) order. `stride` is the number of elements
+    /// between the start of consecutive rows (row-major) or columns (column-major).
+    ///
+    /// The scope is always `Subgroup`.
     ///
     /// # Safety
-    /// `ptr` must be valid for `ROWS * stride` (row-major) or `COLS * stride`
-    /// (column-major) element reads within the subgroup's access pattern.
+    /// - `slice` must point into an array and be valid for all element accesses
+    ///   implied by the matrix dimensions, layout, and stride.
+    /// - All operands must be dynamically uniform within every instance of the
+    ///   subgroup scope.
     #[spirv_std_macros::gpu_only]
     #[doc(alias = "OpCooperativeMatrixLoadKHR")]
     #[inline]
-    pub unsafe fn load(ptr: *const T, layout: u32, stride: u32) -> Self {
+    pub unsafe fn load(slice: &[T], layout: MatrixLayout, stride: u32) -> Self {
         unsafe {
             let mut result = MaybeUninit::<Self>::uninit();
+            let layout_u32 = layout as u32;
+            let ptr = slice.as_ptr();
             asm!(
                 "%u32 = OpTypeInt 32 0",
                 "%layout = OpLoad %u32 {layout}",
@@ -70,24 +101,35 @@ impl<T, const USE: u32, const ROWS: u32, const COLS: u32> CooperativeMatrix<T, U
                 "%result = OpCooperativeMatrixLoadKHR typeof* {out} {ptr} %layout %stride",
                 "OpStore {out} %result",
                 ptr    = in(reg) ptr,
-                layout = in(reg) &layout,
+                layout = in(reg) &layout_u32,
                 stride = in(reg) &stride,
                 out    = in(reg) result.as_mut_ptr(),
             );
             result.assume_init()
         }
     }
 
-    /// Store this cooperative matrix tile to `ptr` using `layout` and `stride`.
+    /// Store a cooperative matrix through a pointer.
+    ///
+    /// `slice` must point into an array. `layout` specifies whether the matrix
+    /// is stored in row-major ([`MatrixLayout::RowMajor`]) or column-major
+    /// ([`MatrixLayout::ColumnMajor`]) order. `stride` is the number of elements
+    /// between the start of consecutive rows (row-major) or columns (column-major).
+    ///
+    /// The scope is always `Subgroup`.
     ///
     /// # Safety
-    /// `ptr` must be valid for `ROWS * stride` (row-major) or `COLS * stride`
-    /// (column-major) element writes within the subgroup's access pattern.
+    /// - `slice` must point into an array and be valid for all element accesses
+    ///   implied by the matrix dimensions, layout, and stride.
+    /// - All operands must be dynamically uniform within every instance of the
+    ///   subgroup scope.
     #[spirv_std_macros::gpu_only]
     #[doc(alias = "OpCooperativeMatrixStoreKHR")]
     #[inline]
-    pub unsafe fn store(self, ptr: *mut T, layout: u32, stride: u32) {
+    pub unsafe fn store(self, slice: &mut [T], layout: MatrixLayout, stride: u32) {
         unsafe {
+            let layout_u32 = layout as u32;
+            let ptr = slice.as_mut_ptr();
             asm!(
                 "%u32 = OpTypeInt 32 0",
                 "%layout = OpLoad %u32 {layout}",
@@ -96,7 +138,7 @@ impl<T, const USE: u32, const ROWS: u32, const COLS: u32> CooperativeMatrix<T, U
                 "OpCooperativeMatrixStoreKHR {ptr} %matrix %layout %stride",
                 ptr    = in(reg) ptr,
                 matrix = in(reg) &self,
-                layout = in(reg) &layout,
+                layout = in(reg) &layout_u32,
                 stride = in(reg) &stride,
             );
         }
@@ -125,26 +167,32 @@ impl<T, const USE: u32, const ROWS: u32, const COLS: u32> CooperativeMatrix<T, U
     }
 }
 
-/// Fused multiply-accumulate: `D = A × B + C`.
+/// Linear-algebraic matrix multiply of `A` by `B` and then component-wise add `C`.
+///
+/// The order of operations is implementation-dependent. All matrices must have the
+/// same scope, which is always subgroup here.
 ///
-/// - `A`: `M × K` matrix with role [`MATRIX_A`]
-/// - `B`: `K × N` matrix with role [`MATRIX_B`]
-/// - `C`: `M × N` matrix with role [`MATRIX_ACCUMULATOR`]
-/// - returns: `M × N` accumulator equal to `A × B + C`
+/// - `A`: `M × K` matrix with use [`MatrixUse::MatrixA`]
+/// - `B`: `K × N` matrix with use [`MatrixUse::MatrixB`]
+/// - `C`: `M × N` matrix with use [`MatrixUse::MatrixAccumulator`]
+/// - returns `D`: `M × N` accumulator equal to `A × B + C`
 ///
+/// All operands must be dynamically uniform within every instance of the subgroup scope.
 ///
 /// # Capability
 /// Requires `CooperativeMatrixKHR` + `SPV_KHR_cooperative_matrix`.
 #[spirv_std_macros::gpu_only]
 #[doc(alias = "OpCooperativeMatrixMulAddKHR")]
 #[inline]
 pub fn mul_add<TA, TB, TC, const M: u32, const N: u32, const K: u32>(
-    a: CooperativeMatrix<TA, MATRIX_A, M, K>,
-    b: CooperativeMatrix<TB, MATRIX_B, K, N>,
-    c: CooperativeMatrix<TC, MATRIX_ACCUMULATOR, M, N>,
-) -> CooperativeMatrix<TC, MATRIX_ACCUMULATOR, M, N> {
+    a: CooperativeMatrix<TA, { MatrixUse::MatrixA as u32 }, M, K>,
+    b: CooperativeMatrix<TB, { MatrixUse::MatrixB as u32 }, K, N>,
+    c: CooperativeMatrix<TC, { MatrixUse::MatrixAccumulator as u32 }, M, N>,
+) -> CooperativeMatrix<TC, { MatrixUse::MatrixAccumulator as u32 }, M, N> {
     unsafe {
-        let mut result = MaybeUninit::<CooperativeMatrix<TC, MATRIX_ACCUMULATOR, M, N>>::uninit();
+        let mut result = MaybeUninit::<
+            CooperativeMatrix<TC, { MatrixUse::MatrixAccumulator as u32 }, M, N>,
+        >::uninit();
         asm!(
             "%a      = OpLoad _ {a}",
             "%b      = OpLoad _ {b}",

tests/compiletests/ui/arch/cooperative_matrix.rs

@@ -17,11 +17,11 @@ pub fn main(
     #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] b: &[f32],
     #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] c: &mut [f32],
 ) {
-    let mat_a = unsafe { MatA::load(a.as_ptr(), ROW_MAJOR, 16) };
-    let mat_b = unsafe { MatB::load(b.as_ptr(), COLUMN_MAJOR, 16) };
-    let mat_c = unsafe { MatAcc::load(c.as_ptr(), ROW_MAJOR, 16) };
+    let mat_a = unsafe { MatA::load(a, ROW_MAJOR, 16) };
+    let mat_b = unsafe { MatB::load(b, COLUMN_MAJOR, 16) };
+    let mat_c = unsafe { MatAcc::load(c, ROW_MAJOR, 16) };
 
     let result = mul_add(mat_a, mat_b, mat_c);
 
-    unsafe { result.store(c.as_mut_ptr(), ROW_MAJOR, 16) };
+    unsafe { result.store(c, ROW_MAJOR, 16) };
 }