cubek_std/tile/ops/
matmul.rs

1//! `Tile::mma` / `Tile::mma_partition` dispatchers and the partition
2//! load/write-back helpers.
3
4use cubecl::prelude::*;
5
6use crate::{
7    StageIdent,
8    stage::Stage,
9    tile::{
10        PartitionScheduler, StageEventListener, Tile, TileExpand, TileKind, TileKindExpand,
11        TileScope, WriteEvent, WriteEventListener,
12    },
13};
14
15#[cube]
16impl<N: Numeric, Sc: TileScope> Tile<N, Sc> {
17    /// `self += lhs · rhs`. For `(Stage, Stage, Partition)` use
18    /// [`Tile::mma_partition`].
19    pub fn mma<L: Numeric, R: Numeric>(&mut self, lhs: &Tile<L, Sc>, rhs: &Tile<R, Sc>) {
20        match (&lhs.kind, &rhs.kind, &mut self.kind) {
21            (TileKind::Cmma(l), TileKind::Cmma(r), TileKind::Cmma(a)) => a.mma(l, r),
22            (TileKind::Cmma(l), TileKind::Cmma(r), TileKind::Bounce(a)) => a.cmma.mma(l, r),
23            (TileKind::Bounce(l), TileKind::Cmma(r), TileKind::Bounce(a)) => a.cmma.mma(&l.cmma, r),
24            (TileKind::Bounce(l), TileKind::Cmma(r), TileKind::Cmma(a)) => a.mma(&l.cmma, r),
25            (TileKind::Mma(l), TileKind::Mma(r), TileKind::Mma(a)) => a.mma(l, r),
26            (TileKind::Register(l), TileKind::Register(r), TileKind::Register(a)) => a.mma(l, r),
27            (TileKind::PlaneVec(l), TileKind::PlaneVec(r), TileKind::PlaneVec(a)) => a.mma(l, r),
28            (TileKind::Interleaved(l), TileKind::Interleaved(r), TileKind::Interleaved(a)) => {
29                a.mma(l, r)
30            }
31            (TileKind::Stage(_), TileKind::Stage(_), TileKind::Partition(_)) => {
32                panic!(
33                    "Tile::mma: (Stage, Stage, Partition) requires extra context — call \
34                     Tile::mma_partition."
35                )
36            }
37            _ => panic!("Unsupported storage combination for mma"),
38        }
39    }
40
41    /// `mma` for `(Stage, Stage, Partition)` operands with rhs fragments
42    /// held under `TileKind::Pipelined`.
43    #[allow(clippy::too_many_arguments)]
44    pub fn mma_partition<
45        LhsS: Numeric,
46        LhsSize: Size,
47        LhsR: Numeric,
48        RhsS: Numeric,
49        RhsSize: Size,
50        RhsR: Numeric,
51        SEL: StageEventListener,
52    >(
53        &mut self,
54        lhs: &Tile<LhsS, Sc>,
55        rhs: &Tile<RhsS, Sc>,
56        a_fragment: &mut Sequence<Tile<LhsR, Sc>>,
57        b_fragments: &mut Tile<RhsR, Sc>,
58        #[comptime] partition_size_k: u32,
59        listener: SEL,
60        scheduler: &PartitionScheduler,
61    ) {
62        match (&lhs.kind, &rhs.kind, &mut self.kind, &mut b_fragments.kind) {
63            (
64                TileKind::Stage(a_stage),
65                TileKind::Stage(b_stage),
66                TileKind::Partition(acc),
67                TileKind::Pipelined(b_frags),
68            ) => acc.execute_with_listener::<LhsS, LhsSize, LhsR, RhsS, RhsSize, RhsR, SEL>(
69                a_stage,
70                b_stage,
71                a_fragment,
72                b_frags,
73                partition_size_k,
74                listener,
75                scheduler,
76            ),
77            _ => panic!(
78                "Tile::mma_partition: requires (lhs, rhs, self, b_fragments) kinds = \
79                 (Stage, Stage, Partition, Pipelined)"
80            ),
81        }
82    }
83}
84
85#[cube]
86/// Fill a partition accumulator from a stage. `None`-kind stage zero-inits.
87pub fn load_partition_from_stage<
88    AccSE: Numeric,
89    AccSS: Size,
90    LhsRE: Numeric,
91    RhsRE: Numeric,
92    AccRE: Numeric,
93    Sc: TileScope,
94    StageAcc: Stage<AccSE>,
95>(
96    stage: &StageAcc,
97    acc: &mut Tile<AccRE, Sc>,
98    scheduler: &PartitionScheduler,
99    #[comptime] partition_size_m: u32,
100    #[comptime] partition_size_n: u32,
101) {
102    let n_iterations = partition_size_n as usize;
103
104    #[unroll]
105    for m in 0..partition_size_m as usize {
106        let m_stage = scheduler.map_m(m as u32);
107
108        #[unroll]
109        for n in 0..n_iterations {
110            let n_stage = scheduler.map_n(n as u32);
111
112            let acc_tile = acc.partition_tile_at_mut(m, n, n_iterations);
113            let tile = StageAcc::tile::<Sc>(stage, (m_stage, n_stage));
114            acc_tile.copy_from::<AccSE, AccSS, LhsRE, RhsRE, AccRE>(&tile, StageIdent::Acc);
115        }
116    }
117}
118
119#[cube]
120#[allow(clippy::too_many_arguments)]
121/// Write a partition accumulator back to an output stage, emitting
122/// `Begin` / `TileStored` / `Finish` events.
123pub fn write_partition_to_stage<
124    OutSE: Numeric,
125    AccSS: Size,
126    LhsRE: Numeric,
127    RhsRE: Numeric,
128    AccRE: Numeric,
129    Sc: TileScope,
130    OutStage: Stage<OutSE>,
131    W: WriteEventListener,
132>(
133    acc: &mut Tile<AccRE, Sc>,
134    out_stage: &mut OutStage,
135    listener: &mut W,
136    scheduler: &PartitionScheduler,
137    #[comptime] partition_size_m: u32,
138    #[comptime] partition_size_n: u32,
139) {
140    let n_iterations = partition_size_n as usize;
141
142    W::on_event(listener, WriteEvent::new_Begin());
143
144    #[unroll]
145    for m_iter in 0..partition_size_m as usize {
146        let m_store = scheduler.map_m(m_iter as u32);
147
148        #[unroll]
149        for n_iter in 0..n_iterations {
150            let n_store = scheduler.map_n(n_iter as u32);
151
152            let tile_accumulator = acc.partition_tile_at_mut(m_iter, n_iter, n_iterations);
153
154            let tile_pos = (m_store, n_store);
155            let mut tile = OutStage::tile::<Sc>(&*out_stage, tile_pos);
156
157            tile.copy_from::<AccRE, AccSS, LhsRE, RhsRE, AccRE>(
158                &*tile_accumulator,
159                StageIdent::Out,
160            );
161
162            W::on_event(listener, WriteEvent::new_TileStored(tile_pos));
163        }
164    }
165
166    W::on_event(listener, WriteEvent::new_Finish());
167}
cubek_std/tile/ops/matmul.rs

cubek_std/tile/ops/
matmul.rs
