rudnn/convolution/components/global/layout/
bias.rs

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

rudnn/convolution/components/global/layout/
bias.rs
