cubecl_core/frontend/container/tensor/
base.rs

1use crate::{
2    ir::{Metadata, Scope},
3    prelude::*,
4    unexpanded,
5};
6use alloc::boxed::Box;
7use core::ops::{Deref, DerefMut};
8use cubecl_ir::VectorSize;
9
10use crate as cubecl;
11
12/// The tensor type is a wrapper around `[T]` that comes with more
13/// metadata such as [stride](Tensor::stride) and [shape](Tensor::shape).
14#[derive(CubeType)]
15pub struct Tensor<T: CubePrimitive> {
16    pub(super) meta: TensorMeta,
17    pub(super) buffer: [T],
18}
19
20#[derive(CubeType, Clone)]
21#[expand(derive(Clone))]
22pub struct OwnedTensor<T: CubePrimitive> {
23    #[allow(unused)]
24    pub(super) meta: TensorMeta,
25    pub(super) buffer: Box<[T]>,
26}
27
28impl<T: CubePrimitive> TensorExpand<T> {
29    /// Expand only because `[T]` can't be passed to a function
30    pub fn __expand_from_parts(meta: TensorMetaExpand, buffer: NativeExpand<[T]>) -> Self {
31        Self { meta, buffer }
32    }
33}
34
35#[cube]
36impl<T: CubePrimitive> OwnedTensor<T> {
37    pub fn from_parts(meta: TensorMeta, buffer: Box<[T]>) -> Self {
38        OwnedTensor::<T> { meta, buffer }
39    }
40}
41
42#[cube]
43impl<T: CubePrimitive> OwnedTensor<T> {
44    pub fn as_slice(&self) -> &[T] {
45        &self.buffer
46    }
47
48    pub fn as_mut_slice(&mut self) -> &mut [T] {
49        &mut self.buffer
50    }
51}
52
53#[cube]
54impl<T: CubePrimitive> Tensor<T> {
55    pub fn as_slice(&self) -> &[T] {
56        &self.buffer
57    }
58
59    pub fn as_mut_slice(&mut self) -> &mut [T] {
60        &mut self.buffer
61    }
62}
63
64/// Module that contains the implementation details of the metadata functions.
65mod metadata {
66    use cubecl_ir::Value;
67
68    use super::*;
69    use crate::ir::{Arithmetic, BinaryOperands, Instruction};
70
71    #[cube]
72    impl<T: CubePrimitive> Tensor<T> {
73        /// Obtain the stride of input at dimension dim
74        pub fn stride(&self, dim: usize) -> usize {
75            intrinsic!(|scope| {
76                let dim: Value = dim.into();
77                let list = self.__extract_list(scope);
78                let out = scope.create_value(usize::__expand_as_type(scope));
79                scope.register(Instruction::new(Metadata::Stride { dim, list }, out));
80                out.into()
81            })
82        }
83
84        /// Obtain the shape of input at dimension dim
85        pub fn shape(&self, dim: usize) -> usize {
86            intrinsic!(|scope| {
87                let dim: Value = dim.into();
88                let list = self.__extract_list(scope);
89                let out = scope.create_value(usize::__expand_as_type(scope));
90                scope.register(Instruction::new(Metadata::Shape { dim, list }, out));
91                out.into()
92            })
93        }
94
95        /// Obtain the coordinate corresponding to the given `index` of the tensor at dimension `dim`.
96        ///
97        /// A coordinate is a list of indices corresponding to the multi-dimensional position of an element in the tensor.
98        /// The `dim` element in a coordinate is the position along the `dim` dimension of the tensor.
99        pub fn coordinate(&self, index: usize, dim: usize) -> usize {
100            intrinsic!(|scope| {
101                let index: Value = index.into();
102                let stride = self.__expand_stride_method(scope, dim.clone());
103                let shape = self.__expand_shape_method(scope, dim.clone());
104
105                // Compute `num_strides = index / stride`.
106                let num_strides = scope.create_value(usize::__expand_as_type(scope));
107                scope.register(Instruction::new(
108                    Arithmetic::Div(BinaryOperands {
109                        lhs: index,
110                        rhs: stride.expand.into(),
111                    }),
112                    num_strides.clone().into(),
113                ));
114
115                // Compute `coordinate = num_strides % shape `.
116                let coordinate = scope.create_value(usize::__expand_as_type(scope));
117                scope.register(Instruction::new(
118                    Arithmetic::Rem(BinaryOperands {
119                        lhs: num_strides,
120                        rhs: shape.expand.into(),
121                    }),
122                    coordinate.clone().into(),
123                ));
124
125                coordinate.into()
126            })
127        }
128
129        /// The number of vectorized elements in the tensor.
130        ///
131        /// # Warning
132        ///
133        /// The length will be affected by the vectorization factor. To obtain the number of elements,
134        /// you should multiply the length by the vectorization factor.
135        #[allow(clippy::len_without_is_empty)]
136        pub fn len(&self) -> usize {
137            self.meta.len
138        }
139
140        /// The length of the buffer representing the tensor in terms of vectorized elements.
141        ///
142        /// # Warning
143        ///
144        /// The buffer length will be affected by the vectorization factor. To obtain the number of
145        /// elements, you should multiply the length by the vectorization factor.
146        #[allow(clippy::len_without_is_empty)]
147        pub fn buffer_len(&self) -> usize {
148            intrinsic!(|scope| { self.__extract_length(scope) })
149        }
150
151        /// Returns the rank of the tensor.
152        pub fn rank(&self) -> usize {
153            self.meta.rank
154        }
155    }
156}
157
158/// Module that contains the implementation details of the `vector_size` function.
159mod vector {
160    use super::*;
161
162    impl<P: Scalar, N: Size> Tensor<Vector<P, N>> {
163        /// Get the size of each vector contained in the tensor.
164        ///
165        /// Same as the following:
166        ///
167        /// ```rust, ignore
168        /// let size = tensor[0].size();
169        /// ```
170        pub fn vector_size(&self) -> VectorSize {
171            N::value()
172        }
173
174        // Expand function of [size](Tensor::vector_size).
175        pub fn __expand_vector_size(
176            expand: <Self as CubeType>::ExpandType,
177            scope: &Scope,
178        ) -> VectorSize {
179            expand.__expand_vector_size_method(scope)
180        }
181    }
182}
183
184impl<'a, E: CubePrimitive> From<&'a OwnedTensorExpand<E>> for &'a TensorExpand<E> {
185    fn from(value: &'a OwnedTensorExpand<E>) -> Self {
186        value.deref()
187    }
188}
189
190impl<'a, E: CubePrimitive> From<&'a mut OwnedTensorExpand<E>> for &'a mut TensorExpand<E> {
191    fn from(value: &'a mut OwnedTensorExpand<E>) -> Self {
192        value.deref_mut()
193    }
194}
195
196impl<'a, E: CubePrimitive> From<&'a TensorExpand<E>> for &'a SliceExpand<E> {
197    fn from(value: &'a TensorExpand<E>) -> Self {
198        value.deref()
199    }
200}
201
202impl<'a, E: CubePrimitive> From<&'a mut TensorExpand<E>> for &'a mut SliceExpand<E> {
203    fn from(value: &'a mut TensorExpand<E>) -> Self {
204        value.deref_mut()
205    }
206}
207
208impl<T: CubePrimitive> SizedContainer<usize> for Tensor<T> {
209    fn len(&self) -> usize {
210        unexpanded!()
211    }
212}
213
214impl<T: CubePrimitive> SizedContainerExpand<usize> for TensorExpand<T> {
215    fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
216        self.__expand_len_method(scope)
217    }
218}
219
220impl<T: CubePrimitive> Iterator for &Tensor<T> {
221    type Item = T;
222
223    fn next(&mut self) -> Option<Self::Item> {
224        unexpanded!()
225    }
226}
227
228impl<T: CubePrimitive> List<T> for Tensor<T> {}
229
230impl<T: CubePrimitive> Deref for Tensor<T> {
231    type Target = [T];
232
233    fn deref(&self) -> &Self::Target {
234        unexpanded!()
235    }
236}
237
238impl<T: CubePrimitive> DerefMut for Tensor<T> {
239    fn deref_mut(&mut self) -> &mut Self::Target {
240        unexpanded!()
241    }
242}
243
244impl<T: CubePrimitive> Deref for TensorExpand<T> {
245    type Target = SliceExpand<T>;
246
247    fn deref(&self) -> &Self::Target {
248        &self.buffer
249    }
250}
251
252impl<T: CubePrimitive> DerefMut for TensorExpand<T> {
253    fn deref_mut(&mut self) -> &mut Self::Target {
254        &mut self.buffer
255    }
256}
257
258impl<T: CubePrimitive> Deref for OwnedTensor<T> {
259    type Target = Tensor<T>;
260
261    fn deref(&self) -> &Self::Target {
262        unexpanded!()
263    }
264}
265
266impl<T: CubePrimitive> DerefMut for OwnedTensor<T> {
267    fn deref_mut(&mut self) -> &mut Self::Target {
268        unexpanded!()
269    }
270}
271
272impl<T: CubePrimitive> Deref for OwnedTensorExpand<T> {
273    type Target = TensorExpand<T>;
274
275    fn deref(&self) -> &Self::Target {
276        // SAFETY: Expand type has compatible layout since the type of the buffer is just a marker
277        unsafe { core::mem::transmute(self) }
278    }
279}
280
281impl<T: CubePrimitive> DerefMut for OwnedTensorExpand<T> {
282    fn deref_mut(&mut self) -> &mut Self::Target {
283        // SAFETY: Expand type has compatible layout since the type of the buffer is just a marker
284        unsafe { core::mem::transmute(self) }
285    }
286}
287
288impl<T: CubePrimitive> ListExpand<T> for TensorExpand<T> {
289    fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
290        Self::__expand_len_method(self, scope)
291    }
292}
293
294impl<T: CubePrimitive> Vectorized for Tensor<T> {}
295impl<T: CubePrimitive> VectorizedExpand for TensorExpand<T> {
296    fn vector_size(&self) -> VectorSize {
297        self.buffer.expand.ty.vector_size()
298    }
299}
cubecl_core/frontend/container/tensor/base.rs

cubecl_core/frontend/container/tensor/
base.rs
