cubecl_tile/stage/pipeline/
slot.rs

1//! [`Slot`]: a payload plus the [`Meeting`] sequencing its fill against its read.
2
3use cubecl::ir::Scope;
4use cubecl::prelude::*;
5use cubecl::unexpanded;
6
7use super::payload::base::{Payload, PayloadExpand};
8use super::payload::pair::OperandPair;
9use crate::*;
10
11pub(crate) const FIRST_SLOT: usize = 0;
12
13/// One slot of a buffered walk: payload `T` and the rendezvous sequencing fill against read.
14#[derive(CubeType)]
15pub struct Slot<T: CubeType> {
16    pub(crate) data: T,
17    pub(crate) pipeline: Meeting,
18    /// When each operand of `T` is filled, in the order `T` holds them.
19    #[cube(comptime)]
20    pub(crate) refills: Vec<Refill>,
21}
22
23#[cube]
24impl<T: CubeType> Slot<T> {
25    /// Wrap an already-built payload and pipeline.
26    pub(crate) fn wrap(data: T, pipeline: Meeting, #[comptime] refills: Vec<Refill>) -> Slot<T> {
27        Slot::<T> {
28            data,
29            pipeline,
30            refills,
31        }
32    }
33
34    /// Whether this slot has any fixed operand.
35    #[allow(dead_code)] // Reached through its expand, from [`StagesFill`].
36    pub(crate) fn has_fixed(&self) -> comptime_type!(bool) {
37        comptime!(self.refills.contains(&Refill::Once))
38    }
39
40    /// Producer acquire: wait until the slot is free; a `Plane` slot waits on `sync_plane`.
41    #[allow(dead_code)] // Reached through its expand, from `Slot::fill` / `Slot::consume`.
42    pub(crate) fn acquire_write(&self) {
43        match &self.pipeline {
44            Meeting::Barrier { empty, writes, .. } => empty.wait_parity(*writes ^ 1),
45            Meeting::Cube => sync_cube(),
46            Meeting::Plane => sync_plane(),
47        }
48    }
49
50    /// Producer release: publish the fill, async copies included ([`Meeting::producers`] say
51    /// who arrives).
52    #[allow(dead_code)] // Reached through its expand, from `Slot::fill` / `Slot::consume`.
53    pub(crate) fn release_write(&mut self) {
54        match &mut self.pipeline {
55            Meeting::Barrier {
56                full,
57                all_publish,
58                commits,
59                elected,
60                writes,
61                ..
62            } => {
63                if *all_publish || UNIT_POS == *elected {
64                    // Handed over before the arrival, so the phase cannot flip on the arrivals
65                    // with this unit's copies still in flight.
66                    if *commits {
67                        full.commit_copy_async();
68                    }
69                    full.arrive();
70                }
71                *writes ^= 1;
72            }
73            Meeting::Cube | Meeting::Plane => {}
74        }
75    }
76
77    /// Consumer acquire: wait for the slot's fill.
78    #[allow(dead_code)] // Reached through its expand, from `Slot::fill` / `Slot::consume`.
79    pub(crate) fn acquire_read(&self) {
80        match &self.pipeline {
81            Meeting::Barrier { full, reads, .. } => full.wait_parity(*reads),
82            Meeting::Cube | Meeting::Plane => {}
83        }
84    }
85
86    /// Consumer release: free the slot.
87    #[allow(dead_code)] // Reached through its expand, from `Slot::fill` / `Slot::consume`.
88    pub(crate) fn release_read(&mut self) {
89        match &mut self.pipeline {
90            Meeting::Barrier { empty, reads, .. } => {
91                empty.arrive();
92                *reads ^= 1;
93            }
94            Meeting::Cube | Meeting::Plane => {}
95        }
96    }
97
98    /// Publish this slot's last fill where no later fill will; call right before
99    /// [`consume`](Slot::consume).
100    pub fn publish(&self) {
101        match &self.pipeline {
102            Meeting::Cube => sync_cube(),
103            Meeting::Plane => sync_plane(),
104            Meeting::Barrier { .. } => {}
105        }
106    }
107}
108
109#[cube]
110impl<P: Payload<P> + CubeType<ExpandType: Clone>> Slot<P> {
111    /// Bring `src`'s window at `region` into this slot, for every operand whose refill is `only`.
112    pub(crate) fn refill(&mut self, src: &P, region: &Region, #[comptime] only: Refill) {
113        let refills = comptime!(self.refills.clone());
114        let writes = comptime!(refills.contains(&only));
115        if comptime!(writes || only == Refill::EveryRegion) {
116            self.acquire_write();
117            self.data
118                .bring(src, &self.pipeline, region, comptime!(refills), only);
119            self.release_write();
120        }
121    }
122}
123
124// `consume` is spelled per payload shape: closure inference cannot resolve `P::ExpandType`.
125impl Slot<OperandPair> {
126    /// Wait for the slot's fill, pass its two tiles to `compute`, then free the slot.
127    pub fn consume(&mut self, _compute: impl FnOnce(&Tile, &Tile)) {
128        unexpanded!()
129    }
130}
131
132impl SlotExpand<OperandPair> {
133    pub fn __expand_consume_method<F>(&mut self, scope: &Scope, compute: F)
134    where
135        F: FnOnce(&Scope, &TileExpand, &TileExpand),
136    {
137        self.__expand_acquire_read_method(scope);
138        compute(scope, &self.data.lhs, &self.data.rhs);
139        self.__expand_release_read_method(scope);
140    }
141}
142
143impl Slot<Tile> {
144    /// [`consume`](Slot::consume) for the sole operand.
145    pub fn consume(&mut self, _compute: impl FnOnce(&Tile)) {
146        unexpanded!()
147    }
148}
149
150impl SlotExpand<Tile> {
151    pub fn __expand_consume_method<F>(&mut self, scope: &Scope, compute: F)
152    where
153        F: FnOnce(&Scope, &TileExpand),
154    {
155        self.__expand_acquire_read_method(scope);
156        compute(scope, &self.data);
157        self.__expand_release_read_method(scope);
158    }
159}
cubecl_tile/stage/pipeline/slot.rs

cubecl_tile/stage/pipeline/
slot.rs
