cubecl_std/tensor/layout/
slice.rs

1use cubecl::prelude::*;
2use cubecl_core::{self as cubecl, intrinsic};
3
4use crate::tensor::layout::{Coordinates, Layout, LayoutExpand};
5
6/// A layout containing a sub-slice of the inner layout.
7#[allow(unused)]
8#[derive(CubeType)]
9pub struct SliceLayout<C: Coordinates> {
10    offset: C,
11    size: C,
12    #[cube(comptime)]
13    checked: bool,
14}
15
16#[cube]
17impl<C: Coordinates> SliceLayout<C> {
18    /// Create a new slice layout.
19    /// `checked` determines whether bounds should be checked, or simply treated as always in bounds.
20    pub fn new(start: C, size: C, #[comptime] checked: bool) -> Self {
21        SliceLayout::<C> {
22            offset: start,
23            size,
24            checked,
25        }
26    }
27
28    fn offset(&self) -> C {
29        intrinsic! {|_| self.offset.clone() }
30    }
31
32    fn size(&self) -> C {
33        intrinsic! {|_| self.size.clone() }
34    }
35}
36
37#[cube]
38impl<C: Coordinates> Layout for SliceLayout<C> {
39    type Coordinates = C;
40    type SourceCoordinates = C;
41
42    fn shape(&self) -> Self::Coordinates {
43        self.size()
44    }
45
46    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
47        C::add(self.offset(), pos)
48    }
49
50    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
51        if comptime![self.checked] {
52            C::is_in_bounds(&pos, &self.size)
53        } else {
54            true.runtime()
55        }
56    }
57
58    #[allow(unused)]
59    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool) {
60        intrinsic!(|scope| {
61            let in_bounds = self
62                .clone()
63                .__expand_is_in_bounds_method(scope, pos.clone());
64            let pos = self.__expand_to_source_pos_method(scope, pos);
65            (pos, in_bounds)
66        })
67    }
68}
cubecl_std/tensor/layout/slice.rs

cubecl_std/tensor/layout/
slice.rs
