cubecl_std/tensor/view/operations/
slice.rs

1use core::ops::Deref;
2
3use super::*;
4use crate::tensor::layout::Coords1d;
5use cubecl::prelude::*;
6use cubecl_core::{self as cubecl, io::read_masked, prelude::barrier::Barrier};
7
8impl<T: CubePrimitive> ViewOperations<T, Coords1d> for Box<[T]> {}
9impl<T: CubePrimitive> ViewOperations<T, Coords1d> for [T] {}
10impl<T: CubePrimitive> ViewOperationsExpand<T, Coords1d> for SliceExpand<T> {
11    fn __expand_read_method(&self, scope: &Scope, pos: NativeExpand<usize>) -> <T>::ExpandType {
12        self.__expand_index_method(scope, pos)
13            .__expand_deref_method(scope)
14    }
15
16    fn __expand_read_checked_method(
17        &self,
18        scope: &Scope,
19        pos: NativeExpand<usize>,
20    ) -> <T>::ExpandType {
21        let len = self.__expand_len_method(scope);
22        let in_bounds = pos.__expand_lt_method(scope, &len);
23        let zero = T::__expand_cast_from(scope, 0.into());
24        read_masked::expand::<T>(scope, in_bounds, self, pos, zero)
25    }
26
27    fn __expand_read_masked_method(
28        &self,
29        scope: &Scope,
30        pos: NativeExpand<usize>,
31        mask_value: <T>::ExpandType,
32    ) -> <T>::ExpandType {
33        let len = self.__expand_len_method(scope);
34        let in_bounds = pos.__expand_lt_method(scope, &len);
35        read_masked::expand::<T>(scope, in_bounds, self, pos, mask_value)
36    }
37
38    fn __expand_read_unchecked_method(
39        &self,
40        scope: &Scope,
41        pos: NativeExpand<usize>,
42    ) -> <T>::ExpandType {
43        unsafe {
44            self.__expand_get_unchecked_method(scope, pos)
45                .__expand_deref_method(scope)
46        }
47    }
48
49    fn __expand_as_linear_slice_method(
50        &self,
51        scope: &Scope,
52        pos: NativeExpand<usize>,
53        end: NativeExpand<usize>,
54    ) -> &SliceExpand<T> {
55        // Convert to exclusive end
56        let end = end.__expand_add_method(scope, 1usize.into_expand(scope));
57        // Handling for shapes that are 0 in at least one dim, ensures the slice is not
58        // negative length.
59        let start = clamp_max::expand(scope, pos, end);
60        <Self as SliceOperatorExpand<T>>::__expand_slice_method(self, scope, start, end)
61    }
62
63    fn __expand_shape_method(&self, scope: &Scope) -> NativeExpand<usize> {
64        <Self as ListExpand<T>>::__expand_len_method(self, scope)
65    }
66
67    fn __expand_is_in_bounds_method(
68        &self,
69        scope: &Scope,
70        pos: NativeExpand<usize>,
71    ) -> NativeExpand<bool> {
72        let len = self.__expand_shape_method(scope);
73        pos.__expand_lt_method(scope, &len)
74    }
75
76    fn __expand_tensor_map_load_method(
77        &self,
78        _scope: &Scope,
79        _barrier: &NativeExpand<Barrier>,
80        _shared_memory: &mut SliceExpand<T>,
81        _pos: NativeExpand<usize>,
82    ) {
83        unimplemented!("Not a tensor map");
84    }
85}
86
87impl<T: CubePrimitive> ViewOperationsExpand<T, Coords1d> for NativeExpand<Box<[T]>> {
88    fn __expand_read_method(&self, scope: &Scope, pos: NativeExpand<usize>) -> <T>::ExpandType {
89        self.deref().__expand_read_method(scope, pos)
90    }
91
92    fn __expand_read_checked_method(
93        &self,
94        scope: &Scope,
95        pos: NativeExpand<usize>,
96    ) -> <T>::ExpandType {
97        self.deref().__expand_read_checked_method(scope, pos)
98    }
99
100    fn __expand_read_masked_method(
101        &self,
102        scope: &Scope,
103        pos: NativeExpand<usize>,
104        value: <T>::ExpandType,
105    ) -> <T>::ExpandType {
106        self.deref().__expand_read_masked_method(scope, pos, value)
107    }
108
109    fn __expand_read_unchecked_method(
110        &self,
111        scope: &Scope,
112        pos: NativeExpand<usize>,
113    ) -> <T>::ExpandType {
114        self.deref().__expand_read_unchecked_method(scope, pos)
115    }
116
117    fn __expand_as_linear_slice_method(
118        &self,
119        scope: &Scope,
120        pos: NativeExpand<usize>,
121        size: NativeExpand<usize>,
122    ) -> &SliceExpand<T> {
123        self.deref()
124            .__expand_as_linear_slice_method(scope, pos, size)
125    }
126
127    fn __expand_shape_method(&self, scope: &Scope) -> NativeExpand<usize> {
128        self.deref().__expand_shape_method(scope)
129    }
130
131    fn __expand_is_in_bounds_method(
132        &self,
133        scope: &Scope,
134        pos: NativeExpand<usize>,
135    ) -> NativeExpand<bool> {
136        self.deref().__expand_is_in_bounds_method(scope, pos)
137    }
138
139    fn __expand_tensor_map_load_method(
140        &self,
141        scope: &Scope,
142        barrier: &NativeExpand<Barrier>,
143        shared_memory: &mut SliceExpand<T>,
144        pos: NativeExpand<usize>,
145    ) {
146        self.deref()
147            .__expand_tensor_map_load_method(scope, barrier, shared_memory, pos);
148    }
149}
150
151impl<T: CubePrimitive> ViewOperationsMut<T, Coords1d> for Box<[T]> {}
152impl<T: CubePrimitive> ViewOperationsMut<T, Coords1d> for [T] {}
153impl<T: CubePrimitive> ViewOperationsMutExpand<T, Coords1d> for SliceExpand<T> {
154    fn __expand_write_method(
155        &self,
156        scope: &Scope,
157        pos: NativeExpand<usize>,
158        value: <T>::ExpandType,
159    ) {
160        self.__expand_as_mut_unchecked_method(scope)
161            .__expand_index_mut_method(scope, pos)
162            .__expand_assign_method(scope, value);
163    }
164
165    fn __expand_write_checked_method(
166        &self,
167        scope: &Scope,
168        pos: NativeExpand<usize>,
169        value: <T>::ExpandType,
170    ) {
171        let len = <Self as ListExpand<T>>::__expand_len_method(self, scope);
172        let in_bounds = pos.__expand_lt_method(scope, &len);
173        if_expand(scope, in_bounds, |scope| {
174            self.__expand_as_mut_unchecked_method(scope)
175                .__expand_index_mut_method(scope, pos)
176                .__expand_assign_method(scope, value)
177        })
178    }
179
180    fn __expand_as_linear_slice_mut_method(
181        &self,
182        scope: &Scope,
183        pos: NativeExpand<usize>,
184        end: NativeExpand<usize>,
185    ) -> &mut SliceExpand<T> {
186        // Convert to exclusive end
187        let end = end.__expand_add_method(scope, 1usize.into_expand(scope));
188        // Handling for shapes that are 0 in at least one dim, ensures the slice is not
189        // negative length.
190        let start = clamp_max::expand(scope, pos, end);
191        <Self as SliceOperatorExpand<T>>::__expand_slice_mut_method(
192            self.__expand_as_mut_unchecked_method(scope),
193            scope,
194            start,
195            end,
196        )
197    }
198
199    fn __expand_tensor_map_store_method(
200        &self,
201        _scope: &Scope,
202        _shared_memory: &SliceExpand<T>,
203        _pos: <Coords1d as CubeType>::ExpandType,
204    ) {
205        unimplemented!("Not a tensor map");
206    }
207}
208
209impl<T: CubePrimitive> ViewOperationsMutExpand<T, Coords1d> for NativeExpand<Box<[T]>> {
210    fn __expand_write_method(
211        &self,
212        scope: &Scope,
213        pos: NativeExpand<usize>,
214        value: <T>::ExpandType,
215    ) {
216        self.deref().__expand_write_method(scope, pos, value);
217    }
218
219    fn __expand_write_checked_method(
220        &self,
221        scope: &Scope,
222        pos: NativeExpand<usize>,
223        value: <T>::ExpandType,
224    ) {
225        self.deref()
226            .__expand_write_checked_method(scope, pos, value);
227    }
228
229    fn __expand_as_linear_slice_mut_method(
230        &self,
231        scope: &Scope,
232        pos: NativeExpand<usize>,
233        size: NativeExpand<usize>,
234    ) -> &mut SliceExpand<T> {
235        self.deref()
236            .__expand_as_linear_slice_mut_method(scope, pos, size)
237    }
238
239    fn __expand_tensor_map_store_method(
240        &self,
241        scope: &Scope,
242        shared_memory: &SliceExpand<T>,
243        pos: <Coords1d as CubeType>::ExpandType,
244    ) {
245        self.deref()
246            .__expand_tensor_map_store_method(scope, shared_memory, pos);
247    }
248}
cubecl_std/tensor/view/operations/slice.rs

cubecl_std/tensor/view/operations/
slice.rs
