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

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

#[derive(CubeType, Debug, Clone, Copy, PartialEq, Eq)]
/// Events that occur during the process of storing tiles to
/// a stage and executing writes
pub enum WriteEvent {
    /// Before any step
    Begin,
    /// After each tile is stored into the stage
    TileStored { tile: Coords2d },
    /// After the last step
    Finish,
}

#[cube]
/// Function that is called at each [WriteEvent]
pub trait WriteEventListener: CubeType {
    fn on_event(this: &mut Self, event: WriteEvent);
}