cubecl_core/frontend/
cmma.rs

1//! This module exposes cooperative matrix-multiply and accumulate operations.
2//!
3//! Most of the functions are actually unsafe, since they mutate their input, even if they are
4//! passed as reference.
5//!
6//! # Example
7//!
8//! This is a basic 16x16x16 matrix multiplication example.
9//!
10//! ```rust, ignore
11//! #[cube(launch)]
12//! pub fn example(lhs: &[f16], rhs: &[f16], out: &mut [f32]) {
13//!     let a = cmma::Matrix::<f16>::new(
14//!         cmma::MatrixIdent::A,
15//!         16,
16//!         16,
17//!         16,
18//!         cmma::MatrixLayout::RowMajor,
19//!     );
20//!     let b = cmma::Matrix::<f16>::new(
21//!         cmma::MatrixIdent::B,
22//!         16,
23//!         16,
24//!         16,
25//!         cmma::MatrixLayout::ColMajor,
26//!     );
27//!     let c = cmma::Matrix::<f32>::new(
28//!         cmma::MatrixIdent::Accumulator,
29//!         16,
30//!         16,
31//!         16,
32//!         cmma::MatrixLayout::Undefined,
33//!     );
34//!     cmma::fill(&c, 0.0);
35//!     cmma::load(&a, lhs.as_slice(), 16);
36//!     cmma::load(&b, rhs.as_slice(), 16);
37//!
38//!     cmma::execute(&a, &b, &c, &c);
39//!
40//!     cmma::store(
41//!         out.as_mut_slice(),
42//!         &c,
43//!         16,
44//!         cmma::MatrixLayout::RowMajor,
45//!     );
46//! }
47//! ```
48
49use super::{CubeDebug, CubePrimitive, CubeType, IntoMut, NativeExpand, SliceExpand};
50use crate::{self as cubecl, prelude::*};
51use crate::{
52    ir::{self, Instruction},
53    unexpanded,
54};
55use core::marker::PhantomData;
56use cubecl_macros::{comptime_type, cube, intrinsic};
57
58use cubecl_ir::{CoopMma, Scope, StorageType, Value, VectorSize};
59pub use ir::{MatrixIdent, MatrixLayout};
60
61#[derive(Clone, Copy)]
62pub struct Plane;
63#[derive(Clone, Copy)]
64pub struct Cube;
65
66pub trait MatrixScope: Copy {
67    const SCOPE: ir::MatrixScope;
68}
69
70impl MatrixScope for Plane {
71    const SCOPE: ir::MatrixScope = ir::MatrixScope::Plane;
72}
73
74impl MatrixScope for Cube {
75    const SCOPE: cubecl_ir::MatrixScope = ir::MatrixScope::Cube;
76}
77
78/// A matrix represent a 2D grid of numbers.
79///
80/// They can either be in a [row major](MatrixLayout::RowMajor) or a
81/// [column major](MatrixLayout::ColMajor) format.
82#[derive(Copy, Clone)]
83pub struct Matrix<C: CubeType, S: MatrixScope = Plane> {
84    _c: PhantomData<C>,
85    _s: PhantomData<S>,
86}
87
88/// Defines a matrix multiplication operation, including the input and output type, and the shape.
89#[derive(Copy, Clone)]
90pub struct MmaDefinition<A: CubeType, B: CubeType, CD: CubeType> {
91    _a: PhantomData<A>,
92    _b: PhantomData<B>,
93    _cd: PhantomData<CD>,
94}
95
96/// Expand type of [Matrix].
97pub struct MatrixExpand<C: CubeType, S: MatrixScope> {
98    elem: Value,
99    ident: MatrixIdent,
100    _c: PhantomData<C>,
101    _s: PhantomData<S>,
102}
103
104/// Expand type of [`MmaDefinition`].
105#[derive(Debug)]
106pub struct MmaDefinitionExpand<A: CubeType, B: CubeType, CD: CubeType> {
107    pub m: usize,
108    pub n: usize,
109    pub k: usize,
110    pub a_type: StorageType,
111    pub b_type: StorageType,
112    pub cd_type: StorageType,
113    pub scales_factor: Option<usize>,
114    pub scales_type: Option<StorageType>,
115    _a: PhantomData<A>,
116    _b: PhantomData<B>,
117    _cd: PhantomData<CD>,
118}
119
120impl<C: CubeType, S: MatrixScope> Clone for MatrixExpand<C, S> {
121    fn clone(&self) -> Self {
122        Self {
123            elem: self.elem,
124            ident: self.ident,
125            _c: self._c,
126            _s: self._s,
127        }
128    }
129}
130
131impl<C: CubeType, S: MatrixScope> ExpandTypeClone for MatrixExpand<C, S> {
132    fn clone_unchecked(&self) -> Self {
133        self.clone()
134    }
135}
136
137impl<A: CubeType, B: CubeType, CD: CubeType> ExpandTypeClone for MmaDefinitionExpand<A, B, CD> {
138    fn clone_unchecked(&self) -> Self {
139        *self
140    }
141}
142
143impl<C: CubeType, S: MatrixScope> AsRefExpand for MatrixExpand<C, S> {
144    fn __expand_ref_method(&self, _scope: &Scope) -> &Self {
145        self
146    }
147}
148impl<C: CubeType, S: MatrixScope> AsMutExpand for MatrixExpand<C, S> {
149    fn __expand_ref_mut_method(&mut self, _scope: &Scope) -> &mut Self {
150        self
151    }
152}
153
154impl<A: CubeType, B: CubeType, CD: CubeType> AsRefExpand for MmaDefinitionExpand<A, B, CD> {
155    fn __expand_ref_method(&self, _scope: &Scope) -> &Self {
156        self
157    }
158}
159impl<A: CubeType, B: CubeType, CD: CubeType> AsMutExpand for MmaDefinitionExpand<A, B, CD> {
160    fn __expand_ref_mut_method(&mut self, _scope: &Scope) -> &mut Self {
161        self
162    }
163}
164
165impl<A: CubeType, B: CubeType, CD: CubeType> Copy for MmaDefinitionExpand<A, B, CD> {}
166impl<A: CubeType, B: CubeType, CD: CubeType> Clone for MmaDefinitionExpand<A, B, CD> {
167    fn clone(&self) -> Self {
168        *self
169    }
170}
171
172impl<C: CubeType, S: MatrixScope> CubeType for Matrix<C, S> {
173    type ExpandType = MatrixExpand<C, S>;
174}
175
176impl<A: CubeType, B: CubeType, CD: CubeType> CubeType for MmaDefinition<A, B, CD> {
177    type ExpandType = MmaDefinitionExpand<A, B, CD>;
178}
179
180impl<C: CubeType, S: MatrixScope> IntoExpand for MatrixExpand<C, S> {
181    type Expand = Self;
182
183    fn into_expand(self, _: &Scope) -> Self::Expand {
184        self
185    }
186}
187
188impl<C: CubeType, S: MatrixScope> IntoMut for MatrixExpand<C, S> {
189    fn into_mut(self, _scope: &Scope) -> Self {
190        self
191    }
192}
193
194impl<C: CubeType, S: MatrixScope> CubeDebug for MatrixExpand<C, S> {
195    fn set_debug_name(&self, scope: &Scope, name: &'static str) {
196        scope.update_value_name(self.elem, name);
197    }
198}
199
200impl<A: CubeType, B: CubeType, CD: CubeType> IntoExpand for MmaDefinitionExpand<A, B, CD> {
201    type Expand = Self;
202
203    fn into_expand(self, _: &Scope) -> Self::Expand {
204        self
205    }
206}
207
208impl<A: CubeType, B: CubeType, CD: CubeType> IntoMut for MmaDefinitionExpand<A, B, CD> {
209    fn into_mut(self, _scope: &Scope) -> Self {
210        self
211    }
212}
213
214impl<A: CubeType, B: CubeType, CD: CubeType> CubeDebug for MmaDefinitionExpand<A, B, CD> {}
215
216#[cube]
217impl<C: CubePrimitive, S: MatrixScope> Matrix<C, S> {
218    /// Create a new uninitialized matrix that is going to be used in the
219    /// [matrix-multiply and accumulate](execute()) function.
220    ///
221    /// # Safety
222    /// Must be initialized with `load` or `fill` before use. Using it without initialization is
223    /// undefined behaviour on CUDA, and completely invalid on Vulkan.
224    ///
225    /// You have to declare the shape used for the execution.
226    /// The shape of the current matrix is determined using the [MatrixIdent].
227    ///
228    /// * [MatrixIdent::A] Shape => (M, K)
229    /// * [`MatrixIdent::B`] Shape => (K, N)
230    /// * [`MatrixIdent::Accumulator`] Shape => (M, N)
231    ///
232    /// Not all shapes are supported, and the permitted shapes depend on the element type.
233    ///
234    /// Refer to [nvidia documentation](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#element-types-and-matrix-sizes).
235    pub unsafe fn uninitialized(
236        #[comptime] ident: MatrixIdent,
237        #[comptime] m: usize,
238        #[comptime] n: usize,
239        #[comptime] k: usize,
240        layout: MatrixLayout,
241    ) -> Self {
242        intrinsic!(|scope| {
243            let elem = C::__expand_as_type(scope).storage_type();
244            let elem = scope.create_local_mut(Type::Matrix(ir::MatrixType::new(
245                ident,
246                m,
247                n,
248                k,
249                elem,
250                layout,
251                S::SCOPE,
252            )));
253            MatrixExpand {
254                elem,
255                ident,
256                _c: PhantomData,
257                _s: PhantomData,
258            }
259        })
260    }
261
262    /// Create a new matrix that is going to be used in the
263    /// [matrix-multiply and accumulate](execute()) function and is filled with `value`.
264    ///
265    /// You have to declare the shape used for the execution.
266    /// The shape of the current matrix is determined using the [MatrixIdent].
267    ///
268    /// * [MatrixIdent::A] Shape => (M, K)
269    /// * [`MatrixIdent::B`] Shape => (K, N)
270    /// * [`MatrixIdent::Accumulator`] Shape => (M, N)
271    ///
272    /// Not all shapes are supported, and the permitted shapes depend on the element type.
273    ///
274    /// Refer to [nvidia documentation](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#element-types-and-matrix-sizes).
275    pub fn from_value(
276        #[comptime] ident: MatrixIdent,
277        #[comptime] m: usize,
278        #[comptime] n: usize,
279        #[comptime] k: usize,
280        layout: MatrixLayout,
281        value: C,
282    ) -> Self
283    where
284        C: Scalar,
285    {
286        let mut mat = unsafe { Self::uninitialized(ident, m, n, k, layout) };
287        fill(&mut mat, value);
288        mat
289    }
290
291    /// Create a new matrix that is going to be used in the
292    /// [matrix-multiply and accumulate](execute()) function and is loaded from `value` with `stride`.
293    ///
294    /// You have to declare the shape used for the execution.
295    /// The shape of the current matrix is determined using the [MatrixIdent].
296    ///
297    /// * [MatrixIdent::A] Shape => (M, K)
298    /// * [`MatrixIdent::B`] Shape => (K, N)
299    /// * [`MatrixIdent::Accumulator`] Shape => (M, N)
300    ///
301    /// Not all shapes are supported, and the permitted shapes depend on the element type.
302    ///
303    /// Refer to [nvidia documentation](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#element-types-and-matrix-sizes).
304    pub fn from_slice(
305        #[comptime] ident: MatrixIdent,
306        #[comptime] m: usize,
307        #[comptime] n: usize,
308        #[comptime] k: usize,
309        layout: MatrixLayout,
310        value: &[C],
311        stride: u32,
312    ) -> Self {
313        let mut mat = unsafe { Self::uninitialized(ident, m, n, k, layout) };
314
315        if comptime![ident == MatrixIdent::Accumulator] {
316            load_with_layout(&mut mat, value, stride, layout);
317        } else {
318            load(&mut mat, value, stride);
319        }
320        mat
321    }
322
323    /// Create a new matrix that is going to be used in the
324    /// [matrix-multiply and accumulate](execute()) function and is loaded from `value` with `stride`.
325    ///
326    /// You have to declare the shape used for the execution.
327    /// The shape of the current matrix is determined using the [MatrixIdent].
328    ///
329    /// * [MatrixIdent::A] Shape => (M, K)
330    /// * [`MatrixIdent::B`] Shape => (K, N)
331    /// * [`MatrixIdent::Accumulator`] Shape => (M, N)
332    ///
333    /// Not all shapes are supported, and the permitted shapes depend on the element type.
334    ///
335    /// Refer to [nvidia documentation](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#element-types-and-matrix-sizes).
336    pub fn from_tensor(
337        #[comptime] ident: MatrixIdent,
338        #[comptime] m: usize,
339        #[comptime] n: usize,
340        #[comptime] k: usize,
341        value: &TensorView<C>,
342    ) -> Self {
343        let mut mat = unsafe { Self::uninitialized(ident, m, n, k, MatrixLayout::Undefined) };
344        load_tensor(&mut mat, value);
345        mat
346    }
347}
348
349#[cube]
350impl<A: Scalar, B: Scalar, CD: Scalar> MmaDefinition<A, B, CD> {
351    /// Create a new matrix definition that is going to be used in the manual
352    /// matrix-multiply and accumulate ``execute_manual_mma()`` function.
353    ///
354    /// You have to declare the shape used for the execution.
355    /// The shape of the current matrix is determined using the [MatrixIdent].
356    ///
357    /// * [MatrixIdent::A] Shape => (M, K)
358    /// * [`MatrixIdent::B`] Shape => (K, N)
359    /// * [`MatrixIdent::Accumulator`] Shape => (M, N)
360    ///
361    /// Not all shapes are supported, and the permitted shapes depend on the element type.
362    /// Layout for manual MMA is determined by the runtime and must be handled manually.
363    /// Use [`Self::vector_layout`] to check the correct data layout for each element.
364    ///
365    /// Refer to [nvidia documentation](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#element-types-and-matrix-sizes).
366    pub fn new(#[comptime] m: usize, #[comptime] n: usize, #[comptime] k: usize) -> Self {
367        intrinsic!(|scope| {
368            let a_type = A::__expand_as_type(scope).storage_type();
369            let b_type = B::__expand_as_type(scope).storage_type();
370            let cd_type = CD::__expand_as_type(scope).storage_type();
371
372            MmaDefinitionExpand {
373                m,
374                n,
375                k,
376                a_type,
377                b_type,
378                cd_type,
379                scales_factor: None,
380                scales_type: None,
381                _a: PhantomData,
382                _b: PhantomData,
383                _cd: PhantomData,
384            }
385        })
386    }
387
388    /// Create a new matrix definition that is going to be used in the manual
389    /// matrix-multiply and accumulate ``execute_manual_mma()`` function.
390    ///
391    /// You have to declare the shape used for the execution.
392    /// The shape of the current matrix is determined using the [MatrixIdent].
393    ///
394    /// * [MatrixIdent::A] Shape => (M, K)
395    /// * [`MatrixIdent::B`] Shape => (K, N)
396    /// * [`MatrixIdent::Accumulator`] Shape => (M, N)
397    ///
398    /// Not all shapes are supported, and the permitted shapes depend on the element type.
399    /// Layout for manual MMA is determined by the runtime and must be handled manually.
400    /// Use [`Self::vector_layout`] to check the correct data layout for each element.
401    ///
402    /// Refer to [nvidia documentation](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#element-types-and-matrix-sizes).
403    pub fn new_scaled<S: CubePrimitive>(
404        #[comptime] m: usize,
405        #[comptime] n: usize,
406        #[comptime] k: usize,
407        #[comptime] scale_factor: usize,
408    ) -> Self {
409        intrinsic!(|scope| {
410            let a_type = A::__expand_as_type(scope).storage_type();
411            let b_type = B::__expand_as_type(scope).storage_type();
412            let cd_type = CD::__expand_as_type(scope).storage_type();
413
414            MmaDefinitionExpand {
415                m,
416                n,
417                k,
418                a_type,
419                b_type,
420                cd_type,
421                scales_factor: Some(scale_factor),
422                scales_type: Some(S::__expand_as_type(scope).storage_type()),
423                _a: PhantomData,
424                _b: PhantomData,
425                _cd: PhantomData,
426            }
427        })
428    }
429
430    /// Number of elements in the matrix
431    #[allow(unused)]
432    pub fn num_elems(&self, #[comptime] ident: MatrixIdent) -> comptime_type!(usize) {
433        intrinsic!(|scope| {
434            match ident {
435                MatrixIdent::A => (self.m * self.k) / self.a_type.packing_factor(),
436                MatrixIdent::B => (self.k * self.n) / self.b_type.packing_factor(),
437                MatrixIdent::Accumulator => (self.m * self.n) / self.cd_type.packing_factor(),
438            }
439        })
440    }
441
442    /// Returns the number of elements handled by each lane. Should be packed into `Vector`s of size
443    /// `vector_size` with [`Self::vector_layout`].
444    ///
445    /// # Note
446    /// "Lane" here refers to the unit relative to a plane, to distinguish it from a unit relative
447    /// to a cube.
448    #[allow(unused)]
449    pub fn elems_per_lane(&self, #[comptime] ident: MatrixIdent) -> comptime_type!(usize) {
450        intrinsic!(|scope| {
451            let elems = self.__expand_num_elems_method(scope, ident);
452            let plane_dim = scope.state().target_properties.mma.const_plane_size as usize;
453            let duplication = match ident {
454                MatrixIdent::A => scope.state().target_properties.mma.register_duplication_a,
455                MatrixIdent::B => scope.state().target_properties.mma.register_duplication_b,
456                MatrixIdent::Accumulator => {
457                    scope.state().target_properties.mma.register_duplication_acc
458                }
459            };
460            (elems * duplication) / plane_dim
461        })
462    }
463
464    /// Returns the number of vectors of size `vector_size` with layout `vector_layout` per lane.
465    ///
466    /// # Note
467    /// "Lane" here refers to the unit relative to a plane, to distinguish it from a unit relative
468    /// to a cube.
469    #[allow(unused)]
470    pub fn vectors_per_lane(&self, #[comptime] ident: MatrixIdent) -> comptime_type!(usize) {
471        intrinsic!(|scope| {
472            let elems = self.clone().__expand_elems_per_lane_method(scope, ident);
473            let vector_size = self.__expand_vector_size_method(scope, ident);
474            elems / vector_size
475        })
476    }
477
478    /// The layout of each vector in this matrix (row major or column major)
479    #[allow(unused)]
480    pub fn vector_layout(&self, #[comptime] ident: MatrixIdent) -> comptime_type!(MatrixLayout) {
481        intrinsic!(|scope| {
482            match ident {
483                MatrixIdent::A => scope.state().target_properties.mma.register_layout_a,
484                MatrixIdent::B => scope.state().target_properties.mma.register_layout_b,
485                MatrixIdent::Accumulator => scope.state().target_properties.mma.register_layout_acc,
486            }
487        })
488    }
489
490    /// Number of elements in each vector passed to the execute function. Represents the maximum
491    /// number of contiguous elements held by the thread.
492    pub fn vector_size(&self, #[comptime] ident: MatrixIdent) -> comptime_type!(VectorSize) {
493        intrinsic!(|scope| {
494            let storage = match ident {
495                MatrixIdent::A => self.a_type,
496                MatrixIdent::B => self.b_type,
497                MatrixIdent::Accumulator => self.cd_type,
498            };
499            let matrix = cubecl_ir::MatrixType {
500                ident,
501                m: self.m,
502                n: self.n,
503                k: self.k,
504                storage: storage,
505                layout: MatrixLayout::ColMajor,
506                scope: ir::MatrixScope::Plane,
507            };
508            scope
509                .state()
510                .target_properties
511                .mma
512                .contiguous_elements
513                .apply(ident, matrix)
514        })
515    }
516
517    /// Returns the coordinates of the `nth` element handled by the `lane_id`
518    /// Each lane contains [`Self::elems_per_lane`] elements in [`Self::vector_size`] chunks.
519    /// Returns (`row_idx`, `col_idx`)
520    ///
521    /// # Note
522    /// "Lane" here refers to the unit relative to a plane, to distinguish it from a unit relative
523    /// to a cube.
524    pub fn position_of_nth(
525        &self,
526        lane_id: u32,
527        elem_idx: u32,
528        #[comptime] ident: MatrixIdent,
529    ) -> (u32, u32) {
530        intrinsic!(|scope| {
531            let lane_id: Value = lane_id.into();
532            let elem_idx: Value = elem_idx.into();
533
534            let ty = match ident {
535                MatrixIdent::A => self.a_type,
536                MatrixIdent::B => self.b_type,
537                MatrixIdent::Accumulator => self.cd_type,
538            };
539            let layout = match ident {
540                MatrixIdent::A => scope.state().target_properties.mma.register_layout_a,
541                MatrixIdent::B => scope.state().target_properties.mma.register_layout_b,
542                MatrixIdent::Accumulator => scope.state().target_properties.mma.register_layout_acc,
543            };
544            let matrix = cubecl_ir::MatrixType {
545                ident,
546                m: self.m,
547                n: self.n,
548                k: self.k,
549                storage: ty,
550                layout,
551                scope: ir::MatrixScope::Plane,
552            };
553
554            let row = scope.create_value(u32::__expand_as_type(scope));
555            let col = scope.create_value(u32::__expand_as_type(scope));
556            scope.register(Instruction::new(
557                CoopMma::RowIndex {
558                    lane_id,
559                    i: elem_idx,
560                    matrix,
561                },
562                row,
563            ));
564            scope.register(Instruction::new(
565                CoopMma::ColIndex {
566                    lane_id,
567                    i: elem_idx,
568                    matrix,
569                },
570                col,
571            ));
572            (row.into(), col.into())
573        })
574    }
575
576    /// Index of the scales for this thread, along the non-major dimension of the matrix.
577    /// Each thread loads all scales in the major direction into a single `Vector`.
578    pub fn scales_index(&self, lane_id: u32, #[comptime] ident: MatrixIdent) -> u32 {
579        // Just do CUDA for now, call an actual intrinsic when HIP gets support
580        let quad_id = lane_id / 4;
581        let t_id = lane_id % 4;
582        match ident {
583            MatrixIdent::A => quad_id + (t_id % 2) * 8,
584            MatrixIdent::B => quad_id,
585            MatrixIdent::Accumulator => panic!("Accumulator doesn't have scales"),
586        }
587    }
588
589    /// Number of scales in each vector (not the vector size!). Vector size may include padding bytes.
590    pub fn scales_count(&self) -> comptime_type!(usize) {
591        // We only have the CUDA version for now, so just use `scales_factor`. The function can
592        // be modified for HIP in the future without having to redo all uses.
593        intrinsic!(|_| {
594            self.scales_factor
595                .expect("Can't retrieve scales count for matrix with no scales")
596        })
597    }
598
599    /// Vector size for the scale factors. May be larger than the total number of scales.
600    pub fn scales_vector_size(&self) -> comptime_type!(VectorSize) {
601        intrinsic!(|scope| {
602            let elem = self
603                .scales_type
604                .expect("Can't retrieve scales vector size for matrix with no scales");
605            scope.state().target_properties.mma.register_size_bits / elem.size_bits()
606        })
607    }
608
609    /// Load one or more matrix register using intrinsic instructions. CUDA only.
610    /// The number of matrices must be 1, 2, or 4. The rows for the nth matrix are passed by the 8
611    /// lanes starting at `n * 8`. All slice starts must be valid, even for non-participating lanes.
612    /// The slice determines the starting address for a 16-byte row loaded by this unit, with
613    /// the row index being `UNIT_POS_PLANE % 8`.
614    /// The number of elements is determined by element size.
615    ///
616    /// # Constraints:
617    /// Address must be aligned to 16 bytes
618    /// Address must be in shared memory
619    pub fn load_matrix<E: CubePrimitive, NO: Size>(
620        &self,
621        row: &[E],
622        #[comptime] ident: MatrixIdent,
623        #[comptime] num_matrices: usize,
624        #[comptime] transpose: bool,
625    ) -> Array<Vector<E::Scalar, NO>> {
626        intrinsic!(|scope| {
627            let slice_vector_size = row.expand.vector_size();
628            let ptr = unsafe { *row.__expand_as_ptr_method(scope) }.expand;
629            let out = Array::__expand_new(scope, num_matrices);
630            scope.register(Instruction::new(
631                CoopMma::LoadMatrix {
632                    ptr,
633                    factor: num_matrices,
634                    transpose,
635                },
636                out.__extract_list(scope),
637            ));
638            out
639        })
640    }
641
642    pub fn load_matrix_inplace<E: Scalar, N: Size>(
643        &self,
644        row: &[E],
645        fragment: &mut Array<Vector<E, N>>,
646        #[comptime] ident: MatrixIdent,
647        #[comptime] num_matrices: usize,
648        #[comptime] transpose: bool,
649    ) {
650        intrinsic!(|scope| {
651            let vector_size = self.__expand_vector_size_method(scope, ident);
652            let slice_vector_size = row.expand.vector_size();
653            let ptr = unsafe { *row.__expand_as_ptr_method(scope) }.expand;
654            let fragment = fragment.__extract_list(scope);
655            scope.register(Instruction::new(
656                CoopMma::LoadMatrix {
657                    ptr,
658                    factor: num_matrices,
659                    transpose,
660                },
661                fragment,
662            ));
663        })
664    }
665
666    /// Store one or more matrix register using intrinsic instructions. CUDA only.
667    /// The number of matrices must be 1, 2, or 4. The rows for the nth matrix are passed by the 8
668    /// lanes starting at `n * 8`. All slice starts must be valid, even for non-participating lanes.
669    /// The slice determines the starting address for a 16-byte row loaded by this unit, with
670    /// the row index being `UNIT_POS_PLANE % 8`.
671    /// The number of elements is determined by element size.
672    ///
673    /// # Constraints:
674    /// Address must be aligned to 16 bytes
675    /// Address must be in shared memory
676    pub fn store_matrix<E: CubePrimitive, N: Size>(
677        &self,
678        row: &mut [E],
679        registers: &Array<Vector<E::Scalar, N>>,
680        #[comptime] ident: MatrixIdent,
681        #[comptime] num_matrices: usize,
682        #[comptime] transpose: bool,
683    ) {
684        intrinsic!(|scope| {
685            let vector_size = self.__expand_vector_size_method(scope, ident);
686
687            let registers = registers.__extract_list(scope);
688            let destination = unsafe { *row.__expand_as_ptr_method(scope) }.expand;
689
690            scope.register(Instruction::no_out(CoopMma::StoreMatrix {
691                registers,
692                destination,
693                factor: num_matrices,
694                transpose,
695            }));
696        })
697    }
698
699    /// Execute a low level `mma` operation with manually managed registers. Register layout
700    /// and index mapping can be retrieved from the [`MmaDefinition`]
701    #[allow(unused)]
702    pub fn execute<NA: Size, NB: Size, NC: Size>(
703        &self,
704        registers_a: &Array<Vector<A, NA>>,
705        registers_b: &Array<Vector<B, NB>>,
706        registers_c: &Array<Vector<CD, NC>>,
707    ) -> Array<Vector<CD, NC>> {
708        intrinsic!(|scope| {
709            let acc_elems = self
710                .clone()
711                .__expand_elems_per_lane_method(scope, MatrixIdent::Accumulator);
712            let acc_vector_size = self
713                .clone()
714                .__expand_vector_size_method(scope, MatrixIdent::Accumulator);
715            let num_registers = acc_elems / acc_vector_size;
716
717            let registers_d = Array::__expand_new(scope, num_registers);
718
719            let registers_a = registers_a.__extract_list(scope);
720            let registers_b = registers_b.__extract_list(scope);
721            let registers_c = registers_c.__extract_list(scope);
722
723            // Only shape is actually used
724            let matrix = cubecl_ir::MatrixType {
725                ident: MatrixIdent::A,
726                m: self.m,
727                n: self.n,
728                k: self.k,
729                storage: self.a_type,
730                layout: MatrixLayout::ColMajor,
731                scope: ir::MatrixScope::Plane,
732            };
733
734            scope.register(Instruction::new(
735                CoopMma::ExecuteManual {
736                    matrix,
737                    registers_a,
738                    registers_b,
739                    registers_c,
740                },
741                registers_d.__extract_list(scope),
742            ));
743
744            registers_d
745        })
746    }
747
748    #[allow(unused)]
749    pub fn execute_inplace<NA: Size, NB: Size, NC: Size>(
750        &self,
751        registers_a: &Array<Vector<A, NA>>,
752        registers_b: &Array<Vector<B, NB>>,
753        registers_c: &mut Array<Vector<CD, NC>>,
754    ) {
755        intrinsic!(|scope| {
756            let acc_elems = self
757                .clone()
758                .__expand_elems_per_lane_method(scope, MatrixIdent::Accumulator);
759            let acc_vector_size = self
760                .clone()
761                .__expand_vector_size_method(scope, MatrixIdent::Accumulator);
762            let num_registers = acc_elems / acc_vector_size;
763
764            let registers_a = registers_a.__extract_list(scope);
765            let registers_b = registers_b.__extract_list(scope);
766            let registers_c = registers_c.__extract_list(scope);
767
768            // Only shape is actually used
769            let matrix = cubecl_ir::MatrixType {
770                ident: MatrixIdent::A,
771                m: self.m,
772                n: self.n,
773                k: self.k,
774                storage: self.a_type,
775                layout: MatrixLayout::ColMajor,
776                scope: ir::MatrixScope::Plane,
777            };
778
779            scope.register(Instruction::new(
780                CoopMma::ExecuteManual {
781                    matrix,
782                    registers_a,
783                    registers_b,
784                    registers_c,
785                },
786                registers_c,
787            ));
788        })
789    }
790
791    /// Execute a low level block scaled `mma` operation with manually managed registers. Register
792    /// layout and index mapping can be retrieved from the [`MmaDefinition`]
793    #[allow(unused)]
794    pub fn execute_scaled<S: Scalar, NA: Size, NB: Size, NC: Size, NS: Size>(
795        &self,
796        registers_a: &Array<Vector<A, NA>>,
797        registers_b: &Array<Vector<B, NB>>,
798        registers_c: &Array<Vector<CD, NC>>,
799        scales_a: Vector<S, NS>,
800        scales_b: Vector<S, NS>,
801    ) -> Array<Vector<CD, NC>> {
802        intrinsic!(|scope| {
803            let acc_elems = self
804                .clone()
805                .__expand_elems_per_lane_method(scope, MatrixIdent::Accumulator);
806            let acc_vector_size = self
807                .clone()
808                .__expand_vector_size_method(scope, MatrixIdent::Accumulator);
809            let num_registers = acc_elems / acc_vector_size;
810
811            let registers_d = Array::__expand_new(scope, num_registers);
812
813            let registers_a = registers_a.__extract_list(scope);
814            let registers_b = registers_b.__extract_list(scope);
815            let registers_c = registers_c.__extract_list(scope);
816
817            // Only shape is actually used
818            let matrix = cubecl_ir::MatrixType {
819                ident: MatrixIdent::A,
820                m: self.m,
821                n: self.n,
822                k: self.k,
823                storage: self.a_type,
824                layout: MatrixLayout::ColMajor,
825                scope: ir::MatrixScope::Plane,
826            };
827
828            scope.register(Instruction::new(
829                CoopMma::ExecuteScaled {
830                    matrix,
831                    registers_a,
832                    registers_b,
833                    registers_c,
834                    scales_a: scales_a.expand,
835                    scales_b: scales_b.expand,
836                    scales_factor: self
837                        .scales_factor
838                        .expect("Can't execute scaled on matrix with no scales"),
839                },
840                registers_d.__extract_list(scope),
841            ));
842
843            registers_d
844        })
845    }
846}
847
848/// Fill the matrix with the provided value.
849#[allow(unused_variables)]
850pub fn fill<C: Scalar, S: MatrixScope>(mat: &mut Matrix<C, S>, value: C) {
851    unexpanded!()
852}
853
854/// Module containing the expand function for [`fill()`].
855pub mod fill {
856    use super::*;
857
858    /// Expand method of [`fill()`].
859    pub fn expand<C: Scalar, S: MatrixScope>(
860        scope: &Scope,
861        mat: &mut MatrixExpand<C, S>,
862        value: NativeExpand<C>,
863    ) {
864        let value: Value = value.into();
865        scope.register(Instruction::new(ir::CoopMma::Fill { value }, mat.elem));
866    }
867}
868
869/// Load the matrix with the provided array using the stride.
870#[allow(unused_variables)]
871pub fn load<C: CubePrimitive, V: CubePrimitive, S: MatrixScope>(
872    mat: &mut Matrix<C, S>,
873    value: &[V],
874    stride: u32,
875) {
876    unexpanded!()
877}
878
879/// Module containing the expand function for [`load()`].
880pub mod load {
881    use super::*;
882
883    /// Expand method of [`load()`].
884    pub fn expand<C: CubePrimitive, V: CubePrimitive, S: MatrixScope>(
885        scope: &Scope,
886        mat: &mut MatrixExpand<C, S>,
887        value: &SliceExpand<V>,
888        stride: NativeExpand<u32>,
889    ) {
890        let stride: Value = stride.into();
891        assert_ne!(
892            mat.ident,
893            MatrixIdent::Accumulator,
894            "Loading accumulator requires explicit layout. Use `load_with_layout` instead."
895        );
896
897        let ptr = unsafe { *value.__expand_as_ptr_method(scope) }.expand;
898
899        scope.register(Instruction::new(
900            ir::CoopMma::Load {
901                ptr,
902                stride,
903                layout: None,
904            },
905            mat.elem,
906        ));
907    }
908}
909
910/// Load the matrix with the provided array using the tensor layout.
911#[allow(unused_variables)]
912pub fn load_tensor<C: CubePrimitive, V: CubePrimitive, S: MatrixScope>(
913    mat: &mut Matrix<C, S>,
914    value: &TensorView<V>,
915) {
916    unexpanded!()
917}
918
919/// Module containing the expand function for [`load_tensor()`].
920pub mod load_tensor {
921    use super::*;
922
923    /// Expand method of [`load()`].
924    pub fn expand<C: CubePrimitive, V: CubePrimitive, S: MatrixScope>(
925        scope: &Scope,
926        mat: &mut MatrixExpand<C, S>,
927        value: &TensorViewExpand<V>,
928    ) {
929        assert_ne!(
930            mat.ident,
931            MatrixIdent::Accumulator,
932            "Loading accumulator requires explicit layout. Use `load_with_layout` instead."
933        );
934        let buffer = value.buffer.__extract_list(scope);
935
936        scope.register(Instruction::new(
937            ir::CoopMma::LoadTensor {
938                buffer,
939                layout: value.layout.expand,
940                view: match &value.view {
941                    ComptimeOptionExpand::None => None,
942                    ComptimeOptionExpand::Some(view) => Some(view.expand),
943                },
944            },
945            mat.elem,
946        ));
947    }
948}
949
950/// Load the matrix with the provided array using the stride with an explicit layout.
951/// Explicit layouts are required when loading accumulators.
952#[allow(unused_variables)]
953pub fn load_with_layout<C: CubePrimitive, V: CubePrimitive, S: MatrixScope>(
954    mat: &mut Matrix<C, S>,
955    value: &[V],
956    stride: u32,
957    layout: MatrixLayout,
958) {
959    unexpanded!()
960}
961
962/// Module containing the expand function for [`load_with_layout()`].
963pub mod load_with_layout {
964    use super::*;
965
966    /// Expand method of [`load_with_layout()`].
967    pub fn expand<C: CubeType, V: CubePrimitive, S: MatrixScope>(
968        scope: &Scope,
969        mat: &mut MatrixExpand<C, S>,
970        value: &SliceExpand<V>,
971        stride: NativeExpand<u32>,
972        layout: MatrixLayout,
973    ) {
974        let stride: Value = stride.into();
975        let ptr = unsafe { *value.__expand_as_ptr_method(scope) }.expand;
976
977        scope.register(Instruction::new(
978            ir::CoopMma::Load {
979                ptr,
980                stride,
981                layout: Some(layout),
982            },
983            mat.elem,
984        ));
985    }
986}
987
988/// Store the matrix in the given array following the given stride and layout.
989#[allow(unused_variables)]
990pub fn store<C: CubePrimitive, O: CubePrimitive, S: MatrixScope>(
991    output: &mut [O],
992    mat: &Matrix<C, S>,
993    stride: u32,
994    layout: MatrixLayout,
995) {
996    unexpanded!()
997}
998
999/// Module containing the expand function for [`store()`].
1000pub mod store {
1001    use super::*;
1002
1003    /// Expand method of [`store()`].
1004    pub fn expand<C: CubePrimitive, O: CubePrimitive, S: MatrixScope>(
1005        scope: &Scope,
1006        output: &mut SliceExpand<O>,
1007        mat: &MatrixExpand<C, S>,
1008        stride: NativeExpand<u32>,
1009        layout: MatrixLayout,
1010    ) {
1011        let stride: Value = stride.into();
1012
1013        let destination = unsafe { *output.__expand_as_ptr_method(scope) }.expand;
1014
1015        scope.register(Instruction::no_out(ir::CoopMma::Store {
1016            mat: mat.elem,
1017            stride,
1018            destination,
1019            layout,
1020        }));
1021    }
1022}
1023
1024/// Store the matrix in the given tensor view.
1025#[allow(unused_variables)]
1026pub fn store_tensor<C: CubePrimitive, O: CubePrimitive, S: MatrixScope>(
1027    output: &mut TensorView<O>,
1028    mat: &Matrix<C, S>,
1029) {
1030    unexpanded!()
1031}
1032
1033/// Module containing the expand function for [`store_tensor()`].
1034pub mod store_tensor {
1035    use super::*;
1036
1037    /// Expand method of [`store()`].
1038    pub fn expand<C: CubePrimitive, O: CubePrimitive, S: MatrixScope>(
1039        scope: &Scope,
1040        output: &mut TensorViewExpand<O>,
1041        mat: &MatrixExpand<C, S>,
1042    ) {
1043        let buffer = output.buffer.__extract_list(scope);
1044        scope.register(Instruction::new(
1045            ir::CoopMma::StoreTensor {
1046                mat: mat.elem,
1047                layout: output.layout.expand,
1048                view: match &output.view {
1049                    ComptimeOptionExpand::None => None,
1050                    ComptimeOptionExpand::Some(view) => Some(view.expand),
1051                },
1052            },
1053            buffer,
1054        ));
1055    }
1056}
1057
1058/// Execute the matrix-multiply and accumulate operation on the given [matrices](Matrix).
1059#[allow(unused_variables)]
1060pub fn execute<
1061    A: CubePrimitive,
1062    B: CubePrimitive,
1063    C: CubePrimitive,
1064    D: CubePrimitive,
1065    S: MatrixScope,
1066>(
1067    mat_a: &Matrix<A, S>,
1068    mat_b: &Matrix<B, S>,
1069    mat_c: &Matrix<C, S>,
1070    mat_d: &Matrix<D, S>,
1071) {
1072    unexpanded!()
1073}
1074
1075/// Module containing the expand function for [`execute()`].
1076pub mod execute {
1077    use super::*;
1078
1079    /// Expand method of [`execute()`].
1080    pub fn expand<
1081        A: CubePrimitive,
1082        B: CubePrimitive,
1083        C: CubePrimitive,
1084        D: CubePrimitive,
1085        S: MatrixScope,
1086    >(
1087        scope: &Scope,
1088        mat_a: &MatrixExpand<A, S>,
1089        mat_b: &MatrixExpand<B, S>,
1090        mat_c: &MatrixExpand<C, S>,
1091        mat_d: &MatrixExpand<D, S>,
1092    ) {
1093        scope.register(Instruction::new(
1094            ir::CoopMma::Execute {
1095                mat_a: mat_a.elem,
1096                mat_b: mat_b.elem,
1097                mat_c: mat_c.elem,
1098            },
1099            mat_d.elem,
1100        ));
1101    }
1102}
1103
1104/// Cast the matrix fragment to a different type
1105#[allow(unused_variables)]
1106pub fn cast<C: CubePrimitive, O: CubePrimitive, S: MatrixScope>(
1107    input: &Matrix<C, S>,
1108) -> Matrix<O, S> {
1109    unexpanded!()
1110}
1111
1112/// Module containing the expand function for [`cast()`].
1113pub mod cast {
1114    use super::*;
1115
1116    /// Expand method of [`cast()`].
1117    pub fn expand<C: CubePrimitive, O: CubePrimitive, S: MatrixScope>(
1118        scope: &Scope,
1119        input: &MatrixExpand<C, S>,
1120    ) -> MatrixExpand<O, S> {
1121        let ident = input.ident;
1122
1123        if core::any::TypeId::of::<C>() == core::any::TypeId::of::<O>() {
1124            return MatrixExpand {
1125                elem: input.elem,
1126                ident,
1127                _c: PhantomData,
1128                _s: PhantomData,
1129            };
1130        }
1131        let input = input.elem;
1132        let input_mat = match input.ty.unwrap_ptr() {
1133            ir::Type::Matrix(mat) => mat,
1134            _ => unreachable!(),
1135        };
1136
1137        let elem = O::__expand_as_type(scope).storage_type();
1138        let elem = scope.create_local_mut(Type::Matrix(ir::MatrixType::new(
1139            ident,
1140            input_mat.m,
1141            input_mat.n,
1142            input_mat.k,
1143            elem,
1144            MatrixLayout::Undefined,
1145            input_mat.scope,
1146        )));
1147
1148        let output = MatrixExpand {
1149            ident,
1150            elem,
1151            _c: PhantomData,
1152            _s: PhantomData,
1153        };
1154        scope.register(Instruction::new(ir::CoopMma::Cast { input }, output.elem));
1155
1156        output
1157    }
1158}
1159
1160/// Cast the matrix fragment to a different type and a different matrix ident.
1161/// This allows casting to otherwise unsupported types, i.e. casting an f32 accumulator to bf16
1162/// (which can't be used as the accumulator type).
1163#[allow(unused_variables)]
1164pub fn cast_with_ident<C: CubePrimitive, O: CubePrimitive, S: MatrixScope>(
1165    input: &Matrix<C, S>,
1166    ident: MatrixIdent,
1167) -> Matrix<O, S> {
1168    unexpanded!()
1169}
1170
1171/// Module containing the expand function for [`cast()`].
1172pub mod cast_with_ident {
1173    use super::*;
1174
1175    /// Expand method of [`cast()`].
1176    pub fn expand<C: CubePrimitive, O: CubePrimitive, S: MatrixScope>(
1177        scope: &Scope,
1178        input: MatrixExpand<C, S>,
1179        ident: MatrixIdent,
1180    ) -> MatrixExpand<O, S> {
1181        if core::any::TypeId::of::<C>() == core::any::TypeId::of::<O>() && ident == input.ident {
1182            return MatrixExpand {
1183                elem: input.elem,
1184                ident,
1185                _c: PhantomData,
1186                _s: PhantomData,
1187            };
1188        }
1189        let input = input.elem;
1190        let input_mat = match input.ty.unwrap_ptr() {
1191            ir::Type::Matrix(mat) => mat,
1192            _ => unreachable!(),
1193        };
1194
1195        let elem = O::__expand_as_type(scope).storage_type();
1196        let elem = scope.create_local_mut(Type::Matrix(ir::MatrixType::new(
1197            ident,
1198            input_mat.m,
1199            input_mat.n,
1200            input_mat.k,
1201            elem,
1202            MatrixLayout::Undefined,
1203            input_mat.scope,
1204        )));
1205
1206        let output = MatrixExpand {
1207            ident,
1208            elem,
1209            _c: PhantomData,
1210            _s: PhantomData,
1211        };
1212        scope.register(Instruction::new(ir::CoopMma::Cast { input }, output.elem));
1213
1214        output
1215    }
1216}
1217
1218impl CubeType for MatrixLayout {
1219    type ExpandType = Self;
1220}
1221
1222impl IntoExpand for MatrixLayout {
1223    type Expand = Self;
1224
1225    fn into_expand(self, _scope: &Scope) -> Self::Expand {
1226        self
1227    }
1228}
1229
1230impl ExpandTypeClone for MatrixLayout {
1231    fn clone_unchecked(&self) -> Self {
1232        *self
1233    }
1234}
1235
1236impl IntoMut for MatrixLayout {
1237    fn into_mut(self, _scope: &Scope) -> Self {
1238        self
1239    }
1240}
1241
1242impl CubeDebug for MatrixLayout {}
1243
1244impl AsRefExpand for MatrixLayout {
1245    fn __expand_ref_method(&self, _: &Scope) -> &Self {
1246        self
1247    }
1248}
1249impl AsMutExpand for MatrixLayout {
1250    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
1251        self
1252    }
1253}
1254
1255/// Execute an elementwise op on the matrix fragment.
1256///
1257/// Function parameters are (row, col, element) -> element
1258#[allow(unused_variables)]
1259pub fn execute_elementwise_op<A: CubePrimitive, S: MatrixScope>(
1260    matrix_in: &Matrix<A, S>,
1261    matrix_out: &Matrix<A, S>,
1262    op: impl Fn(u32, u32, A::Scalar) -> A::Scalar,
1263) {
1264    unexpanded!()
1265}
1266
1267/// Module containing the expand function for [`execute()`].
1268pub mod execute_elementwise_op {
1269    use alloc::vec;
1270
1271    use super::*;
1272
1273    /// Expand method of [`execute()`].
1274    pub fn expand<A: CubePrimitive, S: MatrixScope>(
1275        scope: &Scope,
1276        matrix_in: &MatrixExpand<A, S>,
1277        matrix_out: &MatrixExpand<A, S>,
1278        mut op: impl FnMut(
1279            &Scope,
1280            NativeExpand<u32>,
1281            NativeExpand<u32>,
1282            NativeExpand<A::Scalar>,
1283        ) -> NativeExpand<A::Scalar>,
1284    ) {
1285        let row = scope.create_value(u32::__expand_as_type(scope));
1286        let col = scope.create_value(u32::__expand_as_type(scope));
1287        let elem = scope.create_value(A::Scalar::__expand_as_type(scope));
1288
1289        let mut closure_scope = scope.child();
1290        let return_value = op(&mut closure_scope, row.into(), col.into(), elem.into());
1291        closure_scope.return_value = Some(return_value.expand);
1292
1293        let op = scope.create_function(vec![row, col, elem], closure_scope);
1294
1295        scope.register(Instruction::new(
1296            ir::CoopMma::ExecuteElementwise {
1297                matrix: matrix_in.elem,
1298                op,
1299            },
1300            matrix_out.elem,
1301        ));
1302    }
1303}
cubecl_core/frontend/cmma.rs

cubecl_core/frontend/
cmma.rs
