cubecl_core/frontend/element/
atomic.rs

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

cubecl_core/frontend/element/
atomic.rs
