cubecl-runtime 0.11.0-pre.1

Crate that helps creating high performance async runtimes 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

use cubecl_ir::ElemType;

/// Configuration for a matrix multiplication (CMMA) operation.
#[derive(Eq, PartialEq, Clone, Hash, Debug, Copy)]
#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
pub struct ComputeCmmaConfig {
    /// The data type used to store the running sum.
    pub accumulator_type: AccumulatorType,
    /// The spatial dimensions of the operation.
    pub cmma_dims: CmmaDims,
}

/// The element type of the accumulator.
pub type AccumulatorType = ElemType;

/// The M, N, and K dimensions of a matrix multiplication.
#[derive(Eq, PartialEq, Clone, Hash, Debug, Copy)]
#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
pub struct CmmaDims {
    /// Rows in the output matrix.
    pub m: usize,
    /// Columns in the output matrix.
    pub n: usize,
    /// The shared inner dimension of the input matrices.
    pub k: usize,
}

impl CmmaDims {
    /// Returns the total iteration volume (M * N * K).
    pub fn num_elems(&self) -> usize {
        self.m * self.n * self.k
    }
}