cubek-std 0.3.0-pre.2

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

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
/// Events that occur during the process of loading tiles to
/// registers and executing inner Tile Matmuls
pub enum StageEvent {
    /// Before any step
    Begin,
    /// After loading LHS
    LhsLoaded { current: u32, total: u32 },
    /// After X RHS loads are completed
    RhsLoaded { current: u32, total: u32 },
    /// After X tile matmul operations are completed
    TileMatmulCompleted { current: u32, total: u32 },
    /// After the last step
    Finish,
}

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

#[derive(CubeType)]
/// Use when there is no event listening to do
pub struct NoEvent {}

#[cube]
impl StageEventListener for NoEvent {
    fn on_event(_this: &mut Self, #[comptime] _event: StageEvent) {
        // Nothing to do
    }
}

impl Default for NoEvent {
    fn default() -> Self {
        Self::new()
    }
}

#[cube]
impl NoEvent {
    pub fn new() -> NoEvent {
        NoEvent {}
    }
}

#[derive(CubeType, Debug, Clone, Copy, PartialEq, Eq)]
/// Events emitted while writing per-partition tiles back to an output stage.
pub enum WriteEvent {
    /// Before any step.
    Begin,
    /// After each tile is stored into the stage.
    TileStored { tile: Coords2d },
    /// After the last step.
    Finish,
}

#[cube]
/// Callback invoked at each [`WriteEvent`].
pub trait WriteEventListener: CubeType {
    fn on_event(this: &mut Self, event: WriteEvent);
}