cubecl_core/frontend/container/array/
base.rs

1use core::ops::{Deref, DerefMut};
2
3use cubecl_ir::{Scope, VectorSize};
4
5use crate::frontend::{CubePrimitive, NativeExpand};
6use crate::prelude::*;
7use crate::{self as cubecl};
8use crate::{frontend::CubeType, ir::Type, unexpanded};
9use cubecl_macros::{cube, intrinsic};
10
11/// A contiguous array of elements.
12#[derive(Clone, Copy)]
13pub struct Array<E> {
14    _buffer: [E; 1],
15}
16
17type ArrayExpand<E> = NativeExpand<Array<E>>;
18
19impl<E> AsMutExpand for ArrayExpand<E> {
20    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
21        self
22    }
23}
24
25/// Module that contains the implementation details of the new function.
26mod new {
27    use cubecl_macros::intrinsic;
28
29    use super::*;
30    use crate::frontend::container::slice;
31
32    #[cube]
33    impl<T: CubePrimitive + Clone> Array<T> {
34        /// Create a new array of the given length.
35        pub fn new(#[comptime] length: usize) -> Self {
36            intrinsic!(|scope| {
37                // Allocate as a slice even though it's statically sized, so we can deref to it.
38                // Unlike Rust, we can't construct fat pointers ad-hoc without access to the scope,
39                // so it needs to be prepared in advance.
40                let elem = T::__expand_as_type(scope);
41                let ty = Type::array(elem, length);
42                let buffer = scope.create_local_mut(ty);
43                let slice = slice::from_raw_parts::<T>(
44                    scope,
45                    buffer,
46                    0usize.into_expand(scope),
47                    length.into_expand(scope),
48                );
49                slice.expand.into()
50            })
51        }
52    }
53}
54
55/// Module that contains the implementation details of the `vector_size` function.
56mod vector {
57    use super::*;
58
59    impl<P: CubePrimitive> Array<P> {
60        /// Get the size of each vector contained in the tensor.
61        ///
62        /// Same as the following:
63        ///
64        /// ```rust, ignore
65        /// let size = tensor[0].vector_size();
66        /// ```
67        pub fn vector_size(&self) -> VectorSize {
68            P::vector_size()
69        }
70
71        // Expand function of [size](Tensor::vector_size).
72        pub fn __expand_vector_size(
73            expand: <Self as CubeType>::ExpandType,
74            scope: &Scope,
75        ) -> VectorSize {
76            expand.__expand_vector_size_method(scope)
77        }
78    }
79}
80
81#[cube]
82impl<E: CubePrimitive> Array<E> {
83    /// Obtain the array length
84    #[allow(clippy::len_without_is_empty)]
85    pub fn len(&self) -> comptime_type!(usize) {
86        intrinsic!(|_| self.expand.ty.array_size())
87    }
88}
89
90impl<C: CubePrimitive> Assign for ArrayExpand<C> {
91    fn __expand_assign_method(&mut self, scope: &Scope, value: Self) {
92        let value = value.__extract_list(scope);
93        let arr = self.__extract_list(scope);
94        assert_eq!(
95            value.ty.array_size(),
96            arr.ty.array_size(),
97            "Can't assign differently sized arrays"
98        );
99        assign::expand_element(scope, value, arr);
100    }
101}
102
103impl<C: CubePrimitive> RuntimeAssign for ArrayExpand<C> {
104    fn init_mut(&self, scope: &Scope) -> Self::Expand {
105        Array::__expand_new(scope, self.expand.ty.array_size())
106    }
107}
108
109impl<C: CubeType> CubeType for Array<C> {
110    type ExpandType = NativeExpand<Array<C>>;
111}
112
113impl<C: CubeType> IntoMut for ArrayExpand<C> {
114    fn into_mut(self, _scope: &Scope) -> Self {
115        self
116    }
117}
118
119impl<T: CubePrimitive> SizedContainer<usize> for Array<T> {
120    fn len(&self) -> usize {
121        unexpanded!()
122    }
123}
124
125impl<T: CubePrimitive> SizedContainerExpand<usize> for ArrayExpand<T> {
126    fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
127        self.__expand_len_method(scope).into_expand(scope)
128    }
129}
130
131impl<T: CubeType> Iterator for Array<T> {
132    type Item = T;
133
134    fn next(&mut self) -> Option<Self::Item> {
135        unexpanded!()
136    }
137}
138
139impl<T: CubePrimitive> Deref for Array<T> {
140    type Target = [T];
141
142    fn deref(&self) -> &Self::Target {
143        unexpanded!()
144    }
145}
146
147impl<T: CubePrimitive> DerefMut for Array<T> {
148    fn deref_mut(&mut self) -> &mut Self::Target {
149        unexpanded!()
150    }
151}
152
153impl<T: CubePrimitive> Deref for ArrayExpand<T> {
154    type Target = SliceExpand<T>;
155
156    fn deref(&self) -> &Self::Target {
157        unsafe { self.as_type_ref_unchecked() }
158    }
159}
160
161impl<T: CubePrimitive> DerefMut for ArrayExpand<T> {
162    fn deref_mut(&mut self) -> &mut Self::Target {
163        unsafe { self.as_type_mut_unchecked() }
164    }
165}
166
167impl<T: CubePrimitive> List<T> for Array<T> {}
168impl<T: CubePrimitive> ListExpand<T> for ArrayExpand<T> {
169    fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
170        Array::<T>::__expand_len(scope, self).into_expand(scope)
171    }
172}
173
174impl<T: CubePrimitive> Vectorized for Array<T> {}
175impl<T: CubePrimitive> VectorizedExpand for ArrayExpand<T> {
176    fn vector_size(&self) -> VectorSize {
177        self.expand.ty.vector_size()
178    }
179}
cubecl_core/frontend/container/array/base.rs

cubecl_core/frontend/container/array/
base.rs
