cubecl-tile 0.12.0-pre.1

CubeCL Tile: a space cut into tiles, and the operations on them.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

//! `c.reduce_axis(input, monoid)` and its accumulating twin at a final tile.

use cubecl::prelude::*;

use super::leaf as reduce;
use crate::*;

#[cube]
impl Tile {
    /// `c = reduce(input)`: reduce `input` into `self` across the contracted axes.
    pub fn reduce_axis(&mut self, input: &Tile, #[comptime] monoid: Monoid) {
        let define!(Served) = self.element;
        self.init_identity::<Served>(monoid);
        self.reduce_axis_accumulate(input, monoid);
    }

    /// `self[c] = Σ_r lhs[c, r] · rhs[c, r]`, `r` the axes `lhs` and `rhs` span beyond `self`'s:
    /// each cell the dot product of its run of both, as attention's backward pass sums its output
    /// against the output's gradient along the value dim. `self` is a window of memory a unit or a
    /// plane holds ([`Tile::set_cells`]), every unit of the holder summing every cell; `lhs` and
    /// `rhs` span one space, bound at one width.
    pub fn dot<E: Float>(&mut self, lhs: &Tile, rhs: &Tile) {
        let (width, rhs_width) = (lhs.vector_size(), rhs.vector_size());
        comptime!(assert!(
            lhs.place.space == rhs.place.space && width == rhs_width,
            "Tile::dot: both factors span one space at one width, not {:?} at {width} and {:?} \
             at {rhs_width}",
            lhs.place.space,
            rhs.place.space,
        ));
        let define!(Lhs) = lhs.read_element();
        let define!(Rhs) = rhs.read_element();
        let space = comptime!(self.place.space.clone());
        let cells = comptime!(space.cells());
        let size!(V) = width;
        let values = reduce::dot_cells::<E, Lhs, Rhs, V>(lhs, rhs, space);
        self.set_cells(&values, cells);
    }

    /// `c = reduce(c, input)`: reduce `input` into what `self` holds, which must be seeded.
    pub fn reduce_axis_accumulate(&mut self, input: &Tile, #[comptime] monoid: Monoid) {
        let define!(Served) = self.element;
        let define!(Source) = input.read_element();
        reduce_leaf::<Served, Source>(self, input, monoid)
    }
}

/// Dispatches to the register nest by the accumulator's form.
#[cube]
pub(crate) fn reduce_leaf<Acc: Numeric, In: Numeric>(
    acc: &mut Tile,
    input: &Tile,
    #[comptime] monoid: Monoid,
) {
    let input_space = comptime!(input.place.space.clone());
    let vector_size = input.vector_size();
    comptime!(assert!(
        input_space
            .extent_at(input_space.rank() - 1)
            .is_multiple_of(vector_size),
        "reduce: the input's innermost extent must be divisible by its vector size"
    ));

    let space = comptime!(acc.place.space.clone());
    match &mut acc.kind {
        TileKind::Memory(g) => {
            reduce::memory::<Acc, In>(g, input, space, monoid);
        }
        TileKind::PlaneTile(t) => {
            reduce_plane_tile::<Acc, In>(t, input, space, monoid);
        }
        TileKind::PlanePartition(p) => {
            comptime!(assert!(
                p.m_tiles == 1 && p.n_tiles == 1,
                "reduce_leaf: a multi-tile partition must be contracted at its partition level"
            ));
            let mut t = p.at(0usize, 0usize);
            reduce_plane_tile::<Acc, In>(&mut t, input, space, monoid);
        }
        TileKind::TmaGmem(_) => panic!("reduce: a tma source is not an accumulator sink"),
        TileKind::Procedural(_) | TileKind::Lines(_) => {
            panic!("reduce: a procedural tile and the plane's units are not an accumulator sink")
        }
    }
}

#[cube]
fn reduce_plane_tile<Acc: Numeric, In: Numeric>(
    tile: &mut PlaneTile,
    input: &Tile,
    #[comptime] acc_space: Space,
    #[comptime] monoid: Monoid,
) {
    match tile {
        PlaneTile::Registers(d) => {
            reduce::register_data::<Acc, In>(d, input, acc_space, monoid);
        }
        PlaneTile::Cmma(_) | PlaneTile::Mma(_) | PlaneTile::Wgmma(_) => {
            panic!(
                "reduce: a hardware mma fragment scatters its rows across units in a \
                 layout the elementwise walk cannot address; reduce into a register, \
                 Gmem or Smem accumulator instead"
            );
        }
    }
}