cubek-std 0.3.0

CubeK: Standard Library
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105

use cubecl::prelude::*;

#[derive(Debug, Clone, Copy)]
/// Matrix dimension specifier for matmul operations.
pub enum MatmulDim {
    /// Rows of the output matrix.
    M,
    /// Columns of the output matrix.
    N,
    /// Reduction dimension.
    K,
}

#[macro_export]
macro_rules! define_3d_size_base {
    ($name:ident, $ty:ty) => {
        #[derive(CubeType, Copy, Clone, Debug, Hash, PartialEq, Eq)]
        pub struct $name {
            pub m: $ty,
            pub n: $ty,
            pub k: $ty,
        }

        impl $name {
            pub fn new(m: u32, n: u32, k: u32) -> Self {
                $name {
                    m: <$ty>::try_from(m).unwrap(),
                    n: <$ty>::try_from(n).unwrap(),
                    k: <$ty>::try_from(k).unwrap(),
                }
            }

            pub fn get(&self, dim: $crate::MatmulDim) -> u32 {
                (match dim {
                    $crate::MatmulDim::M => self.m,
                    $crate::MatmulDim::N => self.n,
                    $crate::MatmulDim::K => self.k,
                }) as u32
            }

            pub fn m(&self) -> u32 {
                self.get($crate::MatmulDim::M)
            }

            pub fn n(&self) -> u32 {
                self.get($crate::MatmulDim::N)
            }

            pub fn k(&self) -> u32 {
                self.get($crate::MatmulDim::K)
            }

            pub fn mn(&self) -> u32 {
                self.get($crate::MatmulDim::M) * self.get($crate::MatmulDim::N)
            }

            pub fn mk(&self) -> u32 {
                self.get($crate::MatmulDim::M) * self.get($crate::MatmulDim::K)
            }

            pub fn nk(&self) -> u32 {
                self.get($crate::MatmulDim::N) * self.get($crate::MatmulDim::K)
            }

            pub fn mnk(&self) -> u32 {
                self.get($crate::MatmulDim::M)
                    * self.get($crate::MatmulDim::N)
                    * self.get($crate::MatmulDim::K)
            }
        }
    };
}

#[macro_export]
macro_rules! impl_3d_size_from_tuple {
    ($name:ident, $ty_struct:ty, $ty_tuple:ty) => {
        impl From<($ty_tuple, $ty_tuple, $ty_tuple)> for $name {
            fn from(value: ($ty_tuple, $ty_tuple, $ty_tuple)) -> Self {
                Self {
                    m: value.0 as $ty_struct,
                    n: value.1 as $ty_struct,
                    k: value.2 as $ty_struct,
                }
            }
        }

        impl From<$name> for ($ty_tuple, $ty_tuple, $ty_tuple) {
            fn from(value: $name) -> Self {
                (
                    value.m as $ty_tuple,
                    value.n as $ty_tuple,
                    value.k as $ty_tuple,
                )
            }
        }
    };
}

// Shapes m,n,k of the problem
define_3d_size_base!(MatmulProblemSize, u32);
impl_3d_size_from_tuple!(MatmulProblemSize, u32, u8);
impl_3d_size_from_tuple!(MatmulProblemSize, u32, u32);
impl_3d_size_from_tuple!(MatmulProblemSize, u32, i32);
impl_3d_size_from_tuple!(MatmulProblemSize, u32, u16);
impl_3d_size_from_tuple!(MatmulProblemSize, u32, usize);