cubek_std/tile/
event.rs

1use cubecl::{prelude::*, std::tensor::layout::Coords2d};
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4/// Events that occur during the process of loading tiles to
5/// registers and executing inner Tile Matmuls
6pub enum StageEvent {
7    /// Before any step
8    Begin,
9    /// After loading LHS
10    LhsLoaded { current: u32, total: u32 },
11    /// After X RHS loads are completed
12    RhsLoaded { current: u32, total: u32 },
13    /// After X tile matmul operations are completed
14    TileMatmulCompleted { current: u32, total: u32 },
15    /// After the last step
16    Finish,
17}
18
19#[cube]
20/// Function that is called at each [StageEvent]
21pub trait StageEventListener: CubeType {
22    fn on_event(this: &mut Self, #[comptime] event: StageEvent);
23}
24
25#[derive(CubeType)]
26/// Use when there is no event listening to do
27pub struct NoEvent {}
28
29#[cube]
30impl StageEventListener for NoEvent {
31    fn on_event(_this: &mut Self, #[comptime] _event: StageEvent) {
32        // Nothing to do
33    }
34}
35
36impl Default for NoEvent {
37    fn default() -> Self {
38        Self::new()
39    }
40}
41
42#[cube]
43impl NoEvent {
44    pub fn new() -> NoEvent {
45        NoEvent {}
46    }
47}
48
49#[derive(CubeType, Debug, Clone, Copy, PartialEq, Eq)]
50/// Events emitted while writing per-partition tiles back to an output stage.
51pub enum WriteEvent {
52    /// Before any step.
53    Begin,
54    /// After each tile is stored into the stage.
55    TileStored { tile: Coords2d },
56    /// After the last step.
57    Finish,
58}
59
60#[cube]
61/// Callback invoked at each [`WriteEvent`].
62pub trait WriteEventListener: CubeType {
63    fn on_event(this: &mut Self, event: WriteEvent);
64}
cubek_std/tile/event.rs

cubek_std/tile/
event.rs
