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

use cubecl::prelude::*;
use cubecl_core as cubecl;

use crate::components::tile::{
    StridedTile,
    register::{UNROLL, config::RegisterConfig},
};

/// Writer for the register matmul fragments.
#[derive(CubeType)]
pub struct RegisterStageWriter {}

#[cube]
impl RegisterStageWriter {
    pub fn store_fragment<A: Numeric, E: Numeric>(
        tile: &mut StridedTile<E, ReadWrite>,
        acc: &Array<A>,
        #[comptime] config: RegisterConfig,
    ) {
        let out_line_size = tile.slice.line_size();
        #[unroll(UNROLL)]
        for i in 0..comptime!(config.tile_size.mn() / out_line_size) {
            let mut line = Line::empty(out_line_size);
            #[unroll]
            for j in 0..comptime!(out_line_size) {
                line[j] = acc[i * out_line_size + j];
            }
            tile.slice[i] = Line::cast_from(line);
        }
    }
}