cubecl_tile/ops/matmul/
register_block.rs

1//! The software instruction's execution and unrolling configuration ([`RegisterBlock`]).
2
3/// Execution and unrolling configuration for the software instruction.
4#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
5pub struct RegisterBlock {
6    /// Scalar register budget for the accumulator block; blocks over budget stay rolled.
7    /// Also limits the tap count eligible for gather-loop unrolling.
8    pub budget: usize,
9    /// Whether to generate a fast in-bounds path plus checked fallback for edge tiles.
10    pub split_edge: bool,
11    /// Whether to walk K as (line, component) rather than as a flat scalar walk.
12    pub component_fanout: bool,
13    /// Whether a block contracting an rhs lined along `K` keeps one partial a lane, summed once at
14    /// the drain, rather than summing each run into a one-wide cell every step.
15    pub partials: bool,
16}
17
18impl RegisterBlock {
19    /// A budget with neither specialization turned on.
20    pub const fn new(budget: usize) -> Self {
21        Self {
22            budget,
23            split_edge: false,
24            component_fanout: false,
25            partials: false,
26        }
27    }
28
29    /// Generate the fast in-bounds path plus checked fallback for edge tiles.
30    pub const fn split_edge(self) -> Self {
31        Self {
32            split_edge: true,
33            ..self
34        }
35    }
36
37    /// Walk `K` as (line, component) rather than as a flat scalar walk.
38    pub const fn component_fanout(self) -> Self {
39        Self {
40            component_fanout: true,
41            ..self
42        }
43    }
44
45    /// Hold a run along `K` as one partial a lane: each step of the contraction is one
46    /// multiply-add a cell, `vector_size` values wide, and the lanes are summed once, at the drain,
47    /// rather than every step. Only for a block the drain is the one reader of: a lane holds part
48    /// of a cell, so its cells are not whole values until then.
49    pub const fn partials(self) -> Self {
50        Self {
51            partials: true,
52            ..self
53        }
54    }
55}
cubecl_tile/ops/matmul/register_block.rs

cubecl_tile/ops/matmul/
register_block.rs
