cubecl_tile/ops/reduce/
base.rs

1//! `c.reduce_axis(input, monoid)` and its accumulating twin at a final tile.
2
3use cubecl::prelude::*;
4
5use super::leaf as reduce;
6use crate::*;
7
8#[cube]
9impl Tile {
10    /// `c = reduce(input)`: reduce `input` into `self` across the contracted axes.
11    pub fn reduce_axis(&mut self, input: &Tile, #[comptime] monoid: Monoid) {
12        let define!(Served) = self.element;
13        self.init_identity::<Served>(monoid);
14        self.reduce_axis_accumulate(input, monoid);
15    }
16
17    /// `self[c] = Σ_r lhs[c, r] · rhs[c, r]`, `r` the axes `lhs` and `rhs` span beyond `self`'s:
18    /// each cell the dot product of its run of both, as attention's backward pass sums its output
19    /// against the output's gradient along the value dim. `self` is a window of memory a unit or a
20    /// plane holds ([`Tile::set_cells`]), every unit of the holder summing every cell; `lhs` and
21    /// `rhs` span one space, bound at one width.
22    pub fn dot<E: Float>(&mut self, lhs: &Tile, rhs: &Tile) {
23        let (width, rhs_width) = (lhs.vector_size(), rhs.vector_size());
24        comptime!(assert!(
25            lhs.place.space == rhs.place.space && width == rhs_width,
26            "Tile::dot: both factors span one space at one width, not {:?} at {width} and {:?} \
27             at {rhs_width}",
28            lhs.place.space,
29            rhs.place.space,
30        ));
31        let define!(Lhs) = lhs.read_element();
32        let define!(Rhs) = rhs.read_element();
33        let space = comptime!(self.place.space.clone());
34        let cells = comptime!(space.cells());
35        let size!(V) = width;
36        let values = reduce::dot_cells::<E, Lhs, Rhs, V>(lhs, rhs, space);
37        self.set_cells(&values, cells);
38    }
39
40    /// `c = reduce(c, input)`: reduce `input` into what `self` holds, which must be seeded.
41    pub fn reduce_axis_accumulate(&mut self, input: &Tile, #[comptime] monoid: Monoid) {
42        let define!(Served) = self.element;
43        let define!(Source) = input.read_element();
44        reduce_leaf::<Served, Source>(self, input, monoid)
45    }
46}
47
48/// Dispatches to the register nest by the accumulator's form.
49#[cube]
50pub(crate) fn reduce_leaf<Acc: Numeric, In: Numeric>(
51    acc: &mut Tile,
52    input: &Tile,
53    #[comptime] monoid: Monoid,
54) {
55    let input_space = comptime!(input.place.space.clone());
56    let vector_size = input.vector_size();
57    comptime!(assert!(
58        input_space
59            .extent_at(input_space.rank() - 1)
60            .is_multiple_of(vector_size),
61        "reduce: the input's innermost extent must be divisible by its vector size"
62    ));
63
64    let space = comptime!(acc.place.space.clone());
65    match &mut acc.kind {
66        TileKind::Memory(g) => {
67            reduce::memory::<Acc, In>(g, input, space, monoid);
68        }
69        TileKind::PlaneTile(t) => {
70            reduce_plane_tile::<Acc, In>(t, input, space, monoid);
71        }
72        TileKind::PlanePartition(p) => {
73            comptime!(assert!(
74                p.m_tiles == 1 && p.n_tiles == 1,
75                "reduce_leaf: a multi-tile partition must be contracted at its partition level"
76            ));
77            let mut t = p.at(0usize, 0usize);
78            reduce_plane_tile::<Acc, In>(&mut t, input, space, monoid);
79        }
80        TileKind::TmaGmem(_) => panic!("reduce: a tma source is not an accumulator sink"),
81        TileKind::Procedural(_) | TileKind::Lines(_) => {
82            panic!("reduce: a procedural tile and the plane's units are not an accumulator sink")
83        }
84    }
85}
86
87#[cube]
88fn reduce_plane_tile<Acc: Numeric, In: Numeric>(
89    tile: &mut PlaneTile,
90    input: &Tile,
91    #[comptime] acc_space: Space,
92    #[comptime] monoid: Monoid,
93) {
94    match tile {
95        PlaneTile::Registers(d) => {
96            reduce::register_data::<Acc, In>(d, input, acc_space, monoid);
97        }
98        PlaneTile::Cmma(_) | PlaneTile::Mma(_) | PlaneTile::Wgmma(_) => {
99            panic!(
100                "reduce: a hardware mma fragment scatters its rows across units in a \
101                 layout the elementwise walk cannot address; reduce into a register, \
102                 Gmem or Smem accumulator instead"
103            );
104        }
105    }
106}
cubecl_tile/ops/reduce/base.rs

cubecl_tile/ops/reduce/
base.rs
