ruda_kernel/dsl/frontend/element/
atomic.rs

1use ruda_core::ir::{AtomicOp, ConstantValue, ManagedVariable, StorageType};
2use ruda_kernel_macros::intrinsic;
3
4use super::{NativeAssign, NativeExpand, Numeric};
5use crate::dsl::{
6    frontend::{RudaPrimitive, RudaType},
7    ir::{BinaryOperator, CompareAndSwapOperator, Instruction, Scope, Type, UnaryOperator},
8    prelude::*,
9};
10
11/// An atomic numerical type wrapping a normal numeric primitive. Enables the use of atomic
12/// operations, while disabling normal operations. In WGSL, this is a separate type - on CUDA/SPIR-V
13/// it can theoretically be bitcast to a normal number, but this isn't recommended.
14#[derive(Clone, Copy, Hash, PartialEq, Eq)]
15pub struct Atomic<Inner: RudaPrimitive> {
16    pub val: Inner,
17}
18
19type AtomicExpand<Inner> = NativeExpand<Atomic<Inner>>;
20
21#[ruda]
22impl<Inner: RudaPrimitive<Scalar: Numeric>> Atomic<Inner> {
23    /// Load the value of the atomic.
24    #[allow(unused_variables)]
25    pub fn load(&self) -> Inner {
26        intrinsic!(|scope| {
27            let pointer: ManagedVariable = self.into();
28            let new_var = scope.create_local(Inner::as_type(scope));
29            scope.register(Instruction::new(
30                AtomicOp::Load(UnaryOperator { input: *pointer }),
31                *new_var,
32            ));
33            new_var.into()
34        })
35    }
36
37    /// Store the value of the atomic.
38    #[allow(unused_variables)]
39    pub fn store(&self, value: Inner) {
40        intrinsic!(|scope| {
41            let ptr: ManagedVariable = self.into();
42            let value: ManagedVariable = value.into();
43            scope.register(Instruction::new(
44                AtomicOp::Store(UnaryOperator { input: *value }),
45                *ptr,
46            ));
47        })
48    }
49
50    /// Atomically stores the value into the atomic and returns the old value.
51    #[allow(unused_variables)]
52    pub fn swap(&self, value: Inner) -> Inner {
53        intrinsic!(|scope| {
54            let ptr: ManagedVariable = self.into();
55            let value: ManagedVariable = value.into();
56            let new_var = scope.create_local(Inner::as_type(scope));
57            scope.register(Instruction::new(
58                AtomicOp::Swap(BinaryOperator {
59                    lhs: *ptr,
60                    rhs: *value,
61                }),
62                *new_var,
63            ));
64            new_var.into()
65        })
66    }
67
68    /// Atomically add a number to the atomic variable. Returns the old value.
69    #[allow(unused_variables)]
70    pub fn fetch_add(&self, value: Inner) -> Inner {
71        intrinsic!(|scope| {
72            let ptr: ManagedVariable = self.into();
73            let value: ManagedVariable = value.into();
74            let new_var = scope.create_local(Inner::as_type(scope));
75            scope.register(Instruction::new(
76                AtomicOp::Add(BinaryOperator {
77                    lhs: *ptr,
78                    rhs: *value,
79                }),
80                *new_var,
81            ));
82            new_var.into()
83        })
84    }
85
86    /// Atomically subtracts a number from the atomic variable. Returns the old value.
87    #[allow(unused_variables)]
88    pub fn fetch_sub(&self, value: Inner) -> Inner {
89        intrinsic!(|scope| {
90            let ptr: ManagedVariable = self.into();
91            let value: ManagedVariable = value.into();
92            let new_var = scope.create_local(Inner::as_type(scope));
93            scope.register(Instruction::new(
94                AtomicOp::Sub(BinaryOperator {
95                    lhs: *ptr,
96                    rhs: *value,
97                }),
98                *new_var,
99            ));
100            new_var.into()
101        })
102    }
103
104    /// Atomically sets the value of the atomic variable to `max(current_value, value)`. Returns
105    /// the old value.
106    #[allow(unused_variables)]
107    pub fn fetch_max(&self, value: Inner) -> Inner {
108        intrinsic!(|scope| {
109            let ptr: ManagedVariable = self.into();
110            let value: ManagedVariable = value.into();
111            let new_var = scope.create_local(Inner::as_type(scope));
112            scope.register(Instruction::new(
113                AtomicOp::Max(BinaryOperator {
114                    lhs: *ptr,
115                    rhs: *value,
116                }),
117                *new_var,
118            ));
119            new_var.into()
120        })
121    }
122
123    /// Atomically sets the value of the atomic variable to `min(current_value, value)`. Returns the
124    /// old value.
125    #[allow(unused_variables)]
126    pub fn fetch_min(&self, value: Inner) -> Inner {
127        intrinsic!(|scope| {
128            let ptr: ManagedVariable = self.into();
129            let value: ManagedVariable = value.into();
130            let new_var = scope.create_local(Inner::as_type(scope));
131            scope.register(Instruction::new(
132                AtomicOp::Min(BinaryOperator {
133                    lhs: *ptr,
134                    rhs: *value,
135                }),
136                *new_var,
137            ));
138            new_var.into()
139        })
140    }
141}
142
143#[ruda]
144impl<Inner: RudaPrimitive<Scalar: Int>> Atomic<Inner> {
145    /// Compare the value at `pointer` to `cmp` and set it to `value` only if they are the same.
146    /// Returns the old value of the pointer before the store.
147    ///
148    /// ### Tip
149    /// Compare the returned value to `cmp` to determine whether the store was successful.
150    #[allow(unused_variables)]
151    pub fn compare_exchange_weak(&self, cmp: Inner, value: Inner) -> Inner {
152        intrinsic!(|scope| {
153            let pointer: ManagedVariable = self.into();
154            let cmp: ManagedVariable = cmp.into();
155            let value: ManagedVariable = value.into();
156            let new_var = scope.create_local(Inner::as_type(scope));
157            scope.register(Instruction::new(
158                AtomicOp::CompareAndSwap(CompareAndSwapOperator {
159                    input: *pointer,
160                    cmp: *cmp,
161                    val: *value,
162                }),
163                *new_var,
164            ));
165            new_var.into()
166        })
167    }
168
169    /// Executes an atomic bitwise and operation on the atomic variable. Returns the old value.
170    #[allow(unused_variables)]
171    pub fn fetch_and(&self, value: Inner) -> Inner {
172        intrinsic!(|scope| {
173            let ptr: ManagedVariable = self.into();
174            let value: ManagedVariable = value.into();
175            let new_var = scope.create_local(Inner::as_type(scope));
176            scope.register(Instruction::new(
177                AtomicOp::And(BinaryOperator {
178                    lhs: *ptr,
179                    rhs: *value,
180                }),
181                *new_var,
182            ));
183            new_var.into()
184        })
185    }
186
187    /// Executes an atomic bitwise or operation on the atomic variable. Returns the old value.
188    #[allow(unused_variables)]
189    pub fn fetch_or(&self, value: Inner) -> Inner {
190        intrinsic!(|scope| {
191            let ptr: ManagedVariable = self.into();
192            let value: ManagedVariable = value.into();
193            let new_var = scope.create_local(Inner::as_type(scope));
194            scope.register(Instruction::new(
195                AtomicOp::Or(BinaryOperator {
196                    lhs: *ptr,
197                    rhs: *value,
198                }),
199                *new_var,
200            ));
201            new_var.into()
202        })
203    }
204
205    /// Executes an atomic bitwise xor operation on the atomic variable. Returns the old value.
206    #[allow(unused_variables)]
207    pub fn fetch_xor(&self, value: Inner) -> Inner {
208        intrinsic!(|scope| {
209            let ptr: ManagedVariable = self.into();
210            let value: ManagedVariable = value.into();
211            let new_var = scope.create_local(Inner::as_type(scope));
212            scope.register(Instruction::new(
213                AtomicOp::Xor(BinaryOperator {
214                    lhs: *ptr,
215                    rhs: *value,
216                }),
217                *new_var,
218            ));
219            new_var.into()
220        })
221    }
222}
223
224impl<Inner: RudaPrimitive> RudaType for Atomic<Inner> {
225    type ExpandType = NativeExpand<Self>;
226}
227
228impl<Inner: RudaPrimitive> RudaPrimitive for Atomic<Inner> {
229    type Scalar = Inner::Scalar;
230    type Size = Const<1>;
231    type WithScalar<S: Scalar> = Atomic<S>;
232
233    fn as_type_native() -> Option<Type> {
234        Inner::as_type_native().map(|it| it.with_storage_type(StorageType::Atomic(it.elem_type())))
235    }
236
237    fn as_type(scope: &Scope) -> Type {
238        let inner = Inner::as_type(scope);
239        inner.with_storage_type(StorageType::Atomic(inner.elem_type()))
240    }
241
242    fn as_type_native_unchecked() -> Type {
243        let inner = Inner::as_type_native_unchecked();
244        inner.with_storage_type(StorageType::Atomic(inner.elem_type()))
245    }
246
247    fn size() -> Option<usize> {
248        Inner::size()
249    }
250
251    fn from_expand_elem(elem: ManagedVariable) -> Self::ExpandType {
252        NativeExpand::new(elem)
253    }
254
255    fn from_const_value(_value: ConstantValue) -> Self {
256        panic!("Can't have constant atomic");
257    }
258}
259
260impl<Inner: RudaPrimitive> NativeAssign for Atomic<Inner> {}
ruda_kernel/dsl/frontend/element/atomic.rs

ruda_kernel/dsl/frontend/element/
atomic.rs
