cubecl_core/frontend/element/
atomic.rs

1use cubecl_ir::{
2    ConstantValue, ExpandValue,
3    dialect::{atomic::*, synchronization::SyncScope},
4    types::AtomicType,
5};
6use cubecl_macros::intrinsic;
7use half::{bf16, f16};
8use pliron::{
9    builtin::op_interfaces::OneResultInterface, context::Context, op::Op, r#type::TypeHandle,
10    value::Value,
11};
12
13use super::{NativeAssign, NativeExpand};
14use crate::{
15    self as cubecl,
16    frontend::{CubePrimitive, CubeType},
17    ir::Scope,
18    prelude::*,
19};
20
21/// An atomic numerical type wrapping a normal numeric primitive. Enables the use of atomic
22/// operations, while disabling normal operations. In WGSL, this is a separate type - on CUDA/SPIR-V
23/// it can theoretically be bitcast to a normal number, but this isn't recommended.
24#[derive(Clone, Copy, Hash, PartialEq, Eq)]
25pub struct Atomic<Inner: CubePrimitive> {
26    pub val: Inner,
27}
28
29type AtomicExpand<Inner> = NativeExpand<Atomic<Inner>>;
30
31pub trait AtomicNumeric {
32    fn __expand_fetch_add(
33        scope: &Scope,
34        ptr: ExpandValue,
35        value: ExpandValue,
36        sync: SyncScope,
37    ) -> ExpandValue;
38    fn __expand_fetch_sub(
39        scope: &Scope,
40        ptr: ExpandValue,
41        value: ExpandValue,
42        sync: SyncScope,
43    ) -> ExpandValue;
44    fn __expand_fetch_min(
45        scope: &Scope,
46        ptr: ExpandValue,
47        value: ExpandValue,
48        sync: SyncScope,
49    ) -> ExpandValue;
50    fn __expand_fetch_max(
51        scope: &Scope,
52        ptr: ExpandValue,
53        value: ExpandValue,
54        sync: SyncScope,
55    ) -> ExpandValue;
56}
57
58macro_rules! atomic_numeric {
59    ($($ty: ty),*; $add: ty, $sub: ty, $min: ty, $max: ty) => {
60        $(impl AtomicNumeric for $ty {
61            fn __expand_fetch_add(scope: &Scope, ptr: ExpandValue, value: ExpandValue, sync: SyncScope) -> ExpandValue {
62                atomic_binary_expand(scope, ptr, value, sync, <$add>::new)
63            }
64            fn __expand_fetch_sub(scope: &Scope, ptr: ExpandValue, value: ExpandValue, sync: SyncScope) -> ExpandValue {
65                atomic_binary_expand(scope, ptr, value, sync, <$sub>::new)
66            }
67            fn __expand_fetch_min(scope: &Scope, ptr: ExpandValue, value: ExpandValue, sync: SyncScope) -> ExpandValue {
68                atomic_binary_expand(scope, ptr, value, sync, <$min>::new)
69            }
70            fn __expand_fetch_max(scope: &Scope, ptr: ExpandValue, value: ExpandValue, sync: SyncScope) -> ExpandValue {
71                atomic_binary_expand(scope, ptr, value, sync, <$max>::new)
72            }
73        })*
74    };
75}
76
77atomic_numeric!(i8, i16, i32, i64, isize; AtomicIAddOp, AtomicISubOp, AtomicSMinOp, AtomicSMaxOp);
78atomic_numeric!(u8, u16, u32, u64, usize; AtomicIAddOp, AtomicISubOp, AtomicUMinOp, AtomicUMaxOp);
79atomic_numeric!(f16, bf16, f32, flex32, tf32, f64; AtomicFAddOp, AtomicFSubOp, AtomicFMinOp, AtomicFMaxOp);
80
81/// The scope an atomic on `ptr` synchronizes at: `sync`, or the system's when `ptr` points into a
82/// [`Peer`] buffer, which kernels on other devices reach while this one runs.
83fn reach(scope: &Scope, ptr: Value, sync: SyncScope) -> SyncScope {
84    match scope.is_peer_pointer(ptr) {
85        true => sync.max(SyncScope::System),
86        false => sync,
87    }
88}
89
90fn atomic_binary_expand<F, O>(
91    scope: &Scope,
92    ptr: ExpandValue,
93    value: ExpandValue,
94    sync: SyncScope,
95    func: F,
96) -> ExpandValue
97where
98    F: Fn(&mut Context, Value, Value) -> O,
99    O: Op + OneResultInterface,
100{
101    let ptr = ptr.value(scope);
102    let sync = reach(scope, ptr, sync);
103    let op = func(scope.ctx_mut(), ptr, value.read_value(scope));
104    set_atomic_scope(scope.ctx_mut(), op.get_operation(), sync);
105    scope.register_with_result(&op).into()
106}
107
108#[cube]
109impl<Inner: CubePrimitive<Scalar: AtomicNumeric>> Atomic<Inner> {
110    /// Load the value of the atomic.
111    pub fn load(&self) -> Inner {
112        self.load_at(comptime![SyncScope::Device])
113    }
114
115    /// Store the value of the atomic.
116    pub fn store(&self, value: Inner) {
117        self.store_at(value, comptime![SyncScope::Device])
118    }
119
120    /// Atomically stores the value into the atomic and returns the old value.
121    pub fn exchange(&self, value: Inner) -> Inner {
122        self.exchange_at(value, comptime![SyncScope::Device])
123    }
124
125    /// Atomically add a number to the atomic variable. Returns the old value.
126    pub fn fetch_add(&self, value: Inner) -> Inner {
127        self.fetch_add_at(value, comptime![SyncScope::Device])
128    }
129
130    /// Atomically subtracts a number from the atomic variable. Returns the old value.
131    pub fn fetch_sub(&self, value: Inner) -> Inner {
132        self.fetch_sub_at(value, comptime![SyncScope::Device])
133    }
134
135    /// Atomically sets the value of the atomic variable to `max(current_value, value)`. Returns
136    /// the old value.
137    pub fn fetch_max(&self, value: Inner) -> Inner {
138        self.fetch_max_at(value, comptime![SyncScope::Device])
139    }
140
141    /// Atomically sets the value of the atomic variable to `min(current_value, value)`. Returns the
142    /// old value.
143    pub fn fetch_min(&self, value: Inner) -> Inner {
144        self.fetch_min_at(value, comptime![SyncScope::Device])
145    }
146
147    /// `load`, synchronizing at `sync`.
148    ///
149    /// The plain operations synchronize at `SyncScope::Device`, and at
150    /// `SyncScope::System` on a `Peer` buffer, whatever scope they name. The system scope
151    /// also synchronizes with the host and with kernels on other devices that bind this memory,
152    /// which is what kernels on different devices signal each other through: publish with
153    /// `sync_system()`, then store a flag; spin on the flag, then
154    /// `sync_system()` before reading what it guards.
155    #[allow(unused_variables)]
156    pub fn load_at(&self, #[comptime] sync: SyncScope) -> Inner {
157        intrinsic!(|scope| {
158            let ptr = self.value(scope);
159            let op = AtomicLoadOp::new(scope.ctx_mut(), ptr);
160            let sync = reach(scope, ptr, sync);
161            set_atomic_scope(scope.ctx_mut(), op.get_operation(), sync);
162            scope.register_with_result(&op).into()
163        })
164    }
165
166    /// `store`, synchronizing at `sync` (see `load_at`).
167    #[allow(unused_variables)]
168    pub fn store_at(&self, value: Inner, #[comptime] sync: SyncScope) {
169        intrinsic!(|scope| {
170            let ptr = self.value(scope);
171            let value = value.read_value(scope);
172            let op = AtomicStoreOp::new(scope.ctx_mut(), ptr, value);
173            let sync = reach(scope, ptr, sync);
174            set_atomic_scope(scope.ctx_mut(), op.get_operation(), sync);
175            scope.register(&op);
176        })
177    }
178
179    /// `exchange`, synchronizing at `sync` (see `load_at`).
180    #[allow(unused_variables)]
181    pub fn exchange_at(&self, value: Inner, #[comptime] sync: SyncScope) -> Inner {
182        intrinsic!(|scope| {
183            let ptr = self.value(scope);
184            let value = value.read_value(scope);
185            let op = AtomicExchangeOp::new(scope.ctx_mut(), ptr, value);
186            let sync = reach(scope, ptr, sync);
187            set_atomic_scope(scope.ctx_mut(), op.get_operation(), sync);
188            scope.register_with_result(&op).into()
189        })
190    }
191
192    /// `fetch_add`, synchronizing at `sync` (see `load_at`).
193    #[allow(unused_variables)]
194    pub fn fetch_add_at(&self, value: Inner, #[comptime] sync: SyncScope) -> Inner {
195        intrinsic!(|scope| Inner::Scalar::__expand_fetch_add(
196            scope,
197            self.expand,
198            value.expand,
199            sync
200        )
201        .into())
202    }
203
204    /// `fetch_sub`, synchronizing at `sync` (see `load_at`).
205    #[allow(unused_variables)]
206    pub fn fetch_sub_at(&self, value: Inner, #[comptime] sync: SyncScope) -> Inner {
207        intrinsic!(|scope| Inner::Scalar::__expand_fetch_sub(
208            scope,
209            self.expand,
210            value.expand,
211            sync
212        )
213        .into())
214    }
215
216    /// `fetch_max`, synchronizing at `sync` (see `load_at`).
217    #[allow(unused_variables)]
218    pub fn fetch_max_at(&self, value: Inner, #[comptime] sync: SyncScope) -> Inner {
219        intrinsic!(|scope| Inner::Scalar::__expand_fetch_max(
220            scope,
221            self.expand,
222            value.expand,
223            sync
224        )
225        .into())
226    }
227
228    /// `fetch_min`, synchronizing at `sync` (see `load_at`).
229    #[allow(unused_variables)]
230    pub fn fetch_min_at(&self, value: Inner, #[comptime] sync: SyncScope) -> Inner {
231        intrinsic!(|scope| Inner::Scalar::__expand_fetch_min(
232            scope,
233            self.expand,
234            value.expand,
235            sync
236        )
237        .into())
238    }
239}
240
241#[cube]
242impl<Inner: CubePrimitive<Scalar: Int>> Atomic<Inner> {
243    /// Compare the value at `pointer` to `cmp` and set it to `value` only if they are the same.
244    /// Returns the old value of the pointer before the store.
245    ///
246    /// ### Tip
247    /// Compare the returned value to `cmp` to determine whether the store was successful.
248    pub fn compare_exchange_weak(&self, cmp: Inner, value: Inner) -> Inner {
249        self.compare_exchange_weak_at(cmp, value, comptime![SyncScope::Device])
250    }
251
252    /// Executes an atomic bitwise and operation on the atomic variable. Returns the old value.
253    pub fn fetch_and(&self, value: Inner) -> Inner {
254        self.fetch_and_at(value, comptime![SyncScope::Device])
255    }
256
257    /// Executes an atomic bitwise or operation on the atomic variable. Returns the old value.
258    pub fn fetch_or(&self, value: Inner) -> Inner {
259        self.fetch_or_at(value, comptime![SyncScope::Device])
260    }
261
262    /// Executes an atomic bitwise xor operation on the atomic variable. Returns the old value.
263    pub fn fetch_xor(&self, value: Inner) -> Inner {
264        self.fetch_xor_at(value, comptime![SyncScope::Device])
265    }
266
267    /// `compare_exchange_weak`, synchronizing at `sync` (see
268    /// `load_at`).
269    #[allow(unused_variables)]
270    pub fn compare_exchange_weak_at(
271        &self,
272        cmp: Inner,
273        value: Inner,
274        #[comptime] sync: SyncScope,
275    ) -> Inner {
276        intrinsic!(|scope| {
277            let ptr = self.value(scope);
278            let cmp = cmp.read_value(scope);
279            let value = value.read_value(scope);
280            let op = AtomicCompareExchangeWeakOp::new(scope.ctx_mut(), ptr, cmp, value);
281            let sync = reach(scope, ptr, sync);
282            set_atomic_scope(scope.ctx_mut(), op.get_operation(), sync);
283            scope.register_with_result(&op).into()
284        })
285    }
286
287    /// `fetch_and`, synchronizing at `sync` (see `load_at`).
288    #[allow(unused_variables)]
289    pub fn fetch_and_at(&self, value: Inner, #[comptime] sync: SyncScope) -> Inner {
290        intrinsic!(|scope| atomic_binary_expand(
291            scope,
292            self.expand,
293            value.expand,
294            sync,
295            AtomicAndOp::new
296        )
297        .into())
298    }
299
300    /// `fetch_or`, synchronizing at `sync` (see `load_at`).
301    #[allow(unused_variables)]
302    pub fn fetch_or_at(&self, value: Inner, #[comptime] sync: SyncScope) -> Inner {
303        intrinsic!(|scope| atomic_binary_expand(
304            scope,
305            self.expand,
306            value.expand,
307            sync,
308            AtomicOrOp::new
309        )
310        .into())
311    }
312
313    /// `fetch_xor`, synchronizing at `sync` (see `load_at`).
314    #[allow(unused_variables)]
315    pub fn fetch_xor_at(&self, value: Inner, #[comptime] sync: SyncScope) -> Inner {
316        intrinsic!(|scope| atomic_binary_expand(
317            scope,
318            self.expand,
319            value.expand,
320            sync,
321            AtomicXorOp::new
322        )
323        .into())
324    }
325}
326
327impl<Inner: CubePrimitive> CubeType for Atomic<Inner> {
328    type ExpandType = NativeExpand<Self>;
329}
330
331impl<Inner: CubePrimitive> CubeDebug for Atomic<Inner> {}
332impl<Inner: CubePrimitive> CubePrimitive for Atomic<Inner> {
333    type Scalar = Inner::Scalar;
334    type Size = Const<1>;
335    type WithScalar<S: Scalar> = Atomic<S>;
336
337    fn from_expand_elem(elem: ExpandValue) -> Self::ExpandType {
338        NativeExpand::new(elem)
339    }
340
341    fn from_const_value(_value: ConstantValue) -> Self {
342        panic!("Can't have constant atomic");
343    }
344}
345impl<Inner: CubePrimitive> ConstantTyped for Atomic<Inner> {
346    fn __expand_as_type(scope: &Scope) -> TypeHandle {
347        let inner = Inner::__expand_as_type(scope);
348        AtomicType::get(scope.ctx(), inner).into()
349    }
350}
351
352impl<Inner: CubePrimitive> NativeAssign for Atomic<Inner> {}
cubecl_core/frontend/element/atomic.rs

cubecl_core/frontend/element/
atomic.rs
