cubecl_tile/tile/accumulator/
planes_output.rs

1//! A cube's output as its planes reach it: written by each plane where every plane owns its box of
2//! it, or summed across the planes that share a box before the cube writes it.
3
4use cubecl::prelude::*;
5
6use crate::tile::accumulator::smem_cyclic::SmemCyclicAccumulation;
7use crate::*;
8
9/// The output of one cube, reached by its planes.
10///
11/// Where the partitioning hands every plane a box of the output of its own, a plane's sum drains
12/// into its box. Where it splits an axis the output lacks across the planes, the contracted axis of
13/// a sum they each hold part of, the planes holding the same box drain into a shared-memory
14/// accumulator on the cyclic schedule ([`SmemCyclicAccumulation`]) and [`write`] stores the sum.
15/// The tile says which from the partitioning it was cut by: the caller states nothing.
16///
17/// [`write`]: PlanesOutput::write
18#[derive(CubeType)]
19pub struct PlanesOutput<A: Numeric> {
20    out: Tile,
21    merged: ComptimeOption<SmemCyclicAccumulation<A>>,
22    /// The axes the output lacks that the planes are split along, each with its count of planes.
23    #[cube(comptime)]
24    sharing: Vec<(Axis, usize)>,
25}
26
27#[cube]
28impl Tile {
29    /// This tile, `cube`'s window of an output, as its planes reach it, a sum shared by planes
30    /// added in `A`.
31    pub fn planes_output<A: Numeric>(&self, cube: &Region) -> PlanesOutput<A> {
32        let space = comptime!(self.place.space.clone());
33        let sharing = comptime!(cube.path.planes_along(|axis| !space.contains(axis)));
34        let writers = comptime!(sharing.iter().map(|&(_, count)| count).product::<usize>());
35        let merged = if comptime!(writers == 1) {
36            ComptimeOption::new_None()
37        } else {
38            ComptimeOption::new_Some(self.smem_cyclic_accumulation::<A>(writers))
39        };
40        PlanesOutput::<A> {
41            out: self.clone(),
42            merged,
43            sharing,
44        }
45    }
46}
47
48#[cube]
49impl<A: Numeric> PlanesOutput<A> {
50    /// Drain `partial`, the sum `plane` holds over its box of the output. Where planes share the
51    /// box it is a meeting of the cube: every unit of every plane drains, each with its plane's
52    /// partial.
53    pub fn drain(&self, partial: &Tile, plane: &Region) {
54        #[comptime]
55        match &self.merged {
56            ComptimeOption::None => partial.drained_into(&self.out.at(plane)),
57            ComptimeOption::Some(merged) => {
58                merged.drain_at(partial, plane, PlanesOutput::<A>::writer(self, plane))
59            }
60        }
61    }
62
63    /// Store the sum the planes sharing a box drained into it, once every plane has; nothing where
64    /// each plane drained into its own box.
65    pub fn write(&mut self) {
66        #[comptime]
67        match &self.merged {
68            ComptimeOption::None => {}
69            ComptimeOption::Some(merged) => self.out.copy_from(&merged.source),
70        }
71    }
72
73    /// `plane`'s place among the planes sharing its box.
74    fn writer(&self, plane: &Region) -> usize {
75        let mut writer = 0usize;
76        #[unroll]
77        for i in 0..comptime!(self.sharing.len()) {
78            let (axis, count) = comptime!(self.sharing[i]);
79            writer = writer * count + plane.coord(axis);
80        }
81        writer
82    }
83}
cubecl_tile/tile/accumulator/planes_output.rs

cubecl_tile/tile/accumulator/
planes_output.rs
