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

ruda_kernel/dsl/frontend/container/
pointer.rs
