cubek-tile 0.3.0

CubeK: Tile
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

//! `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<Acc: Numeric> Tile<Acc> {
    /// `c = fold(input)`: reduce `input` into `self` across the contracted axes.
    pub fn reduce_axis<In: Numeric>(&mut self, input: &Tile<In>, #[comptime] monoid: Monoid) {
        self.init_identity(monoid);
        self.reduce_axis_accumulate(input, monoid);
    }

    /// `c = fold(c, input)`: fold `input` into what `self` holds, which must be seeded.
    pub fn reduce_axis_accumulate<In: Numeric>(
        &mut self,
        input: &Tile<In>,
        #[comptime] monoid: Monoid,
    ) {
        reduce_leaf(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<Acc>,
    input: &Tile<In>,
    #[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(g, input, space, monoid);
        }
        TileKind::PlaneTile(t) => {
            reduce_plane_tile(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(&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<Acc>,
    input: &Tile<In>,
    #[comptime] acc_space: Space,
    #[comptime] monoid: Monoid,
) {
    match tile {
        PlaneTile::Registers(d) => {
            reduce::register_data(d, input, acc_space, monoid);
        }
        PlaneTile::Cmma(_) | PlaneTile::Mma(_) => {
            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"
            );
        }
    }
}