cubecl_cpp/cuda/mma/
manual.rs

1use cubecl_core::{
2    cmma::{MatrixIdent, MatrixType},
3    ir::{
4        ElemType, FloatKind, IntKind, UIntKind,
5        dialect::matrix::{ColIndexOp, RowIndexOp},
6        features::{MmaConfig, ScaledMmaConfig},
7        interfaces::TypedExt,
8        prelude::*,
9    },
10    prelude::*,
11};
12use itertools::Itertools;
13use pliron::r#type::TypedHandle;
14
15use cubecl_core::prelude::polyfills::mma::{col_index, row_index};
16
17use crate::{
18    cuda::arch::CudaArchitecture,
19    shared::{
20        Architecture, SupportedMmaCombinations, SupportedScaledMmaCombinations, lowering::LowerOp,
21    },
22    target::Cuda,
23};
24
25#[op_interface_impl]
26impl LowerOp<Cuda> for RowIndexOp {
27    fn lower(&self, scope: &Scope) -> Vec<Value> {
28        let matrix = *self.matrix_ty(scope.ctx()).deref(scope.ctx());
29        let elems_per_reg = 32 / matrix.unpacked_elem_size_bits(scope.ctx());
30        let lane_id = self.lane_id(scope.ctx());
31        let i = self.i(scope.ctx());
32        let out = row_index::expand(scope, lane_id.into(), i.into(), elems_per_reg, matrix.ident);
33        vec![out.value(scope)]
34    }
35}
36
37#[op_interface_impl]
38impl LowerOp<Cuda> for ColIndexOp {
39    fn lower(&self, scope: &Scope) -> Vec<Value> {
40        let matrix = *self.matrix_ty(scope.ctx()).deref(scope.ctx());
41        let elems_per_reg = 32 / matrix.unpacked_elem_size_bits(scope.ctx());
42        let lane_id = self.lane_id(scope.ctx());
43        let i = self.i(scope.ctx());
44        let out = col_index::expand(scope, lane_id.into(), i.into(), elems_per_reg, matrix.ident);
45        vec![out.value(scope)]
46    }
47}
48
49pub fn supported_mma_combinations(arch: &CudaArchitecture) -> SupportedMmaCombinations {
50    if !arch.tensor_cores {
51        return vec![];
52    }
53    let mut result: SupportedMmaCombinations = vec![];
54    // Higher than WMMA because we only support the newest shapes. Other shapes would make things
55    // very complicated.
56    // Also only use f32 accumulators for now
57    if arch.get_version() >= 80 {
58        result.extend([
59            MmaConfig {
60                a_type: ElemType::Float(FloatKind::F16),  // a
61                b_type: ElemType::Float(FloatKind::F16),  // b
62                cd_type: ElemType::Float(FloatKind::F32), // cd
63                m: 16,
64                n: 8,
65                k: 16,
66            },
67            MmaConfig {
68                a_type: ElemType::Float(FloatKind::BF16),
69                b_type: ElemType::Float(FloatKind::BF16),
70                cd_type: ElemType::Float(FloatKind::F32),
71                m: 16,
72                n: 8,
73                k: 16,
74            },
75            MmaConfig {
76                a_type: ElemType::Float(FloatKind::TF32),
77                b_type: ElemType::Float(FloatKind::TF32),
78                cd_type: ElemType::Float(FloatKind::F32),
79                m: 16,
80                n: 8,
81                k: 8,
82            },
83            MmaConfig {
84                a_type: ElemType::Int(IntKind::I8),
85                b_type: ElemType::Int(IntKind::I8),
86                cd_type: ElemType::Int(IntKind::I32),
87                m: 16,
88                n: 8,
89                k: 32,
90            },
91            MmaConfig {
92                a_type: ElemType::UInt(UIntKind::U8),
93                b_type: ElemType::UInt(UIntKind::U8),
94                cd_type: ElemType::Int(IntKind::I32),
95                m: 16,
96                n: 8,
97                k: 32,
98            },
99            MmaConfig {
100                a_type: ElemType::Int(IntKind::I8),
101                b_type: ElemType::UInt(UIntKind::U8),
102                cd_type: ElemType::Int(IntKind::I32),
103                m: 16,
104                n: 8,
105                k: 32,
106            },
107            MmaConfig {
108                a_type: ElemType::UInt(UIntKind::U8),
109                b_type: ElemType::Int(IntKind::I8),
110                cd_type: ElemType::Int(IntKind::I32),
111                m: 16,
112                n: 8,
113                k: 32,
114            },
115            // TODO: u4/i4/b1, there's no types for them yet
116        ]);
117    }
118    if arch.get_version() >= 89 {
119        let f8f6f4_types = [
120            FloatKind::E4M3,
121            FloatKind::E5M2,
122            FloatKind::E3M2,
123            FloatKind::E2M3,
124            FloatKind::E2M1,
125        ];
126        let combinations = f8f6f4_types.iter().cartesian_product(f8f6f4_types.iter());
127        result.extend(combinations.map(|(t1, t2)| MmaConfig {
128            a_type: ElemType::Float(*t1),
129            b_type: ElemType::Float(*t2),
130            cd_type: ElemType::Float(FloatKind::F32),
131            m: 16,
132            n: 8,
133            k: 32,
134        }));
135    }
136    // Warning: this likely does not follow the same layout pattern as those after 80
137    if arch.get_version() >= 70 && arch.get_version() < 80 {
138        result.push(MmaConfig {
139            a_type: ElemType::Float(FloatKind::F16),
140            b_type: ElemType::Float(FloatKind::F16),
141            cd_type: ElemType::Float(FloatKind::F32),
142            m: 16,
143            n: 8,
144            k: 8,
145        });
146    }
147    result
148}
149
150pub fn supported_scaled_mma_combinations(
151    arch: &CudaArchitecture,
152) -> SupportedScaledMmaCombinations {
153    if !arch.tensor_cores {
154        return vec![];
155    }
156    let mut result: SupportedScaledMmaCombinations = vec![];
157    // sm_120f
158    if arch.get_version() >= 120 && arch.get_version() < 130 {
159        let f8f6f4_types = [
160            FloatKind::E4M3,
161            FloatKind::E5M2,
162            FloatKind::E3M2,
163            FloatKind::E2M3,
164            FloatKind::E2M1,
165        ];
166        let combinations = f8f6f4_types
167            .iter()
168            .flat_map(|t1| f8f6f4_types.iter().map(move |t2| (t1, t2)));
169
170        result.extend(combinations.map(|(t1, t2)| ScaledMmaConfig {
171            a_type: ElemType::Float(*t1),
172            b_type: ElemType::Float(*t2),
173            cd_type: ElemType::Float(FloatKind::F32),
174            scales_type: ElemType::Float(FloatKind::UE8M0),
175            m: 16,
176            n: 8,
177            k: 32,
178            scales_factor: 1,
179        }));
180
181        result.extend([
182            ScaledMmaConfig {
183                a_type: ElemType::Float(FloatKind::E2M1x2),
184                b_type: ElemType::Float(FloatKind::E2M1x2),
185                cd_type: ElemType::Float(FloatKind::F32),
186                scales_type: ElemType::Float(FloatKind::UE8M0),
187                m: 16,
188                n: 8,
189                k: 64,
190                scales_factor: 2,
191            },
192            // Sign of scales is ignored
193            ScaledMmaConfig {
194                a_type: ElemType::Float(FloatKind::E2M1x2),
195                b_type: ElemType::Float(FloatKind::E2M1x2),
196                cd_type: ElemType::Float(FloatKind::F32),
197                scales_type: ElemType::Float(FloatKind::E4M3),
198                m: 16,
199                n: 8,
200                k: 64,
201                scales_factor: 4,
202            },
203        ]);
204    }
205    result
206}
207
208pub fn contiguous_elements_cuda(
209    ctx: &Context,
210    ident: MatrixIdent,
211    matrix: TypedHandle<MatrixType>,
212) -> usize {
213    let elem = matrix.deref(ctx).elem_ty;
214    match ident {
215        MatrixIdent::A | MatrixIdent::B => 32 / elem.size_bits(ctx),
216        MatrixIdent::Accumulator => 2,
217    }
218}
219
220#[cfg(test)]
221mod tests {
222    use super::supported_mma_combinations;
223    use crate::cuda::arch::CudaArchitecture;
224
225    /// A die with no tensor cores offers no MMA at all, so nothing downstream can pick a tile
226    /// and measure the FP16 pipeline in units that claim tensor hardware.
227    #[test]
228    fn a_turing_die_without_tensor_cores_offers_no_mma() {
229        let turing = |name: &str| CudaArchitecture {
230            version: 75,
231            tensor_cores: CudaArchitecture::has_tensor_cores(75, name),
232        };
233
234        assert!(supported_mma_combinations(&turing("NVIDIA GeForce GTX 1660 SUPER")).is_empty());
235        assert!(!supported_mma_combinations(&turing("NVIDIA GeForce RTX 2060")).is_empty());
236    }
237}
cubecl_cpp/cuda/mma/manual.rs

cubecl_cpp/cuda/mma/
manual.rs
