cubek-matmul 0.2.0

CubeK: Matrix Multiplication Kernels
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

use cubecl::{prelude::*, std::tensor::layout::*};

#[derive(CubeType, Clone, Copy)]
pub struct VecLayout {
    batch: usize,
    shape: Coords1d,
}

#[cube]
impl VecLayout {
    pub fn new(batch: usize, shape: Coords1d) -> Self {
        VecLayout { batch, shape }
    }
}

#[cube]
impl Layout for VecLayout {
    type Coordinates = Coords1d;
    type SourceCoordinates = (usize, u32, u32);

    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
        (self.batch, 0, pos as u32)
    }

    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
        pos < self.shape
    }

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

    fn shape(&self) -> Self::Coordinates {
        self.shape
    }
}

#[derive(CubeType, Clone, Copy)]
pub struct MatLayout {
    batch: usize,
    shape: Coords2d,
}

#[cube]
impl MatLayout {
    pub fn new(batch: usize, shape: Coords2d) -> Self {
        MatLayout { batch, shape }
    }
}

#[cube]
impl Layout for MatLayout {
    type Coordinates = Coords2d;
    type SourceCoordinates = (usize, u32, u32);

    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
        (self.batch, pos.0, pos.1)
    }

    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
        pos.0 < self.shape.0 && pos.1 < self.shape.1
    }

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

    fn shape(&self) -> Self::Coordinates {
        self.shape
    }
}