rublas/kernel_ir/definition/
spec.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::ir::StorageType;
3use ruda_kernel::dsl::prelude::*;
4use half::{bf16, f16};
5
6use crate::kernel_ir::definition::MatmulIdent;
7
8/// Matrix multiplication precisions.
9pub trait MatmulPrecision: Send + Sync + Copy + 'static {
10    /// Element type of lhs input tensor of the kernel.
11    type Lhs: MatrixPrecision;
12    /// Element type of rhs input tensor of the kernel.
13    type Rhs: MatrixPrecision;
14    /// Element type of acc input tensor of the kernel.
15    type Acc: MatrixPrecision;
16}
17
18pub trait MatrixPrecision: Send + Sync + Copy + 'static {
19    /// Element type of input tensor in global memory
20    type Global: Numeric;
21    /// Element type once stored in shared memory
22    type Stage: Numeric;
23    /// Element type once in registers for computation
24    type Register: Numeric;
25}
26
27impl<EG: Numeric, ES: Numeric> MatrixPrecision for (EG, ES) {
28    type Global = EG;
29    type Stage = ES;
30    type Register = ES;
31}
32
33impl<EG: Numeric, ES: Numeric, ER: Numeric> MatrixPrecision for (EG, ES, ER) {
34    type Global = EG;
35    type Stage = ES;
36    type Register = ER;
37}
38
39/// Matrix multiplication precisions.
40pub trait MatmulTypes: Send + Sync + Copy + 'static {
41    /// Element type of lhs input tensor of the kernel.
42    type Lhs: MatrixTypes;
43    /// Element type of rhs input tensor of the kernel.
44    type Rhs: MatrixTypes;
45    /// Element type of acc input tensor of the kernel.
46    type Acc: MatrixTypes;
47}
48
49pub trait MatrixTypes: Send + Sync + Copy + 'static {
50    /// Element type of input tensor in global memory
51    type Global: Numeric;
52    type GlobalSize: Size;
53    /// Element type once stored in shared memory
54    type Stage: Numeric;
55    type StageSize: Size;
56    /// Element type once in registers for computation
57    type Register: Numeric;
58    type RegisterSize: Size;
59}
60
61impl<EG: Numeric, SG: Size, ES: Numeric, SS: Size, ER: Numeric, SR: Size> MatrixTypes
62    for (EG, SG, ES, SS, ER, SR)
63{
64    type Global = EG;
65    type GlobalSize = SG;
66    type Stage = ES;
67    type StageSize = SS;
68    type Register = ER;
69    type RegisterSize = SR;
70}
71
72impl MatmulPrecision for f16 {
73    type Lhs = (f16, f16);
74    type Rhs = (f16, f16);
75    #[cfg(target_os = "macos")]
76    type Acc = (f16, f16);
77    #[cfg(not(target_os = "macos"))]
78    type Acc = (f16, f32);
79}
80
81impl MatmulPrecision for flex32 {
82    type Lhs = (f32, f16);
83    type Rhs = (f32, f16);
84    type Acc = (f32, f32);
85}
86
87impl MatmulPrecision for bf16 {
88    type Lhs = (bf16, bf16);
89    type Rhs = (bf16, bf16);
90    #[cfg(target_os = "macos")]
91    type Acc = (bf16, bf16);
92    #[cfg(not(target_os = "macos"))]
93    type Acc = (bf16, f32);
94}
95
96impl MatmulPrecision for f32 {
97    type Lhs = (f32, f32);
98    type Rhs = (f32, f32);
99    type Acc = (f32, f32);
100}
101
102impl MatmulPrecision for f64 {
103    type Lhs = (f64, f32);
104    type Rhs = (f64, f32);
105    type Acc = (f64, f32);
106}
107
108impl MatmulPrecision for u8 {
109    type Lhs = (u8, u8);
110    type Rhs = (u8, u8);
111    type Acc = (i32, i32);
112}
113
114impl MatmulPrecision for u16 {
115    type Lhs = (u16, u16);
116    type Rhs = (u16, u16);
117    type Acc = (i32, i32);
118}
119
120impl MatmulPrecision for u32 {
121    type Lhs = (u32, u32);
122    type Rhs = (u32, u32);
123    type Acc = (u32, u32);
124}
125
126impl MatmulPrecision for u64 {
127    type Lhs = (u64, u64);
128    type Rhs = (u64, u64);
129    type Acc = (u64, u64);
130}
131
132impl MatmulPrecision for i8 {
133    type Lhs = (i8, i8);
134    type Rhs = (i8, i8);
135    type Acc = (i32, i32);
136}
137
138impl MatmulPrecision for i16 {
139    type Lhs = (i16, i16);
140    type Rhs = (i16, i16);
141    type Acc = (i32, i32);
142}
143
144impl MatmulPrecision for i32 {
145    type Lhs = (i32, i32);
146    type Rhs = (i32, i32);
147    type Acc = (i32, i32);
148}
149
150impl MatmulPrecision for i64 {
151    type Lhs = (i64, i64);
152    type Rhs = (i64, i64);
153    type Acc = (i64, i64);
154}
155
156impl<Lhs: MatrixPrecision, Rhs: MatrixPrecision, Acc: MatrixPrecision> MatmulPrecision
157    for (Lhs, Rhs, Acc)
158{
159    type Lhs = Lhs;
160    type Rhs = Rhs;
161    type Acc = Acc;
162}
163
164impl<Lhs: MatrixTypes, Rhs: MatrixTypes, Acc: MatrixTypes> MatmulTypes for (Lhs, Rhs, Acc) {
165    type Lhs = Lhs;
166    type Rhs = Rhs;
167    type Acc = Acc;
168}
169
170pub type Lhs<MT> = <MT as MatmulTypes>::Lhs;
171pub type Rhs<MT> = <MT as MatmulTypes>::Rhs;
172pub type Acc<MT> = <MT as MatmulTypes>::Acc;
173
174pub type Global<MT> = <MT as MatrixTypes>::Global;
175pub type GlobalSize<MT> = <MT as MatrixTypes>::GlobalSize;
176
177pub type Stage<MT> = <MT as MatrixTypes>::Stage;
178pub type StageSize<MT> = <MT as MatrixTypes>::StageSize;
179
180pub type Register<MT> = <MT as MatrixTypes>::Register;
181pub type RegisterSize<MT> = <MT as MatrixTypes>::RegisterSize;
182
183// ==================== LHS ====================
184
185// Vector forms
186pub type LhsG<MT> = Vector<Global<Lhs<MT>>, GlobalSize<Lhs<MT>>>;
187pub type LhsS<MT> = Vector<Stage<Lhs<MT>>, StageSize<Lhs<MT>>>;
188pub type LhsR<MT> = Vector<Register<Lhs<MT>>, RegisterSize<Lhs<MT>>>;
189
190// Element / Size splits
191pub type LhsGE<MT> = <Lhs<MT> as MatrixTypes>::Global;
192pub type LhsGS<MT> = <Lhs<MT> as MatrixTypes>::GlobalSize;
193
194pub type LhsSE<MT> = <Lhs<MT> as MatrixTypes>::Stage;
195pub type LhsSS<MT> = <Lhs<MT> as MatrixTypes>::StageSize;
196
197pub type LhsRE<MT> = <Lhs<MT> as MatrixTypes>::Register;
198pub type LhsRS<MT> = <Lhs<MT> as MatrixTypes>::RegisterSize;
199
200// ==================== RHS ====================
201
202// Vector forms
203pub type RhsG<MT> = Vector<Global<Rhs<MT>>, GlobalSize<Rhs<MT>>>;
204pub type RhsS<MT> = Vector<Stage<Rhs<MT>>, StageSize<Rhs<MT>>>;
205pub type RhsR<MT> = Vector<Register<Rhs<MT>>, RegisterSize<Rhs<MT>>>;
206
207// Element / Size splits
208pub type RhsGE<MT> = <Rhs<MT> as MatrixTypes>::Global;
209pub type RhsGS<MT> = <Rhs<MT> as MatrixTypes>::GlobalSize;
210
211pub type RhsSE<MT> = <Rhs<MT> as MatrixTypes>::Stage;
212pub type RhsSS<MT> = <Rhs<MT> as MatrixTypes>::StageSize;
213
214pub type RhsRE<MT> = <Rhs<MT> as MatrixTypes>::Register;
215pub type RhsRS<MT> = <Rhs<MT> as MatrixTypes>::RegisterSize;
216
217// ==================== ACC ====================
218
219// Vector forms
220pub type AccG<MT> = Vector<Global<Acc<MT>>, GlobalSize<Acc<MT>>>;
221pub type AccS<MT> = Vector<Stage<Acc<MT>>, StageSize<Acc<MT>>>;
222pub type AccR<MT> = Vector<Register<Acc<MT>>, RegisterSize<Acc<MT>>>;
223
224// Element / Size splits
225pub type AccGE<MT> = <Acc<MT> as MatrixTypes>::Global;
226pub type AccGS<MT> = <Acc<MT> as MatrixTypes>::GlobalSize;
227
228pub type AccSE<MT> = <Acc<MT> as MatrixTypes>::Stage;
229pub type AccSS<MT> = <Acc<MT> as MatrixTypes>::StageSize;
230
231pub type AccRE<MT> = <Acc<MT> as MatrixTypes>::Register;
232pub type AccRS<MT> = <Acc<MT> as MatrixTypes>::RegisterSize;
233
234/// Arithmetic permitted for F32 inputs; other dtypes retain their own precision rules.
235#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)]
236pub enum F32MathMode {
237    /// Do not automatically reduce F32 operands to a lower-precision format.
238    #[default]
239    Strict,
240    /// Permit TF32 multiplication with F32 accumulation; F32 fallback is allowed.
241    AllowTf32,
242}
243
244#[derive(Debug, Clone, Eq, PartialEq, Hash)]
245pub struct MatmulElems {
246    pub f32_math: F32MathMode,
247    pub lhs_global: StorageType,
248    pub rhs_global: StorageType,
249    pub acc_global: StorageType,
250    pub lhs_stage: StorageType,
251    pub rhs_stage: StorageType,
252    pub acc_stage: StorageType,
253    pub lhs_register: StorageType,
254    pub rhs_register: StorageType,
255    pub acc_register: StorageType,
256}
257
258#[derive(Clone, Debug)]
259pub struct MatmulGlobalElems {
260    pub f32_math: F32MathMode,
261    pub lhs: StorageType,
262    pub rhs: StorageType,
263    pub out: StorageType,
264}
265
266impl MatmulElems {
267    pub fn new_deprecated<MP: MatmulPrecision>() -> Self {
268        Self {
269            f32_math: F32MathMode::Strict,
270            lhs_global: <MP::Lhs as MatrixPrecision>::Global::as_type_native_unchecked()
271                .storage_type(),
272            rhs_global: <MP::Rhs as MatrixPrecision>::Global::as_type_native_unchecked()
273                .storage_type(),
274            acc_global: <MP::Acc as MatrixPrecision>::Global::as_type_native_unchecked()
275                .storage_type(),
276            lhs_stage: <MP::Lhs as MatrixPrecision>::Stage::as_type_native_unchecked()
277                .storage_type(),
278            rhs_stage: <MP::Rhs as MatrixPrecision>::Stage::as_type_native_unchecked()
279                .storage_type(),
280            acc_stage: <MP::Acc as MatrixPrecision>::Stage::as_type_native_unchecked()
281                .storage_type(),
282            lhs_register: <MP::Lhs as MatrixPrecision>::Register::as_type_native_unchecked()
283                .storage_type(),
284            rhs_register: <MP::Rhs as MatrixPrecision>::Register::as_type_native_unchecked()
285                .storage_type(),
286            acc_register: <MP::Acc as MatrixPrecision>::Register::as_type_native_unchecked()
287                .storage_type(),
288        }
289    }
290
291    pub fn from_globals(global_elems: &MatmulGlobalElems) -> Self {
292        let acc_type = if global_elems.out == half::f16::as_type_native_unchecked().storage_type()
293            || global_elems.out == half::bf16::as_type_native_unchecked().storage_type()
294        {
295            f32::as_type_native_unchecked().storage_type()
296        } else {
297            global_elems.out
298        };
299
300        Self {
301            f32_math: global_elems.f32_math,
302            lhs_global: global_elems.lhs,
303            rhs_global: global_elems.rhs,
304            acc_global: global_elems.out,
305            lhs_stage: global_elems.lhs,
306            rhs_stage: global_elems.rhs,
307            acc_stage: acc_type,
308            lhs_register: global_elems.lhs,
309            rhs_register: global_elems.rhs,
310            acc_register: acc_type,
311        }
312    }
313
314    pub fn from_single_dtype(dtype: Type) -> Self {
315        let dtype = dtype.storage_type();
316        Self {
317            f32_math: F32MathMode::Strict,
318            lhs_global: dtype,
319            rhs_global: dtype,
320            acc_global: dtype,
321            lhs_stage: dtype,
322            rhs_stage: dtype,
323            acc_stage: dtype,
324            lhs_register: dtype,
325            rhs_register: dtype,
326            acc_register: dtype,
327        }
328    }
329
330    pub fn global(&self, ident: MatmulIdent) -> StorageType {
331        match ident {
332            MatmulIdent::Lhs => self.lhs_global,
333            MatmulIdent::Rhs => self.rhs_global,
334            MatmulIdent::Out => self.acc_global,
335        }
336    }
337
338    pub fn stage(&self, ident: MatmulIdent) -> StorageType {
339        match ident {
340            MatmulIdent::Lhs => self.lhs_stage,
341            MatmulIdent::Rhs => self.rhs_stage,
342            MatmulIdent::Out => self.acc_stage,
343        }
344    }
345
346    pub fn register(&self, ident: MatmulIdent) -> StorageType {
347        match ident {
348            MatmulIdent::Lhs => self.lhs_register,
349            MatmulIdent::Rhs => self.rhs_register,
350            MatmulIdent::Out => self.acc_register,
351        }
352    }
353
354    pub fn as_global_elems(&self) -> MatmulGlobalElems {
355        MatmulGlobalElems {
356            f32_math: self.f32_math,
357            lhs: self.lhs_global,
358            rhs: self.rhs_global,
359            out: self.acc_global,
360        }
361    }
362
363    /// Prefer output type for stage because it's the same size at best, but often smaller.
364    /// Having stage == global also enables things like TMA, and an f16 stage for output enables
365    /// using `stmatrix` on the registers after casting.
366    pub fn adjust_stage_dtypes(&mut self) {
367        self.lhs_stage = self.lhs_global;
368        self.rhs_stage = self.rhs_global;
369        self.acc_stage = self.acc_global;
370    }
371}
rublas/kernel_ir/definition/spec.rs

rublas/kernel_ir/definition/
spec.rs
