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

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

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