cubecl_tile/stage/pipeline/
base.rs

1//! The buffered walk ([`Stages`]): `depth` [`Slot`]s driven as a circular software pipeline,
2//! `depth - 1` regions in flight while one computes; built, filled and consumed here, driven by
3//! the [`schedule`](super::schedule).
4
5use cubecl::ir::Scope;
6use cubecl::prelude::*;
7use cubecl::unexpanded;
8
9use super::payload::base::{Payload, PayloadExpand, StageSpec};
10use super::payload::pair::OperandPair;
11use super::plan::StagePlan;
12use super::schedule::Prefetch;
13use crate::*;
14
15/// What a plane does with the stages' slots.
16#[derive(CubeType, CubeTypeMut, IntoRuntime)]
17#[cube(runtime_variants)]
18pub enum Role {
19    /// Fills the slots and takes no tile.
20    Fill,
21    /// Reads the slots and computes.
22    Compute,
23}
24
25/// The `depth` slots of one buffered walk and the operands they are filled from.
26#[derive(CubeType)]
27pub struct Stages<T: CubeType> {
28    pub(crate) slots: Sequence<Slot<T>>,
29    pub(crate) sources: T,
30    #[allow(dead_code)]
31    #[cube(comptime)]
32    pub(crate) depth: usize,
33    /// Planes that only fill these slots.
34    #[cube(comptime)]
35    pub(crate) fillers: usize,
36}
37
38#[cube]
39impl<T: CubeType> Stages<T> {
40    /// Wrap already-built slots over the operands they stage.
41    pub(crate) fn wrap(
42        slots: Sequence<Slot<T>>,
43        sources: T,
44        #[comptime] depth: usize,
45        #[comptime] fillers: usize,
46    ) -> Stages<T> {
47        comptime!(assert!(
48            depth > 0,
49            "Stages: a pipeline needs at least one slot"
50        ));
51        Stages::<T> {
52            slots,
53            sources,
54            depth,
55            fillers,
56        }
57    }
58
59    /// What this plane does with the stages' slots; filler planes sit at the end of the cube.
60    pub fn role(&self) -> Role {
61        if comptime!(self.fillers == 0) {
62            Role::new_Compute()
63        } else if UNIT_POS >= Meeting::consumers(comptime!(self.fillers)) {
64            Role::new_Fill()
65        } else {
66            Role::new_Compute()
67        }
68    }
69
70    /// Slot `index`, known at comptime.
71    pub fn slot_mut(&mut self, #[comptime] index: usize) -> &mut Slot<T> {
72        self.slots.index_mut(index)
73    }
74}
75
76#[cube]
77impl<P: Payload<P> + Clone + CubeType<ExpandType: Clone>> Stages<P> {
78    /// Build the `depth` slots staging `sources` into shared memory laid out as `storage`.
79    ///
80    /// `width` overrides the stage's line width; the operand must then be scalar and unpacked.
81    /// `decoded` is the element a stage whose fill decodes its operand holds, `f32` where `None`.
82    #[allow(dead_code)]
83    pub(crate) fn new(
84        walk: &Walk,
85        sources: &P,
86        #[comptime] storage: StageStorage,
87        #[comptime] width: Option<usize>,
88        #[comptime] depth: usize,
89        #[comptime] decoded: Option<ElemType>,
90    ) -> Stages<P> {
91        let operands = sources.operands();
92        let owner = walk.stage_owner();
93        let plan = comptime!(StagePlan::new(&operands, &walk.space, &walk.level, owner));
94        let spec = comptime!(StageSpec {
95            level: walk.level.clone(),
96            depth: walk.depth(),
97            storage,
98            width,
99            decoded,
100            owner: plan.owner(),
101        });
102
103        // Later slots share buffers the first slot owns, so it is built first.
104        let first = sources.staged(
105            sources,
106            comptime!(spec.clone()),
107            comptime!(plan.refills(FIRST_SLOT)),
108        );
109        let mut slots = Sequence::<Slot<P>>::new();
110        #[unroll]
111        for slot in 0..depth {
112            let data = if comptime!(slot == FIRST_SLOT) {
113                first.clone()
114            } else {
115                sources.staged(
116                    &first,
117                    comptime!(spec.clone()),
118                    comptime!(plan.refills(slot)),
119                )
120            };
121            slots.push(Slot::wrap(
122                data,
123                Meeting::new(
124                    comptime!(plan.sync()),
125                    comptime!(plan.collective_full()),
126                    comptime!(plan.commits()),
127                    comptime!(plan.fences()),
128                    comptime!(plan.fillers()),
129                ),
130                comptime!(plan.refills(slot)),
131            ));
132        }
133        Stages::wrap(slots, sources.clone(), depth, comptime!(plan.fillers()))
134    }
135
136    /// Free the shared memory every slot holds, for what the kernel declares after it to take:
137    /// once the walk is over, nothing reads a stage again. Every unit of the cube reaches it, and
138    /// a cube barrier comes before any unit touches what is declared in its place, as
139    /// [`Tile::smem_accumulation`] and [`Tile::smem_cyclic_accumulation`] open with.
140    pub fn free(&self) {
141        #[unroll]
142        for slot in 0..comptime!(self.depth) {
143            self.slots.index(slot).data.free();
144        }
145    }
146
147    /// Fill slot `slot` from `region`'s window: one step of a [`Fill`](Role::Fill) plane's walk.
148    pub fn fill(&mut self, #[comptime] slot: usize, region: &Region) {
149        let sources = self.sources.clone();
150        self.slot_mut(slot)
151            .refill(&sources, region, comptime!(Refill::EveryRegion));
152    }
153}
154
155#[cube]
156impl Stages<OperandPair> {
157    /// The `depth` slots staging both operands of a contraction into shared memory laid out as
158    /// `storage`.
159    ///
160    /// A stage whose fill decodes its operand, from scales, a codebook or fields as wide as their
161    /// word, holds the values at `decoded`, or at `f32` where `None`. An operand staged as it lies
162    /// keeps its own element.
163    pub fn smem(
164        walk: &Walk,
165        lhs: &Tile,
166        rhs: &Tile,
167        #[comptime] storage: StageStorage,
168        #[comptime] depth: usize,
169        #[comptime] decoded: Option<ElemType>,
170    ) -> Stages<OperandPair> {
171        let sources = OperandPair {
172            lhs: lhs.clone(),
173            rhs: rhs.clone(),
174        };
175        Stages::new(walk, &sources, storage, comptime!(None), depth, decoded)
176    }
177}
178
179#[cube]
180impl Stages<Tile> {
181    /// [`smem`](Stages::smem) for the sole operand `input`.
182    pub fn smem_single(
183        walk: &Walk,
184        input: &Tile,
185        #[comptime] storage: StageStorage,
186        #[comptime] depth: usize,
187        #[comptime] decoded: Option<ElemType>,
188    ) -> Stages<Tile> {
189        Stages::new(walk, input, storage, comptime!(None), depth, decoded)
190    }
191
192    /// [`smem_single`](Stages::smem_single) with the stage served in `width`-wide lines.
193    pub fn smem_single_at(
194        walk: &Walk,
195        input: &Tile,
196        #[comptime] storage: StageStorage,
197        #[comptime] width: Option<usize>,
198        #[comptime] depth: usize,
199        #[comptime] decoded: Option<ElemType>,
200    ) -> Stages<Tile> {
201        Stages::new(walk, input, storage, width, depth, decoded)
202    }
203}
204
205// Line widths of a register-staged fill, bound by `fetch_buffers`.
206define_size!(pub(crate) FL);
207define_size!(pub(crate) FR);
208
209/// Registers holding one unit's register-staged fill of both operands.
210pub(crate) type FetchBuffers<Lhs, Rhs> = (Array<Vector<Lhs, FL>>, Array<Vector<Rhs, FR>>);
211
212/// Bind the fetched line widths `FL` and `FR` for the rest of the kernel's scope.
213#[cube]
214fn register_fetched_widths(#[comptime] lhs: usize, #[comptime] rhs: usize) {
215    intrinsic!(|scope| {
216        scope.register_size::<FL>(lhs);
217        scope.register_size::<FR>(rhs);
218    });
219}
220
221#[cube]
222impl Stages<OperandPair> {
223    /// Registers holding one unit's fill of both operands across a contraction.
224    #[allow(dead_code)] // Reached through its expand, from `prefetched`.
225    pub(crate) fn fetch_buffers<Lhs: Numeric, Rhs: Numeric>(&self) -> FetchBuffers<Lhs, Rhs> {
226        let staged = self.slots.index(FIRST_SLOT);
227        let lhs = staged.data.lhs.fetched_scalars();
228        let rhs = staged.data.rhs.fetched_scalars();
229        comptime!(assert!(
230            Prefetch::fits(lhs + rhs),
231            "Stages: a fill fetched into registers holds {lhs} + {rhs} scalars a unit, past the \
232             {MOST_FETCHED_SCALARS} a unit keeps beside its contraction"
233        ));
234        let lhs_width = staged.data.lhs.vector_size();
235        let rhs_width = staged.data.rhs.vector_size();
236        register_fetched_widths(lhs_width, rhs_width);
237        (
238            Array::<Vector<Lhs, FL>>::new(comptime!(lhs / lhs_width)),
239            Array::<Vector<Rhs, FR>>::new(comptime!(rhs / rhs_width)),
240        )
241    }
242
243    /// Read this unit's share of filling slot `slot` with `sources` at `region` into `lhs` and
244    /// `rhs`.
245    #[allow(dead_code)] // Reached through its expand, from `prefetched`.
246    pub(crate) fn fetch<Lhs: Numeric, Rhs: Numeric>(
247        &self,
248        #[comptime] slot: usize,
249        sources: &OperandPair,
250        region: &Region,
251        lhs: &mut Array<Vector<Lhs, FL>>,
252        rhs: &mut Array<Vector<Rhs, FR>>,
253    ) {
254        let staged = self.slots.index(slot);
255        staged.data.lhs.fetch_from(&sources.lhs.at(region), lhs);
256        staged.data.rhs.fetch_from(&sources.rhs.at(region), rhs);
257    }
258
259    /// Write what [`fetch`](Stages::fetch) read into slot `slot`; the caller owns the rendezvous.
260    #[allow(dead_code)] // Reached through its expand, from `prefetched`.
261    pub(crate) fn store<Lhs: Numeric, Rhs: Numeric>(
262        &mut self,
263        #[comptime] slot: usize,
264        lhs: &Array<Vector<Lhs, FL>>,
265        rhs: &Array<Vector<Rhs, FR>>,
266    ) {
267        let slot = self.slot_mut(slot);
268        slot.data.lhs.store_fetched(lhs);
269        slot.data.rhs.store_fetched(rhs);
270    }
271
272    /// Refuse stages a register-staged schedule cannot drive.
273    #[allow(dead_code)] // Reached through its expand, from `prefetched`.
274    pub(crate) fn assert_copied_by_every_unit(&self) {
275        let lhs = self.sources.lhs.delivery();
276        let rhs = self.sources.rhs.delivery();
277        comptime!(assert!(
278            self.fillers == 0 && lhs == Delivery::SyncPerUnit && rhs == Delivery::SyncPerUnit,
279            "Stages: a register-staged schedule fills its slots with every unit's own copy of \
280             memory; these stages are filled by {} plane(s) of their own, or their sources are \
281             delivered {lhs:?} and {rhs:?}",
282            self.fillers
283        ));
284    }
285}
286
287// `consume` is spelled per payload shape, as [`Slot::consume`] is.
288impl Stages<OperandPair> {
289    /// Consume slot `slot`: one step of a [`Compute`](Role::Compute) plane's walk.
290    pub fn consume(&mut self, _slot: usize, _compute: impl FnOnce(&Tile, &Tile)) {
291        unexpanded!()
292    }
293}
294
295impl StagesExpand<OperandPair> {
296    pub fn __expand_consume_method<F>(&mut self, scope: &Scope, slot: usize, compute: F)
297    where
298        F: FnOnce(&Scope, &TileExpand, &TileExpand),
299    {
300        self.__expand_slot_mut_method(scope, slot)
301            .__expand_consume_method(scope, compute);
302    }
303}
304
305impl Stages<Tile> {
306    /// [`consume`](Stages::consume) for the sole operand.
307    pub fn consume(&mut self, _slot: usize, _compute: impl FnOnce(&Tile)) {
308        unexpanded!()
309    }
310}
311
312impl StagesExpand<Tile> {
313    pub fn __expand_consume_method<F>(&mut self, scope: &Scope, slot: usize, compute: F)
314    where
315        F: FnOnce(&Scope, &TileExpand),
316    {
317        self.__expand_slot_mut_method(scope, slot)
318            .__expand_consume_method(scope, compute);
319    }
320}
cubecl_tile/stage/pipeline/base.rs

cubecl_tile/stage/pipeline/
base.rs
