cubecl_tile/ops/matmul/
instruction.rs

1//! What one plane contracts through ([`Instruction`]).
2
3use crate::{MmaIo, RegisterBlock};
4
5/// What one plane contracts through: a software register block or a hardware fragment form.
6#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
7pub enum Instruction {
8    /// A register block, run under `config`. Asks nothing of the hardware.
9    Registers { config: RegisterBlock },
10    /// The vendor's fragment API: `wmma` on CUDA, rocWMMA on HIP.
11    Cmma,
12    /// This crate's own fragment transport over the matrix intrinsics, with per-role `io`, over
13    /// `operands` of one form.
14    Mma { io: MmaIo, operands: MmaOperands },
15    /// Hopper's warpgroup MMA: a plane group of four contracts its `64 × n` tile together,
16    /// reading both operands out of shared memory, asynchronously.
17    Wgmma,
18}
19
20impl Instruction {
21    /// Whether a sum this instruction contracts is scaled where it is held: a register block's
22    /// cells and a manual-mma fragment's are each unit's own registers. A cmma fragment hides
23    /// which unit holds a cell, so it is scaled only through a round trip in shared memory, and a
24    /// warpgroup's sum only once its MMAs land.
25    pub fn scales_its_sum_in_place(self) -> bool {
26        match self {
27            Self::Registers { .. } | Self::Mma { .. } => true,
28            Self::Cmma | Self::Wgmma => false,
29        }
30    }
31}
32
33/// What a manual-mma instruction's operands are: the blueprint's choice, which every leaf reads, so
34/// the operands a leaf is handed never decide it.
35#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, Default)]
36pub enum MmaOperands {
37    /// Values as the instruction takes them. A scaled or packed operand is decoded before it
38    /// reaches one, or read through lane runs where its fragment is deeper than the instruction.
39    #[default]
40    Plain,
41    /// `e2m1` values under `e4m3` block scales, each read as stored by the device's block-scaled
42    /// instruction.
43    BlockScaled,
44}
cubecl_tile/ops/matmul/instruction.rs

cubecl_tile/ops/matmul/
instruction.rs
