cubek_matmul/definition/
precision.rs

1//! The element types a matmul runs at, as a type-level mapping: one `MatmulPrecision`
2//! names the global/stage/register element of each operand. Both architectures launch at
3//! a precision, so it sits beside [`MatmulElems`](crate::definition::MatmulElems).
4
5use cubecl::prelude::*;
6use half::{bf16, f16};
7
8use crate::definition::MatmulElems;
9
10/// Matrix multiplication precisions.
11pub trait MatmulPrecision: Send + Sync + Copy + 'static {
12    /// Element type of lhs input tensor of the kernel.
13    type Lhs: MatrixPrecision;
14    /// Element type of rhs input tensor of the kernel.
15    type Rhs: MatrixPrecision;
16    /// Element type of acc input tensor of the kernel.
17    type Acc: MatrixPrecision;
18}
19
20pub trait MatrixPrecision: Send + Sync + Copy + 'static {
21    /// Element type of input tensor in global memory
22    type Global: Numeric;
23    /// Element type once stored in shared memory
24    type Stage: Numeric;
25    /// Element type once in registers for computation
26    type Register: Numeric;
27}
28
29impl<EG: Numeric, ES: Numeric> MatrixPrecision for (EG, ES) {
30    type Global = EG;
31    type Stage = ES;
32    type Register = ES;
33}
34
35impl<EG: Numeric, ES: Numeric, ER: Numeric> MatrixPrecision for (EG, ES, ER) {
36    type Global = EG;
37    type Stage = ES;
38    type Register = ER;
39}
40
41impl MatmulPrecision for f16 {
42    type Lhs = (f16, f16);
43    type Rhs = (f16, f16);
44    #[cfg(target_os = "macos")]
45    type Acc = (f16, f16);
46    #[cfg(not(target_os = "macos"))]
47    type Acc = (f16, f32);
48}
49
50impl MatmulPrecision for flex32 {
51    type Lhs = (f32, f16);
52    type Rhs = (f32, f16);
53    type Acc = (f32, f32);
54}
55
56impl MatmulPrecision for bf16 {
57    type Lhs = (bf16, bf16);
58    type Rhs = (bf16, bf16);
59    #[cfg(target_os = "macos")]
60    type Acc = (bf16, bf16);
61    #[cfg(not(target_os = "macos"))]
62    type Acc = (bf16, f32);
63}
64
65impl MatmulPrecision for f32 {
66    type Lhs = (f32, f32);
67    type Rhs = (f32, f32);
68    type Acc = (f32, f32);
69}
70
71impl MatmulPrecision for f64 {
72    type Lhs = (f64, f32);
73    type Rhs = (f64, f32);
74    type Acc = (f64, f32);
75}
76
77impl MatmulPrecision for u8 {
78    type Lhs = (u8, u8);
79    type Rhs = (u8, u8);
80    type Acc = (i32, i32);
81}
82
83impl MatmulPrecision for u16 {
84    type Lhs = (u16, u16);
85    type Rhs = (u16, u16);
86    type Acc = (i32, i32);
87}
88
89impl MatmulPrecision for u32 {
90    type Lhs = (u32, u32);
91    type Rhs = (u32, u32);
92    type Acc = (u32, u32);
93}
94
95impl MatmulPrecision for u64 {
96    type Lhs = (u64, u64);
97    type Rhs = (u64, u64);
98    type Acc = (u64, u64);
99}
100
101impl MatmulPrecision for i8 {
102    type Lhs = (i8, i8);
103    type Rhs = (i8, i8);
104    type Acc = (i32, i32);
105}
106
107impl MatmulPrecision for i16 {
108    type Lhs = (i16, i16);
109    type Rhs = (i16, i16);
110    type Acc = (i32, i32);
111}
112
113impl MatmulPrecision for i32 {
114    type Lhs = (i32, i32);
115    type Rhs = (i32, i32);
116    type Acc = (i32, i32);
117}
118
119impl MatmulPrecision for i64 {
120    type Lhs = (i64, i64);
121    type Rhs = (i64, i64);
122    type Acc = (i64, i64);
123}
124
125impl<Lhs: MatrixPrecision, Rhs: MatrixPrecision, Acc: MatrixPrecision> MatmulPrecision
126    for (Lhs, Rhs, Acc)
127{
128    type Lhs = Lhs;
129    type Rhs = Rhs;
130    type Acc = Acc;
131}
132
133impl MatmulElems {
134    pub fn new_deprecated<MP: MatmulPrecision>() -> Self {
135        Self {
136            lhs_global: <MP::Lhs as MatrixPrecision>::Global::elem_type_native(),
137            rhs_global: <MP::Rhs as MatrixPrecision>::Global::elem_type_native(),
138            acc_global: <MP::Acc as MatrixPrecision>::Global::elem_type_native(),
139            lhs_stage: <MP::Lhs as MatrixPrecision>::Stage::elem_type_native(),
140            rhs_stage: <MP::Rhs as MatrixPrecision>::Stage::elem_type_native(),
141            acc_stage: <MP::Acc as MatrixPrecision>::Stage::elem_type_native(),
142            lhs_register: <MP::Lhs as MatrixPrecision>::Register::elem_type_native(),
143            rhs_register: <MP::Rhs as MatrixPrecision>::Register::elem_type_native(),
144            acc_register: <MP::Acc as MatrixPrecision>::Register::elem_type_native(),
145        }
146    }
147}
cubek_matmul/definition/precision.rs

cubek_matmul/definition/
precision.rs
