spirv_std/cooperative_matrix.rs
1//! Cooperative matrix types and operations (`SPV_KHR_cooperative_matrix`).
2//!
3//! Requires the `CooperativeMatrixKHR` capability and `SPV_KHR_cooperative_matrix` extension:
4//! ```text
5//! -C target-feature=+CooperativeMatrixKHR,+ext:SPV_KHR_cooperative_matrix
6//! ```
7//!
8//! See the [SPV_KHR_cooperative_matrix specification] for full details.
9//!
10//! [SPV_KHR_cooperative_matrix specification]: https://github.khronos.org/SPIRV-Registry/extensions/KHR/SPV_KHR_cooperative_matrix.html
11#[cfg(target_arch = "spirv")]
12use core::arch::asm;
13use core::marker::PhantomData;
14use core::mem::MaybeUninit;
15
16/// Matrix role in a cooperative multiply-accumulate operation (`D = A × B + C`).
17#[derive(Copy, Clone, Debug, PartialEq, Eq)]
18#[repr(u32)]
19pub enum MatrixUse {
20 /// Input operand A.
21 MatrixA = 0,
22 /// Input operand B.
23 MatrixB = 1,
24 /// Accumulator / result.
25 MatrixAccumulator = 2,
26}
27
28/// Memory layout for cooperative matrix load/store operations.
29#[derive(Copy, Clone, Debug, PartialEq, Eq)]
30#[repr(u32)]
31pub enum MatrixLayout {
32 /// Rows are stored contiguously.
33 RowMajor = 0,
34 /// Columns are stored contiguously.
35 ColumnMajor = 1,
36}
37
38/// Matrix role: input operand A in D = A × B + C.
39pub const MATRIX_A: u32 = MatrixUse::MatrixA as u32;
40/// Matrix role: input operand B in D = A × B + C.
41pub const MATRIX_B: u32 = MatrixUse::MatrixB as u32;
42/// Matrix role: accumulator / result in D = A × B + C.
43pub const MATRIX_ACCUMULATOR: u32 = MatrixUse::MatrixAccumulator as u32;
44
45/// Memory layout: rows are stored contiguously.
46pub const ROW_MAJOR: MatrixLayout = MatrixLayout::RowMajor;
47/// Memory layout: columns are stored contiguously.
48pub const COLUMN_MAJOR: MatrixLayout = MatrixLayout::ColumnMajor;
49
50/// A cooperative matrix distributed across the subgroup.
51///
52/// Each invocation holds a fragment of the full `ROWS × COLS` matrix.
53/// The hardware maps elements to invocations automatically.
54///
55/// # Type parameters
56/// - `T`: element type (`f32`, `f64`, `i32`, `u32`, `i8`, `u8`, etc.)
57/// - `USE`: matrix role — one of [`MatrixUse::MatrixA`], [`MatrixUse::MatrixB`], [`MatrixUse::MatrixAccumulator`] cast to `u32`
58/// - `ROWS`: number of rows
59/// - `COLS`: number of columns
60///
61/// # Capability
62/// Requires `CooperativeMatrixKHR` + `SPV_KHR_cooperative_matrix`.
63#[spirv(cooperative_matrix)]
64#[derive(Copy, Clone)]
65#[repr(C)]
66pub struct CooperativeMatrix<T, const USE: u32, const ROWS: u32, const COLS: u32> {
67 // HACK: keeps the Rust layout non-ZST so #[spirv(cooperative_matrix)] can
68 // special-case it before it gets elided.
69 _anti_zst_padding: MaybeUninit<u32>,
70 _phantom: PhantomData<T>,
71}
72
73impl<T, const USE: u32, const ROWS: u32, const COLS: u32> CooperativeMatrix<T, USE, ROWS, COLS> {
74 /// Load a cooperative matrix through a pointer.
75 ///
76 /// `slice` must point into an array. `layout` specifies whether the matrix
77 /// is stored in row-major ([`MatrixLayout::RowMajor`]) or column-major
78 /// ([`MatrixLayout::ColumnMajor`]) order. `stride` is the number of elements
79 /// between the start of consecutive rows (row-major) or columns (column-major).
80 ///
81 /// The scope is always `Subgroup`.
82 ///
83 /// # Safety
84 /// - `slice` must point into an array and be valid for all element accesses
85 /// implied by the matrix dimensions, layout, and stride.
86 /// - All operands must be dynamically uniform within every instance of the
87 /// subgroup scope.
88 #[spirv_std_macros::gpu_only]
89 #[doc(alias = "OpCooperativeMatrixLoadKHR")]
90 #[inline]
91 pub unsafe fn load(slice: &[T], layout: MatrixLayout, stride: u32) -> Self {
92 unsafe {
93 let mut result = MaybeUninit::<Self>::uninit();
94 let layout_u32 = layout as u32;
95 let ptr = slice.as_ptr();
96 asm!(
97 "%u32 = OpTypeInt 32 0",
98 "%layout = OpLoad %u32 {layout}",
99 "%stride = OpLoad %u32 {stride}",
100 // Use typeof* to get the cooperative matrix type from the output pointer.
101 "%result = OpCooperativeMatrixLoadKHR typeof* {out} {ptr} %layout %stride",
102 "OpStore {out} %result",
103 ptr = in(reg) ptr,
104 layout = in(reg) &layout_u32,
105 stride = in(reg) &stride,
106 out = in(reg) result.as_mut_ptr(),
107 );
108 result.assume_init()
109 }
110 }
111
112 /// Store a cooperative matrix through a pointer.
113 ///
114 /// `slice` must point into an array. `layout` specifies whether the matrix
115 /// is stored in row-major ([`MatrixLayout::RowMajor`]) or column-major
116 /// ([`MatrixLayout::ColumnMajor`]) order. `stride` is the number of elements
117 /// between the start of consecutive rows (row-major) or columns (column-major).
118 ///
119 /// The scope is always `Subgroup`.
120 ///
121 /// # Safety
122 /// - `slice` must point into an array and be valid for all element accesses
123 /// implied by the matrix dimensions, layout, and stride.
124 /// - All operands must be dynamically uniform within every instance of the
125 /// subgroup scope.
126 #[spirv_std_macros::gpu_only]
127 #[doc(alias = "OpCooperativeMatrixStoreKHR")]
128 #[inline]
129 pub unsafe fn store(self, slice: &mut [T], layout: MatrixLayout, stride: u32) {
130 unsafe {
131 let layout_u32 = layout as u32;
132 let ptr = slice.as_mut_ptr();
133 asm!(
134 "%u32 = OpTypeInt 32 0",
135 "%layout = OpLoad %u32 {layout}",
136 "%stride = OpLoad %u32 {stride}",
137 "%matrix = OpLoad _ {matrix}",
138 "OpCooperativeMatrixStoreKHR {ptr} %matrix %layout %stride",
139 ptr = in(reg) ptr,
140 matrix = in(reg) &self,
141 layout = in(reg) &layout_u32,
142 stride = in(reg) &stride,
143 );
144 }
145 }
146
147 /// Returns the number of matrix components this invocation is responsible for.
148 ///
149 /// The sum across all invocations in the subgroup equals `ROWS * COLS`.
150 #[spirv_std_macros::gpu_only]
151 #[doc(alias = "OpCooperativeMatrixLengthKHR")]
152 #[inline]
153 pub fn length(&self) -> u32 {
154 unsafe {
155 let mut result: u32 = 0;
156 asm!(
157 "%u32 = OpTypeInt 32 0",
158 // typeof* {self_ptr} resolves to the CooperativeMatrix type (pointee of &self).
159 "%coop_ty = typeof* {self_ptr}",
160 "%result = OpCooperativeMatrixLengthKHR %u32 %coop_ty",
161 "OpStore {out} %result",
162 self_ptr = in(reg) self,
163 out = in(reg) &mut result,
164 );
165 result
166 }
167 }
168}
169
170/// Linear-algebraic matrix multiply of `A` by `B` and then component-wise add `C`.
171///
172/// The order of operations is implementation-dependent. All matrices must have the
173/// same scope, which is always subgroup here.
174///
175/// - `A`: `M × K` matrix with use [`MatrixUse::MatrixA`]
176/// - `B`: `K × N` matrix with use [`MatrixUse::MatrixB`]
177/// - `C`: `M × N` matrix with use [`MatrixUse::MatrixAccumulator`]
178/// - returns `D`: `M × N` accumulator equal to `A × B + C`
179///
180/// All operands must be dynamically uniform within every instance of the subgroup scope.
181///
182/// # Capability
183/// Requires `CooperativeMatrixKHR` + `SPV_KHR_cooperative_matrix`.
184#[spirv_std_macros::gpu_only]
185#[doc(alias = "OpCooperativeMatrixMulAddKHR")]
186#[inline]
187pub fn mul_add<TA, TB, TC, const M: u32, const N: u32, const K: u32>(
188 a: CooperativeMatrix<TA, { MatrixUse::MatrixA as u32 }, M, K>,
189 b: CooperativeMatrix<TB, { MatrixUse::MatrixB as u32 }, K, N>,
190 c: CooperativeMatrix<TC, { MatrixUse::MatrixAccumulator as u32 }, M, N>,
191) -> CooperativeMatrix<TC, { MatrixUse::MatrixAccumulator as u32 }, M, N> {
192 unsafe {
193 let mut result = MaybeUninit::<
194 CooperativeMatrix<TC, { MatrixUse::MatrixAccumulator as u32 }, M, N>,
195 >::uninit();
196 asm!(
197 "%a = OpLoad _ {a}",
198 "%b = OpLoad _ {b}",
199 "%c = OpLoad _ {c}",
200 "%result = OpCooperativeMatrixMulAddKHR _ %a %b %c",
201 "OpStore {out} %result",
202 a = in(reg) &a,
203 b = in(reg) &b,
204 c = in(reg) &c,
205 out = in(reg) result.as_mut_ptr(),
206 );
207 result.assume_init()
208 }
209}