cubecl_convolution/components/stage/
reader.rs

1use cubecl::prelude::*;
2use cubecl_core as cubecl;
3use cubecl_matmul::components::{
4    InvalidConfigError, MatrixLayout,
5    global::memory::GlobalMemoryConfig,
6    stage::{
7        StageMemoryConfig, StridedStageMemory, TilingLayout, TilingLayoutEnum, TilingValidation,
8    },
9    tile::StridedTile,
10};
11use cubecl_std::tensor::layout::Coords2d;
12
13#[derive(Clone, Copy)]
14/// Tiling layout specific for bias, which is one-dimensional with stride 0
15pub struct BiasTilingLayout {}
16
17#[cube]
18impl TilingLayout for BiasTilingLayout {
19    fn get_tile<ES: Numeric>(
20        stage: &StridedStageMemory<ES, Self>,
21        tile: Coords2d,
22        #[comptime] config: StageMemoryConfig,
23    ) -> StridedTile<ES> {
24        if comptime!(config.num_stages > 1) {
25            unimplemented!()
26        }
27
28        let (_, col) = tile;
29
30        let stage_line_size = config.stage_line_size;
31        let tile_size_col = config.elements_in_tile_col / stage_line_size;
32
33        let length = tile_size_col;
34        let start = col * tile_size_col;
35
36        StridedTile::new_strided(
37            stage.as_slice(stage_line_size),
38            start,
39            start + length,
40            0,
41            stage.swizzle,
42            MatrixLayout::RowMajor,
43        )
44    }
45
46    fn to_enum() -> comptime_type!(TilingLayoutEnum) {
47        comptime![TilingLayoutEnum::Other]
48    }
49}
50
51impl TilingValidation for BiasTilingLayout {
52    fn check(config: GlobalMemoryConfig) -> Result<(), InvalidConfigError> {
53        let stage_width = config.elements_in_stage_col();
54        if config.line_size() > stage_width {
55            return Err(Box::new("Invalid line size"));
56        }
57        Ok(())
58    }
59}
cubecl_convolution/components/stage/reader.rs

cubecl_convolution/components/stage/
reader.rs
