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

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

use crate::components::{as_cmma_layout, tile::StridedTile};

/// Writer using the cmma store function.
#[derive(CubeType)]
pub struct CmmaStageWriter {}

#[cube]
impl CmmaStageWriter {
    pub fn store_fragment<E: Numeric, V: Numeric>(
        tile: &mut StridedTile<V, ReadWrite>,
        fragment: &cmma::Matrix<E>,
    ) {
        let layout = as_cmma_layout(tile.layout);
        let (mut slice, stride) = tile.as_unlined();
        cmma::store(&mut slice, fragment, stride, layout);
    }
}