cubecl_convolution/components/global/layout/
bias.rs

1use cubecl::prelude::*;
2use cubecl_core as cubecl;
3use cubecl_std::tensor::layout::*;
4
5#[derive(CubeType, CubeLaunch)]
6pub struct BiasLayout {
7    shape: u32,
8    #[cube(comptime)]
9    line_size: u32,
10}
11
12#[cube]
13impl Layout for BiasLayout {
14    type Coordinates = Coords3d;
15    type SourceCoordinates = Coords1d;
16
17    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
18        let (_, _, n) = pos;
19        n / self.line_size
20    }
21
22    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
23        let (_, _, n) = pos;
24        n < self.shape
25    }
26
27    fn shape(&self) -> Self::Coordinates {
28        (1, 1, self.shape)
29    }
30
31    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool) {
32        (self.to_source_pos(pos), self.is_in_bounds(pos))
33    }
34}