cubek-convolution 0.2.0-pre.5

CubeK: Convolution Kernels
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

use cubecl::{
    std::tensor::{launch::BufferArg, layout::*},
    {prelude::*, std::tensor::launch::ViewLayoutLaunchArg},
};
use cubek_matmul::launch::BatchedCoords;

#[derive(CubeType)]
pub struct BiasLayout {
    shape: u32,
    #[cube(comptime)]
    vector_size: u32,
}

#[cube]
impl Layout for BiasLayout {
    type Coordinates = BatchedCoords;
    type SourceCoordinates = Coords1d;

    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
        let (_, _, n) = pos;
        (n / self.vector_size) as usize
    }

    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
        let (_, _, n) = pos;
        n < self.shape
    }

    fn shape(&self) -> Self::Coordinates {
        (1, 1, self.shape)
    }

    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool) {
        (self.to_source_pos(pos), self.is_in_bounds(pos))
    }
}

impl ViewLayoutLaunchArg for BiasLayout {
    type RuntimeArg<R: Runtime> = ();
    type CompilationArg = ();

    fn register<R: Runtime, B: BufferArg>(
        _: Self::RuntimeArg<R>,
        buffer: &B,
        _: Type,
        launcher: &mut KernelLauncher<R>,
    ) {
        let shape = buffer.len();
        <u32 as LaunchArg>::register(shape as u32, launcher);
    }

    fn expand(
        _: &Self::CompilationArg,
        ty: Type,
        builder: &mut KernelBuilder,
    ) -> <Self as CubeType>::ExpandType {
        BiasLayoutExpand {
            shape: <u32 as LaunchArg>::expand(&(), builder),
            vector_size: ty.vector_size() as u32,
        }
    }
}