Skip to main content

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

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::prelude::*;
3use ruda_kernel::library::tensor::layout::Coords2d;
4use ruda_kernel::library::tensor::layout::Layout;
5use ruda_kernel::library::tensor::layout::LayoutExpand;
6use rublas::kernel_ir::launch::BatchedCoords;
7
8use crate::convolution::components::ConvolutionProblem;
9
10/// Weight backwards needs a consolidated layout to work properly across the combined `k` dimension.
11/// Padding to an even tile shape on width isn't valid, because `im2col` doesn't do this.
12/// Wouldn't be necessary with `im2colWide`, should investigate at some point.
13#[derive(RudaType, RudaLaunch)]
14pub struct TmaOutGradLayout {
15    rows: u32,
16    cols: u32,
17}
18
19#[ruda]
20impl Layout for TmaOutGradLayout {
21    type Coordinates = BatchedCoords;
22    type SourceCoordinates = Coords2d;
23
24    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
25        let (_, row, col) = pos;
26        (row, col)
27    }
28
29    fn is_in_bounds(&self, _pos: Self::Coordinates) -> bool {
30        true.runtime()
31    }
32
33    fn shape(&self) -> Self::Coordinates {
34        (1, self.rows, self.cols)
35    }
36
37    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool) {
38        (self.to_source_pos(pos), self.is_in_bounds(pos))
39    }
40}
41
42impl<R: Runtime> TmaOutGradLayoutLaunch<R> {
43    pub fn from_problem(problem: &ConvolutionProblem) -> Self {
44        TmaOutGradLayoutLaunch::new(problem.k as u32, problem.m as u32)
45    }
46}