ruda_kernel/dsl/frontend/container/slice/
operator.rs

1use super::{ReadOnly, ReadWrite, Slice, SliceExpand, SliceOriginExpand, SliceVisibility};
2use crate::dsl::{ir::Scope, prelude::*, unexpanded};
3use ruda_core::tf32;
4use ruda_core::ir::ManagedVariable;
5
6pub(crate) fn is_tf32<C: RudaPrimitive, T: RudaPrimitive>(scope: &mut Scope) -> bool {
7    let ty_c = C::as_type(scope).storage_type();
8    let ty_t = T::as_type(scope).storage_type();
9    let ty_f32 = f32::as_type(scope).storage_type();
10    let ty_tf32 = tf32::as_type(scope).storage_type();
11
12    (ty_c == ty_f32 && ty_t == ty_tf32) || (ty_c == ty_tf32 && ty_t == ty_f32)
13}
14
15impl<E: RudaPrimitive> SliceOperator<E> for SharedMemory<E> {}
16impl<E: RudaPrimitive> SliceOperatorExpand<E> for NativeExpand<SharedMemory<E>> {
17    fn __expand_slice_method(
18        &self,
19        scope: &mut Scope,
20        start: NativeExpand<usize>,
21        end: NativeExpand<usize>,
22    ) -> SliceExpand<E, ReadOnly> {
23        Slice::__expand_new(
24            scope,
25            SliceOriginExpand::SharedMemory(self.clone()),
26            start,
27            end,
28        )
29    }
30
31    fn __expand_to_slice_method(&self, scope: &mut Scope) -> SliceExpand<E, ReadOnly> {
32        let len = expand_length_native(scope, *self.expand);
33
34        Slice::__expand_new(
35            scope,
36            SliceOriginExpand::SharedMemory(self.clone()),
37            NativeExpand::from_lit(scope, 0usize),
38            ManagedVariable::Plain(len).into(),
39        )
40    }
41}
42
43impl<E: RudaPrimitive> SliceMutOperator<E> for SharedMemory<E> {}
44impl<E: RudaPrimitive> SliceMutOperatorExpand<E> for NativeExpand<SharedMemory<E>> {
45    fn __expand_slice_mut_method(
46        &self,
47        scope: &mut Scope,
48        start: NativeExpand<usize>,
49        end: NativeExpand<usize>,
50    ) -> SliceExpand<E, ReadWrite> {
51        Slice::__expand_new(
52            scope,
53            SliceOriginExpand::SharedMemory(self.clone()),
54            start,
55            end,
56        )
57    }
58
59    fn __expand_to_slice_mut_method(&self, scope: &mut Scope) -> SliceExpand<E, ReadWrite> {
60        let len = expand_length_native(scope, *self.expand);
61
62        Slice::__expand_new(
63            scope,
64            SliceOriginExpand::SharedMemory(self.clone()),
65            NativeExpand::from_lit(scope, 0usize),
66            ManagedVariable::Plain(len).into(),
67        )
68    }
69}
70
71impl<E: RudaPrimitive> SliceOperator<E> for Tensor<E> {}
72impl<E: RudaPrimitive> SliceOperatorExpand<E> for NativeExpand<Tensor<E>> {
73    fn __expand_slice_method(
74        &self,
75        scope: &mut Scope,
76        start: NativeExpand<usize>,
77        end: NativeExpand<usize>,
78    ) -> SliceExpand<E, ReadOnly> {
79        Slice::__expand_new(scope, SliceOriginExpand::Tensor(self.clone()), start, end)
80    }
81
82    fn __expand_to_slice_method(&self, scope: &mut Scope) -> SliceExpand<E, ReadOnly> {
83        let len = self.clone().__expand_len_method(scope);
84        Slice::__expand_new(
85            scope,
86            SliceOriginExpand::Tensor(self.clone()),
87            NativeExpand::from_lit(scope, 0usize),
88            len,
89        )
90    }
91}
92
93impl<E: RudaPrimitive> SliceMutOperator<E> for Tensor<E> {}
94impl<E: RudaPrimitive> SliceMutOperatorExpand<E> for NativeExpand<Tensor<E>> {
95    fn __expand_slice_mut_method(
96        &self,
97        scope: &mut Scope,
98        start: NativeExpand<usize>,
99        end: NativeExpand<usize>,
100    ) -> SliceExpand<E, ReadWrite> {
101        Slice::__expand_new(scope, SliceOriginExpand::Tensor(self.clone()), start, end)
102    }
103
104    fn __expand_to_slice_mut_method(&self, scope: &mut Scope) -> SliceExpand<E, ReadWrite> {
105        let len = self.clone().__expand_len_method(scope);
106        Slice::__expand_new(
107            scope,
108            SliceOriginExpand::Tensor(self.clone()),
109            NativeExpand::from_lit(scope, 0usize),
110            len,
111        )
112    }
113}
114
115impl<E: RudaPrimitive> SliceOperator<E> for Array<E> {}
116impl<E: RudaPrimitive> SliceOperatorExpand<E> for NativeExpand<Array<E>> {
117    fn __expand_slice_method(
118        &self,
119        scope: &mut Scope,
120        start: NativeExpand<usize>,
121        end: NativeExpand<usize>,
122    ) -> SliceExpand<E, ReadOnly> {
123        Slice::__expand_new(scope, SliceOriginExpand::Array(self.clone()), start, end)
124    }
125
126    fn __expand_to_slice_method(&self, scope: &mut Scope) -> SliceExpand<E, ReadOnly> {
127        let len = self.clone().__expand_len_method(scope);
128        Slice::__expand_new(
129            scope,
130            SliceOriginExpand::Array(self.clone()),
131            NativeExpand::from_lit(scope, 0usize),
132            len,
133        )
134    }
135}
136
137impl<E: RudaPrimitive> SliceMutOperator<E> for Array<E> {}
138impl<E: RudaPrimitive> SliceMutOperatorExpand<E> for NativeExpand<Array<E>> {
139    fn __expand_slice_mut_method(
140        &self,
141        scope: &mut Scope,
142        start: NativeExpand<usize>,
143        end: NativeExpand<usize>,
144    ) -> SliceExpand<E, ReadWrite> {
145        Slice::__expand_new(scope, SliceOriginExpand::Array(self.clone()), start, end)
146    }
147
148    fn __expand_to_slice_mut_method(&self, scope: &mut Scope) -> SliceExpand<E, ReadWrite> {
149        let len = self.clone().__expand_len_method(scope);
150        Slice::__expand_new(
151            scope,
152            SliceOriginExpand::Array(self.clone()),
153            NativeExpand::from_lit(scope, 0usize),
154            len,
155        )
156    }
157}
158
159impl<E: RudaPrimitive, IO: SliceVisibility> SliceOperator<E> for Slice<E, IO> {}
160impl<E: RudaPrimitive, IO: SliceVisibility> SliceOperatorExpand<E> for SliceExpand<E, IO> {
161    fn __expand_slice_method(
162        &self,
163        scope: &mut Scope,
164        start: NativeExpand<usize>,
165        end: NativeExpand<usize>,
166    ) -> SliceExpand<E, ReadOnly> {
167        let length = crate::dsl::frontend::sub::expand(scope, end, start.clone());
168        let offset = crate::dsl::frontend::add::expand(scope, start, self.offset.clone());
169
170        SliceExpand {
171            origin: self.origin.clone(),
172            io: core::marker::PhantomData,
173            offset,
174            length,
175            vector_size: self.vector_size,
176        }
177    }
178
179    fn __expand_to_slice_method(&self, _scope: &mut Scope) -> SliceExpand<E, ReadOnly> {
180        SliceExpand {
181            origin: self.origin.clone(),
182            io: core::marker::PhantomData,
183            offset: self.offset.clone(),
184            length: self.length.clone(),
185            vector_size: self.vector_size,
186        }
187    }
188}
189
190impl<E: RudaPrimitive> SliceMutOperator<E> for Slice<E, ReadWrite> {}
191impl<E: RudaPrimitive> SliceMutOperatorExpand<E> for SliceExpand<E, ReadWrite> {
192    fn __expand_slice_mut_method(
193        &self,
194        scope: &mut Scope,
195        start: NativeExpand<usize>,
196        end: NativeExpand<usize>,
197    ) -> SliceExpand<E, ReadWrite> {
198        let length = crate::dsl::frontend::sub::expand(scope, end, start.clone());
199        let offset = crate::dsl::frontend::add::expand(scope, start, self.offset.clone());
200
201        SliceExpand {
202            origin: self.origin.clone(),
203            io: core::marker::PhantomData,
204            offset,
205            length,
206            vector_size: self.vector_size,
207        }
208    }
209
210    fn __expand_to_slice_mut_method(&self, _scope: &mut Scope) -> SliceExpand<E, ReadWrite> {
211        SliceExpand {
212            origin: self.origin.clone(),
213            io: core::marker::PhantomData,
214            offset: self.offset.clone(),
215            length: self.length.clone(),
216            vector_size: self.vector_size,
217        }
218    }
219}
220
221#[ruda(self_type = "ref")]
222pub trait SliceOperator<E: RudaPrimitive> {
223    /// Return a read-only view of all elements comprise between the `start` and `end` indices.
224    /// In `checked` mode, if the `end` index is out-of-bound, it is replaced by
225    /// the length of `self`.
226    #[allow(unused_variables)]
227    fn slice(&self, start: usize, end: usize) -> Slice<E, ReadOnly> {
228        unexpanded!()
229    }
230
231    /// Reinterprete the current type as a read-only slice.
232    #[allow(unused_variables)]
233    fn to_slice(&self) -> Slice<E, ReadOnly> {
234        unexpanded!()
235    }
236}
237
238#[ruda(self_type = "ref")]
239pub trait SliceMutOperator<E: RudaPrimitive> {
240    /// Return a read-write view of all elements comprise between the `start` and `end` indices.
241    /// In `checked` mode, if the `end` index is out-of-bound, it is replaced by
242    /// the length of `self`.
243    #[allow(unused_variables)]
244    fn slice_mut(&mut self, start: usize, end: usize) -> Slice<E, ReadWrite> {
245        unexpanded!()
246    }
247
248    /// Reinterprete the current type as a read-write slice.
249    #[allow(unused_variables)]
250    fn to_slice_mut(&mut self) -> Slice<E, ReadWrite> {
251        unexpanded!()
252    }
253}
254
255// Automatic implementation for references to SliceOperator.
256impl<'a, T: RudaPrimitive, L: SliceOperator<T>> SliceOperator<T> for &'a L where
257    &'a L: RudaType<ExpandType = L::ExpandType>
258{
259}
260
261// Automatic implementation for mutable references to SliceOperator.
262impl<'a, T: RudaPrimitive, L: SliceOperator<T>> SliceOperator<T> for &'a mut L where
263    &'a mut L: RudaType<ExpandType = L::ExpandType>
264{
265}
266
267// Automatic implementation for references to SliceMutOperator.
268impl<'a, T: RudaPrimitive, L: SliceMutOperator<T>> SliceMutOperator<T> for &'a L where
269    &'a L: RudaType<ExpandType = L::ExpandType>
270{
271}
272
273// Automatic implementation for mutable references to SliceMutOperator.
274impl<'a, T: RudaPrimitive, L: SliceMutOperator<T>> SliceMutOperator<T> for &'a mut L where
275    &'a mut L: RudaType<ExpandType = L::ExpandType>
276{
277}
ruda_kernel/dsl/frontend/container/slice/operator.rs

ruda_kernel/dsl/frontend/container/slice/
operator.rs
