cubecl_ir/types/
matrix.rs

1use cubecl_macros_internal::TypeHash;
2use derive_more::Display;
3use derive_new::new;
4use pliron::derive::{format, pliron_type, type_interface_impl};
5
6use crate::{
7    interfaces::{AlignedType, HasElementType, TypedExt},
8    prelude::*,
9};
10
11#[allow(missing_docs)]
12#[derive(new, Debug, Clone, Copy, PartialEq, Eq, Hash)]
13#[pliron_type(
14    name = "cube.matrix",
15    format = "$ident `<` $scope `, ` $elem_ty `, ` $shape `, ` $layout `>`",
16    generate_get = true,
17    verifier = "succ"
18)]
19pub struct MatrixType {
20    pub ident: MatrixIdent,
21    pub shape: MatrixShape,
22    pub elem_ty: TypeHandle,
23    pub layout: MatrixLayout,
24    pub scope: MatrixScope,
25}
26
27impl MatrixType {
28    /// Size of the unpacked matrix elements, in bits
29    pub fn unpacked_elem_size_bits(&self, ctx: &Context) -> usize {
30        let size_bits = self.elem_ty.size(ctx) * 8;
31        size_bits / self.elem_ty.packing_factor(ctx)
32    }
33}
34
35#[type_interface_impl]
36impl AlignedType for MatrixType {
37    fn align(&self, ctx: &Context) -> usize {
38        self.elem_ty.align(ctx)
39    }
40}
41
42#[type_interface_impl]
43impl HasElementType for MatrixType {
44    fn element_type(&self, ctx: &Context) -> Option<TypeHandle> {
45        type_cast::<dyn HasElementType>(&*self.elem_ty.deref(ctx))?.element_type(ctx)
46    }
47}
48
49#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
50#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
51#[format("$m `x` $n `x` $k")]
52pub struct MatrixShape {
53    pub m: usize,
54    pub n: usize,
55    pub k: usize,
56}
57
58impl MatrixShape {
59    pub fn num_elems(&self, ident: MatrixIdent) -> usize {
60        match ident {
61            MatrixIdent::A => self.m * self.k,
62            MatrixIdent::B => self.k * self.n,
63            MatrixIdent::Accumulator => self.m * self.n,
64        }
65    }
66}
67
68impl From<(usize, usize, usize)> for MatrixShape {
69    fn from((m, n, k): (usize, usize, usize)) -> Self {
70        Self { m, n, k }
71    }
72}
73
74#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
75#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord, Display)]
76#[format]
77#[allow(missing_docs)]
78pub enum MatrixIdent {
79    #[display("IdentA")]
80    A,
81    #[display("IdentB")]
82    B,
83    #[display("IdentAcc")]
84    Accumulator,
85}
86
87#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
88#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord, Display)]
89#[display(rename_all = "snake_case")]
90#[format]
91#[allow(missing_docs)]
92pub enum MatrixLayout {
93    ColMajor,
94    RowMajor,
95    Undefined,
96}
97
98#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
99#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord, Display)]
100#[display(rename_all = "snake_case")]
101#[format]
102#[allow(missing_docs)]
103pub enum MatrixScope {
104    Plane,
105    Cube,
106}
cubecl_ir/types/matrix.rs

cubecl_ir/types/
matrix.rs
