cubecl_std/tensor/view/operations/
slice.rs

1use super::*;
2use crate::tensor::layout::Coords1d;
3use cubecl::prelude::*;
4use cubecl_core::{self as cubecl, io::read_masked, prelude::barrier::BarrierExpand};
5
6impl<T: CubePrimitive, IO: SliceVisibility> ViewOperations<T, Coords1d> for Slice<T, IO> {}
7impl<T: CubePrimitive, IO: SliceVisibility> ViewOperationsExpand<T, Coords1d>
8    for SliceExpand<T, IO>
9{
10    fn __expand_read_method(
11        &self,
12        scope: &mut Scope,
13        pos: ExpandElementTyped<u32>,
14    ) -> <T>::ExpandType {
15        <Self as ListExpand<T>>::__expand_read_method(self, scope, pos)
16    }
17
18    fn __expand_read_checked_method(
19        &self,
20        scope: &mut Scope,
21        pos: ExpandElementTyped<u32>,
22    ) -> <T>::ExpandType {
23        let len = self.__expand_len_method(scope);
24        let in_bounds = lt::expand(scope, pos.clone(), len);
25        let slice = self.clone().__expand_to_slice_method(scope);
26        let zero = T::__expand_cast_from(scope, 0.into());
27        read_masked::expand::<T>(scope, in_bounds, slice, pos, zero)
28    }
29
30    fn __expand_read_masked_method(
31        &self,
32        scope: &mut Scope,
33        pos: ExpandElementTyped<u32>,
34        mask_value: <T>::ExpandType,
35    ) -> <T>::ExpandType {
36        let len = self.__expand_len_method(scope);
37        let in_bounds = lt::expand(scope, pos.clone(), len);
38        let slice = self.clone().__expand_to_slice_method(scope);
39        read_masked::expand::<T>(scope, in_bounds, slice, pos, mask_value)
40    }
41
42    fn __expand_read_unchecked_method(
43        &self,
44        scope: &mut Scope,
45        pos: ExpandElementTyped<u32>,
46    ) -> <T>::ExpandType {
47        <Self as ListExpand<T>>::__expand_read_unchecked_method(self, scope, pos)
48    }
49
50    fn __expand_to_linear_slice_method(
51        &self,
52        scope: &mut Scope,
53        pos: ExpandElementTyped<u32>,
54        end: ExpandElementTyped<u32>,
55    ) -> SliceExpand<T, ReadOnly> {
56        // Convert to exclusive end
57        let end = add::expand(scope, end, 1u32.into());
58        // Handling for shapes that are 0 in at least one dim, ensures the slice is not
59        // negative length.
60        let start = Min::__expand_min(scope, pos, end.clone());
61        <Self as SliceOperatorExpand<T>>::__expand_slice_method(self, scope, start, end)
62    }
63
64    fn __expand_shape_method(&self, scope: &mut Scope) -> ExpandElementTyped<u32> {
65        <Self as ListExpand<T>>::__expand_len_method(self, scope)
66    }
67
68    fn __expand_is_in_bounds_method(
69        &self,
70        scope: &mut Scope,
71        pos: ExpandElementTyped<u32>,
72    ) -> ExpandElementTyped<bool> {
73        let len = self.__expand_shape_method(scope);
74        lt::expand(scope, pos, len)
75    }
76
77    fn __expand_tensor_map_load_method(
78        &self,
79        _scope: &mut Scope,
80        _barrier: BarrierExpand,
81        _shared_memory: SliceExpand<T, ReadWrite>,
82        _pos: ExpandElementTyped<u32>,
83    ) {
84        unimplemented!("Not a tensor map");
85    }
86}
87
88impl<T: CubePrimitive> ViewOperationsMut<T, Coords1d> for Slice<T, ReadWrite> {}
89impl<T: CubePrimitive> ViewOperationsMutExpand<T, Coords1d> for SliceExpand<T, ReadWrite> {
90    fn __expand_write_method(
91        &self,
92        scope: &mut Scope,
93        pos: ExpandElementTyped<u32>,
94        value: <T>::ExpandType,
95    ) {
96        <Self as ListMutExpand<T>>::__expand_write_method(self, scope, pos, value)
97    }
98
99    fn __expand_write_checked_method(
100        &self,
101        scope: &mut Scope,
102        pos: ExpandElementTyped<u32>,
103        value: <T>::ExpandType,
104    ) {
105        let len = <Self as ListExpand<T>>::__expand_len_method(self, scope);
106        let in_bounds = lt::expand(scope, pos.clone(), len);
107        if_expand(scope, in_bounds.into(), |scope| {
108            <Self as ListMutExpand<T>>::__expand_write_method(self, scope, pos, value)
109        })
110    }
111
112    fn __expand_to_linear_slice_mut_method(
113        &self,
114        scope: &mut Scope,
115        pos: ExpandElementTyped<u32>,
116        end: ExpandElementTyped<u32>,
117    ) -> SliceExpand<T, ReadWrite> {
118        // Convert to exclusive end
119        let end = add::expand(scope, end, 1u32.into());
120        // Handling for shapes that are 0 in at least one dim, ensures the slice is not
121        // negative length.
122        let start = Min::__expand_min(scope, pos, end.clone());
123        <Self as SliceMutOperatorExpand<T>>::__expand_slice_mut_method(self, scope, start, end)
124    }
125
126    fn __expand_tensor_map_store_method(
127        &self,
128        _scope: &mut Scope,
129        _shared_memory: SliceExpand<T, ReadOnly>,
130        _pos: <Coords1d as CubeType>::ExpandType,
131    ) {
132        unimplemented!("Not a tensor map");
133    }
134}
cubecl_std/tensor/view/operations/slice.rs

cubecl_std/tensor/view/operations/
slice.rs
