cubecl-matmul 0.8.1

CubeCL Matrix Multiplication Kernels Engine
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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

use cubecl_core as cubecl;
use cubecl_core::prelude::*;
use cubecl_std::tensor::{View, layout::Coords2d};

use crate::components::{
    MatrixPrecision, StageIdent,
    global::{
        GlobalWriter, GlobalWriterFamily, PartitionedStage, PartitionedStageFamily, WriteEvent,
        WriteEventExpand, WriteEventListener,
        memory::GlobalMemoryConfig,
        read::tiled::{TiledCoords, TiledLayout},
    },
    stage::{StageConfig, StageMemoryConfig, StagePartitioner, UnitPartitioner},
};

#[derive(CubeType)]
/// Writes tiles from out shared memory to output global memory
/// using a unit for each tile
pub struct UnitWriter<IP: MatrixPrecision> {
    global: View<Line<IP::Global>, TiledCoords, ReadWrite>,
    stage: PartitionedStage<IP::Stage>,

    #[cube(comptime)]
    config: GlobalMemoryConfig,
}

#[cube]
impl<IP: MatrixPrecision> UnitWriter<IP> {
    pub fn new<S: StageConfig>(
        global: View<Line<IP::Global>, Coords2d, ReadWrite>,
        #[comptime] global_config: GlobalMemoryConfig,
        #[comptime] stage_config: S,
    ) -> Self {
        let stage_mem_config = comptime![stage_memory_config(stage_config)];
        let stage = PartitionedStage::new(tile_pos::<S>(stage_config), stage_mem_config);

        UnitWriter::<IP> {
            global: global.view_mut(TiledLayout::new(global_config)),
            stage,
            config: global_config,
        }
    }

    fn write(&mut self, tile: Coords2d) {
        let smem_tile = &self.stage.unit_tile;
        let config = comptime![self.config];

        let tile_size = config.elements_in_tile_row * config.elements_in_tile_col;
        let output_line_size = self.global.line_size();
        let out_smem_slice = smem_tile.slice.with_line_size(output_line_size);

        let num_lines = tile_size / output_line_size;

        for i in 0..num_lines {
            let value = out_smem_slice[i];
            self.global
                .write_checked((tile, i * output_line_size), Line::cast_from(value));
        }
    }
}

#[cube]
fn tile_pos<S: StageConfig>(#[comptime] config: S) -> (u32, u32) {
    UnitPartitioner::coordinates::<S>(config)
}

fn stage_memory_config<S: StageConfig>(config: S) -> StageMemoryConfig {
    let units = config.num_main_flow_planes() * config.plane_dim();
    let size_n = config.tiling_scheme().stage_partitions_in_stage_n();
    let base = config.stage_memory_config(StageIdent::Acc);
    StageMemoryConfig {
        tiles_in_stage_row: units / size_n,
        tiles_in_stage_col: size_n,
        ..base
    }
}

#[cube]
impl<IP: MatrixPrecision> WriteEventListener for UnitWriter<IP> {
    fn on_event(this: &mut Self, event: super::WriteEvent) {
        #[allow(clippy::single_match)]
        match event {
            WriteEvent::TileStored { tile } => this.write(tile),
            _ => {}
        }
    }
}

#[cube]
impl<IP: MatrixPrecision> GlobalWriter<IP> for UnitWriter<IP> {
    type Stage = PartitionedStage<IP::Stage>;

    fn init<S: StageConfig>(
        tensor: View<Line<IP::Global>, Coords2d, ReadWrite>,
        #[comptime] config: GlobalMemoryConfig,
        #[comptime] stage_config: S,
    ) -> Self {
        Self::new::<S>(tensor, config, stage_config)
    }

    fn stage(this: &Self) -> Self::Stage {
        this.stage
    }
}

pub struct UnitWriterFamily;

impl GlobalWriterFamily for UnitWriterFamily {
    type Stage = PartitionedStageFamily;
    type Writer<IP: MatrixPrecision> = UnitWriter<IP>;
}