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> {}
55
56#[cube]
57impl<E: CubePrimitive> Array<E> {
58    pub fn as_slice(&self) -> &[E] {
59        intrinsic!(|_| self.deref())
60    }
61
62    pub fn as_mut_slice(&mut self) -> &mut [E] {
63        intrinsic!(|_| self.deref_mut())
64    }
65}
66
67#[cube]
68impl<E: CubePrimitive> Shared<[E]> {
69    pub fn as_slice(&self) -> &[E] {
70        intrinsic!(|_| self.deref())
71    }
72
73    pub fn as_mut_slice(&mut self) -> &mut [E] {
74        intrinsic!(|_| self.deref_mut())
75    }
76}
77
78impl<E: CubePrimitive> SliceOperator<E> for [E] {}
79impl<E: CubePrimitive> SliceOperatorExpand<E> for SliceExpand<E> {
80    fn __expand_slice_method(
81        &self,
82        scope: &Scope,
83        start: NativeExpand<usize>,
84        end: NativeExpand<usize>,
85    ) -> &SliceExpand<E> {
86        let length = end.__expand_sub_method(scope, start);
87        let list = self.__extract_list(scope);
88        let offset = self.__extract_offset(scope);
89        let offset = start.__expand_add_method(scope, offset);
90        let slice = slice::from_raw_parts(scope, list, offset, length);
91        scope.create_kernel_ref(slice)
92    }
93
94    fn __expand_slice_mut_method(
95        &mut self,
96        scope: &Scope,
97        start: NativeExpand<usize>,
98        end: NativeExpand<usize>,
99    ) -> &mut SliceExpand<E> {
100        let length = end.__expand_sub_method(scope, start);
101        let list = self.__extract_list(scope);
102        let offset = self.__extract_offset(scope);
103        let offset = start.__expand_add_method(scope, offset);
104
105        let slice = slice::from_raw_parts(scope, list, offset, length);
106        scope.create_kernel_ref(slice)
107    }
108}
109
110#[cube]
111pub trait SliceOperator<E: CubePrimitive> {
112    /// Return a read-only view of all elements comprise between the `start` and `end` indices.
113    /// In `checked` mode, if the `end` index is out-of-bound, it is replaced by
114    /// the length of `self`.
115    #[allow(unused_variables)]
116    fn slice(&self, start: usize, end: usize) -> &[E] {
117        unexpanded!()
118    }
119
120    /// Return a read-write view of all elements comprise between the `start` and `end` indices.
121    /// In `checked` mode, if the `end` index is out-of-bound, it is replaced by
122    /// the length of `self`.
123    #[allow(unused_variables)]
124    fn slice_mut(&mut self, start: usize, end: usize) -> &mut [E] {
125        unexpanded!()
126    }
127}
128
129// Simple heuristic
130const MEMCPY_UNROLL_LIMIT: usize = 8;
131
132impl<E: CubePrimitive> SliceExpand<E> {
133    pub fn __expand_copy_from_slice_method(&mut self, scope: &Scope, source: &SliceExpand<E>) {
134        let len = source.__expand_len_method(scope);
135        let unroll = source
136            .const_len()
137            .is_some_and(|it| it <= MEMCPY_UNROLL_LIMIT);
138        for_expand(
139            scope,
140            range::expand(scope, 0usize.into_expand(scope), len),
141            unroll,
142            |scope, idx| {
143                copy::expand(
144                    scope,
145                    source.__expand_index_method(scope, idx),
146                    self.__expand_index_mut_method(scope, idx),
147                );
148            },
149        );
150    }
151}
cubecl_core/frontend/container/slice/operator.rs

cubecl_core/frontend/container/slice/
operator.rs
