cubecl_core/frontend/container/array/
base.rs

1use core::{
2    fmt::Debug,
3    ops::{Deref, DerefMut},
4};
5
6use alloc::vec::Vec;
7use cubecl_ir::{
8    Scope, VectorSize,
9    interfaces::MaybeVectorizedType,
10    pliron::{
11        r#type::{Typed, TypedHandle},
12        value::Value,
13    },
14    prelude::TypeHandle,
15    read_value,
16    types::{ArrayType, PointerType, aggregate::SliceType},
17};
18
19use crate::frontend::{
20    CubePrimitive, NativeExpand,
21    container::iter::{SizedContainer, SizedContainerExpand},
22};
23use crate::prelude::*;
24use crate::{self as cubecl};
25use crate::{frontend::CubeType, unexpanded};
26use cubecl_macros::{cube, intrinsic};
27
28/// A contiguous array of elements.
29#[derive(Clone, Copy)]
30pub struct Array<E> {
31    _buffer: [E; 1],
32}
33
34type ArrayExpand<E> = NativeExpand<Array<E>>;
35
36impl<E> AsMutExpand for ArrayExpand<E> {
37    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
38        self
39    }
40}
41
42impl<E, const N: usize> AsMutExpand for NativeExpand<[E; N]> {
43    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
44        self
45    }
46}
47
48/// Module that contains the implementation details of the new function.
49mod new {
50    use cubecl_ir::{
51        attributes::ZeroAttr, dialect::vector::CompositeConstructOp, types::ArrayType,
52    };
53    use cubecl_macros::intrinsic;
54
55    use super::*;
56    use crate::frontend::container::slice;
57
58    #[cube]
59    impl<T: CubePrimitive> Array<T> {
60        /// Create a new array of the given length.
61        pub fn new(#[comptime] length: usize) -> Self {
62            intrinsic!(|scope| {
63                // Allocate as a slice even though it's statically sized, so we can deref to it.
64                // Unlike Rust, we can't construct fat pointers ad-hoc without access to the scope,
65                // so it needs to be prepared in advance.
66                let elem = T::__expand_as_type(scope);
67                let ty = ArrayType::get(scope.ctx(), elem, length);
68                let null = ZeroAttr::new(ty);
69                let buffer = scope.create_local_mut(ty, Some(null.into()));
70                let slice = slice::from_raw_parts::<T>(
71                    scope,
72                    buffer,
73                    0usize.into_expand(scope),
74                    length.into_expand(scope),
75                );
76                slice.expand.into()
77            })
78        }
79    }
80
81    impl<T: CubePrimitive> ArrayExpand<T> {
82        pub fn __expand_clone_method(&self, scope: &Scope) -> Self {
83            Self::into_mut(*self, scope)
84        }
85    }
86
87    #[doc(hidden)]
88    impl<T: CubePrimitive, const N: usize> NativeExpand<[T; N]> {
89        pub fn __expand_zeroed(scope: &Scope) -> Self {
90            let elem = T::__expand_as_type(scope);
91            let ty = ArrayType::get(scope.ctx(), elem, N);
92            let null = ZeroAttr::new(ty);
93            let buffer = scope.create_local_mut(ty, Some(null.into()));
94            let slice = slice::from_raw_parts::<T>(
95                scope,
96                buffer,
97                0usize.into_expand(scope),
98                N.into_expand(scope),
99            );
100            slice.expand.into()
101        }
102
103        pub fn __expand_construct(scope: &Scope, values: [NativeExpand<T>; N]) -> Self {
104            let length = values.len();
105            let values = values.into_iter().map(|it| it.read_value(scope)).collect();
106            let elem = T::__expand_as_type(scope);
107            let ty = ArrayType::get(scope.ctx(), elem, length).to_handle();
108            let value =
109                scope.register_with_result(&CompositeConstructOp::new(scope.ctx_mut(), ty, values));
110            let uninit = Self::__expand_zeroed(scope);
111            assign::expand_element(scope, value.into(), uninit.__extract_list(scope).into());
112            uninit
113        }
114
115        pub fn __expand_clone_method(&self, scope: &Scope) -> Self {
116            Self::into_mut(*self, scope)
117        }
118    }
119}
120
121/// Module that contains the implementation details of the `vector_size` function.
122mod vector {
123    use cubecl_ir::{interfaces::TypedExt, read_value};
124
125    use super::*;
126
127    impl<P: CubePrimitive> Array<P> {
128        /// Get the size of each vector contained in the array.
129        ///
130        /// Same as the following:
131        ///
132        /// ```rust, ignore
133        /// let size = tensor[0].vector_size();
134        /// ```
135        pub fn vector_size(&self) -> VectorSize {
136            P::vector_size()
137        }
138
139        // Expand function of [size](Tensor::vector_size).
140        pub fn __expand_vector_size(
141            expand: <Self as CubeType>::ExpandType,
142            scope: &Scope,
143        ) -> VectorSize {
144            expand.__expand_vector_size_method(scope)
145        }
146    }
147
148    #[cube]
149    impl<P: CubePrimitive> Array<P> {
150        pub fn into_vector<N: Size>(self) -> Vector<P::Scalar, N> {
151            intrinsic!(|scope| {
152                let arr = read_value(scope, self.__extract_list(scope));
153                let vec_ty = Vector::<P::Scalar, N>::__expand_as_type(scope);
154                reinterpret_value(scope, arr, vec_ty).into()
155            })
156        }
157
158        pub fn from_vector<S: Scalar, N: Size>(vector: Vector<S, N>) -> Array<P> {
159            intrinsic!(|scope| {
160                let vec = vector.read_value(scope);
161                let vec_p = P::__expand_vector_size(scope);
162                let len = vec.vector_size(scope.ctx()) / vec_p;
163                let arr_ty =
164                    ArrayType::get(scope.ctx(), P::__expand_as_type(scope), len).to_handle();
165                reinterpret_value(scope, vec, arr_ty).into()
166            })
167        }
168    }
169}
170
171#[cube]
172impl<E: CubePrimitive> Array<E> {
173    /// Obtain the array length
174    #[allow(clippy::len_without_is_empty)]
175    pub fn len(&self) -> comptime_type!(usize) {
176        intrinsic!(|scope| {
177            let ty = inner_array_ty(scope, self.value(scope));
178            ty.deref(scope.ctx()).length
179        })
180    }
181}
182
183impl<T: CubePrimitive, const N: usize> NativeExpand<[T; N]> {
184    pub fn __expand_len_method(&self, _: &Scope) -> usize {
185        N
186    }
187}
188
189impl<C: CubePrimitive> Assign for ArrayExpand<C> {
190    fn __expand_assign_method(&mut self, scope: &Scope, value: Self) {
191        let value = value.read_value(scope);
192        let arr = self.__extract_list(scope);
193        assign::expand_element(scope, value.into(), arr.into());
194    }
195}
196
197impl<C: CubePrimitive, const N: usize> Assign for NativeExpand<[C; N]> {
198    fn __expand_assign_method(&mut self, scope: &Scope, value: Self) {
199        let value = value.read_value(scope);
200        let arr = self.__extract_list(scope);
201        assign::expand_element(scope, value.into(), arr.into());
202    }
203}
204
205impl<C: CubePrimitive> RuntimeAssign for ArrayExpand<C> {
206    fn init_mut(&self, scope: &Scope) -> Self::Expand {
207        let ty = inner_array_ty(scope, self.value(scope));
208        let length = ty.deref(scope.ctx()).length;
209        Array::__expand_new(scope, length)
210    }
211}
212
213impl<C: CubePrimitive, const N: usize> RuntimeAssign for NativeExpand<[C; N]> {
214    fn init_mut(&self, scope: &Scope) -> Self::Expand {
215        Self::__expand_zeroed(scope)
216    }
217}
218impl<C: CubePrimitive, const N: usize> ConstRuntimeAssign for NativeExpand<[C; N]> {
219    fn const_init_mut(scope: &Scope) -> Self::Expand {
220        Self::__expand_zeroed(scope)
221    }
222}
223
224impl<C: CubePrimitive> CubeType for Array<C> {
225    type ExpandType = NativeExpand<Array<C>>;
226}
227
228impl<C: CubePrimitive, const N: usize> CubeType for [C; N] {
229    type ExpandType = NativeExpand<[C; N]>;
230}
231
232impl<T: CubePrimitive, const N: usize> ConstantTyped for [T; N] {
233    fn __expand_as_type(scope: &Scope) -> TypeHandle {
234        let inner = T::__expand_as_type(scope);
235        ArrayType::get(scope.ctx(), inner, N).to_handle()
236    }
237}
238
239impl<T: CubePrimitive> ReadValue for NativeExpand<Array<T>> {
240    fn read_value(&self, scope: &Scope) -> Value {
241        read_value(scope, self.__extract_list(scope))
242    }
243}
244
245impl<T: CubePrimitive, const N: usize> ReadValue for NativeExpand<[T; N]> {
246    fn read_value(&self, scope: &Scope) -> Value {
247        read_value(scope, self.__extract_list(scope))
248    }
249}
250
251impl<C: CubePrimitive> IntoMut for ArrayExpand<C> {
252    fn into_mut(self, scope: &Scope) -> Self {
253        let mut out = Array::__expand_new(scope, self.__expand_len_method(scope));
254        out.__expand_assign_method(scope, self);
255        out
256    }
257}
258
259impl<C: CubePrimitive, const N: usize> IntoMut for NativeExpand<[C; N]> {
260    fn into_mut(self, scope: &Scope) -> Self {
261        let mut uninit = Self::__expand_zeroed(scope);
262        uninit.__expand_assign_method(scope, self);
263        uninit
264    }
265}
266
267impl<T: CubePrimitive> SizedContainer<usize> for Array<T> {
268    fn len(&self) -> usize {
269        unexpanded!()
270    }
271}
272
273impl<T: CubePrimitive> SizedContainerExpand<usize> for ArrayExpand<T> {
274    fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
275        self.__expand_len_method(scope).into_expand(scope)
276    }
277}
278
279impl<T: CubePrimitive, const N: usize> SizedContainer<usize> for [T; N] {
280    fn len(&self) -> usize {
281        N
282    }
283}
284
285impl<T: CubePrimitive, const N: usize> SizedContainerExpand<usize> for NativeExpand<[T; N]> {
286    fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
287        N.into_expand(scope)
288    }
289}
290
291impl<T: CubeType> Iterator for Array<T> {
292    type Item = T;
293
294    fn next(&mut self) -> Option<Self::Item> {
295        unexpanded!()
296    }
297}
298
299impl<T: CubePrimitive> Deref for Array<T> {
300    type Target = [T];
301
302    fn deref(&self) -> &Self::Target {
303        unexpanded!()
304    }
305}
306
307impl<T: CubePrimitive> DerefMut for Array<T> {
308    fn deref_mut(&mut self) -> &mut Self::Target {
309        unexpanded!()
310    }
311}
312
313impl<T: CubePrimitive> Deref for ArrayExpand<T> {
314    type Target = SliceExpand<T>;
315
316    fn deref(&self) -> &Self::Target {
317        unsafe { self.as_type_ref_unchecked() }
318    }
319}
320
321impl<T: CubePrimitive> DerefMut for ArrayExpand<T> {
322    fn deref_mut(&mut self) -> &mut Self::Target {
323        unsafe { self.as_type_mut_unchecked() }
324    }
325}
326
327impl<T: CubePrimitive, const N: usize> Deref for NativeExpand<[T; N]> {
328    type Target = SliceExpand<T>;
329
330    fn deref(&self) -> &Self::Target {
331        unsafe { self.as_type_ref_unchecked() }
332    }
333}
334
335impl<T: CubePrimitive, const N: usize> DerefMut for NativeExpand<[T; N]> {
336    fn deref_mut(&mut self) -> &mut Self::Target {
337        unsafe { self.as_type_mut_unchecked() }
338    }
339}
340
341impl<T: CubePrimitive> List<T> for Array<T> {}
342impl<T: CubePrimitive> ListExpand<T> for ArrayExpand<T> {
343    fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
344        Array::<T>::__expand_len(scope, self).into_expand(scope)
345    }
346}
347
348impl<T: CubePrimitive> Vectorized for Array<T> {}
349impl<T: CubePrimitive> VectorizedExpand for ArrayExpand<T> {
350    fn __expand_vector_size_method(&self, scope: &Scope) -> VectorSize {
351        let ty = inner_array_ty(scope, self.value(scope));
352        ty.deref(scope.ctx()).vector_size(scope.ctx())
353    }
354}
355
356// `len` needs to be implemented because array doesn't implement it directly and this one takes
357// priority over the one on slice
358impl<T: CubePrimitive, const N: usize> List<T> for [T; N] {
359    fn len(&self) -> usize {
360        N
361    }
362}
363impl<T: CubePrimitive, const N: usize> ListExpand<T> for NativeExpand<[T; N]> {
364    fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
365        self.__expand_len_method(scope).into_expand(scope)
366    }
367}
368
369impl<T: CubePrimitive, const N: usize> Vectorized for [T; N] {}
370impl<T: CubePrimitive, const N: usize> VectorizedExpand for NativeExpand<[T; N]> {
371    fn __expand_vector_size_method(&self, scope: &Scope) -> VectorSize {
372        let ty = inner_array_ty(scope, self.value(scope));
373        ty.deref(scope.ctx()).vector_size(scope.ctx())
374    }
375}
376
377pub(crate) fn inner_array_ty(scope: &Scope, value: Value) -> TypedHandle<ArrayType> {
378    let ctx = scope.ctx();
379    let ty = value.get_type(ctx).deref(ctx);
380    let SliceType { base_ty, .. } = *ty.downcast_ref().unwrap();
381    let base_ty = base_ty.deref(ctx);
382    let PointerType { inner, .. } = base_ty.downcast_ref().unwrap();
383    TypedHandle::from_handle(*inner, ctx).unwrap()
384}
385
386impl<T: LaunchArg + CubePrimitive + Debug + 'static, const N: usize> LaunchArg for [T; N] {
387    type RuntimeArg = [T::RuntimeArg; N];
388    type CompilationArg = [T::CompilationArg; N];
389
390    fn register(arg: Self::RuntimeArg, launcher: &mut KernelLauncher) -> Self::CompilationArg {
391        arg.into_iter()
392            .map(|arg| T::register(arg, launcher))
393            .collect::<Vec<_>>()
394            .try_into()
395            .unwrap()
396    }
397
398    fn expand(arg: &Self::CompilationArg, builder: &mut KernelBuilder) -> NativeExpand<[T; N]> {
399        let values = arg
400            .iter()
401            .map(|value| T::expand(value, builder))
402            .collect::<Vec<_>>()
403            .try_into()
404            .unwrap();
405
406        NativeExpand::<[T; N]>::__expand_construct(builder, values)
407    }
408}
cubecl_core/frontend/container/array/base.rs

cubecl_core/frontend/container/array/
base.rs
