cubecl_tile/ops/matmul/
mma_io.rs

1//! Which transports a manual-mma tile loads through ([`MmaIo`]).
2
3use cubecl::cmma::MatrixIdent;
4
5/// Which transports a manual-mma tile may load through; every unit stores its own cells, since
6/// `stmatrix` does not reach a tile's memory window. `LoadMatrix` lets an operand
7/// load through `ldmatrix` wherever the device offers it for the element and the stage serves it,
8/// which the tile decides as it expands; `Manual` keeps every unit loading its own cells.
9/// [`Default`] is what a caller with no reason to choose takes: `ldmatrix` for both operands where
10/// it serves.
11#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
12pub struct MmaIo {
13    pub lhs_load_method: LoadMethod,
14    pub rhs_load_method: LoadMethod,
15    pub acc_load_method: LoadMethod,
16}
17
18#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
19pub enum LoadMethod {
20    Manual,
21    LoadMatrix,
22}
23
24/// Both operands through `ldmatrix` where the device and the stage let them, the accumulator each
25/// unit's own.
26impl Default for MmaIo {
27    fn default() -> Self {
28        Self {
29            lhs_load_method: LoadMethod::LoadMatrix,
30            rhs_load_method: LoadMethod::LoadMatrix,
31            ..Self::manual()
32        }
33    }
34}
35
36impl MmaIo {
37    /// A config forcing the manual path for every role: what a test of that path asks for.
38    pub fn manual() -> Self {
39        Self {
40            lhs_load_method: LoadMethod::Manual,
41            rhs_load_method: LoadMethod::Manual,
42            acc_load_method: LoadMethod::Manual,
43        }
44    }
45
46    /// The transport `ident` may load through.
47    pub fn load_method(&self, ident: MatrixIdent) -> LoadMethod {
48        match ident {
49            MatrixIdent::A => self.lhs_load_method,
50            MatrixIdent::B => self.rhs_load_method,
51            MatrixIdent::Accumulator => self.acc_load_method,
52        }
53    }
54}
cubecl_tile/ops/matmul/mma_io.rs

cubecl_tile/ops/matmul/
mma_io.rs
