cubek_convolution/components/global/layout/
tma_out_grad.rs

1use cubecl::{
2    prelude::*,
3    std::tensor::layout::{Coords2d, Layout, LayoutExpand},
4};
5use cubek_matmul::launch::BatchedCoords;
6
7/// Weight backwards needs a consolidated layout to work properly across the combined `k` dimension.
8/// Padding to an even tile shape on width isn't valid, because `im2col` doesn't do this.
9/// Wouldn't be necessary with `im2colWide`, should investigate at some point.
10#[derive(CubeType, CubeLaunch)]
11pub struct TmaOutGradLayout {}
12
13#[cube]
14impl Layout for TmaOutGradLayout {
15    type Coordinates = BatchedCoords;
16    type SourceCoordinates = Coords2d;
17
18    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
19        let (_, row, col) = pos;
20        (row, col)
21    }
22
23    fn is_in_bounds(&self, _pos: Self::Coordinates) -> bool {
24        true.runtime()
25    }
26
27    fn shape(&self) -> Self::Coordinates {
28        (u32::MAX as usize, u32::MAX, u32::MAX).runtime()
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}