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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89

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

#[derive(Clone, CubeType)]
/// An iterator over global memory, advancing along k.
pub struct GlobalIterator<EI: CubePrimitive> {
    global_view: View<EI, Coords2d>,
    offset: RuntimeCell<u32>,
    /// The amount to advance by on each iteration
    step: u32,
    view_size: Coords2d,
    #[cube(comptime)]
    view_direction: ViewDirection,
    #[cube(comptime)]
    checked: bool,
}

unsafe impl<EG: CubePrimitive> Sync for GlobalIterator<EG> {}
unsafe impl<EG: CubePrimitive> Send for GlobalIterator<EG> {}

#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq, Default)]
pub enum ViewDirection {
    Row,
    Col,
    /// Cannot advance if direction is none
    #[default]
    None,
}

#[cube]
impl<EG: CubePrimitive> GlobalIterator<EG> {
    /// Instantiate a read iterator over the given global view, which should be sliced to the size
    /// of one `m`/`n` stage and the full range of `k` handled by this matmul instance.
    ///
    /// `step` is the amount advanced in `view_direction` each iteration.
    /// `checked` determines whether the slices should be created as checked or unchecked.
    pub fn new(
        global_view: View<EG, Coords2d>,
        step: u32,
        #[comptime] view_direction: ViewDirection,
        #[comptime] checked: bool,
    ) -> Self {
        let (size_row, size_col) = global_view.shape();
        let view_size = match view_direction {
            ViewDirection::Row => (step, size_col),
            ViewDirection::Col => (size_row, step),
            ViewDirection::None => (size_row, size_col),
        };

        GlobalIterator::<EG> {
            global_view,
            offset: RuntimeCell::new(0),
            step,
            view_size,
            view_direction,
            checked,
        }
    }

    /// Advance the view along the k dimension by a specified offset, `k_offset`.
    pub fn advance(&self) {
        self.offset.store(self.offset.read() + self.step);
    }

    /// Returns the current view slice of the iterator
    pub fn view(&self) -> View<EG, Coords2d> {
        let offset = match self.view_direction.comptime() {
            ViewDirection::Row => (self.offset.read(), 0u32),
            ViewDirection::Col => (0u32, self.offset.read()),
            ViewDirection::None => (0u32, 0u32).runtime(),
        };
        if self.checked.comptime() {
            self.global_view.slice(offset, self.view_size)
        } else {
            self.global_view.slice_unchecked(offset, self.view_size)
        }
    }

    /// Returns the vector size of the global view
    pub fn vector_size(&self) -> comptime_type!(VectorSize) {
        self.global_view.vector_size()
    }

    pub fn offset(&self) -> u32 {
        self.offset.read()
    }
}