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

cubecl_core/frontend/
cmma.rs
