cubek-tile 0.3.0

CubeK: Tile
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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

//! Host-side configuration: manual-mma load/store selection ([`MmaIo`]) and the software
//! instruction's execution config ([`RegisterBlock`]).

use cubecl::{
    cmma::MatrixIdent,
    ir::{DeviceProperties, ElemType},
};

/// Which transports a manual-mma tile may load and store through. `LoadMatrix` lets an operand
/// load through `ldmatrix` wherever the device offers it for the element and the stage serves it,
/// which the tile decides as it expands; `Manual` keeps every unit loading its own cells.
/// [`Default`] is what a caller with no reason to choose takes: `ldmatrix` for both operands where
/// it serves.
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
pub struct MmaIo {
    pub lhs_load_method: LoadMethod,
    pub rhs_load_method: LoadMethod,
    pub acc_load_method: LoadMethod,
    pub store_method: StoreMethod,
}

#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub enum LoadMethod {
    Manual,
    LoadMatrix,
}

#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub enum StoreMethod {
    Manual,
    StoreMatrix,
}

/// Both operands through `ldmatrix` where the device and the stage let them, the accumulator and
/// the store each unit's own: `stmatrix` does not reach a tile's memory window.
impl Default for MmaIo {
    fn default() -> Self {
        Self {
            lhs_load_method: LoadMethod::LoadMatrix,
            rhs_load_method: LoadMethod::LoadMatrix,
            ..Self::manual()
        }
    }
}

impl MmaIo {
    /// Select each role's transport from the device's `ldmatrix`/`stmatrix` support.
    pub fn new(
        device_props: &DeviceProperties,
        lhs_stage: ElemType,
        rhs_stage: ElemType,
        acc_stage: ElemType,
    ) -> Self {
        Self {
            lhs_load_method: load_method(device_props, lhs_stage),
            rhs_load_method: load_method(device_props, rhs_stage),
            acc_load_method: load_method(device_props, acc_stage),
            store_method: store_method(device_props, acc_stage),
        }
    }

    /// A config forcing the manual path for every role: what a test of that path asks for.
    pub fn manual() -> Self {
        Self {
            lhs_load_method: LoadMethod::Manual,
            rhs_load_method: LoadMethod::Manual,
            acc_load_method: LoadMethod::Manual,
            store_method: StoreMethod::Manual,
        }
    }

    /// The transport `ident` may load through.
    pub fn load_method(&self, ident: MatrixIdent) -> LoadMethod {
        match ident {
            MatrixIdent::A => self.lhs_load_method,
            MatrixIdent::B => self.rhs_load_method,
            MatrixIdent::Accumulator => self.acc_load_method,
        }
    }

    pub fn store_method(&self) -> StoreMethod {
        self.store_method
    }
}

fn load_method(device_props: &DeviceProperties, dtype: ElemType) -> LoadMethod {
    if device_props.features.matmul.ldmatrix.contains(&dtype) {
        LoadMethod::LoadMatrix
    } else {
        LoadMethod::Manual
    }
}

fn store_method(device_props: &DeviceProperties, dtype: ElemType) -> StoreMethod {
    if device_props.features.matmul.stmatrix.contains(&dtype) {
        StoreMethod::StoreMatrix
    } else {
        StoreMethod::Manual
    }
}

/// Execution and unrolling configuration for the software instruction.
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
pub struct RegisterBlock {
    /// Scalar register budget for the accumulator block; blocks over budget stay rolled.
    pub budget: usize,
    /// Whether to generate a fast in-bounds path plus checked fallback for edge tiles.
    pub split_edge: bool,
    /// Whether to walk K as (line, component) rather than as a flat scalar walk.
    pub component_fanout: bool,
}

impl RegisterBlock {
    /// A budget with neither specialization turned on.
    pub const fn new(budget: usize) -> Self {
        Self {
            budget,
            split_edge: false,
            component_fanout: false,
        }
    }

    /// Generate the fast in-bounds path plus checked fallback for edge tiles.
    pub const fn split_edge(self) -> Self {
        Self {
            split_edge: true,
            ..self
        }
    }

    /// Walk `K` as (line, component) rather than as a flat scalar walk.
    pub const fn component_fanout(self) -> Self {
        Self {
            component_fanout: true,
            ..self
        }
    }
}