ruda_kernel/dsl/frontend/container/
shared_memory.rs

1use core::marker::PhantomData;
2use core::ops::{Deref, DerefMut};
3
4use crate::dsl::{
5    prelude::{Vectorized, VectorizedExpand},
6    unexpanded,
7};
8use ruda_core::ir::{Marker, VariableKind, VectorSize};
9use ruda_kernel_macros::{ruda, intrinsic};
10
11use crate::dsl::{
12    frontend::{RudaPrimitive, RudaType, IntoMut, NativeExpand},
13    ir::Scope,
14    prelude::*,
15};
16
17pub type SharedMemoryExpand<T> = NativeExpand<SharedMemory<T>>;
18pub type SharedExpand<T> = NativeExpand<Shared<T>>;
19
20#[derive(Clone, Copy)]
21pub struct Shared<E: RudaPrimitive> {
22    _val: PhantomData<E>,
23}
24
25#[derive(Clone, Copy)]
26pub struct SharedMemory<E: RudaPrimitive> {
27    _val: PhantomData<E>,
28}
29
30impl<T: RudaPrimitive> IntoMut for NativeExpand<SharedMemory<T>> {
31    fn into_mut(self, _scope: &mut Scope) -> Self {
32        self
33    }
34}
35
36impl<T: RudaPrimitive> RudaType for SharedMemory<T> {
37    type ExpandType = NativeExpand<SharedMemory<T>>;
38}
39
40impl<T: RudaPrimitive> IntoMut for NativeExpand<Shared<T>> {
41    fn into_mut(self, _scope: &mut Scope) -> Self {
42        self
43    }
44}
45
46impl<T: RudaPrimitive> RudaType for Shared<T> {
47    type ExpandType = NativeExpand<Shared<T>>;
48}
49
50#[ruda]
51impl<T: RudaPrimitive + Clone> SharedMemory<T> {
52    #[allow(unused_variables)]
53    pub fn new(#[comptime] size: usize) -> Self {
54        intrinsic!(|scope| {
55            scope
56                .create_shared_array(T::as_type(scope), size, None)
57                .into()
58        })
59    }
60
61    #[allow(clippy::len_without_is_empty)]
62    pub fn len(&self) -> usize {
63        intrinsic!(|_| len_static(&self))
64    }
65
66    pub fn buffer_len(&self) -> usize {
67        self.len()
68    }
69}
70
71#[ruda]
72impl<T: RudaPrimitive> Shared<T> {
73    pub fn new() -> Self {
74        intrinsic!(|scope| {
75            let var = scope.create_shared(T::as_type(scope));
76            NativeExpand::new(var)
77        })
78    }
79}
80
81pub trait AsRefExpand<T: RudaType> {
82    /// Converts this type into a shared reference of the (usually inferred) input type.
83    fn __expand_as_ref_method(self, scope: &mut Scope) -> T::ExpandType;
84}
85impl<T: RudaPrimitive> AsRefExpand<T> for NativeExpand<T> {
86    fn __expand_as_ref_method(self, _scope: &mut Scope) -> NativeExpand<T> {
87        self
88    }
89}
90pub trait AsMutExpand<T: RudaType> {
91    /// Converts this type into a shared reference of the (usually inferred) input type.
92    fn __expand_as_mut_method(self, scope: &mut Scope) -> T::ExpandType;
93}
94impl<T: RudaPrimitive> AsMutExpand<T> for NativeExpand<T> {
95    fn __expand_as_mut_method(self, _scope: &mut Scope) -> <T as RudaType>::ExpandType {
96        self
97    }
98}
99
100/// Type inference won't allow things like assign to work normally, so we need to manually call
101/// `as_ref` or `as_mut` for those. Things like barrier ops should take `AsRef` so the conversion
102/// is automatic.
103impl<T: RudaPrimitive> AsRef<T> for Shared<T> {
104    fn as_ref(&self) -> &T {
105        unexpanded!()
106    }
107}
108impl<T: RudaPrimitive> AsRefExpand<T> for SharedExpand<T> {
109    fn __expand_as_ref_method(self, _scope: &mut Scope) -> <T as RudaType>::ExpandType {
110        self.expand.into()
111    }
112}
113
114impl<T: RudaPrimitive> AsMut<T> for Shared<T> {
115    fn as_mut(&mut self) -> &mut T {
116        unexpanded!()
117    }
118}
119impl<T: RudaPrimitive> AsMutExpand<T> for SharedExpand<T> {
120    fn __expand_as_mut_method(self, _scope: &mut Scope) -> <T as RudaType>::ExpandType {
121        self.expand.into()
122    }
123}
124
125impl<T: RudaPrimitive> Default for Shared<T> {
126    fn default() -> Self {
127        Self::new()
128    }
129}
130impl<T: RudaPrimitive> Shared<T> {
131    pub fn __expand_default(scope: &mut Scope) -> <Self as RudaType>::ExpandType {
132        Self::__expand_new(scope)
133    }
134}
135
136#[ruda]
137impl<T: RudaPrimitive + Clone> SharedMemory<T> {
138    #[allow(unused_variables)]
139    pub fn new_aligned(#[comptime] size: usize, #[comptime] alignment: usize) -> SharedMemory<T> {
140        intrinsic!(|scope| {
141            let var = scope.create_shared_array(T::as_type(scope), size, Some(alignment));
142            NativeExpand::new(var)
143        })
144    }
145
146    /// Frees the shared memory for reuse, if possible on the target runtime.
147    ///
148    /// # Safety
149    /// *Must* be used in uniform control flow
150    /// *Must not* have any dangling references to this shared memory
151    pub unsafe fn free(self) {
152        intrinsic!(|scope| { scope.register(Marker::Free(*self.expand)) })
153    }
154}
155
156fn len_static<T: RudaPrimitive>(shared: &NativeExpand<SharedMemory<T>>) -> NativeExpand<usize> {
157    let VariableKind::SharedArray { length, .. } = shared.expand.kind else {
158        unreachable!("Kind of shared memory is always shared memory")
159    };
160    length.into()
161}
162
163/// Module that contains the implementation details of the index functions.
164mod indexation {
165    use ruda_core::ir::{IndexAssignOperator, IndexOperator, Operator};
166
167    use crate::dsl::ir::Instruction;
168
169    use super::*;
170
171    type SharedMemoryExpand<E> = NativeExpand<SharedMemory<E>>;
172
173    #[ruda]
174    impl<E: RudaPrimitive> SharedMemory<E> {
175        /// Perform an unchecked index into the array
176        ///
177        /// # Safety
178        /// Out of bounds indexing causes undefined behaviour and may segfault. Ensure index is
179        /// always in bounds
180        #[allow(unused_variables)]
181        pub unsafe fn index_unchecked(&self, i: usize) -> &E {
182            intrinsic!(|scope| {
183                let out = scope.create_local(self.expand.ty);
184                scope.register(Instruction::new(
185                    Operator::UncheckedIndex(IndexOperator {
186                        list: *self.expand,
187                        index: i.expand.consume(),
188                        vector_size: 0,
189                        unroll_factor: 1,
190                    }),
191                    *out,
192                ));
193                out.into()
194            })
195        }
196
197        /// Perform an unchecked index assignment 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_assign_unchecked(&mut self, i: usize, value: E) {
204            intrinsic!(|scope| {
205                scope.register(Instruction::new(
206                    Operator::UncheckedIndexAssign(IndexAssignOperator {
207                        index: i.expand.consume(),
208                        value: value.expand.consume(),
209                        vector_size: 0,
210                        unroll_factor: 1,
211                    }),
212                    *self.expand,
213                ));
214            })
215        }
216    }
217}
218
219impl<T: RudaPrimitive> List<T> for SharedMemory<T> {
220    fn __expand_read(
221        scope: &mut Scope,
222        this: NativeExpand<SharedMemory<T>>,
223        idx: NativeExpand<usize>,
224    ) -> NativeExpand<T> {
225        index::expand(scope, this, idx)
226    }
227}
228
229impl<T: RudaPrimitive> Deref for SharedMemory<T> {
230    type Target = [T];
231
232    fn deref(&self) -> &Self::Target {
233        unexpanded!()
234    }
235}
236
237impl<T: RudaPrimitive> DerefMut for SharedMemory<T> {
238    fn deref_mut(&mut self) -> &mut Self::Target {
239        unexpanded!()
240    }
241}
242
243impl<T: RudaPrimitive> ListExpand<T> for NativeExpand<SharedMemory<T>> {
244    fn __expand_read_method(&self, scope: &mut Scope, idx: NativeExpand<usize>) -> NativeExpand<T> {
245        index::expand(scope, self.clone(), idx)
246    }
247    fn __expand_read_unchecked_method(
248        &self,
249        scope: &mut Scope,
250        idx: NativeExpand<usize>,
251    ) -> NativeExpand<T> {
252        index_unchecked::expand(scope, self.clone(), idx)
253    }
254
255    fn __expand_len_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
256        Self::__expand_len_method(self.clone(), scope)
257    }
258}
259
260impl<T: RudaPrimitive> Vectorized for SharedMemory<T> {}
261impl<T: RudaPrimitive> VectorizedExpand for NativeExpand<SharedMemory<T>> {
262    fn vector_size(&self) -> VectorSize {
263        self.expand.ty.vector_size()
264    }
265}
266
267impl<T: RudaPrimitive> ListMut<T> for SharedMemory<T> {
268    fn __expand_write(
269        scope: &mut Scope,
270        this: NativeExpand<SharedMemory<T>>,
271        idx: NativeExpand<usize>,
272        value: NativeExpand<T>,
273    ) {
274        index_assign::expand(scope, this, idx, value);
275    }
276}
277
278impl<T: RudaPrimitive> ListMutExpand<T> for NativeExpand<SharedMemory<T>> {
279    fn __expand_write_method(
280        &self,
281        scope: &mut Scope,
282        idx: NativeExpand<usize>,
283        value: NativeExpand<T>,
284    ) {
285        index_assign::expand(scope, self.clone(), idx, value);
286    }
287}
ruda_kernel/dsl/frontend/container/shared_memory.rs

ruda_kernel/dsl/frontend/container/
shared_memory.rs
