cubecl_ir/dialect/
atomic.rs

1use cubecl_macros_internal::{cube_op, op_traits};
2
3use crate::{CanMaterialize, dialect::ptr_value_ty, prelude::*};
4
5macro_rules! atomic_binop {
6    ($name: literal, $ty: ident) => {
7        #[cube_op(name = $name)]
8        #[result_ty(same_as = value)]
9        #[op_traits(CanMaterialize)]
10        pub struct $ty {
11            #[operand(ptr_read, ptr_write)]
12            pub ptr: Value,
13            pub value: Value,
14        }
15    };
16}
17
18atomic_binop!("atomic.exchange", AtomicExchangeOp);
19atomic_binop!("atomic.i_add", AtomicIAddOp);
20atomic_binop!("atomic.f_add", AtomicFAddOp);
21atomic_binop!("atomic.i_sub", AtomicISubOp);
22atomic_binop!("atomic.f_sub", AtomicFSubOp);
23atomic_binop!("atomic.s_min", AtomicSMinOp);
24atomic_binop!("atomic.u_min", AtomicUMinOp);
25atomic_binop!("atomic.f_min", AtomicFMinOp);
26atomic_binop!("atomic.s_max", AtomicSMaxOp);
27atomic_binop!("atomic.u_max", AtomicUMaxOp);
28atomic_binop!("atomic.f_max", AtomicFMaxOp);
29atomic_binop!("atomic.and", AtomicAndOp);
30atomic_binop!("atomic.or", AtomicOrOp);
31atomic_binop!("atomic.xor", AtomicXorOp);
32
33#[cube_op(name = "atomic.load")]
34#[result_ty(from_inputs = ptr_value_ty)]
35#[op_traits(CanMaterialize)]
36pub struct AtomicLoadOp {
37    #[operand(ptr_read)]
38    pub ptr: Value,
39}
40
41#[cube_op(name = "atomic.store")]
42#[result_ty(none)]
43#[op_traits(CanMaterialize)]
44pub struct AtomicStoreOp {
45    #[operand(ptr_write)]
46    pub ptr: Value,
47    pub value: Value,
48}
49
50#[cube_op(name = "atomic.compare_exchange_weak")]
51#[result_ty(same_as = value)]
52#[op_traits(CanMaterialize)]
53pub struct AtomicCompareExchangeWeakOp {
54    #[operand(ptr_read, ptr_write)]
55    pub ptr: Value,
56    pub cmp: Value,
57    pub value: Value,
58}
cubecl_ir/dialect/atomic.rs

cubecl_ir/dialect/
atomic.rs
