cubek_convolution/components/global/layout/
bias.rs

1use cubecl::{
2    prelude::*,
3    std::tensor::{
4        launch::{MemoryArg, ViewLayoutLaunchArg},
5        layout::*,
6    },
7};
8use cubek_matmul::args::BatchedCoords;
9
10#[derive(CubeType)]
11pub struct BiasLayout {
12    shape: u32,
13    #[cube(comptime)]
14    vector_size: u32,
15}
16
17#[cube]
18impl Layout for BiasLayout {
19    type Coordinates = BatchedCoords;
20    type SourceCoordinates = Coords1d;
21
22    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
23        let (_, _, n) = pos;
24        (n / self.vector_size) as usize
25    }
26
27    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
28        let (_, _, n) = pos;
29        n < self.shape
30    }
31
32    fn shape(&self) -> Self::Coordinates {
33        (1, 1, self.shape)
34    }
35
36    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool) {
37        (self.to_source_pos(pos), self.is_in_bounds(pos))
38    }
39}
40
41impl ViewLayoutLaunchArg for BiasLayout {
42    type RuntimeArg<R: Runtime> = ();
43    type CompilationArg = ();
44
45    fn register<R: Runtime, B: MemoryArg>(
46        _: Self::RuntimeArg<R>,
47        buffer: &B,
48        _: Type,
49        launcher: &mut KernelLauncher<R>,
50    ) {
51        let shape = buffer.len();
52        <u32 as LaunchArg>::register(shape as u32, launcher);
53    }
54
55    fn expand(
56        _: &Self::CompilationArg,
57        ty: Type,
58        builder: &mut KernelBuilder,
59    ) -> <Self as CubeType>::ExpandType {
60        BiasLayoutExpand {
61            shape: <u32 as LaunchArg>::expand(&(), builder),
62            vector_size: ty.vector_size() as u32,
63        }
64    }
65}
cubek_convolution/components/global/layout/bias.rs

cubek_convolution/components/global/layout/
bias.rs
