cubecl_tile/launch/arg/
accumulate.rs

1//! An output several instances accumulate into, as a single launch argument.
2
3use cubecl::define_size;
4use cubecl::ir::{ExpandValue, VectorSize};
5use cubecl::prelude::*;
6use cubecl::unexpanded;
7use cubecl_std::tensor::{
8    ErasedTensor, ErasedTensorExpand, ErasedTensorOperationsExpand, WriteOnly, WritesLines,
9};
10
11use crate::*;
12
13/// An output several instances accumulate into, as a single launch argument, in lines of
14/// `width`. The buffer must arrive holding zero: the atomic only ever adds.
15#[derive(CubeType, CubeLaunch)]
16pub struct AccumulateArg<'a, E: Numeric> {
17    pub tensor: &'a Tensor<Atomic<E>>,
18    /// The width a line is added in: the buffer's own elements are scalar.
19    #[cube(comptime)]
20    pub width: usize,
21    #[cube(comptime)]
22    pub spec: TileSpec,
23}
24
25#[cube]
26impl<'a, E: Numeric> AccumulateArg<'a, E> {
27    /// Serve the output as a [`Tile`] that accumulates into it in lines of its width.
28    pub fn tile(&self, partitioning: &Partitioning) -> Tile {
29        let space = comptime!(partitioning.clone());
30        // An atomic element is scalar, so these strides are already in scalars.
31        let geometry = RuntimeGeometry::of_tensor::<Atomic<E>>(
32            self.tensor,
33            comptime!(self.spec.projection.physical_rank()),
34        );
35        let sink = atomic_sink::<E>(self.tensor, comptime!(self.width));
36        GlobalOperand::sink(
37            sink,
38            geometry,
39            comptime!(self.width),
40            comptime!(space.space().clone()),
41            comptime!(self.spec.clone()),
42            Write::Accumulate,
43        )
44        .tile::<E>(comptime!(space.levels().to_vec()))
45    }
46}
47
48/// The erased tensor over an atomic buffer, accumulating in lines of `width`.
49fn atomic_sink<E: Numeric>(
50    _values: &Tensor<Atomic<E>>,
51    _width: usize,
52) -> ErasedTensor<E, WriteOnly> {
53    unexpanded!()
54}
55
56mod atomic_sink {
57    use super::*;
58
59    pub(super) fn expand<E: Numeric>(
60        _scope: &Scope,
61        values: &<Tensor<Atomic<E>> as CubeType>::ExpandType,
62        width: usize,
63    ) -> ErasedTensorExpand<E, WriteOnly> {
64        ErasedTensorExpand::new(AtomicAccumulate::<E> {
65            values: ExpandTypeClone::clone_unchecked(values),
66            width,
67        })
68    }
69}
70
71/// A write-only backing that accumulates into an `Atomic<E>` buffer in lines of `width`.
72struct AtomicAccumulate<E: Numeric> {
73    values: <Tensor<Atomic<E>> as CubeType>::ExpandType,
74    width: usize,
75}
76
77// The line a write hands over, typed at the width it is added in: registered where the write
78// expands and read nowhere after.
79define_size!(Line);
80
81impl<E: Numeric> ErasedTensorOperationsExpand<E> for AtomicAccumulate<E> {
82    fn __expand_vector_size_method(&self, _scope: &Scope) -> VectorSize {
83        self.width
84    }
85
86    /// In lines of `width`, as the trait counts them, off a buffer whose own elements are scalar:
87    /// the buffer's, not the tensor's, since a pitched buffer holds cells past its shape's product.
88    fn __expand_lines_method(&self, scope: &Scope) -> NativeExpand<usize> {
89        let scalars = self.values.__expand_buffer_len_method(scope);
90        let width = self.width.__expand_runtime_method(scope);
91        scalars.__expand_div_method(scope, width)
92    }
93
94    fn __expand_write_line_method(
95        &mut self,
96        scope: &Scope,
97        index: NativeExpand<usize>,
98        value: ExpandValue,
99    ) {
100        scope.register_size::<Line>(self.width);
101        accumulate_line::expand::<E, Line>(scope, &self.values, index, value.into());
102    }
103}
104
105impl<E: Numeric> WritesLines<E> for AtomicAccumulate<E> {}
106
107/// Accumulate one line into the buffer: `N` scalar adds at the line's own offset.
108#[cube]
109fn accumulate_line<E: Numeric, N: Size>(
110    values: &Tensor<Atomic<E>>,
111    index: usize,
112    value: Vector<E, N>,
113) {
114    let base = index * N::value();
115    #[unroll]
116    for k in 0..N::value() {
117        values[base + k].fetch_add(value.extract(k));
118    }
119}
cubecl_tile/launch/arg/accumulate.rs

cubecl_tile/launch/arg/
accumulate.rs
