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

cubecl_core/frontend/
cmma.rs
