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

cubecl_core/frontend/
cmma.rs
