cubek_convolution/components/stage/
reader.rs

1use cubecl;
2use cubecl::{prelude::*, std::tensor::layout::Coords2d};
3use cubek_std::{
4    stage::StageMemoryConfig,
5    tile::{StridedTile, TilingValidation},
6    {InvalidConfigError, MatrixLayout},
7};
8
9use crate::components::stage::bias_stage::BiasStageMemory;
10
11#[derive(Clone, Copy)]
12/// Tiling layout specific for bias, which is one-dimensional with stride 0
13pub struct BiasTilingLayout {}
14
15#[cube]
16impl BiasTilingLayout {
17    pub fn get_tile<ES: Numeric, NS: Size>(
18        stage: &BiasStageMemory<ES, NS>,
19        tile: Coords2d,
20        #[comptime] config: StageMemoryConfig,
21    ) -> StridedTile<ES, NS> {
22        if config.num_stages > 1 {
23            unimplemented!()
24        }
25
26        let (_, col) = tile;
27
28        let stage_vector_size = config.vector_size;
29        let tile_size_col = config.elements_per_tile_along_col / stage_vector_size;
30
31        let length = tile_size_col;
32        let start = col * tile_size_col;
33
34        StridedTile::new_strided(
35            stage.as_slice(),
36            start,
37            start + length,
38            0,
39            stage.swizzle,
40            MatrixLayout::RowMajor,
41        )
42    }
43}
44
45impl TilingValidation for BiasTilingLayout {
46    fn check(config: StageMemoryConfig) -> Result<(), InvalidConfigError> {
47        let stage_width = config.elements_per_stage_along_col();
48        if config.vector_size > stage_width {
49            return Err(Box::new(format!(
50                "Invalid vector size. Got {:?} which should not be >{:?}",
51                config.vector_size, stage_width,
52            )));
53        }
54        Ok(())
55    }
56}
cubek_convolution/components/stage/reader.rs

cubek_convolution/components/stage/
reader.rs
