cubecl_cpp/shared/
atomic.rs

1use cubecl_core::{
2    self as cubecl,
3    ir::{dialect::atomic::*, interfaces::TypedExt, prelude::*},
4    prelude::*,
5};
6use num_traits::One;
7
8use crate::{
9    shared::{
10        binary::lower_binop, lowering::LowerOp, scoped_block, shared_op_with_out, ty::TypeExtCPP,
11    },
12    target::{CtxTarget, Target},
13};
14
15#[cube]
16fn atomic_i_sub<T: Numeric + CubeNot, N: Size>(
17    ptr: Atomic<Vector<T, N>>,
18    value: Vector<T, N>,
19) -> Vector<T, N> {
20    ptr.fetch_add(!value + Vector::one())
21}
22
23#[cube]
24fn atomic_f_sub<T: Numeric + CubeNeg, N: Size>(
25    ptr: Atomic<Vector<T, N>>,
26    value: Vector<T, N>,
27) -> Vector<T, N> {
28    ptr.fetch_add(-value)
29}
30
31#[cube]
32fn atomic_store<T: Numeric + CubeNeg, N: Size>(ptr: Atomic<Vector<T, N>>, value: Vector<T, N>) {
33    ptr.exchange(value);
34}
35
36#[op_interface_impl]
37impl LowerOp for AtomicStoreOp {
38    fn lower(&self, scope: &Scope) -> Vec<Value> {
39        define_scalar!(T);
40        define_size!(S);
41        let ptr = self.ptr(scope.ctx());
42        let value = self.value(scope.ctx());
43        scope.register_value_type::<T, S>(value);
44        atomic_store::expand::<T, S>(scope, ptr.into(), value.into());
45        vec![]
46    }
47}
48
49lower_binop!(AtomicISubOp, atomic_i_sub, |_, ctx| {
50    ctx.target() != Target::Metal
51});
52lower_binop!(AtomicFSubOp, atomic_f_sub, |_, ctx| {
53    ctx.target() != Target::Metal
54});
55
56shared_op_with_out!(AtomicLoadOp, |op, ctx| {
57    let ptr = op.ptr(ctx).name(ctx);
58    let out_ty = op.get_result(ctx).get_type(ctx);
59    let uint_ty = match out_ty.size(ctx) {
60        1 => "uint8_t",
61        2 => "uint16_t",
62        4 => "uint32_t",
63        8 => "uint64_t",
64        16 => {
65            return scoped_block! {
66                format!("{} tmp;", out_ty.to_cpp(ctx))
67                format!("__nv_atomic_load({ptr}, &tmp, __NV_ATOMIC_RELAXED);")
68                "return tmp;"
69            };
70        }
71        _ => unreachable!(),
72    };
73    scoped_block! {
74        format!("volatile {uint_ty} const* tmp = reinterpret_cast<volatile {uint_ty} const*>({ptr});")
75        format!("const {uint_ty} tmp_2 = *tmp;")
76        format!("return reinterpret_cast<const {}&>(tmp_2);", out_ty.to_cpp(ctx))
77    }
78});
79
80shared_op_with_out!(AtomicExchangeOp, |op, ctx| {
81    let ptr = op.ptr(ctx).name(ctx);
82    let value = op.value(ctx).name(ctx);
83    let out_ty = op.get_result(ctx).get_type(ctx);
84    let uint_ty = match out_ty.size(ctx) {
85        1 => "uint8_t",
86        2 => "uint16_t",
87        4 => "uint32_t",
88        8 => "uint64_t",
89        16 => {
90            return format!("atomicExch({ptr}, {value})");
91        }
92        _ => unreachable!(),
93    };
94    let ptr = format!("reinterpret_cast<{uint_ty}*>({ptr})");
95    let value = format!("reinterpret_cast<const {uint_ty}&>({value})");
96    scoped_block! {
97        format!("const {uint_ty} tmp = atomicExch({ptr}, {value});")
98        format!("return reinterpret_cast<const {}&>(tmp);", out_ty.to_cpp(ctx))
99    }
100});
101
102shared_op_with_out!(AtomicCompareExchangeWeakOp, |op, ctx| {
103    let ptr = op.ptr(ctx).name(ctx);
104    let cmp = op.cmp(ctx).name(ctx);
105    let value = op.value(ctx).name(ctx);
106    let out_ty = op.get_result(ctx).get_type(ctx);
107    let uint_ty = match out_ty.size(ctx) {
108        1 => "uint8_t",
109        2 => "uint16_t",
110        4 => "uint32_t",
111        8 => "uint64_t",
112        16 => {
113            return format!("atomicCAS({ptr}, {cmp}, {value})");
114        }
115        _ => unreachable!(),
116    };
117    let ptr = format!("reinterpret_cast<{uint_ty}*>({ptr})");
118    let cmp = format!("reinterpret_cast<const {uint_ty}&>({cmp})");
119    let value = format!("reinterpret_cast<const {uint_ty}&>({value})");
120    scoped_block! {
121        format!("const {uint_ty} tmp = atomicCAS({ptr}, {cmp}, {value});")
122        format!("return reinterpret_cast<const {}&>(tmp);", out_ty.to_cpp(ctx))
123    }
124});
125
126shared_op_with_out!(AtomicSMinOp, |op, ctx| {
127    let ptr = op.ptr(ctx).name(ctx);
128    let value = op.value(ctx).name(ctx);
129    format!("atomicMin({ptr}, {value})")
130});
131shared_op_with_out!(AtomicUMinOp, |op, ctx| {
132    let ptr = op.ptr(ctx).name(ctx);
133    let value = op.value(ctx).name(ctx);
134    format!("atomicMin({ptr}, {value})")
135});
136shared_op_with_out!(AtomicFMinOp, |op, ctx| {
137    let ptr = op.ptr(ctx).name(ctx);
138    let value = op.value(ctx).name(ctx);
139    format!("atomicMin({ptr}, {value})")
140});
141
142shared_op_with_out!(AtomicSMaxOp, |op, ctx| {
143    let ptr = op.ptr(ctx).name(ctx);
144    let value = op.value(ctx).name(ctx);
145    format!("atomicMax({ptr}, {value})")
146});
147shared_op_with_out!(AtomicUMaxOp, |op, ctx| {
148    let ptr = op.ptr(ctx).name(ctx);
149    let value = op.value(ctx).name(ctx);
150    format!("atomicMax({ptr}, {value})")
151});
152shared_op_with_out!(AtomicFMaxOp, |op, ctx| {
153    let ptr = op.ptr(ctx).name(ctx);
154    let value = op.value(ctx).name(ctx);
155    format!("atomicMax({ptr}, {value})")
156});
157
158shared_op_with_out!(AtomicAndOp, |op, ctx| {
159    let ptr = op.ptr(ctx).name(ctx);
160    let value = op.value(ctx).name(ctx);
161    format!("atomicAnd({ptr}, {value})")
162});
163
164shared_op_with_out!(AtomicOrOp, |op, ctx| {
165    let ptr = op.ptr(ctx).name(ctx);
166    let value = op.value(ctx).name(ctx);
167    format!("atomicOr({ptr}, {value})")
168});
169
170shared_op_with_out!(AtomicXorOp, |op, ctx| {
171    let ptr = op.ptr(ctx).name(ctx);
172    let value = op.value(ctx).name(ctx);
173    format!("atomicXor({ptr}, {value})")
174});
cubecl_cpp/shared/atomic.rs

cubecl_cpp/shared/
atomic.rs
