ruda_kernel/dsl/frontend/container/tensor/
base.rs

1use crate::dsl::{
2    frontend::{RudaPrimitive, RudaType, NativeExpand, SizedContainer},
3    ir::{Metadata, Scope},
4    prelude::*,
5    unexpanded,
6};
7use core::{
8    marker::PhantomData,
9    ops::{Deref, DerefMut},
10};
11use ruda_core::ir::VectorSize;
12use ruda_kernel_macros::{ruda, intrinsic};
13
14
15/// The tensor type is similar to the [array type](crate::dsl::prelude::Array), however it comes with more
16/// metadata such as [stride](Tensor::stride) and [shape](Tensor::shape).
17#[derive(new, Clone, Copy)]
18pub struct Tensor<T: RudaType> {
19    _val: PhantomData<T>,
20}
21
22type TensorExpand<T> = NativeExpand<Tensor<T>>;
23
24/// Module that contains the implementation details of the metadata functions.
25mod metadata {
26    use ruda_core::ir::ManagedVariable;
27
28    use super::*;
29    use crate::dsl::{
30        ir::{Arithmetic, BinaryOperator, Instruction},
31        prelude::Array,
32    };
33
34    #[ruda]
35    impl<T: RudaType> Tensor<T> {
36        /// Obtain the stride of input at dimension dim
37        #[allow(unused_variables)]
38        pub fn stride(&self, dim: usize) -> usize {
39            intrinsic!(|scope| {
40                let dim: ManagedVariable = dim.into();
41                let out = scope.create_local(usize::as_type(scope));
42                scope.register(Instruction::new(
43                    Metadata::Stride {
44                        dim: *dim,
45                        var: self.expand.into(),
46                    },
47                    out.clone().into(),
48                ));
49                out.into()
50            })
51        }
52
53        /// Obtain the shape of input at dimension dim
54        #[allow(unused_variables)]
55        pub fn shape(&self, dim: usize) -> usize {
56            intrinsic!(|scope| {
57                let dim: ManagedVariable = dim.into();
58                let out = scope.create_local(usize::as_type(scope));
59                scope.register(Instruction::new(
60                    Metadata::Shape {
61                        dim: *dim,
62                        var: self.expand.into(),
63                    },
64                    out.clone().into(),
65                ));
66                out.into()
67            })
68        }
69
70        /// Obtain the coordinate corresponding to the given `index` of the tensor at dimension `dim`.
71        ///
72        /// A coordinate is a list of indices corresponding to the multi-dimensional position of an element in the tensor.
73        /// The `dim` element in a coordinate is the position along the `dim` dimension of the tensor.
74        #[allow(unused_variables)]
75        pub fn coordinate(&self, index: usize, dim: usize) -> usize {
76            intrinsic!(|scope| {
77                let index: ManagedVariable = index.into();
78                let stride = self.clone().__expand_stride_method(scope, dim.clone());
79                let shape = self.clone().__expand_shape_method(scope, dim.clone());
80
81                // Compute `num_strides = index / stride`.
82                let num_strides = scope.create_local(usize::as_type(scope));
83                scope.register(Instruction::new(
84                    Arithmetic::Div(BinaryOperator {
85                        lhs: *index,
86                        rhs: stride.expand.into(),
87                    }),
88                    num_strides.clone().into(),
89                ));
90
91                // Compute `coordinate = num_strides % shape `.
92                let coordinate = scope.create_local(usize::as_type(scope));
93                scope.register(Instruction::new(
94                    Arithmetic::Modulo(BinaryOperator {
95                        lhs: *num_strides,
96                        rhs: shape.expand.into(),
97                    }),
98                    coordinate.clone().into(),
99                ));
100
101                coordinate.into()
102            })
103        }
104
105        /// The number of vectorized elements in the tensor.
106        ///
107        /// # Warning
108        ///
109        /// The length will be affected by the vectorization factor. To obtain the number of elements,
110        /// you should multiply the length by the vectorization factor.
111        #[allow(clippy::len_without_is_empty)]
112        pub fn len(&self) -> usize {
113            intrinsic!(|scope| {
114                let elem: NativeExpand<Array<u32>> = self.expand.into();
115                elem.__expand_len_method(scope)
116            })
117        }
118
119        /// The length of the buffer representing the tensor in terms of vectorized elements.
120        ///
121        /// # Warning
122        ///
123        /// The buffer length will be affected by the vectorization factor. To obtain the number of
124        /// elements, you should multiply the length by the vectorization factor.
125        #[allow(clippy::len_without_is_empty)]
126        pub fn buffer_len(&self) -> usize {
127            intrinsic!(|scope| {
128                let elem: NativeExpand<Array<u32>> = self.expand.into();
129                elem.__expand_buffer_len_method(scope)
130            })
131        }
132
133        /// Returns the rank of the tensor.
134        pub fn rank(&self) -> usize {
135            intrinsic!(|scope| {
136                let out = scope.create_local(usize::as_type(scope));
137                scope.register(Instruction::new(Metadata::Rank { var: *self.expand }, *out));
138                out.into()
139            })
140        }
141    }
142}
143
144/// Module that contains the implementation details of the index functions.
145mod indexation {
146    use ruda_core::ir::{IndexAssignOperator, IndexOperator, Operator};
147
148    use crate::dsl::ir::Instruction;
149
150    use super::*;
151
152    #[ruda]
153    impl<E: RudaPrimitive> Tensor<E> {
154        /// Perform an unchecked index into the array
155        ///
156        /// # Safety
157        /// Out of bounds indexing causes undefined behaviour and may segfault. Ensure index is
158        /// always in bounds
159        #[allow(unused_variables)]
160        pub unsafe fn index_unchecked(&self, i: usize) -> &E {
161            intrinsic!(|scope| {
162                let out = scope.create_local(self.expand.ty);
163                scope.register(Instruction::new(
164                    Operator::UncheckedIndex(IndexOperator {
165                        list: *self.expand,
166                        index: i.expand.consume(),
167                        vector_size: 0,
168                        unroll_factor: 1,
169                    }),
170                    *out,
171                ));
172                out.into()
173            })
174        }
175
176        /// Perform an unchecked index assignment into the array
177        ///
178        /// # Safety
179        /// Out of bounds indexing causes undefined behaviour and may segfault. Ensure index is
180        /// always in bounds
181        #[allow(unused_variables)]
182        pub unsafe fn index_assign_unchecked(&mut self, i: usize, value: E) {
183            intrinsic!(|scope| {
184                scope.register(Instruction::new(
185                    Operator::UncheckedIndexAssign(IndexAssignOperator {
186                        index: i.expand.consume(),
187                        value: value.expand.consume(),
188                        vector_size: 0,
189                        unroll_factor: 1,
190                    }),
191                    *self.expand,
192                ));
193            })
194        }
195    }
196}
197
198/// Module that contains the implementation details of the `vector_size` function.
199mod vector {
200    use super::*;
201
202    impl<P: Scalar, N: Size> Tensor<Vector<P, N>> {
203        /// Get the size of each vector contained in the tensor.
204        ///
205        /// Same as the following:
206        ///
207        /// ```rust, ignore
208        /// let size = tensor[0].size();
209        /// ```
210        pub fn vector_size(&self) -> VectorSize {
211            N::value()
212        }
213
214        // Expand function of [size](Tensor::vector_size).
215        pub fn __expand_vector_size(
216            expand: <Self as RudaType>::ExpandType,
217            scope: &mut Scope,
218        ) -> VectorSize {
219            expand.__expand_vector_size_method(scope)
220        }
221    }
222}
223
224impl<T: RudaPrimitive> SizedContainer for Tensor<T> {
225    type Item = T;
226}
227
228impl<T: RudaType> Iterator for &Tensor<T> {
229    type Item = T;
230
231    fn next(&mut self) -> Option<Self::Item> {
232        unexpanded!()
233    }
234}
235
236impl<T: RudaType> RudaType for Tensor<T> {
237    type ExpandType = NativeExpand<Tensor<T>>;
238}
239
240impl<T: RudaType> RudaType for *const Tensor<T> {
241    type ExpandType = NativeExpand<Tensor<T>>;
242}
243
244impl<T: RudaType> RudaType for *mut Tensor<T> {
245    type ExpandType = NativeExpand<Tensor<T>>;
246}
247
248impl<T: RudaType> RudaType for &mut Tensor<T> {
249    type ExpandType = NativeExpand<Tensor<T>>;
250}
251
252impl<T: RudaType> RudaType for &Tensor<T> {
253    type ExpandType = NativeExpand<Tensor<T>>;
254}
255
256impl<C: RudaType> IntoMut for NativeExpand<Tensor<C>> {
257    fn into_mut(self, _scope: &mut Scope) -> Self {
258        self
259    }
260}
261
262impl<T: RudaPrimitive> List<T> for Tensor<T> {
263    fn __expand_read(
264        scope: &mut Scope,
265        this: NativeExpand<Tensor<T>>,
266        idx: NativeExpand<usize>,
267    ) -> NativeExpand<T> {
268        index::expand(scope, this, idx)
269    }
270}
271
272impl<T: RudaPrimitive> Deref for Tensor<T> {
273    type Target = [T];
274
275    fn deref(&self) -> &Self::Target {
276        unexpanded!()
277    }
278}
279
280impl<T: RudaPrimitive> DerefMut for Tensor<T> {
281    fn deref_mut(&mut self) -> &mut Self::Target {
282        unexpanded!()
283    }
284}
285
286impl<T: RudaPrimitive> ListExpand<T> for NativeExpand<Tensor<T>> {
287    fn __expand_read_method(&self, scope: &mut Scope, idx: NativeExpand<usize>) -> NativeExpand<T> {
288        index::expand(scope, self.clone(), idx)
289    }
290    fn __expand_read_unchecked_method(
291        &self,
292        scope: &mut Scope,
293        idx: NativeExpand<usize>,
294    ) -> NativeExpand<T> {
295        index_unchecked::expand(scope, self.clone(), idx)
296    }
297
298    fn __expand_len_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
299        Self::__expand_len(scope, self.clone())
300    }
301}
302
303impl<T: RudaPrimitive> Vectorized for Tensor<T> {}
304impl<T: RudaPrimitive> VectorizedExpand for NativeExpand<Tensor<T>> {
305    fn vector_size(&self) -> VectorSize {
306        self.expand.ty.vector_size()
307    }
308}
309
310impl<T: RudaPrimitive> ListMut<T> for Tensor<T> {
311    fn __expand_write(
312        scope: &mut Scope,
313        this: NativeExpand<Tensor<T>>,
314        idx: NativeExpand<usize>,
315        value: NativeExpand<T>,
316    ) {
317        index_assign::expand(scope, this, idx, value);
318    }
319}
320
321impl<T: RudaPrimitive> ListMutExpand<T> for NativeExpand<Tensor<T>> {
322    fn __expand_write_method(
323        &self,
324        scope: &mut Scope,
325        idx: NativeExpand<usize>,
326        value: NativeExpand<T>,
327    ) {
328        index_assign::expand(scope, self.clone(), idx, value);
329    }
330}
ruda_kernel/dsl/frontend/container/tensor/base.rs

ruda_kernel/dsl/frontend/container/tensor/
base.rs
