cubecl_core/frontend/container/slice/
base.rs

1use alloc::vec;
2use core::ops::{Deref, DerefMut};
3
4use alloc::boxed::Box;
5
6use crate::{self as cubecl, unexpanded};
7use cubecl::prelude::*;
8use cubecl_ir::{
9    AggregateKind, Branch, ElemType, FloatKind, Instruction, MetadataKind, Operation, RangeLoop,
10    SliceMetadata, Value, VectorSize,
11};
12
13pub type SliceExpand<T> = NativeExpand<[T]>;
14
15#[derive(Clone, Copy)]
16pub struct ReadOnly;
17#[derive(Clone, Copy)]
18pub struct ReadWrite;
19
20pub trait SliceVisibility: Clone + Copy + Send + Sync + 'static {}
21
22impl SliceVisibility for ReadOnly {}
23
24impl SliceVisibility for ReadWrite {}
25
26impl<E: CubePrimitive> SliceExpand<E> {
27    pub fn __extract_list(&self, scope: &Scope) -> Value {
28        let Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) = self.expand.ty else {
29            unreachable!("Should be slice aggregate")
30        };
31        scope.extract_field(self.expand, *inner_ty, SliceMetadata::LIST)
32    }
33
34    pub fn __extract_offset(&self, scope: &Scope) -> NativeExpand<usize> {
35        let ty = usize::__expand_as_type(scope);
36        let field = scope.extract_field(self.expand, ty, SliceMetadata::OFFSET);
37        field.into()
38    }
39
40    pub fn __extract_length(&self, scope: &Scope) -> NativeExpand<usize> {
41        let ty = usize::__expand_as_type(scope);
42        let field = scope.extract_field(self.expand, ty, SliceMetadata::LENGTH);
43        field.into()
44    }
45}
46
47pub trait SliceVectorExt<E: Scalar, N: Size> {
48    fn with_vector_size<N2: Size>(&self) -> &[Vector<E, N2>] {
49        unexpanded!()
50    }
51    fn with_vector_size_mut<N2: Size>(&mut self) -> &mut [Vector<E, N2>] {
52        unexpanded!()
53    }
54    fn __expand_with_vector_size<'infer, N2: Size>(
55        scope: &Scope,
56        this: &'infer SliceExpand<Vector<E, N>>,
57    ) -> &'infer SliceExpand<Vector<E, N2>> {
58        this.__expand_with_vector_size_method(scope)
59    }
60    fn __expand_with_vector_size_mut<'infer, N2: Size>(
61        scope: &Scope,
62        this: &'infer mut SliceExpand<Vector<E, N>>,
63    ) -> &'infer mut SliceExpand<Vector<E, N2>> {
64        this.__expand_with_vector_size_mut_method(scope)
65    }
66}
67
68impl<E: Scalar, N: Size> SliceVectorExt<E, N> for [Vector<E, N>] {}
69impl<E: Scalar, N: Size> SliceExpand<Vector<E, N>> {
70    pub fn __expand_with_vector_size_method<'infer, N2: Size>(
71        &'infer self,
72        scope: &Scope,
73    ) -> &'infer SliceExpand<Vector<E, N2>> {
74        let slice = self.with_vector_size_inner::<N2>(scope);
75        scope.create_kernel_ref(slice)
76    }
77
78    pub fn __expand_with_vector_size_mut_method<'infer, N2: Size>(
79        &'infer mut self,
80        scope: &Scope,
81    ) -> &'infer mut SliceExpand<Vector<E, N2>> {
82        let slice = self.with_vector_size_inner::<N2>(scope);
83        scope.create_kernel_ref(slice)
84    }
85}
86
87impl<E: Scalar, N: Size> SliceExpand<Vector<E, N>> {
88    fn with_vector_size_inner<N2: Size>(&self, scope: &Scope) -> SliceExpand<Vector<E, N2>> {
89        let vector_size = N2::__expand_value(scope);
90        let list = self.__extract_list(scope);
91        let item = list.value_type();
92
93        let length = self.__extract_length(scope);
94        let offset = self.__extract_offset(scope);
95
96        let current = list.ty.vector_size();
97
98        if vector_size == item.vector_size() {
99            return self.expand.into();
100        }
101
102        let mut new_ptr = list;
103        new_ptr.ty = new_ptr.ty.with_vector_size(vector_size);
104
105        if current < vector_size {
106            let ratio = vector_size / current;
107            let offset = offset.__expand_div_method(scope, ratio.into());
108            let length = length.__expand_div_method(scope, ratio.into());
109            from_raw_parts(scope, new_ptr, offset, length)
110        } else {
111            let ratio = current / vector_size;
112            let offset = offset.__expand_mul_method(scope, ratio.into());
113            let length = length.__expand_mul_method(scope, ratio.into());
114            from_raw_parts(scope, new_ptr, offset, length)
115        }
116    }
117}
118
119pub trait SliceExt<E: CubePrimitive> {
120    /// Returns the same slice, but with the type reinterpreted as `Vector`.
121    /// Preserves existing vector size of the primitive.
122    fn as_vectorized(&self) -> &[Vector<E::Scalar, E::Size>] {
123        unexpanded!()
124    }
125
126    /// Returns the same slice, but with the type reinterpreted as `Vector`.
127    /// Preserves existing vector size of the primitive.
128    fn as_vectorized_mut(&mut self) -> &mut [Vector<E::Scalar, E::Size>] {
129        unexpanded!()
130    }
131
132    /// Downcast the slice to the given type and panic if the type isn't the same.
133    ///
134    /// This function should only be used to satisfy the Rust type system, when two generic
135    /// types are supposed to be the same.
136    fn downcast<T: CubePrimitive>(&self) -> &[T] {
137        unexpanded!()
138    }
139
140    /// Downcast the slice to the given type and panic if the type isn't the same.
141    ///
142    /// This function should only be used to satisfy the Rust type system, when two generic
143    /// types are supposed to be the same.
144    fn downcast_mut<T: CubePrimitive>(&mut self) -> &mut [T] {
145        unexpanded!()
146    }
147
148    /// Unsafely downcast the slice to the given type and panic if the type isn't the same.
149    ///
150    /// # Safety
151    /// This function converts unsafely, and should only be used for temporary storage with a dummy
152    /// type (i.e. `ReinterpretSlice`)
153    unsafe fn downcast_unchecked<T: CubePrimitive>(&self) -> &[T] {
154        unexpanded!()
155    }
156
157    /// Unsafely downcast the slice to the given type and panic if the type isn't the same.
158    ///
159    /// # Safety
160    /// This function converts unsafely, and should only be used for temporary storage with a dummy
161    /// type (i.e. `ReinterpretSlice`)
162    unsafe fn downcast_mut_unchecked<T: CubePrimitive>(&mut self) -> &mut [T] {
163        unexpanded!()
164    }
165
166    /// Unsafely cast an immutable slice to a mutable one.
167    ///
168    /// # Safety
169    /// This is safe in practice, but breaks semantics. Should only be used if absolutely necessary.
170    /// May cause problems if an immutable input is reinterpreted as mutable.
171    #[allow(clippy::mut_from_ref)]
172    unsafe fn as_mut_unchecked(&self) -> &mut [E] {
173        unexpanded!()
174    }
175
176    /// Convert the slice to a start pointer, without any bounds checks.
177    ///
178    /// # Safety
179    /// See [`get_unchecked`]([E]::get_unchecked)
180    unsafe fn as_ptr_unchecked(&self) -> *const E {
181        unexpanded!()
182    }
183
184    /// Convert the slice to a mutable start pointer, without any bounds checks.
185    ///
186    /// # Safety
187    /// See [`get_unchecked_mut`]([E]::get_unchecked_mut)
188    unsafe fn as_mut_ptr_unchecked(&mut self) -> *mut E {
189        unexpanded!()
190    }
191
192    /// Convert to an owned boxed slice. This is very unsafe as it completely erases the lifetime.
193    /// Only use it for global kernel inputs, which have a static lifetime.
194    ///
195    /// # Safety
196    /// Erases the lifetime. Only use when an owned representation is absolutely needed.
197    unsafe fn as_boxed_unchecked(&self) -> Box<[E]> {
198        unexpanded!()
199    }
200
201    fn __expand_as_vectorized<'infer>(
202        scope: &Scope,
203        this: &'infer SliceExpand<E>,
204    ) -> &'infer SliceExpand<Vector<E::Scalar, E::Size>>;
205
206    fn __expand_as_vectorized_mut<'infer>(
207        scope: &Scope,
208        this: &'infer mut SliceExpand<E>,
209    ) -> &'infer mut SliceExpand<Vector<E::Scalar, E::Size>>;
210
211    fn __expand_downcast<'infer, T: CubePrimitive>(
212        scope: &Scope,
213        this: &'infer SliceExpand<E>,
214    ) -> &'infer SliceExpand<T>;
215
216    fn __expand_downcast_mut<'infer, T: CubePrimitive>(
217        scope: &Scope,
218        this: &'infer mut SliceExpand<E>,
219    ) -> &'infer mut SliceExpand<T>;
220
221    fn __expand_downcast_unchecked<'infer, T: CubePrimitive>(
222        scope: &Scope,
223        this: &'infer SliceExpand<E>,
224    ) -> &'infer SliceExpand<T>;
225
226    fn __expand_downcast_mut_unchecked<'infer, T: CubePrimitive>(
227        scope: &Scope,
228        this: &'infer mut SliceExpand<E>,
229    ) -> &'infer mut SliceExpand<T>;
230
231    fn __expand_as_ptr(scope: &Scope, this: &SliceExpand<E>) -> *const NativeExpand<E>;
232
233    fn __expand_as_mut_ptr(scope: &Scope, this: &mut SliceExpand<E>) -> *mut NativeExpand<E>;
234
235    #[doc(hidden)]
236    unsafe fn __expand_as_ptr_unchecked(
237        scope: &Scope,
238        this: &SliceExpand<E>,
239    ) -> *const NativeExpand<E>;
240
241    #[doc(hidden)]
242    unsafe fn __expand_as_mut_ptr_unchecked(
243        scope: &Scope,
244        this: &mut SliceExpand<E>,
245    ) -> *mut NativeExpand<E>;
246
247    #[allow(clippy::mut_from_ref)]
248    #[doc(hidden)]
249    unsafe fn __expand_as_mut_unchecked<'infer>(
250        scope: &Scope,
251        this: &'infer SliceExpand<E>,
252    ) -> &'infer mut SliceExpand<E>;
253
254    #[doc(hidden)]
255    unsafe fn __expand_as_boxed_unchecked(
256        scope: &Scope,
257        this: &SliceExpand<E>,
258    ) -> NativeExpand<Box<[E]>>;
259}
260
261impl<E: CubePrimitive> SliceExt<E> for [E] {
262    fn __expand_as_vectorized<'infer>(
263        scope: &Scope,
264        this: &'infer SliceExpand<E>,
265    ) -> &'infer SliceExpand<Vector<E::Scalar, E::Size>> {
266        this.__expand_as_vectorized_method(scope)
267    }
268
269    fn __expand_as_vectorized_mut<'infer>(
270        scope: &Scope,
271        this: &'infer mut SliceExpand<E>,
272    ) -> &'infer mut SliceExpand<Vector<E::Scalar, E::Size>> {
273        this.__expand_as_vectorized_mut_method(scope)
274    }
275
276    fn __expand_downcast<'infer, T: CubePrimitive>(
277        scope: &Scope,
278        this: &'infer SliceExpand<E>,
279    ) -> &'infer SliceExpand<T> {
280        this.__expand_downcast_method::<T>(scope)
281    }
282
283    fn __expand_downcast_mut<'infer, T: CubePrimitive>(
284        scope: &Scope,
285        this: &'infer mut SliceExpand<E>,
286    ) -> &'infer mut SliceExpand<T> {
287        this.__expand_downcast_mut_method::<T>(scope)
288    }
289
290    fn __expand_downcast_unchecked<'infer, T: CubePrimitive>(
291        scope: &Scope,
292        this: &'infer SliceExpand<E>,
293    ) -> &'infer SliceExpand<T> {
294        this.__expand_downcast_unchecked_method::<T>(scope)
295    }
296
297    fn __expand_downcast_mut_unchecked<'infer, T: CubePrimitive>(
298        scope: &Scope,
299        this: &'infer mut SliceExpand<E>,
300    ) -> &'infer mut SliceExpand<T> {
301        this.__expand_downcast_mut_unchecked_method::<T>(scope)
302    }
303
304    fn __expand_as_ptr(scope: &Scope, this: &SliceExpand<E>) -> *const NativeExpand<E> {
305        this.__expand_as_ptr_method(scope)
306    }
307
308    fn __expand_as_mut_ptr(scope: &Scope, this: &mut SliceExpand<E>) -> *mut NativeExpand<E> {
309        this.__expand_as_mut_ptr_method(scope)
310    }
311
312    unsafe fn __expand_as_ptr_unchecked(
313        scope: &Scope,
314        this: &SliceExpand<E>,
315    ) -> *const NativeExpand<E> {
316        unsafe { this.__expand_as_ptr_unchecked_method(scope) }
317    }
318
319    unsafe fn __expand_as_mut_ptr_unchecked(
320        scope: &Scope,
321        this: &mut SliceExpand<E>,
322    ) -> *mut NativeExpand<E> {
323        unsafe { this.__expand_as_mut_ptr_unchecked_method(scope) }
324    }
325
326    unsafe fn __expand_as_mut_unchecked<'infer>(
327        scope: &Scope,
328        this: &'infer SliceExpand<E>,
329    ) -> &'infer mut SliceExpand<E> {
330        this.__expand_as_mut_unchecked_method(scope)
331    }
332
333    unsafe fn __expand_as_boxed_unchecked(
334        scope: &Scope,
335        this: &SliceExpand<E>,
336    ) -> NativeExpand<Box<[E]>> {
337        unsafe { this.__expand_as_boxed_unchecked_method(scope) }
338    }
339}
340
341impl<E: CubePrimitive> SliceExpand<E> {
342    pub fn __expand_as_vectorized_method(
343        &self,
344        _: &Scope,
345    ) -> &SliceExpand<Vector<E::Scalar, E::Size>> {
346        unsafe { self.as_type_ref_unchecked() }
347    }
348
349    pub fn __expand_as_vectorized_mut_method(
350        &mut self,
351        _: &Scope,
352    ) -> &mut SliceExpand<Vector<E::Scalar, E::Size>> {
353        unsafe { self.as_type_mut_unchecked() }
354    }
355
356    pub fn __expand_downcast_method<T: CubePrimitive>(&self, scope: &Scope) -> &SliceExpand<T> {
357        if T::__expand_as_type(scope) != E::__expand_as_type(scope) && !is_tf32::<E, T>(scope) {
358            let elems = [
359                T::__expand_as_type(scope).elem_type(),
360                E::__expand_as_type(scope).elem_type(),
361            ];
362            let is_flex32_cast = elems.contains(&ElemType::Float(FloatKind::F32))
363                && elems.contains(&ElemType::Float(FloatKind::Flex32));
364
365            if !is_flex32_cast {
366                panic!(
367                    "Downcast should only be used to satisfy the Rust type system.
368Expected types to be the same, got [{}, {}]",
369                    elems[0], elems[1]
370                )
371            }
372        }
373
374        self.__expand_downcast_unchecked_method(scope)
375    }
376
377    pub fn __expand_downcast_mut_method<T: CubePrimitive>(
378        &mut self,
379        scope: &Scope,
380    ) -> &mut SliceExpand<T> {
381        if T::__expand_as_type(scope) != E::__expand_as_type(scope) && !is_tf32::<E, T>(scope) {
382            let elems = [
383                T::__expand_as_type(scope).elem_type(),
384                E::__expand_as_type(scope).elem_type(),
385            ];
386            let is_flex32_cast = elems.contains(&ElemType::Float(FloatKind::F32))
387                && elems.contains(&ElemType::Float(FloatKind::Flex32));
388
389            if !is_flex32_cast {
390                panic!(
391                    "Downcast should only be used to satisfy the Rust type system.
392Expected types to be the same, got [{}, {}]",
393                    elems[0], elems[1]
394                )
395            }
396        }
397
398        self.__expand_downcast_mut_unchecked_method(scope)
399    }
400
401    #[doc(hidden)]
402    pub fn __expand_downcast_unchecked_method<T: CubePrimitive>(
403        &self,
404        _: &Scope,
405    ) -> &SliceExpand<T> {
406        unsafe { self.as_type_ref_unchecked() }
407    }
408
409    #[doc(hidden)]
410    pub fn __expand_downcast_mut_unchecked_method<T: CubePrimitive>(
411        &mut self,
412        _: &Scope,
413    ) -> &mut SliceExpand<T> {
414        unsafe { self.as_type_mut_unchecked() }
415    }
416
417    pub fn __expand_as_ptr_method(&self, scope: &Scope) -> *const NativeExpand<E> {
418        self.__expand_index_method(scope, NativeExpand::<usize>::from_lit(scope, 0))
419    }
420
421    pub fn __expand_as_mut_ptr_method(&mut self, scope: &Scope) -> *mut NativeExpand<E> {
422        self.__expand_index_mut_method(scope, NativeExpand::<usize>::from_lit(scope, 0))
423    }
424
425    #[doc(hidden)]
426    pub unsafe fn __expand_as_ptr_unchecked_method(&self, scope: &Scope) -> *const NativeExpand<E> {
427        unsafe {
428            self.__expand_get_unchecked_method(scope, NativeExpand::<usize>::from_lit(scope, 0))
429        }
430    }
431
432    #[doc(hidden)]
433    pub unsafe fn __expand_as_mut_ptr_unchecked_method(
434        &mut self,
435        scope: &Scope,
436    ) -> *mut NativeExpand<E> {
437        unsafe {
438            self.__expand_get_unchecked_mut_method(scope, NativeExpand::<usize>::from_lit(scope, 0))
439        }
440    }
441
442    pub fn __expand_as_mut_unchecked_method(&self, scope: &Scope) -> &mut SliceExpand<E> {
443        scope.create_kernel_ref(self.expand.into())
444    }
445
446    #[doc(hidden)]
447    pub unsafe fn __expand_as_boxed_unchecked_method(&self, _: &Scope) -> NativeExpand<Box<[E]>> {
448        self.expand.into()
449    }
450}
451
452pub fn from_raw_parts<E: CubePrimitive>(
453    scope: &Scope,
454    list: Value,
455    offset: NativeExpand<usize>,
456    length: NativeExpand<usize>,
457) -> SliceExpand<E> {
458    let ty = Type::Aggregate(AggregateKind::ptr(list.ty, MetadataKind::Slice));
459    let out = scope.create_value(ty);
460    scope.register(Instruction::new(
461        Operation::ConstructAggregate(vec![list, offset.expand, length.expand]),
462        out,
463    ));
464    out.into()
465}
466
467impl<E: CubePrimitive> SliceExpand<E> {
468    /// Get the length of the slice.
469    pub fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
470        self.__extract_length(scope)
471    }
472    /// Returns true if the slice is empty.
473    pub fn is_empty(&self, scope: &Scope) -> NativeExpand<bool> {
474        self.__extract_length(scope)
475            .__expand_eq_method(scope, &0usize.into_expand(scope))
476    }
477}
478
479impl<E: CubePrimitive> CubeType for [E] {
480    type ExpandType = SliceExpand<E>;
481}
482
483impl<E: CubePrimitive> CubeType for Box<[E]> {
484    type ExpandType = NativeExpand<Box<[E]>>;
485}
486
487impl<E> Deref for NativeExpand<Box<[E]>> {
488    type Target = NativeExpand<[E]>;
489
490    fn deref(&self) -> &Self::Target {
491        unsafe { self.as_type_ref_unchecked() }
492    }
493}
494
495impl<E> DerefMut for NativeExpand<Box<[E]>> {
496    fn deref_mut(&mut self) -> &mut Self::Target {
497        unsafe { self.as_type_mut_unchecked() }
498    }
499}
500
501macro_rules! impl_expand_traits {
502    ($generic: ident, $ty: ty) => {
503        impl<$generic: CubePrimitive> AsMutExpand for $ty {
504            fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
505                self
506            }
507        }
508
509        impl<$generic: CubePrimitive> IntoMut for $ty {
510            fn into_mut(self, _scope: &Scope) -> Self {
511                self
512            }
513        }
514    };
515}
516
517impl_expand_traits!(E, SliceExpand<E>);
518impl_expand_traits!(E, NativeExpand<Box<[E]>>);
519
520impl<'a, E: CubePrimitive> From<&'a NativeExpand<Box<[E]>>> for &'a NativeExpand<[E]> {
521    fn from(value: &'a NativeExpand<Box<[E]>>) -> Self {
522        unsafe { value.as_type_ref_unchecked() }
523    }
524}
525
526impl<'a, E: CubePrimitive> From<&'a mut NativeExpand<Box<[E]>>> for &'a mut NativeExpand<[E]> {
527    fn from(value: &'a mut NativeExpand<Box<[E]>>) -> Self {
528        unsafe { value.as_type_mut_unchecked() }
529    }
530}
531
532impl<E: CubePrimitive> SizedContainer<usize> for [E] {
533    fn len(&self) -> usize {
534        unexpanded!()
535    }
536}
537
538impl<E: CubePrimitive> SizedContainerExpand<usize> for SliceExpand<E> {
539    fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
540        self.__expand_len_method(scope)
541    }
542}
543
544impl<E: CubePrimitive> Iterable for SliceExpand<E> {
545    type Item = E::ExpandType;
546
547    fn expand(self, scope: &Scope, mut body: impl FnMut(&Scope, Self::Item)) {
548        let index_ty = usize::__expand_as_type(scope);
549        let len = self.__extract_length(scope).expand;
550
551        let child = scope.child();
552        let i = scope.create_local_mut(index_ty);
553
554        let index = NativeExpand::new(i);
555        let item = self
556            .__expand_index_method(&child, index)
557            .__expand_deref_method(&child);
558        body(&child, item);
559
560        scope.register(Branch::RangeLoop(Box::new(RangeLoop {
561            i,
562            start: 0usize.into(),
563            end: len,
564            step: None,
565            inclusive: false,
566            scope: child,
567        })));
568    }
569
570    fn expand_unroll(self, _scope: &Scope, _body: impl FnMut(&Scope, Self::Item)) {
571        unimplemented!("Can't unroll slice iterator")
572    }
573}
574
575impl<'a, E: CubePrimitive> Iterable for &'a SliceExpand<E> {
576    type Item = &'a E::ExpandType;
577
578    fn expand(self, scope: &Scope, mut body: impl FnMut(&Scope, Self::Item)) {
579        let index_ty = usize::__expand_as_type(scope);
580        let len = self.__extract_length(scope).expand;
581
582        let child = scope.child();
583        let i = scope.create_local_mut(index_ty);
584
585        let index = NativeExpand::new(i);
586        let item = self.__expand_index_method(&child, index);
587        body(&child, item);
588
589        scope.register(Branch::RangeLoop(Box::new(RangeLoop {
590            i,
591            start: 0usize.into(),
592            end: len,
593            step: None,
594            inclusive: false,
595            scope: child,
596        })));
597    }
598
599    fn expand_unroll(self, _scope: &Scope, _body: impl FnMut(&Scope, Self::Item)) {
600        unimplemented!("Can't unroll slice iterator")
601    }
602}
603
604impl<'a, E: CubePrimitive> Iterable for &'a mut SliceExpand<E> {
605    type Item = &'a mut E::ExpandType;
606
607    fn expand(self, scope: &Scope, mut body: impl FnMut(&Scope, Self::Item)) {
608        let index_ty = usize::__expand_as_type(scope);
609        let len = self.__extract_length(scope).expand;
610
611        let child = scope.child();
612        let i = scope.create_local_mut(index_ty);
613
614        let index = NativeExpand::new(i);
615        let item = self.__expand_index_mut_method(&child, index);
616        body(&child, item);
617
618        scope.register(Branch::RangeLoop(Box::new(RangeLoop {
619            i,
620            start: 0usize.into(),
621            end: len,
622            step: None,
623            inclusive: false,
624            scope: child,
625        })));
626    }
627
628    fn expand_unroll(self, _scope: &Scope, _body: impl FnMut(&Scope, Self::Item)) {
629        unimplemented!("Can't unroll slice iterator")
630    }
631}
632
633impl<E: CubePrimitive> SliceExpand<E> {
634    #[doc(hidden)]
635    pub unsafe fn __expand_get_unchecked_method(
636        &self,
637        scope: &Scope,
638        index: NativeExpand<usize>,
639    ) -> &NativeExpand<E> {
640        read_offset::expand::<E>(scope, self, index, false)
641    }
642
643    #[doc(hidden)]
644    pub unsafe fn __expand_get_unchecked_mut_method(
645        &mut self,
646        scope: &Scope,
647        index: NativeExpand<usize>,
648    ) -> &mut NativeExpand<E> {
649        write_offset::expand::<E>(scope, self, index, false)
650    }
651}
652
653impl<E: CubePrimitive> IndexExpand<NativeExpand<usize>> for SliceExpand<E> {
654    type Output = E::ExpandType;
655
656    fn __expand_index_method(&self, scope: &Scope, index: NativeExpand<usize>) -> &Self::Output {
657        read_offset::expand::<E>(scope, self, index, true)
658    }
659}
660
661impl<E: CubePrimitive> IndexMutExpand<NativeExpand<usize>> for SliceExpand<E> {
662    fn __expand_index_mut_method(
663        &mut self,
664        scope: &Scope,
665        index: NativeExpand<usize>,
666    ) -> &mut Self::Output {
667        write_offset::expand::<E>(scope, self, index, true)
668    }
669}
670
671impl_slice_ranges!(SliceExpand<E>);
672
673impl<E: CubePrimitive> List<E> for [E] {}
674impl<E: CubePrimitive> ListExpand<E> for SliceExpand<E> {
675    fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
676        self.__expand_len_method(scope)
677    }
678}
679
680impl<E: CubePrimitive> Vectorized for Box<[E]> {}
681impl<E: CubePrimitive> Vectorized for [E] {}
682impl<E: CubePrimitive> VectorizedExpand for SliceExpand<E> {
683    fn vector_size(&self) -> VectorSize {
684        self.expand.vector_size()
685    }
686}
687impl<E: CubePrimitive> VectorizedExpand for NativeExpand<Box<[E]>> {
688    fn vector_size(&self) -> VectorSize {
689        self.expand.vector_size()
690    }
691}
692
693mod read_offset {
694    use super::*;
695
696    pub fn expand<'a, E: CubePrimitive>(
697        scope: &Scope,
698        slice: &SliceExpand<E>,
699        index: NativeExpand<usize>,
700        checked: bool,
701    ) -> &'a <E as cubecl::prelude::CubeType>::ExpandType {
702        let list = slice.__extract_list(scope);
703        let offset = slice.__extract_offset(scope);
704        let index = offset.__expand_add_method(scope, index);
705
706        expand_index_native(scope, list, index, checked)
707    }
708}
709
710mod write_offset {
711    use super::*;
712
713    pub fn expand<'a, E: CubePrimitive>(
714        scope: &Scope,
715        slice: &SliceExpand<E>,
716        index: <usize as CubeType>::ExpandType,
717        checked: bool,
718    ) -> &'a mut E::ExpandType {
719        let list = slice.__extract_list(scope);
720        let offset = slice.__extract_offset(scope);
721        let index = offset.__expand_add_method(scope, index);
722
723        expand_index_mut_native(scope, list, index, checked)
724    }
725}
cubecl_core/frontend/container/slice/base.rs

cubecl_core/frontend/container/slice/
base.rs
