cubek_std/microkernel/
estimate.rs

1use cubecl::tune::Work;
2
3/// A kernel's performance, estimated: what the kernel a routine prepared asks of a device,
4/// counted from its own walk, before any device is known.
5///
6/// One vocabulary for every routine of every family, so kernels of different routines are
7/// compared on the same quantities. It starts from the kernel's grid (its `From<&Grid>`), and
8/// each part has one owner, which sets it whole: the problem its [`demand`](Self::demand), the
9/// component placing the stage's reads its [`staging`](Self::staging), the microkernel holding
10/// the work in its units its [`work`](Self::work). A routine sets what it owns and hands the
11/// estimate to its microkernel, which does the same, as the kernel hands down its walk.
12///
13/// Counts only: every rate, latency and capacity is the device's. A part left unset is a routine
14/// that forgot its own.
15#[derive(Debug, Clone, PartialEq)]
16pub struct KernelEstimate {
17    /// Cubes the grid spawns; `None` where the count is only known at launch.
18    pub cubes: Option<u64>,
19    /// What one cube holds of its worker.
20    pub cube: CubeFootprint,
21    /// What the problem asks for, whichever routine runs it: its arithmetic, and the bytes its
22    /// operands hold, what a kernel reading each byte once would move. The same [`Work`] an
23    /// autotune table's roofline is drawn from.
24    pub demand: Option<Work>,
25    pub staging: Option<CubeStaging>,
26    pub work: Option<UnitWork>,
27}
28
29/// What one cube holds of its worker, from the grid.
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub struct CubeFootprint {
32    pub units: u32,
33    /// Bytes of shared memory.
34    pub shared_memory: usize,
35}
36
37/// How one cube walks its work through global memory: in stages, each loaded while the stages
38/// before it compute, then drained.
39#[derive(Debug, Clone, PartialEq)]
40pub struct CubeStaging {
41    /// Arithmetic operations one cube performs, the cells it covers past the problem's edge
42    /// included.
43    pub operations: f64,
44    /// What one cube reads over its walk.
45    pub reads: Vec<GlobalStream>,
46    /// What one cube writes once its walk ends.
47    pub writes: Vec<GlobalStream>,
48    /// Stages one cube walks.
49    pub stages: u64,
50    /// Stages one cube holds at once: one where it loads a stage only once the last is computed,
51    /// more where it loads ahead.
52    pub stages_in_flight: u32,
53    /// Atomic adds into shared memory one cube issues as it drains.
54    pub drain_atomics: u64,
55}
56
57/// Bytes one cube moves through one operand, and the stretch of them a plane touches contiguously.
58#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59pub struct GlobalStream {
60    pub bytes: u64,
61    /// Consecutive bytes a plane touches over the cube's walk: what a memory transaction is filled
62    /// from, a line a later stage finishes reading counted as filled.
63    pub run_bytes: usize,
64}
65
66/// What one unit holds of the work and how it computes it.
67#[derive(Debug, Clone, Copy, PartialEq, Eq)]
68pub struct UnitWork {
69    /// The instructions it runs on, which decide the rate.
70    pub instruction: ArithmeticInstruction,
71    /// 32-bit registers it keeps its work in: a lower bound, counting what the microkernel keeps
72    /// live and not what the compiler adds around it.
73    pub registers: u32,
74}
75
76impl UnitWork {
77    /// This work and `other` held by one unit at once, as a microkernel built from two runs them
78    /// one after the other within each step: their registers added, on the instruction both run.
79    pub fn beside(self, other: UnitWork) -> UnitWork {
80        debug_assert_eq!(
81            self.instruction, other.instruction,
82            "work held beside other work runs on the same instruction"
83        );
84        UnitWork {
85            instruction: self.instruction,
86            registers: self.registers + other.registers,
87        }
88    }
89}
90
91/// Which of the device's arithmetic rates a kernel runs at.
92#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
93pub enum ArithmeticInstruction {
94    /// Vector fused multiply-adds, one unit at a time.
95    Vector,
96    /// A matrix instruction a plane issues together.
97    Matrix,
98}
cubek_std/microkernel/estimate.rs

cubek_std/microkernel/
estimate.rs
