cubecl_core/frontend/container/slice/
base.rs

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

cubecl_core/frontend/container/slice/
base.rs
