cubecl_std/tensor/layout/
dynamic_rank_strided.rs

1use cubecl::prelude::*;
2use cubecl_core as cubecl;
3
4use crate::tensor::layout::{Coords1d, CoordsDyn, Layout, LayoutExpand};
5
6/// Maps multi-dimensional physical coordinates to a linear buffer offset using
7/// per-axis strides. The physical layout's rank can be any value.
8#[derive(CubeType, Clone)]
9pub struct DynamicRankStridedLayout {
10    shape: CoordsDyn,
11    strides: CoordsDyn,
12}
13
14#[cube]
15impl DynamicRankStridedLayout {
16    pub fn new(shape: CoordsDyn, strides: CoordsDyn) -> DynamicRankStridedLayout {
17        DynamicRankStridedLayout { shape, strides }
18    }
19}
20
21#[cube]
22impl Layout for DynamicRankStridedLayout {
23    type Coordinates = CoordsDyn;
24    type SourceCoordinates = Coords1d;
25
26    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
27        let mut offset = 0u32;
28        #[unroll]
29        for i in 0..self.strides.len() {
30            offset += pos[i] * self.strides[i];
31        }
32        offset as usize
33    }
34
35    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool) {
36        (self.to_source_pos(pos.clone()), self.is_in_bounds(pos))
37    }
38
39    fn shape(&self) -> Self::Coordinates {
40        self.shape.clone()
41    }
42
43    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
44        let mut is_valid = true;
45        #[unroll]
46        for i in 0..self.shape.len() {
47            is_valid = is_valid && pos[i] < self.shape[i];
48        }
49        is_valid
50    }
51}
cubecl_std/tensor/layout/dynamic_rank_strided.rs

cubecl_std/tensor/layout/
dynamic_rank_strided.rs
