cubecl_spirv/
cmma.rs

1use crate::{
2    SpirvCompiler, SpirvTarget,
3    item::{Elem, Item},
4    lookups::Matrix,
5    value::Value,
6};
7use cubecl_core::ir::{self as core, CoopMma, ElemType, Id, MatrixLayout, MatrixScope};
8use rspirv::{
9    dr::Operand,
10    spirv::{
11        self, Capability, CooperativeMatrixLayout, CooperativeMatrixOperands, CooperativeMatrixUse,
12        MemoryAccess, TensorAddressingOperands,
13    },
14};
15
16impl<T: SpirvTarget> SpirvCompiler<T> {
17    pub fn compile_cmma(&mut self, cmma: CoopMma, out: Option<core::Value>) {
18        self.capabilities.insert(Capability::CooperativeMatrixKHR);
19
20        if let Some(out) = out {
21            if out.elem_type() == ElemType::Float(core::FloatKind::BF16) {
22                self.capabilities
23                    .insert(Capability::BFloat16CooperativeMatrixKHR);
24            }
25            if matches!(
26                out.elem_type(),
27                ElemType::Float(core::FloatKind::E5M2 | core::FloatKind::E4M3)
28            ) {
29                self.capabilities
30                    .insert(Capability::Float8CooperativeMatrixEXT);
31            }
32        }
33
34        match cmma {
35            CoopMma::Fill { value } => self.compile_fill(out.unwrap(), value),
36            CoopMma::Load {
37                ptr,
38                stride,
39                layout,
40            } => self.compile_load(out.unwrap(), ptr, stride, layout),
41            CoopMma::LoadTensor {
42                buffer,
43                layout,
44                view,
45            } => self.compile_load_tensor(out.unwrap(), buffer, layout, view),
46            CoopMma::Execute {
47                mat_a,
48                mat_b,
49                mat_c,
50            } => self.compile_execute(mat_a, mat_b, mat_c, out.unwrap()),
51            CoopMma::ExecuteElementwise { matrix, op } => {
52                self.compile_elementwise_op(matrix, op, out.unwrap());
53            }
54            CoopMma::Store {
55                mat,
56                stride,
57                layout,
58                destination,
59            } => self.compile_store(mat, stride, destination, layout),
60            CoopMma::StoreTensor { mat, layout, view } => {
61                self.compile_store_tensor(mat, out.unwrap(), layout, view)
62            }
63            CoopMma::Cast { input } => self.compile_cast(input, out.unwrap()),
64            CoopMma::RowIndex { .. }
65            | CoopMma::ColIndex { .. }
66            | CoopMma::LoadMatrix { .. }
67            | CoopMma::StoreMatrix { .. }
68            | CoopMma::ExecuteManual { .. }
69            | CoopMma::ExecuteScaled { .. } => {
70                panic!("Manual register management not currently supported in SPIR-V")
71            }
72        }
73    }
74
75    fn compile_load(
76        &mut self,
77        mat: core::Value,
78        ptr: core::Value,
79        stride: core::Value,
80        layout: Option<MatrixLayout>,
81    ) {
82        let mat = self.compile_value(mat);
83        let write_id = self.write_id_cmma(&mat);
84
85        let ptr = self.compile_value(ptr);
86        let stride = self.compile_value(stride);
87        let stride_item = stride.item();
88        let mut stride = self.read(&stride);
89
90        let value_ty = ptr.item().value_type();
91        let align = value_ty.size();
92
93        if let Item::Vector(_, vector_size) = value_ty {
94            let shift = stride_item.const_u32(self, vector_size.trailing_zeros());
95            let stride_ty = stride_item.id(self);
96            stride = self
97                .shift_right_logical(stride_ty, None, stride, shift)
98                .unwrap();
99        }
100
101        let layout = layout
102            .and_then(compile_layout)
103            .or(matrix_layout(&mat))
104            .unwrap_or(CooperativeMatrixLayout::RowMajorKHR);
105        let memory_layout = self.const_u32(layout as u32);
106
107        let ptr = self.read(&ptr);
108        let out_ty = mat.item().unwrap_ptr();
109        let ty = out_ty.id(self);
110
111        let mat_id = self
112            .cooperative_matrix_load_khr(
113                ty,
114                Some(write_id),
115                ptr,
116                memory_layout,
117                Some(stride),
118                Some(MemoryAccess::ALIGNED),
119                [align.into()],
120            )
121            .unwrap();
122
123        self.write_cmma(&mat, mat_id);
124    }
125
126    fn compile_load_tensor(
127        &mut self,
128        mat: core::Value,
129        buffer: core::Value,
130        layout: core::Value,
131        view: Option<core::Value>,
132    ) {
133        self.capabilities
134            .insert(Capability::CooperativeMatrixTensorAddressingNV);
135
136        let mat = self.compile_value(mat);
137        let write_id = self.write_id_cmma(&mat);
138
139        let buffer = self.compile_value(buffer);
140        let layout = self.compile_value(layout);
141        let view = view.map(|view| self.compile_value(view));
142        let layout = self.read(&layout);
143        let view = view.map(|view| self.read(&view));
144
145        let ptr = buffer.id(self);
146        let out_ty = mat.item().unwrap_ptr();
147        let align = buffer.item().value_type().size();
148        let ty = out_ty.id(self);
149
150        let zero = Item::Scalar(mat.elem()).const_u32(self, 0);
151        let clipped_fallback = self.composite_construct(ty, None, [zero]).unwrap();
152
153        let (operands, extra_args) = match view {
154            Some(view) => (
155                TensorAddressingOperands::TENSOR_VIEW,
156                vec![Operand::IdRef(view)],
157            ),
158            None => (TensorAddressingOperands::NONE, vec![]),
159        };
160
161        let mat_id = self
162            .cooperative_matrix_load_tensor_nv(
163                ty,
164                Some(write_id),
165                ptr,
166                clipped_fallback,
167                layout,
168                MemoryAccess::ALIGNED,
169                [align.into()],
170                operands,
171                extra_args,
172            )
173            .unwrap();
174
175        self.write_cmma(&mat, mat_id);
176    }
177
178    fn compile_fill(&mut self, mat: core::Value, value: core::Value) {
179        let mat = self.compile_value(mat);
180        let value = self.compile_value(value);
181        let mat_id = self.write_id_cmma(&mat);
182
183        let item = mat.item().unwrap_ptr();
184        let ty = item.id(self);
185        let mat_id = match value {
186            Value::Constant(id, _, _) => self.constant_composite(ty, vec![id]),
187            val => {
188                let val = self.read(&val);
189                self.composite_construct(ty, Some(mat_id), vec![val])
190                    .unwrap()
191            }
192        };
193
194        self.write_cmma(&mat, mat_id);
195    }
196
197    fn compile_store(
198        &mut self,
199        mat: core::Value,
200        stride: core::Value,
201        destination: core::Value,
202        layout: MatrixLayout,
203    ) {
204        let mat = self.compile_value(mat);
205        let mat_obj = self.read(&mat);
206        //assert_ne!(mat_obj, 0, "Can't store uninitialized matrix");
207
208        let ptr = self.compile_value(destination);
209        let stride = self.compile_value(stride);
210        let value_ty = ptr.item().value_type();
211
212        let stride_item = stride.item();
213        let mut stride = self.read(&stride);
214        let layout = compile_layout(layout).unwrap_or(CooperativeMatrixLayout::RowMajorKHR);
215        let memory_layout = self.const_u32(layout as u32);
216
217        let ptr = self.read(&ptr);
218
219        let align = value_ty.size();
220
221        if let Item::Vector(_, vector_size) = value_ty {
222            let shift = stride_item.const_u32(self, vector_size.trailing_zeros());
223            let stride_ty = stride_item.id(self);
224            stride = self
225                .shift_right_logical(stride_ty, None, stride, shift)
226                .unwrap();
227        }
228
229        self.cooperative_matrix_store_khr(
230            ptr,
231            mat_obj,
232            memory_layout,
233            Some(stride),
234            Some(MemoryAccess::ALIGNED),
235            [align.into()],
236        )
237        .unwrap();
238    }
239
240    fn compile_store_tensor(
241        &mut self,
242        mat: core::Value,
243        out: core::Value,
244        layout: core::Value,
245        view: Option<core::Value>,
246    ) {
247        self.capabilities
248            .insert(Capability::CooperativeMatrixTensorAddressingNV);
249
250        let mat = self.compile_value(mat);
251        let mat_obj = self.read(&mat);
252        //assert_ne!(mat_obj, 0, "Can't store uninitialized matrix");
253
254        let out = self.compile_value(out);
255        let layout = self.compile_value(layout);
256        let view = view.map(|view| self.compile_value(view));
257
258        let layout = self.read(&layout);
259        let view = view.map(|view| self.read(&view));
260
261        let align = out.item().value_type().size();
262        let ptr = out.id(self);
263
264        let (operands, extra_args) = match view {
265            Some(view) => (
266                TensorAddressingOperands::TENSOR_VIEW,
267                vec![Operand::IdRef(view)],
268            ),
269            None => (TensorAddressingOperands::NONE, vec![]),
270        };
271
272        self.cooperative_matrix_store_tensor_nv(
273            ptr,
274            mat_obj,
275            layout,
276            MemoryAccess::ALIGNED,
277            [align.into()],
278            operands,
279            extra_args,
280        )
281        .unwrap();
282    }
283
284    fn compile_execute(
285        &mut self,
286        mat_a: core::Value,
287        mat_b: core::Value,
288        mat_c: core::Value,
289        mat_d: core::Value,
290    ) {
291        let mat_a = self.compile_value(mat_a);
292        let mat_b = self.compile_value(mat_b);
293        let mat_c = self.compile_value(mat_c);
294        let mat_d = self.compile_value(mat_d);
295
296        let mat_a_id = self.read(&mat_a);
297        let mat_b_id = self.read(&mat_b);
298        let mat_c_id = self.read(&mat_c);
299        let mat_d_id = self.write_id_cmma(&mat_d);
300
301        let ty = mat_d.item().unwrap_ptr().id(self);
302
303        let mut operands = CooperativeMatrixOperands::NONE_KHR;
304        if matches!(mat_a.elem(), Elem::Int(_, true)) {
305            operands |= CooperativeMatrixOperands::MATRIX_A_SIGNED_COMPONENTS_KHR;
306        }
307        if matches!(mat_b.elem(), Elem::Int(_, true)) {
308            operands |= CooperativeMatrixOperands::MATRIX_B_SIGNED_COMPONENTS_KHR;
309        }
310        if matches!(mat_c.elem(), Elem::Int(_, true)) {
311            operands |= CooperativeMatrixOperands::MATRIX_C_SIGNED_COMPONENTS_KHR;
312        }
313        if matches!(mat_d.elem(), Elem::Int(_, true)) {
314            operands |= CooperativeMatrixOperands::MATRIX_RESULT_SIGNED_COMPONENTS_KHR;
315        }
316
317        self.cooperative_matrix_mul_add_khr(
318            ty,
319            Some(mat_d_id),
320            mat_a_id,
321            mat_b_id,
322            mat_c_id,
323            Some(operands),
324        )
325        .unwrap();
326
327        self.write_cmma(&mat_d, mat_d_id);
328    }
329
330    fn compile_elementwise_op(&mut self, matrix: core::Value, op: Id, output: core::Value) {
331        self.capabilities
332            .insert(Capability::CooperativeMatrixPerElementOperationsNV);
333
334        let matrix = self.compile_value(matrix);
335        let output = self.compile_value(output);
336
337        let matrix_ty = matrix.item().unwrap_ptr().id(self);
338        let matrix_id = self.read(&matrix);
339        let output_id = self.write_id_cmma(&output);
340
341        let captures = self.opt.global_state.extra_functions[&op]
342            .implicit_params
343            .clone();
344        let captures = captures
345            .into_iter()
346            .map(|val| self.compile_value(val))
347            .collect::<Vec<_>>();
348        let captures = captures
349            .iter()
350            .map(|val| self.read(val))
351            .collect::<Vec<_>>();
352        let func = self.state.extra_funcs[&op].id;
353
354        self.cooperative_matrix_per_element_op_nv(
355            matrix_ty,
356            Some(output_id),
357            matrix_id,
358            func,
359            captures,
360        )
361        .unwrap();
362
363        self.write_cmma(&output, output_id);
364    }
365
366    fn compile_cast(&mut self, input: core::Value, output: core::Value) {
367        let input = self.compile_value(input);
368        let output = self.compile_value(output);
369
370        let input_ident = matrix_ident(&input);
371        let output_ident = matrix_ident(&output);
372
373        if input_ident != output_ident {
374            self.capabilities
375                .insert(Capability::CooperativeMatrixConversionsNV);
376        }
377
378        let input_ty = input.item();
379        let output_ty = output.item().unwrap_ptr();
380
381        let output_ty_id = output_ty.id(self);
382
383        let input_id = self.read(&input);
384        let output_id = self.write_id_cmma(&output);
385
386        if input_ty == output_ty && input_ident != output_ident {
387            self.cooperative_matrix_convert_nv(output_ty_id, Some(output_id), input_id)
388                .unwrap()
389        } else {
390            input_ty.cast_to(self, Some(output_id), input_id, &output_ty)
391        };
392
393        self.write_cmma(&output, output_id);
394    }
395
396    pub fn matrix_rows(&mut self, mat: &Matrix) -> u32 {
397        let rows = match mat.ident {
398            CooperativeMatrixUse::MatrixAKHR => mat.m,
399            CooperativeMatrixUse::MatrixBKHR => mat.k,
400            CooperativeMatrixUse::MatrixAccumulatorKHR => mat.m,
401        };
402        self.const_u32(rows)
403    }
404
405    pub fn matrix_columns(&mut self, mat: &Matrix) -> u32 {
406        let columns = match mat.ident {
407            CooperativeMatrixUse::MatrixAKHR => mat.k,
408            CooperativeMatrixUse::MatrixBKHR => mat.n,
409            CooperativeMatrixUse::MatrixAccumulatorKHR => mat.n,
410        };
411        self.const_u32(columns)
412    }
413
414    pub fn compile_matrix(&mut self, mat: &core::MatrixType) -> Matrix {
415        let elem = self.compile_type(core::Type::new(mat.storage)).elem();
416        let ident = match mat.ident {
417            core::MatrixIdent::A => CooperativeMatrixUse::MatrixAKHR,
418            core::MatrixIdent::B => CooperativeMatrixUse::MatrixBKHR,
419            core::MatrixIdent::Accumulator => CooperativeMatrixUse::MatrixAccumulatorKHR,
420        };
421        let layout = compile_layout(mat.layout);
422        let scope = compile_scope(mat.scope);
423
424        Matrix {
425            id: 0,
426            ident,
427            m: mat.m as u32,
428            n: mat.n as u32,
429            k: mat.k as u32,
430            elem,
431            layout,
432            scope,
433        }
434    }
435}
436
437fn matrix_ident(val: &Value) -> CooperativeMatrixUse {
438    let Item::CoopMatrix { ident, .. } = val.item().unwrap_ptr() else {
439        unreachable!()
440    };
441    ident
442}
443
444fn matrix_layout(val: &Value) -> Option<CooperativeMatrixLayout> {
445    let Item::CoopMatrix { layout, .. } = val.item().unwrap_ptr() else {
446        unreachable!()
447    };
448    layout
449}
450
451fn compile_layout(layout: MatrixLayout) -> Option<CooperativeMatrixLayout> {
452    match layout {
453        core::MatrixLayout::ColMajor => Some(CooperativeMatrixLayout::ColumnMajorKHR),
454        core::MatrixLayout::RowMajor => Some(CooperativeMatrixLayout::RowMajorKHR),
455        core::MatrixLayout::Undefined => None,
456    }
457}
458
459fn compile_scope(scope: MatrixScope) -> spirv::Scope {
460    match scope {
461        MatrixScope::Plane => spirv::Scope::Subgroup,
462        MatrixScope::Cube => spirv::Scope::Workgroup,
463    }
464}
cubecl_spirv/cmma.rs

cubecl_spirv/
cmma.rs
