rudnn/convolution/components/stage/
reader.rs

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

rudnn/convolution/components/stage/
reader.rs
