t4a-cubecl-std 0.11.0

CubeCL Standard Library.
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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

use cubecl_core::ir::{ElemType, FloatKind, features::Features};
use cubecl_runtime::{client::Client, throughput::ComputeCmmaConfig};

use crate::throughput::compute_direct;

/// The operand types and vector widths an arithmetic ceiling is measured in.
pub(super) struct Arithmetic;

impl Arithmetic {
    /// Without an fma of its own for `dtype` a device emulates one per
    /// operation, and a kernel converts and accumulates in f32 rather than pay
    /// that, so both are measured and the faster kept.
    pub(super) fn dtypes(client: &Client, dtype: ElemType) -> alloc::vec::Vec<ElemType> {
        let mut dtypes = alloc::vec![dtype];

        if let Some(accumulator) = Self::promoted(dtype)
            && client.properties().features.supports_type(accumulator)
        {
            dtypes.push(accumulator);
        }

        dtypes
    }

    /// The probe issues no loads, so it sweeps down from the widest vector kept in registers.
    pub(super) fn widths(client: &Client, dtype: ElemType) -> alloc::vec::Vec<usize> {
        client
            .properties()
            .vector_sizes_in_registers(dtype.size(), compute_direct::LIVE_VECTORS)
            .collect()
    }

    /// What a kernel with `dtype` operands accumulates in when the device has no
    /// arithmetic of its own for them, or `None` where `dtype` is already the
    /// widest of the two.
    fn promoted(dtype: ElemType) -> Option<ElemType> {
        let accumulator = ElemType::Float(FloatKind::F32);
        let narrower_float =
            matches!(dtype, ElemType::Float(_)) && dtype.size() < accumulator.size();

        narrower_float.then_some(accumulator)
    }
}

/// The cooperative matrix a compute-cmma probe launches.
pub(super) struct CooperativeMatrix;

impl CooperativeMatrix {
    /// A non-empty capability list says the device has tensor hardware, not this
    /// shape of it. `mma` is not consulted: the probe issues `cmma::execute`.
    pub(super) fn implemented(
        features: &Features,
        dtype: ElemType,
        config: ComputeCmmaConfig,
    ) -> bool {
        features.matmul.cmma.iter().any(|it| {
            it.a_type == dtype
                && it.b_type == dtype
                && it.cd_type == config.accumulator_type
                && it.m as usize == config.cmma_dims.m
                && it.n as usize == config.cmma_dims.n
                && it.k as usize == config.cmma_dims.k
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use cubecl_core::ir::{IntKind, UIntKind};

    const F32: ElemType = ElemType::Float(FloatKind::F32);

    /// Where the device does the arithmetic itself, its rate is the ceiling and
    /// a second measurement of the same thing costs a probe for nothing.
    #[test]
    fn nothing_as_wide_as_the_accumulator_is_promoted() {
        assert_eq!(Arithmetic::promoted(F32), None);
        assert_eq!(Arithmetic::promoted(ElemType::Float(FloatKind::F64)), None);
    }

    /// The probe retires a multiply for an integer, which no float rate bounds.
    #[test]
    fn an_integer_is_not_promoted() {
        assert_eq!(Arithmetic::promoted(ElemType::Int(IntKind::I8)), None);
        assert_eq!(Arithmetic::promoted(ElemType::UInt(UIntKind::U16)), None);
    }

    #[test]
    fn every_float_narrower_than_the_accumulator_is_promoted() {
        for dtype in [FloatKind::F16, FloatKind::BF16, FloatKind::E4M3] {
            assert_eq!(Arithmetic::promoted(ElemType::Float(dtype)), Some(F32));
        }
    }
}