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

cubecl_core/frontend/container/slice/
operator.rs
