cubecl_tile/tile/accumulator/
smem_cyclic.rs

1//! A shared-memory destination the planes of a cube write exclusively, on the cyclic schedule, and
2//! a tile that reads the sum back.
3//!
4//! [`SmemAccumulation`](super::smem_accumulation::SmemAccumulation)'s writers add into every cell
5//! at once through atomics, in whatever order they arrive, into a buffer zeroed first. Here each
6//! writer's fragments are dealt into as many chunks as there are writers, by their place in the
7//! fragment grid, and the drain runs in as many rounds: in round `r`, writer `w` lands its
8//! fragments of chunk `(w + r) mod writers` alone, replacing their cells in round zero and adding
9//! into them after, with a plain read and write. Every writer works every round, on cells no other
10//! writer touches, a cube barrier between rounds, and drains each fragment once: no atomics, no
11//! zeroing, and each cell takes its partials in a fixed order, so the sum is the same bits from run
12//! to run.
13
14use core::marker::PhantomData;
15
16use cubecl::ir::{ExpandValue, VectorSize};
17use cubecl::prelude::*;
18use cubecl::unexpanded;
19use cubecl_std::tensor::{
20    ErasedTensor, ErasedTensorExpand, ErasedTensorOperationsExpand, ReadsLines, WriteOnly,
21    WritesLines,
22};
23
24use crate::*;
25
26/// A cube's shared-memory accumulator over one box, summing in `A` and read back at the element
27/// of the tile it was opened over, that `writers` writers reach exclusively, on the
28/// [`Cyclic`](Schedule::Cyclic) schedule.
29///
30/// A writer drains its partial with [`drain`](SmemCyclicAccumulation::drain) (or
31/// [`drain_at`](SmemCyclicAccumulation::drain_at) for a partial over one window), which runs every
32/// round and synchronizes the cube after each; once every writer has, [`source`] holds the sum.
33/// Every unit of the cube reaches the drain, each plane with the partial and window it holds.
34///
35/// [`source`]: SmemCyclicAccumulation::source
36#[derive(CubeType)]
37pub struct SmemCyclicAccumulation<A: Numeric> {
38    values: Shared<[A]>,
39    /// The box the accumulator was opened over, which every round's sink is shaped like.
40    like: Tile,
41    /// Reads the sum back at the opened tile's element; valid only once every writer has drained.
42    pub source: Tile,
43    #[cube(comptime)]
44    writers: usize,
45}
46
47#[cube]
48impl Tile {
49    /// A shared-memory accumulator over this tile's box, summing in `A`, that `writers` writers
50    /// reach exclusively, one chunk each a round. Nothing is zeroed: every cell's first writer
51    /// replaces it. It is served in lines as wide as this tile's, so a partial drains into it and
52    /// its sum lands in this tile line for line.
53    pub fn smem_cyclic_accumulation<A: Numeric>(
54        &self,
55        #[comptime] writers: usize,
56    ) -> SmemCyclicAccumulation<A> {
57        let define!(Served) = self.element;
58        let space = comptime!(self.place.space.clone());
59        let width = self.vector_size();
60        let size!(W) = width;
61        let form = comptime!(StageForm::dense(
62            &space,
63            width,
64            StageStorage::Strided,
65            LineBytes(A::size().comptime() * width)
66        ));
67        let values = Shared::<[A]>::new_slice(comptime!(form.cells() * width));
68        // The previous open from this call site may still be read; let it finish first.
69        sync_cube();
70
71        let units = self.units();
72        let source = Memory::smem_backed::<Served>(
73            space,
74            width,
75            comptime!(FillUnits::cube(units)),
76            Backing::new_ReadCall(RawCall::of::<Served>(
77                ErasedTensor::<Served, ReadOnly>::of_smem_plain_load::<A, W>(&values),
78            )),
79            comptime!(Packing::Plain),
80            comptime!(form.clone()),
81            RuntimeMap::integral(comptime!(form.physical_rank())),
82            ComptimeOption::new_None(),
83            comptime!(Write::Replace),
84            comptime!(Overhang::Never),
85        );
86        SmemCyclicAccumulation::<A> {
87            values,
88            like: self.clone(),
89            source: source.nested_like(self),
90            writers,
91        }
92    }
93}
94
95#[cube]
96impl<A: Numeric> SmemCyclicAccumulation<A> {
97    /// The sink every writer drains into in `round`: it replaces each cell it is handed in round
98    /// zero and adds into it after.
99    fn sink(&self, round: usize) -> Tile {
100        let space = comptime!(self.like.place.space.clone());
101        let width = self.like.vector_size();
102        let size!(W) = width;
103        let form = comptime!(StageForm::dense(
104            &space,
105            width,
106            StageStorage::Strided,
107            LineBytes(A::size().comptime() * width)
108        ));
109        let units = self.like.units();
110        let sink = Memory::smem_backed::<A>(
111            space,
112            width,
113            comptime!(FillUnits::cube(units)),
114            Backing::new_WriteCall(RawCall::of::<A>(
115                ErasedTensor::<A, WriteOnly>::of_smem_cyclic::<W>(&self.values, round == 0),
116            )),
117            comptime!(Packing::Plain),
118            comptime!(form.clone()),
119            RuntimeMap::integral(comptime!(form.physical_rank())),
120            ComptimeOption::new_None(),
121            comptime!(Write::Exclusive(Schedule::Cyclic)),
122            comptime!(Overhang::Masked),
123        );
124        sink.nested_like(&self.like)
125    }
126
127    /// Drain `partial`, an accumulator over the whole box, every round of the schedule, the cube
128    /// synchronized after each: each unit lands its plane's fragments of the round's chunk as
129    /// `writer`, its plane's place among the writers. Once every writer of the cube has, the sum is
130    /// in [`source`](SmemCyclicAccumulation::source).
131    pub fn drain(&self, partial: &Tile, writer: usize) {
132        let define!(Served) = partial.element;
133        #[unroll]
134        for round in 0..self.writers {
135            let turn = (writer + round) % comptime!(self.writers);
136            partial.drained_chunk_into::<Served, A>(
137                &self.sink(round),
138                turn,
139                comptime!(self.writers),
140            );
141            sync_cube();
142        }
143    }
144
145    /// [`drain`](SmemCyclicAccumulation::drain) for a partial over one `window` of the box, as a
146    /// kernel that walks its planes itself holds one.
147    pub fn drain_at(&self, partial: &Tile, window: &Region, writer: usize) {
148        let define!(Served) = partial.element;
149        #[unroll]
150        for round in 0..self.writers {
151            let turn = (writer + round) % comptime!(self.writers);
152            partial.drained_chunk_into::<Served, A>(
153                &self.sink(round).at(window),
154                turn,
155                comptime!(self.writers),
156            );
157            sync_cube();
158        }
159    }
160}
161
162/// The erased tensors over a shared buffer of plain values.
163pub(crate) trait SmemCyclicSink<E: Numeric> {
164    /// The sink that writes into `values` in lines `W` scalars wide, replacing each cell where
165    /// `first` and adding into it otherwise.
166    fn of_smem_cyclic<W: Size>(_values: &Shared<[E]>, _first: bool) -> ErasedTensor<E, WriteOnly> {
167        unexpanded!()
168    }
169
170    fn __expand_of_smem_cyclic<W: Size>(
171        _scope: &Scope,
172        values: &<Shared<[E]> as CubeType>::ExpandType,
173        first: NativeExpand<bool>,
174    ) -> ErasedTensorExpand<E, WriteOnly> {
175        ErasedTensorExpand::new(SmemCyclic::<E, W> {
176            values: ExpandTypeClone::clone_unchecked(values),
177            first,
178            _e: PhantomData,
179        })
180    }
181}
182
183impl<E: Numeric> SmemCyclicSink<E> for ErasedTensor<E, WriteOnly> {}
184
185/// The source that loads out of a shared buffer of plain `A`, served as `T`.
186pub(crate) trait SmemPlainLoadSource<T: Numeric> {
187    /// The source that loads `values`, cast to `T`, in lines `W` scalars wide.
188    fn of_smem_plain_load<A: Numeric, W: Size>(_values: &Shared<[A]>) -> ErasedTensor<T, ReadOnly> {
189        unexpanded!()
190    }
191
192    fn __expand_of_smem_plain_load<A: Numeric, W: Size>(
193        _scope: &Scope,
194        values: &<Shared<[A]> as CubeType>::ExpandType,
195    ) -> ErasedTensorExpand<T, ReadOnly> {
196        ErasedTensorExpand::new(SmemPlainLoad::<A, T, W> {
197            values: ExpandTypeClone::clone_unchecked(values),
198            _t: PhantomData,
199        })
200    }
201}
202
203impl<T: Numeric> SmemPlainLoadSource<T> for ErasedTensor<T, ReadOnly> {}
204
205/// A backing that replaces each cell it is handed, or adds into it, in a shared buffer read in
206/// lines `W` scalars wide: a read and a write, not an atomic, so it declares [`WritesLines`]
207/// alone. The rounds keep one writer at a cell at a time.
208struct SmemCyclic<E: Numeric, W: Size> {
209    values: <Shared<[E]> as CubeType>::ExpandType,
210    first: NativeExpand<bool>,
211    _e: PhantomData<(E, W)>,
212}
213
214impl<E: Numeric, W: Size> ErasedTensorOperationsExpand<E> for SmemCyclic<E, W> {
215    fn __expand_vector_size_method(&self, scope: &Scope) -> VectorSize {
216        W::__expand_value(scope)
217    }
218
219    fn __expand_lines_method(&self, scope: &Scope) -> NativeExpand<usize> {
220        lines_of::expand::<E, W>(scope, &self.values)
221    }
222
223    fn __expand_write_line_method(
224        &mut self,
225        scope: &Scope,
226        index: NativeExpand<usize>,
227        value: ExpandValue,
228    ) {
229        cyclic_line::expand::<E, W>(scope, &mut self.values, index, value.into(), self.first);
230    }
231}
232
233impl<E: Numeric, W: Size> WritesLines<E> for SmemCyclic<E, W> {}
234
235/// A backing that loads out of a shared buffer of plain `A` in lines `W` scalars wide, casting
236/// to `T`.
237struct SmemPlainLoad<A: Numeric, T: Numeric, W: Size> {
238    values: <Shared<[A]> as CubeType>::ExpandType,
239    _t: PhantomData<(T, W)>,
240}
241
242impl<A: Numeric, T: Numeric, W: Size> ErasedTensorOperationsExpand<T> for SmemPlainLoad<A, T, W> {
243    fn __expand_vector_size_method(&self, scope: &Scope) -> VectorSize {
244        W::__expand_value(scope)
245    }
246
247    fn __expand_lines_method(&self, scope: &Scope) -> NativeExpand<usize> {
248        lines_of::expand::<A, W>(scope, &self.values)
249    }
250
251    fn __expand_read_line_method(&self, scope: &Scope, index: NativeExpand<usize>) -> ExpandValue {
252        load_plain_line::expand::<A, T, W>(scope, &self.values, index).expand
253    }
254}
255
256impl<A: Numeric, T: Numeric, W: Size> ReadsLines<T> for SmemPlainLoad<A, T, W> {}
257
258/// The lines `W` scalars wide `values` holds.
259#[cube]
260fn lines_of<E: Numeric, W: Size>(values: &Shared<[E]>) -> usize {
261    values.len() / W::value()
262}
263
264/// Land line `index`: each of its cells the value where `first`, the cell plus the value
265/// otherwise.
266#[cube]
267fn cyclic_line<E: Numeric, W: Size>(
268    values: &mut Shared<[E]>,
269    index: usize,
270    value: Vector<E, W>,
271    first: bool,
272) {
273    #[unroll]
274    for j in 0..W::value() {
275        let cell = index * W::value() + j;
276        if first {
277            values[cell] = value.extract(j);
278        } else {
279            values[cell] += value.extract(j);
280        }
281    }
282}
283
284/// Load line `index`, cast to `T`.
285#[cube]
286fn load_plain_line<A: Numeric, T: Numeric, W: Size>(
287    values: &Shared<[A]>,
288    index: usize,
289) -> Vector<T, W> {
290    let mut line = Vector::<T, W>::empty();
291    #[unroll]
292    for j in 0..W::value() {
293        line.insert(j, T::cast_from(values[index * W::value() + j]));
294    }
295    line
296}
cubecl_tile/tile/accumulator/smem_cyclic.rs

cubecl_tile/tile/accumulator/
smem_cyclic.rs
