cubecl_tile/tile/accumulator/
smem_slots.rs

1//! A shared-memory destination where each writer of a cube lands its partial in a slot of its own,
2//! and a tile that reads the slots' sum back.
3//!
4//! Every writer replaces its own copy of the box at once, one cube barrier follows, and a read adds
5//! the copies cell by cell, in writer order. No rounds, no atomics, no zeroing, and each cell takes
6//! its partials in a fixed order, so the sum is the same bits from run to run. It holds a copy of
7//! the box a writer: little where the box is, as a decode's output is, against the rounds a
8//! [`SmemCyclicAccumulation`] synchronizes, which a box of fewer fragments than writers lands one
9//! writer at a time.
10
11use core::marker::PhantomData;
12
13use cubecl::ir::{ExpandValue, VectorSize};
14use cubecl::prelude::*;
15use cubecl::unexpanded;
16use cubecl_std::tensor::{
17    ErasedTensor, ErasedTensorExpand, ErasedTensorOperationsExpand, ReadsLines, WriteOnly,
18    WritesLines,
19};
20
21use crate::*;
22
23/// A cube's shared-memory accumulator over one box, summing in `A` and read back as `T`, with a
24/// slot of its own for each of `writers` writers.
25///
26/// A writer drains its partial with [`drain`](SmemSlotsAccumulation::drain), which synchronizes
27/// the cube once; once every writer has, [`source`] reads the sum. Every unit of the cube reaches
28/// the drain, each plane with the partial it holds.
29///
30/// [`source`]: SmemSlotsAccumulation::source
31#[derive(CubeType)]
32pub struct SmemSlotsAccumulation<A: Numeric> {
33    values: Shared<[A]>,
34    /// The box the accumulator was opened over, which every writer's sink is shaped like.
35    like: Tile,
36    /// Reads the sum of the slots back at the opened tile's element; valid only once every
37    /// writer has drained.
38    pub source: Tile,
39}
40
41#[cube]
42impl Tile {
43    /// A shared-memory accumulator over this tile's box, summing in `A`, with a slot for each of
44    /// `writers` writers. Nothing is zeroed: every writer replaces its whole slot. It is served in
45    /// lines as wide as this tile's, so a partial drains into it and its sum lands in this tile
46    /// line for line.
47    pub fn smem_slots_accumulation<A: Numeric>(
48        &self,
49        #[comptime] writers: usize,
50    ) -> SmemSlotsAccumulation<A> {
51        let define!(Served) = self.element;
52        let space = comptime!(self.place.space.clone());
53        let width = self.vector_size();
54        let size!(W) = width;
55        let form = comptime!(StageForm::dense(
56            &space,
57            width,
58            StageStorage::Strided,
59            LineBytes(A::size().comptime() * width)
60        ));
61        let values = Shared::<[A]>::new_slice(comptime!(form.cells() * width * writers));
62        // The previous open from this call site may still be read; let it finish first.
63        sync_cube();
64
65        let units = self.units();
66        let source = Memory::smem_backed::<Served>(
67            space,
68            width,
69            comptime!(FillUnits::cube(units)),
70            Backing::new_ReadCall(RawCall::of::<Served>(
71                ErasedTensor::<Served, ReadOnly>::of_smem_slots_sum::<A, W>(&values, writers),
72            )),
73            comptime!(Packing::Plain),
74            comptime!(form.clone()),
75            RuntimeMap::integral(comptime!(form.physical_rank())),
76            ComptimeOption::new_None(),
77            comptime!(Write::Replace),
78            comptime!(Overhang::Never),
79        );
80        SmemSlotsAccumulation::<A> {
81            values,
82            like: self.clone(),
83            source: source.nested_like(self),
84        }
85    }
86}
87
88#[cube]
89impl<A: Numeric> SmemSlotsAccumulation<A> {
90    /// The sink `writer` drains into: its own slot, each cell it is handed replaced.
91    fn sink(&self, writer: usize) -> Tile {
92        let space = comptime!(self.like.place.space.clone());
93        let width = self.like.vector_size();
94        let size!(W) = width;
95        let form = comptime!(StageForm::dense(
96            &space,
97            width,
98            StageStorage::Strided,
99            LineBytes(A::size().comptime() * width)
100        ));
101        let slot = comptime!(form.cells() * width);
102        let units = self.like.units();
103        let sink = Memory::smem_backed::<A>(
104            space,
105            width,
106            comptime!(FillUnits::cube(units)),
107            Backing::new_WriteCall(RawCall::of::<A>(
108                ErasedTensor::<A, WriteOnly>::of_smem_slot::<W>(&self.values, writer * slot, slot),
109            )),
110            comptime!(Packing::Plain),
111            comptime!(form.clone()),
112            RuntimeMap::integral(comptime!(form.physical_rank())),
113            ComptimeOption::new_None(),
114            comptime!(Write::Replace),
115            comptime!(Overhang::Masked),
116        );
117        sink.nested_like(&self.like)
118    }
119
120    /// Drain `partial`, an accumulator over the whole box, into `writer`'s slot, the cube
121    /// synchronized after. Once every writer of the cube has, the sum is in
122    /// [`source`](SmemSlotsAccumulation::source).
123    pub fn drain(&self, partial: &Tile, writer: usize) {
124        partial.drained_into(&self.sink(writer));
125        sync_cube();
126    }
127
128    /// Replace line `index` of the slot starting at `start` with `value`.
129    #[allow(dead_code)] // Reached through its expand, from [`SmemSlot`].
130    fn store_line<W: Size>(
131        values: &mut Shared<[A]>,
132        start: usize,
133        index: usize,
134        value: Vector<A, W>,
135    ) {
136        #[unroll]
137        for j in 0..W::value() {
138            values[start + index * W::value() + j] = value.extract(j);
139        }
140    }
141
142    /// Line `index` summed over the `writers` slots, in writer order, cast to `T`.
143    #[allow(dead_code)] // Reached through its expand, from [`SmemSlotsSum`].
144    fn sum_line<T: Numeric, W: Size>(
145        values: &Shared<[A]>,
146        #[comptime] writers: usize,
147        index: usize,
148    ) -> Vector<T, W> {
149        let slot = values.len() / writers;
150        let mut line = Vector::<T, W>::empty();
151        #[unroll]
152        for j in 0..W::value() {
153            let cell = index * W::value() + j;
154            let mut sum = values[cell];
155            #[unroll]
156            for writer in 1..writers {
157                sum += values[writer * slot + cell];
158            }
159            line.insert(j, T::cast_from(sum));
160        }
161        line
162    }
163}
164
165/// The erased tensors over a shared buffer of slots.
166pub(crate) trait SmemSlotsSink<E: Numeric> {
167    /// The sink that replaces cells of the slot `len` scalars long starting at `start` in
168    /// `values`, in lines `W` scalars wide.
169    fn of_smem_slot<W: Size>(
170        _values: &Shared<[E]>,
171        _start: usize,
172        _len: usize,
173    ) -> ErasedTensor<E, WriteOnly> {
174        unexpanded!()
175    }
176
177    fn __expand_of_smem_slot<W: Size>(
178        _scope: &Scope,
179        values: &<Shared<[E]> as CubeType>::ExpandType,
180        start: NativeExpand<usize>,
181        len: usize,
182    ) -> ErasedTensorExpand<E, WriteOnly> {
183        ErasedTensorExpand::new(SmemSlot::<E, W> {
184            values: ExpandTypeClone::clone_unchecked(values),
185            start,
186            len,
187            _e: PhantomData,
188        })
189    }
190}
191
192impl<E: Numeric> SmemSlotsSink<E> for ErasedTensor<E, WriteOnly> {}
193
194/// The source that sums the slots of a shared buffer of plain `A`, served as `T`.
195pub(crate) trait SmemSlotsSumSource<T: Numeric> {
196    /// The source that sums the `writers` slots of `values`, cast to `T`, in lines `W` scalars
197    /// wide.
198    fn of_smem_slots_sum<A: Numeric, W: Size>(
199        _values: &Shared<[A]>,
200        _writers: usize,
201    ) -> ErasedTensor<T, ReadOnly> {
202        unexpanded!()
203    }
204
205    fn __expand_of_smem_slots_sum<A: Numeric, W: Size>(
206        _scope: &Scope,
207        values: &<Shared<[A]> as CubeType>::ExpandType,
208        writers: usize,
209    ) -> ErasedTensorExpand<T, ReadOnly> {
210        ErasedTensorExpand::new(SmemSlotsSum::<A, T, W> {
211            values: ExpandTypeClone::clone_unchecked(values),
212            writers,
213            _t: PhantomData,
214        })
215    }
216}
217
218impl<T: Numeric> SmemSlotsSumSource<T> for ErasedTensor<T, ReadOnly> {}
219
220/// A backing that replaces cells of one slot of a shared buffer, in lines `W` scalars wide.
221struct SmemSlot<E: Numeric, W: Size> {
222    values: <Shared<[E]> as CubeType>::ExpandType,
223    start: NativeExpand<usize>,
224    /// Scalars a slot holds.
225    len: usize,
226    _e: PhantomData<(E, W)>,
227}
228
229impl<E: Numeric, W: Size> ErasedTensorOperationsExpand<E> for SmemSlot<E, W> {
230    fn __expand_vector_size_method(&self, scope: &Scope) -> VectorSize {
231        W::__expand_value(scope)
232    }
233
234    fn __expand_lines_method(&self, scope: &Scope) -> NativeExpand<usize> {
235        (self.len / W::__expand_value(scope)).__expand_runtime_method(scope)
236    }
237
238    fn __expand_write_line_method(
239        &mut self,
240        scope: &Scope,
241        index: NativeExpand<usize>,
242        value: ExpandValue,
243    ) {
244        SmemSlotsAccumulation::<E>::__expand_store_line::<W>(
245            scope,
246            &mut self.values,
247            self.start,
248            index,
249            value.into(),
250        );
251    }
252}
253
254impl<E: Numeric, W: Size> WritesLines<E> for SmemSlot<E, W> {}
255
256/// A backing that sums the slots of a shared buffer of plain `A` in lines `W` scalars wide,
257/// casting to `T`.
258struct SmemSlotsSum<A: Numeric, T: Numeric, W: Size> {
259    values: <Shared<[A]> as CubeType>::ExpandType,
260    writers: usize,
261    _t: PhantomData<(T, W)>,
262}
263
264impl<A: Numeric, T: Numeric, W: Size> ErasedTensorOperationsExpand<T> for SmemSlotsSum<A, T, W> {
265    fn __expand_vector_size_method(&self, scope: &Scope) -> VectorSize {
266        W::__expand_value(scope)
267    }
268
269    fn __expand_lines_method(&self, scope: &Scope) -> NativeExpand<usize> {
270        let scalars = self.values.__expand_len_method(scope);
271        let per_line = (self.writers * W::__expand_value(scope)).__expand_runtime_method(scope);
272        scalars.__expand_div_method(scope, per_line)
273    }
274
275    fn __expand_read_line_method(&self, scope: &Scope, index: NativeExpand<usize>) -> ExpandValue {
276        SmemSlotsAccumulation::<A>::__expand_sum_line::<T, W>(
277            scope,
278            &self.values,
279            self.writers,
280            index,
281        )
282        .expand
283    }
284}
285
286impl<A: Numeric, T: Numeric, W: Size> ReadsLines<T> for SmemSlotsSum<A, T, W> {}
cubecl_tile/tile/accumulator/smem_slots.rs

cubecl_tile/tile/accumulator/
smem_slots.rs
