cubecl_core/frontend/container/slice/
operator.rs

1use core::ops::{Deref, DerefMut};
2
3use super::SliceExpand;
4use crate::{
5    self as cubecl,
6    frontend::{container::slice, ranges::range},
7};
8use crate::{ir::Scope, prelude::*, unexpanded};
9use cubecl_common::tf32;
10
11pub(crate) fn is_tf32<C: CubePrimitive, T: CubePrimitive>(scope: &Scope) -> bool {
12    let ty_c = C::__expand_as_type(scope).storage_type();
13    let ty_t = T::__expand_as_type(scope).storage_type();
14    let ty_f32 = f32::__expand_as_type(scope).storage_type();
15    let ty_tf32 = tf32::__expand_as_type(scope).storage_type();
16
17    (ty_c == ty_f32 && ty_t == ty_tf32) || (ty_c == ty_tf32 && ty_t == ty_f32)
18}
19
20type ArrayExpand<E> = NativeExpand<Array<E>>;
21
22impl<E: CubePrimitive, T: Deref<Target = SliceExpand<E>> + DerefMut> SliceOperatorExpand<E> for T {
23    fn __expand_slice_method(
24        &self,
25        scope: &Scope,
26        start: NativeExpand<usize>,
27        end: NativeExpand<usize>,
28    ) -> &SliceExpand<E> {
29        self.deref().__expand_slice_method(scope, start, end)
30    }
31
32    fn __expand_slice_mut_method(
33        &mut self,
34        scope: &Scope,
35        start: NativeExpand<usize>,
36        end: NativeExpand<usize>,
37    ) -> &mut SliceExpand<E> {
38        self.deref_mut()
39            .__expand_slice_mut_method(scope, start, end)
40    }
41}
42
43impl<E: CubePrimitive> SliceOperator<E> for Shared<[E]> {}
44impl<E: CubePrimitive> SliceOperator<E> for Tensor<E> {}
45impl<E: CubePrimitive> SliceOperator<E> for Array<E> {}
46
47#[cube]
48impl<E: CubePrimitive> Array<E> {
49    pub fn as_slice(&self) -> &[E] {
50        intrinsic!(|_| self.deref())
51    }
52
53    pub fn as_mut_slice(&mut self) -> &mut [E] {
54        intrinsic!(|_| self.deref_mut())
55    }
56}
57
58#[cube]
59impl<E: CubePrimitive> Shared<[E]> {
60    pub fn as_slice(&self) -> &[E] {
61        intrinsic!(|_| self.deref())
62    }
63
64    pub fn as_mut_slice(&mut self) -> &mut [E] {
65        intrinsic!(|_| self.deref_mut())
66    }
67}
68
69impl<E: CubePrimitive> SliceOperator<E> for [E] {}
70impl<E: CubePrimitive> SliceOperatorExpand<E> for SliceExpand<E> {
71    fn __expand_slice_method(
72        &self,
73        scope: &Scope,
74        start: NativeExpand<usize>,
75        end: NativeExpand<usize>,
76    ) -> &SliceExpand<E> {
77        let length = end.__expand_sub_method(scope, start);
78        let list = self.__extract_list(scope);
79        let offset = self.__extract_offset(scope);
80        let offset = start.__expand_add_method(scope, offset);
81        let slice = slice::from_raw_parts(scope, list, offset, length);
82        scope.create_kernel_ref(slice)
83    }
84
85    fn __expand_slice_mut_method(
86        &mut self,
87        scope: &Scope,
88        start: NativeExpand<usize>,
89        end: NativeExpand<usize>,
90    ) -> &mut SliceExpand<E> {
91        let length = end.__expand_sub_method(scope, start);
92        let list = self.__extract_list(scope);
93        let offset = self.__extract_offset(scope);
94        let offset = start.__expand_add_method(scope, offset);
95
96        let slice = slice::from_raw_parts(scope, list, offset, length);
97        scope.create_kernel_ref(slice)
98    }
99}
100
101#[cube]
102pub trait SliceOperator<E: CubePrimitive> {
103    /// Return a read-only view of all elements comprise between the `start` and `end` indices.
104    /// In `checked` mode, if the `end` index is out-of-bound, it is replaced by
105    /// the length of `self`.
106    #[allow(unused_variables)]
107    fn slice(&self, start: usize, end: usize) -> &[E] {
108        unexpanded!()
109    }
110
111    /// Return a read-write view of all elements comprise between the `start` and `end` indices.
112    /// In `checked` mode, if the `end` index is out-of-bound, it is replaced by
113    /// the length of `self`.
114    #[allow(unused_variables)]
115    fn slice_mut(&mut self, start: usize, end: usize) -> &mut [E] {
116        unexpanded!()
117    }
118}
119
120// Simple heuristic
121const MEMCPY_UNROLL_LIMIT: usize = 8;
122
123impl<E: CubePrimitive> SliceExpand<E> {
124    pub fn __expand_copy_from_slice_method(&mut self, scope: &Scope, source: &SliceExpand<E>) {
125        let len = source.__expand_len_method(scope);
126        let unroll = source
127            .const_len()
128            .is_some_and(|it| it <= MEMCPY_UNROLL_LIMIT);
129        for_expand(
130            scope,
131            range::expand(scope, 0usize.into_expand(scope), len),
132            unroll,
133            |scope, idx| {
134                copy::expand(
135                    scope,
136                    source.__expand_index_method(scope, idx),
137                    self.__expand_index_mut_method(scope, idx),
138                );
139            },
140        );
141    }
142}
cubecl_core/frontend/container/slice/operator.rs

cubecl_core/frontend/container/slice/
operator.rs
