cubek-std 0.2.0-pre.5

CubeK: Standard Library
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 crate::{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, N: Size>(
        tile: &mut StridedTile<V, N, ReadWrite>,
        fragment: &cmma::Matrix<E>,
    ) {
        let layout = as_cmma_layout(tile.layout);
        let stride = tile.unvectorized_stride();
        let mut slice = tile.as_slice_mut();
        cmma::store(&mut slice, fragment, stride, layout);
    }
}