ruda_kernel/dsl/frontend/container/
pointer.rs

1use crate::dsl::prelude::*;
2use alloc::boxed::Box;
3use core::marker::PhantomData;
4use ruda_core::ir::{
5    AddressSpace, Branch, ElemType, IfElse, Instruction, ManagedVariable, PointerAccess,
6    PointerType, TypedMemory, TypedPointer, Variable,
7};
8use ruda_kernel_macros::{intrinsic, ruda};
9
10#[cfg(test)]
11mod tests;
12
13mod sealed {
14    pub trait Sealed {}
15}
16
17pub trait PointerSpace: sealed::Sealed + Clone + Copy + Send + Sync + 'static {
18    const SPACE: AddressSpace;
19}
20pub trait PointerPermission: sealed::Sealed + SliceVisibility {
21    const ACCESS: PointerAccess;
22}
23
24#[derive(Clone, Copy)]
25pub struct GlobalSpace;
26#[derive(Clone, Copy)]
27pub struct Workgroup;
28#[derive(Clone, Copy)]
29pub struct Local;
30
31impl sealed::Sealed for GlobalSpace {}
32impl sealed::Sealed for Workgroup {}
33impl sealed::Sealed for Local {}
34impl sealed::Sealed for ReadOnly {}
35impl sealed::Sealed for ReadWrite {}
36impl PointerSpace for GlobalSpace {
37    const SPACE: AddressSpace = AddressSpace::Global;
38}
39impl PointerSpace for Workgroup {
40    const SPACE: AddressSpace = AddressSpace::Shared;
41}
42impl PointerSpace for Local {
43    const SPACE: AddressSpace = AddressSpace::Local;
44}
45impl PointerPermission for ReadOnly {
46    const ACCESS: PointerAccess = PointerAccess::ReadOnly;
47}
48impl PointerPermission for ReadWrite {
49    const ACCESS: PointerAccess = PointerAccess::ReadWrite;
50}
51
52#[derive(Clone)]
53enum PointerNode {
54    Element(TypedPointer),
55    Select {
56        condition: Variable,
57        then: Box<PointerNode>,
58        otherwise: Box<PointerNode>,
59    },
60}
61
62/// A logical element pointer whose storage space and permissions are part of its Rust type.
63///
64/// Read-only pointers cannot store:
65/// ```compile_fail
66/// use ruda_kernel::dsl::prelude::{GlobalSpace, Ptr};
67/// fn write(pointer: Ptr<f32, GlobalSpace>) { pointer.store(1.0); }
68/// ```
69/// Address spaces cannot be interchanged:
70/// ```compile_fail
71/// use ruda_kernel::dsl::prelude::{GlobalSpace, Ptr, Workgroup};
72/// fn global(_: Ptr<f32, GlobalSpace>) {}
73/// fn shared(pointer: Ptr<f32, Workgroup>) { global(pointer); }
74/// ```
75/// Read-only permissions cannot be upgraded:
76/// ```compile_fail
77/// use ruda_kernel::dsl::prelude::{GlobalSpace, Ptr, ReadWrite};
78/// fn upgrade(pointer: Ptr<f32, GlobalSpace>) -> Ptr<f32, GlobalSpace, ReadWrite> { pointer }
79/// ```
80#[derive(Clone, Copy)]
81pub struct Ptr<E: RudaPrimitive, S: PointerSpace, A: PointerPermission = ReadOnly> {
82    marker: PhantomData<(E, S, A)>,
83}
84
85pub struct PtrExpand<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> {
86    node: PointerNode,
87    marker: PhantomData<(E, S, A)>,
88}
89
90impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> Clone for PtrExpand<E, S, A> {
91    fn clone(&self) -> Self {
92        Self {
93            node: self.node.clone(),
94            marker: PhantomData,
95        }
96    }
97}
98impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> IntoMut for PtrExpand<E, S, A> {
99    fn into_mut(self, _: &mut Scope) -> Self {
100        self
101    }
102}
103impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> RudaDebug for PtrExpand<E, S, A> {}
104impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> RudaType for Ptr<E, S, A> {
105    type ExpandType = PtrExpand<E, S, A>;
106}
107impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> RudaType for &Ptr<E, S, A> {
108    type ExpandType = PtrExpand<E, S, A>;
109}
110
111impl PointerNode {
112    fn pointee(&self) -> Type {
113        match self {
114            Self::Element(pointer) => pointer.ty.pointee,
115            Self::Select { then, .. } => then.pointee(),
116        }
117    }
118
119    fn offset(&self, scope: &mut Scope, offset: NativeExpand<usize>) -> Self {
120        match self {
121            Self::Element(pointer) => {
122                let index = add::expand(
123                    scope,
124                    NativeExpand::<usize>::new(ManagedVariable::Plain(pointer.index)),
125                    offset,
126                );
127                Self::Element(TypedPointer {
128                    index: index.expand.consume(),
129                    ..*pointer
130                })
131            }
132            Self::Select {
133                condition,
134                then,
135                otherwise,
136            } => Self::Select {
137                condition: *condition,
138                then: Box::new(then.offset(scope, offset.clone())),
139                otherwise: Box::new(otherwise.offset(scope, offset)),
140            },
141        }
142    }
143
144    fn read_only(&self) -> Self {
145        match self {
146            Self::Element(pointer) => Self::Element(pointer.read_only()),
147            Self::Select {
148                condition,
149                then,
150                otherwise,
151            } => Self::Select {
152                condition: *condition,
153                then: Box::new(then.read_only()),
154                otherwise: Box::new(otherwise.read_only()),
155            },
156        }
157    }
158
159    fn emit(&self, scope: &mut Scope, value: Variable, store: bool) {
160        match self {
161            Self::Element(pointer) => {
162                let (op, out) = if store {
163                    (
164                        TypedMemory::Store {
165                            pointer: *pointer,
166                            value,
167                        },
168                        pointer.base,
169                    )
170                } else {
171                    (TypedMemory::Load(*pointer), value)
172                };
173                op.validate(Some(out)).expect("valid typed pointer access");
174                scope.register(Instruction::new(op, out));
175            }
176            Self::Select {
177                condition,
178                then,
179                otherwise,
180            } => {
181                let mut scope_if = scope.child();
182                let mut scope_else = scope.child();
183                then.emit(&mut scope_if, value, store);
184                otherwise.emit(&mut scope_else, value, store);
185                scope.register(Instruction::no_out(Branch::IfElse(Box::new(IfElse {
186                    cond: *condition,
187                    scope_if,
188                    scope_else,
189                }))));
190            }
191        }
192    }
193}
194
195#[ruda]
196impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> Ptr<E, S, A> {
197    pub fn offset(&self, offset: usize) -> Self {
198        intrinsic!(|scope| {
199            PtrExpand {
200                node: self.node.offset(scope, offset),
201                marker: PhantomData,
202            }
203        })
204    }
205
206    pub fn read_only(&self) -> Ptr<E, S, ReadOnly> {
207        intrinsic!(|_scope| {
208            PtrExpand {
209                node: self.node.read_only(),
210                marker: PhantomData,
211            }
212        })
213    }
214
215    pub fn select(&self, condition: bool, otherwise: Self) -> Self {
216        intrinsic!(|_scope| {
217            assert_eq!(
218                self.node.pointee(),
219                otherwise.node.pointee(),
220                "typed pointer selection requires equal pointee types"
221            );
222            assert_eq!(
223                condition.expand.ty,
224                Type::scalar(ElemType::Bool),
225                "typed pointer selection requires a scalar condition"
226            );
227            PtrExpand {
228                node: PointerNode::Select {
229                    condition: condition.expand.consume(),
230                    then: Box::new(self.node),
231                    otherwise: Box::new(otherwise.node),
232                },
233                marker: PhantomData,
234            }
235        })
236    }
237
238    pub fn load(&self) -> E {
239        intrinsic!(|scope| {
240            let out = if matches!(&self.node, PointerNode::Select { .. }) {
241                scope.create_local_mut(self.node.pointee())
242            } else {
243                scope.create_local(self.node.pointee())
244            };
245            self.node.emit(scope, *out, false);
246            NativeExpand::<E>::new(out)
247        })
248    }
249}
250
251#[ruda]
252impl<E: RudaPrimitive, S: PointerSpace> Ptr<E, S, ReadWrite> {
253    pub fn store(&self, value: E) {
254        intrinsic!(|scope| {
255            self.node.emit(scope, value.expand.consume(), true);
256        })
257    }
258}
259
260pub struct TypedSlice<E: RudaPrimitive, S: PointerSpace, A: PointerPermission = ReadOnly> {
261    marker: PhantomData<(E, S, A)>,
262}
263pub struct TypedSliceExpand<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> {
264    slice: SliceExpand<E, A>,
265    marker: PhantomData<S>,
266}
267impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> Clone for TypedSliceExpand<E, S, A> {
268    fn clone(&self) -> Self {
269        Self {
270            slice: self.slice.clone(),
271            marker: PhantomData,
272        }
273    }
274}
275impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> IntoMut
276    for TypedSliceExpand<E, S, A>
277{
278    fn into_mut(self, _: &mut Scope) -> Self {
279        self
280    }
281}
282impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> RudaDebug
283    for TypedSliceExpand<E, S, A>
284{
285}
286impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> RudaType for TypedSlice<E, S, A> {
287    type ExpandType = TypedSliceExpand<E, S, A>;
288}
289impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> RudaType for &TypedSlice<E, S, A> {
290    type ExpandType = TypedSliceExpand<E, S, A>;
291}
292
293#[ruda]
294impl<E: RudaPrimitive, A: PointerPermission> Slice<E, A> {
295    pub fn typed<S: PointerSpace>(&self) -> TypedSlice<E, S, A> {
296        intrinsic!(|_scope| {
297            let (base, index) = self.__to_raw_parts();
298            let pointee = base
299                .ty
300                .with_vector_size(self.vector_size.unwrap_or(base.ty.vector_size()));
301            TypedPointer {
302                base,
303                index,
304                ty: PointerType {
305                    pointee,
306                    address_space: S::SPACE,
307                    access: A::ACCESS,
308                },
309                checked: true,
310                unroll_factor: 1,
311            }
312            .validate()
313            .expect("valid typed slice conversion");
314            TypedSliceExpand {
315                slice: self,
316                marker: PhantomData,
317            }
318        })
319    }
320}
321
322#[ruda]
323impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> TypedSlice<E, S, A> {
324    pub fn ptr(&self, index: usize) -> Ptr<E, S, A> {
325        intrinsic!(|scope| {
326            let (base, start) = self.slice.__to_raw_parts();
327            let index = add::expand(
328                scope,
329                NativeExpand::<usize>::new(ManagedVariable::Plain(start)),
330                index,
331            );
332            let pointer = TypedPointer {
333                base,
334                index: index.expand.consume(),
335                ty: PointerType {
336                    pointee: base
337                        .ty
338                        .with_vector_size(self.slice.vector_size.unwrap_or(base.ty.vector_size())),
339                    address_space: S::SPACE,
340                    access: A::ACCESS,
341                },
342                checked: true,
343                unroll_factor: 1,
344            };
345            pointer.validate().expect("valid typed pointer conversion");
346            PtrExpand {
347                node: PointerNode::Element(pointer),
348                marker: PhantomData,
349            }
350        })
351    }
352
353    pub fn len(&self) -> usize {
354        intrinsic!(|_scope| { self.slice.length })
355    }
356
357    pub fn slice(&self, start: usize, end: usize) -> Self {
358        intrinsic!(|scope| {
359            let offset = add::expand(scope, self.slice.offset, start.clone());
360            let length = sub::expand(scope, end, start);
361            TypedSliceExpand {
362                slice: SliceExpand {
363                    offset,
364                    length,
365                    ..self.slice
366                },
367                marker: PhantomData,
368            }
369        })
370    }
371}
ruda_kernel/dsl/frontend/container/pointer.rs

ruda_kernel/dsl/frontend/container/
pointer.rs
