ruda_kernel/dsl/frontend/container/pointer/
value.rs

1use super::*;
2use alloc::{vec, vec::Vec};
3use ruda::runtime::backend::Runtime;
4use ruda_core::ir::{AtomicOpCode, BinaryOperator, PointerOp, Select, UnaryOperator};
5
6/// A first-class pointer value. Its address space and permissions are statically distinct.
7///
8/// ```compile_fail
9/// use ruda_kernel::dsl::prelude::{GlobalSpace, ValuePtr};
10/// fn write(pointer: ValuePtr<f32, GlobalSpace>) { pointer.store(1.0); }
11/// ```
12/// ```compile_fail
13/// use ruda_kernel::dsl::prelude::{GlobalSpace, ValuePtr, Workgroup};
14/// fn global(_: ValuePtr<f32, GlobalSpace>) {}
15/// fn shared(pointer: ValuePtr<f32, Workgroup>) { global(pointer); }
16/// ```
17/// ```compile_fail
18/// use ruda_kernel::dsl::prelude::{Atomic, GlobalSpace, ValuePtr};
19/// fn mutate(pointer: ValuePtr<Atomic<u32>, GlobalSpace>) { pointer.fetch_add(1); }
20/// ```
21#[derive(Clone, Copy)]
22pub struct ValuePtr<E: RudaPrimitive, S: PointerSpace, A: PointerPermission = ReadOnly> {
23    marker: PhantomData<(E, S, A)>,
24}
25
26pub type ValuePtrExpand<E, S, A = ReadOnly> = NativeExpand<ValuePtr<E, S, A>>;
27
28impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> RudaType for ValuePtr<E, S, A> {
29    type ExpandType = NativeExpand<Self>;
30}
31impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> RudaType for &ValuePtr<E, S, A> {
32    type ExpandType = NativeExpand<ValuePtr<E, S, A>>;
33}
34impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> RudaType for &mut ValuePtr<E, S, A> {
35    type ExpandType = NativeExpand<ValuePtr<E, S, A>>;
36}
37impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> NativeAssign for ValuePtr<E, S, A> {}
38impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> Assign for ValuePtr<E, S, A> {
39    fn expand_assign(&mut self, _: &mut Scope, value: Self) {
40        *self = value;
41    }
42    fn init_mut(&self, _: &mut Scope) -> Self {
43        *self
44    }
45}
46
47fn emit(scope: &mut Scope, operation: PointerOp, ty: Type) -> ManagedVariable {
48    let out = scope.create_local(ty);
49    operation
50        .validate(Some(*out))
51        .expect("valid first-class pointer operation");
52    scope.register(Instruction::new(operation, *out));
53    out
54}
55
56#[ruda]
57impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> ValuePtr<E, S, A> {
58    pub fn offset(&self, offset: usize) -> Self {
59        intrinsic!(|scope| {
60            let ty = self.expand.ty;
61            NativeExpand::new(emit(
62                scope,
63                PointerOp::Offset(BinaryOperator {
64                    lhs: *self.expand,
65                    rhs: *offset.expand,
66                }),
67                ty,
68            ))
69        })
70    }
71
72    pub fn read_only(&self) -> ValuePtr<E, S, ReadOnly> {
73        intrinsic!(|scope| {
74            let mut ty = self.expand.ty.pointer_type().expect("pointer value type");
75            ty.access = PointerAccess::ReadOnly;
76            NativeExpand::new(emit(
77                scope,
78                PointerOp::ReadOnly(UnaryOperator {
79                    input: *self.expand,
80                }),
81                Type::Pointer(ty),
82            ))
83        })
84    }
85
86    pub fn select(&self, condition: bool, otherwise: Self) -> Self {
87        intrinsic!(|scope| {
88            let ty = self.expand.ty;
89            NativeExpand::new(emit(
90                scope,
91                PointerOp::Select(Select {
92                    cond: *condition.expand,
93                    then: *self.expand,
94                    or_else: *otherwise.expand,
95                }),
96                ty,
97            ))
98        })
99    }
100
101    pub fn load(&self) -> E {
102        intrinsic!(|scope| {
103            let ty = self
104                .expand
105                .ty
106                .pointer_type()
107                .expect("pointer value type")
108                .pointee();
109            NativeExpand::new(emit(
110                scope,
111                PointerOp::Load {
112                    pointer: *self.expand,
113                    checked: true,
114                },
115                ty,
116            ))
117        })
118    }
119
120    /// Export a byte address on backends that support native-address operations.
121    pub fn native_address(&self) -> u64 {
122        intrinsic!(|scope| {
123            let ty = u64::as_type(scope);
124            NativeExpand::new(emit(
125                scope,
126                PointerOp::NativeAddress(UnaryOperator {
127                    input: *self.expand,
128                }),
129                ty,
130            ))
131        })
132    }
133}
134
135#[ruda]
136impl<E: RudaPrimitive, S: PointerSpace> ValuePtr<E, S, ReadWrite> {
137    pub fn store(&self, value: E) {
138        intrinsic!(|scope| {
139            let op = PointerOp::Store {
140                pointer: *self.expand,
141                value: *value.expand,
142                checked: true,
143            };
144            op.validate(None).expect("valid writable pointer store");
145            scope.register(Instruction::no_out(op));
146        })
147    }
148}
149
150#[ruda]
151impl<E: RudaPrimitive, S: PointerSpace, A: PointerPermission> ValuePtr<E, S, A> {
152    /// Import a valid, aligned byte address in `S`. The allocation must outlive kernel execution;
153    /// its address space, permissions and length are the caller's responsibility. Native backends only.
154    pub fn from_native(address: u64) -> Self {
155        intrinsic!(|scope| {
156            let ty = PointerType {
157                pointee: E::as_type(scope),
158                address_space: S::SPACE,
159                access: A::ACCESS,
160            }
161            .value_type();
162            NativeExpand::new(emit(
163                scope,
164                PointerOp::FromNative {
165                    address: *address.expand,
166                    ty: ty.pointer_type().unwrap(),
167                },
168                ty,
169            ))
170        })
171    }
172}
173
174fn atomic(
175    scope: &mut Scope,
176    pointer: Variable,
177    operation: AtomicOpCode,
178    values: Vec<Variable>,
179    ty: Type,
180) -> ManagedVariable {
181    emit(
182        scope,
183        PointerOp::Atomic {
184            pointer,
185            operation,
186            values,
187        },
188        ty,
189    )
190}
191
192#[ruda]
193impl<E: RudaPrimitive<Scalar: Numeric>, S: PointerSpace, A: PointerPermission>
194    ValuePtr<Atomic<E>, S, A>
195{
196    pub fn atomic_load(&self) -> E {
197        intrinsic!(|scope| {
198            let ty = E::as_type(scope);
199            NativeExpand::new(atomic(scope, *self.expand, AtomicOpCode::Load, vec![], ty))
200        })
201    }
202}
203
204macro_rules! numeric_atomic {
205    ($name:ident, $opcode:ident) => {
206        #[ruda]
207        impl<E: RudaPrimitive<Scalar: Numeric>, S: PointerSpace> ValuePtr<Atomic<E>, S, ReadWrite> {
208            pub fn $name(&self, value: E) -> E {
209                intrinsic!(|scope| {
210                    let ty = E::as_type(scope);
211                    NativeExpand::new(atomic(
212                        scope,
213                        *self.expand,
214                        AtomicOpCode::$opcode,
215                        vec![*value.expand],
216                        ty,
217                    ))
218                })
219            }
220        }
221    };
222}
223numeric_atomic!(swap, Swap);
224numeric_atomic!(fetch_add, Add);
225numeric_atomic!(fetch_sub, Sub);
226numeric_atomic!(fetch_min, Min);
227numeric_atomic!(fetch_max, Max);
228
229#[ruda]
230impl<E: RudaPrimitive<Scalar: Numeric>, S: PointerSpace> ValuePtr<Atomic<E>, S, ReadWrite> {
231    pub fn atomic_store(&self, value: E) {
232        intrinsic!(|scope| {
233            let op = PointerOp::Atomic {
234                pointer: *self.expand,
235                operation: AtomicOpCode::Store,
236                values: vec![*value.expand],
237            };
238            op.validate(None).expect("valid writable atomic store");
239            scope.register(Instruction::no_out(op));
240        })
241    }
242}
243
244macro_rules! integer_atomic {
245    ($name:ident, $opcode:ident) => {
246        #[ruda]
247        impl<E: RudaPrimitive<Scalar: Int>, S: PointerSpace> ValuePtr<Atomic<E>, S, ReadWrite> {
248            pub fn $name(&self, value: E) -> E {
249                intrinsic!(|scope| {
250                    let ty = E::as_type(scope);
251                    NativeExpand::new(atomic(
252                        scope,
253                        *self.expand,
254                        AtomicOpCode::$opcode,
255                        vec![*value.expand],
256                        ty,
257                    ))
258                })
259            }
260        }
261    };
262}
263integer_atomic!(fetch_and, And);
264integer_atomic!(fetch_or, Or);
265integer_atomic!(fetch_xor, Xor);
266
267#[ruda]
268impl<E: RudaPrimitive<Scalar: Int>, S: PointerSpace> ValuePtr<Atomic<E>, S, ReadWrite> {
269    pub fn compare_exchange_weak(&self, compare: E, value: E) -> E {
270        intrinsic!(|scope| {
271            let ty = E::as_type(scope);
272            NativeExpand::new(atomic(
273                scope,
274                *self.expand,
275                AtomicOpCode::CompareAndSwap,
276                vec![*compare.expand, *value.expand],
277                ty,
278            ))
279        })
280    }
281}
282
283impl<E: RudaPrimitive, A: PointerPermission> LaunchArg for ValuePtr<E, GlobalSpace, A> {
284    type RuntimeArg<R: Runtime> = <Array<E> as LaunchArg>::RuntimeArg<R>;
285    type CompilationArg = <Array<E> as LaunchArg>::CompilationArg;
286
287    fn register<R: Runtime>(
288        arg: Self::RuntimeArg<R>,
289        launcher: &mut KernelLauncher<R>,
290    ) -> Self::CompilationArg {
291        <Array<E> as LaunchArg>::register(arg, launcher)
292    }
293    fn expand(arg: &Self::CompilationArg, builder: &mut KernelBuilder) -> Self::ExpandType {
294        let array = if A::ACCESS == PointerAccess::ReadWrite {
295            <Array<E> as LaunchArg>::expand_output(arg, builder)
296        } else {
297            <Array<E> as LaunchArg>::expand(arg, builder)
298        };
299        let ty = PointerType {
300            pointee: array.expand.ty,
301            address_space: AddressSpace::Global,
302            access: A::ACCESS,
303        };
304        let index = usize::as_type(&builder.scope).constant(ruda_core::ir::ConstantValue::UInt(0));
305        let op = PointerOp::Address(TypedPointer {
306            base: *array.expand,
307            index,
308            ty,
309            checked: true,
310            unroll_factor: 1,
311        });
312        NativeExpand::new(emit(&mut builder.scope, op, ty.value_type()))
313    }
314    fn expand_output(arg: &Self::CompilationArg, builder: &mut KernelBuilder) -> Self::ExpandType {
315        let array = <Array<E> as LaunchArg>::expand_output(arg, builder);
316        let ty = PointerType {
317            pointee: array.expand.ty,
318            address_space: AddressSpace::Global,
319            access: A::ACCESS,
320        };
321        let index = usize::as_type(&builder.scope).constant(ruda_core::ir::ConstantValue::UInt(0));
322        NativeExpand::new(emit(
323            &mut builder.scope,
324            PointerOp::Address(TypedPointer {
325                base: *array.expand,
326                index,
327                ty,
328                checked: true,
329                unroll_factor: 1,
330            }),
331            ty.value_type(),
332        ))
333    }
334}
ruda_kernel/dsl/frontend/container/pointer/value.rs

ruda_kernel/dsl/frontend/container/pointer/
value.rs
