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        comment!("is_in_bounds");
52        let x = if comptime![self.checked] {
53            C::is_in_bounds(&pos, &self.size)
54        } else {
55            true.runtime()
56        };
57        comment!("is_in_bounds_end");
58        x
59    }
60
61    #[allow(unused)]
62    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool) {
63        intrinsic!(|scope| {
64            let in_bounds = self
65                .clone()
66                .__expand_is_in_bounds_method(scope, pos.clone());
67            let pos = self.__expand_to_source_pos_method(scope, pos);
68            (pos, in_bounds)
69        })
70    }
71}
cubecl_std/tensor/layout/slice.rs

cubecl_std/tensor/layout/
slice.rs
