cubecl-matmul 0.8.1

CubeCL Matrix Multiplication Kernels Engine
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

use cubecl::prelude::*;
use cubecl_core as cubecl;
use cubecl_std::tensor::layout::{Coords2d, Layout, LayoutExpand};

use crate::components::{MatrixLayout, global::memory::GlobalMemoryConfig};

pub type TiledCoords = (Coords2d, u32);

/// Tiling mapping on a 2D layout. Unit offset is translated to a 2D offset within the tile.
#[derive(CubeType)]
pub struct TiledLayout {
    #[cube(comptime)]
    config: GlobalMemoryConfig,
}

#[cube]
impl TiledLayout {
    pub fn new(#[comptime] config: GlobalMemoryConfig) -> Self {
        TiledLayout { config }
    }
}

#[cube]
impl Layout for TiledLayout {
    type Coordinates = TiledCoords;
    type SourceCoordinates = Coords2d;

    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
        let (tile, unit_pos) = pos;
        let (tile_row, tile_col) = tile;

        let tile_size_row = comptime![self.config.elements_in_tile_row];
        let tile_size_col = comptime![self.config.elements_in_tile_col];

        let view_tile_row = tile_row * tile_size_row;
        let view_tile_col = tile_col * tile_size_col;

        let (unit_row, unit_col) = match comptime![self.config.matrix_layout] {
            MatrixLayout::RowMajor => (unit_pos / tile_size_col, unit_pos % tile_size_col),
            MatrixLayout::ColMajor => (unit_pos % tile_size_row, unit_pos / tile_size_row),
        };

        (view_tile_row + unit_row, view_tile_col + unit_col)
    }

    fn shape(&self) -> Self::Coordinates {
        let tile_size_row = comptime![self.config.elements_in_tile_row];
        let tile_size_col = comptime![self.config.elements_in_tile_col];

        let tiles_row = comptime![self.config.elements_in_stage_row / tile_size_row];
        let tiles_col = comptime![self.config.elements_in_stage_col / tile_size_col];
        let tile_size = tile_size_row * tile_size_col;
        ((tiles_row, tiles_col), tile_size).runtime()
    }

    fn is_in_bounds(&self, _pos: Self::Coordinates) -> bool {
        // Bounds checking should be handled by underlying layout
        true.runtime()
    }

    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool) {
        (self.to_source_pos(pos), self.is_in_bounds(pos))
    }
}