cubecl-cpp 0.11.0-pre.4

CPP transpiler for CubeCL
Documentation 
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

use super::WMMA_MINIMUM_VERSION;
use crate::{
    cuda::arch::CudaArchitecture,
    shared::{Architecture, SupportedMmaCombinations},
};
use cubecl_core::ir::{ElemType, FloatKind, IntKind, UIntKind, features::MmaConfig};
use itertools::Itertools;

pub(super) fn supported_cmma_combinations_wmma(
    arch: &CudaArchitecture,
) -> SupportedMmaCombinations {
    let mut result: SupportedMmaCombinations = vec![];
    if arch.get_version() >= WMMA_MINIMUM_VERSION && arch.tensor_cores {
        let tdims = vec![(16, 16, 16), (32, 8, 16), (8, 32, 16)];
        // Types fully supported.
        let types = vec![
            (
                ElemType::Float(FloatKind::F16), // m
                ElemType::Float(FloatKind::F16), // n
                ElemType::Float(FloatKind::F16), // k
            ),
            (
                ElemType::Float(FloatKind::F16),
                ElemType::Float(FloatKind::F16),
                ElemType::Float(FloatKind::F32),
            ),
            (
                ElemType::Float(FloatKind::BF16),
                ElemType::Float(FloatKind::BF16),
                ElemType::Float(FloatKind::F32),
            ),
            (
                ElemType::Int(IntKind::I8),
                ElemType::Int(IntKind::I8),
                ElemType::Int(IntKind::I32),
            ),
            (
                ElemType::UInt(UIntKind::U8),
                ElemType::UInt(UIntKind::U8),
                ElemType::Int(IntKind::I32),
            ),
        ];
        let combinations: SupportedMmaCombinations = types
            .into_iter()
            .cartesian_product(tdims)
            .map(|((a, b, c), (m, n, k))| MmaConfig {
                a_type: a,
                b_type: b,
                cd_type: c,
                m,
                n,
                k,
            })
            .collect();
        result.extend(combinations);
        if arch.get_version() >= 80 {
            result.push(MmaConfig {
                a_type: ElemType::Float(FloatKind::TF32),
                b_type: ElemType::Float(FloatKind::TF32),
                cd_type: ElemType::Float(FloatKind::F32),
                m: 16,
                n: 16,
                k: 8,
            });
        }
    }
    result
}