cubecl_cpp/cuda/mma/
cuda_compiler.rs

1use super::WMMA_MINIMUM_VERSION;
2use crate::{
3    cuda::arch::CudaArchitecture,
4    shared::{Architecture, SupportedMmaCombinations},
5};
6use cubecl_core::ir::{ElemType, FloatKind, IntKind, UIntKind, features::MmaConfig};
7use itertools::Itertools;
8
9pub(super) fn supported_cmma_combinations_wmma(
10    arch: &CudaArchitecture,
11) -> SupportedMmaCombinations {
12    let mut result: SupportedMmaCombinations = vec![];
13    if arch.get_version() >= WMMA_MINIMUM_VERSION && arch.tensor_cores {
14        let tdims = vec![(16, 16, 16), (32, 8, 16), (8, 32, 16)];
15        // Types fully supported.
16        let types = vec![
17            (
18                ElemType::Float(FloatKind::F16), // m
19                ElemType::Float(FloatKind::F16), // n
20                ElemType::Float(FloatKind::F16), // k
21            ),
22            (
23                ElemType::Float(FloatKind::F16),
24                ElemType::Float(FloatKind::F16),
25                ElemType::Float(FloatKind::F32),
26            ),
27            (
28                ElemType::Float(FloatKind::BF16),
29                ElemType::Float(FloatKind::BF16),
30                ElemType::Float(FloatKind::F32),
31            ),
32            (
33                ElemType::Int(IntKind::I8),
34                ElemType::Int(IntKind::I8),
35                ElemType::Int(IntKind::I32),
36            ),
37            (
38                ElemType::UInt(UIntKind::U8),
39                ElemType::UInt(UIntKind::U8),
40                ElemType::Int(IntKind::I32),
41            ),
42        ];
43        let combinations: SupportedMmaCombinations = types
44            .into_iter()
45            .cartesian_product(tdims)
46            .map(|((a, b, c), (m, n, k))| MmaConfig {
47                a_type: a,
48                b_type: b,
49                cd_type: c,
50                m,
51                n,
52                k,
53            })
54            .collect();
55        result.extend(combinations);
56        if arch.get_version() >= 80 {
57            result.push(MmaConfig {
58                a_type: ElemType::Float(FloatKind::TF32),
59                b_type: ElemType::Float(FloatKind::TF32),
60                cd_type: ElemType::Float(FloatKind::F32),
61                m: 16,
62                n: 16,
63                k: 8,
64            });
65        }
66    }
67    result
68}
cubecl_cpp/cuda/mma/cuda_compiler.rs

cubecl_cpp/cuda/mma/
cuda_compiler.rs
