cubek_std/layout/
row_major.rs

1use cubecl::{
2    prelude::*,
3    std::tensor::layout::{Coords1d, Layout, LayoutExpand},
4};
5
6#[derive(CubeType, Clone, Copy)]
7pub struct RowMajorLayout {
8    width: usize,
9    height: usize,
10    vector_size: usize,
11}
12
13#[cube]
14impl RowMajorLayout {
15    pub fn new(width: usize, height: usize, vector_size: usize) -> Self {
16        RowMajorLayout {
17            width,
18            height,
19            vector_size,
20        }
21    }
22}
23
24#[cube]
25impl Layout for RowMajorLayout {
26    type Coordinates = (usize, usize);
27    type SourceCoordinates = Coords1d;
28
29    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
30        (self.width * pos.0 + pos.1) / self.vector_size
31    }
32
33    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool) {
34        let is_valid = pos.0 < self.height && pos.1 < self.width;
35        (self.to_source_pos(pos), is_valid)
36    }
37
38    fn shape(&self) -> Self::Coordinates {
39        (self.width, self.height)
40    }
41
42    fn is_in_bounds(&self, _pos: Self::Coordinates) -> bool {
43        true.runtime()
44    }
45}
cubek_std/layout/row_major.rs

cubek_std/layout/
row_major.rs
