cubecl_ir/
atomic.rs

1use core::fmt::Display;
2
3use crate::{AtomicBinaryOperands, StoreOperands, TypeHash};
4
5use crate::{OperationArgs, OperationReflect, Value};
6
7/// Operations that operate on atomics
8#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
9#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationReflect)]
10#[operation(opcode_name = AtomicOpCode)]
11pub enum AtomicOp {
12    Load(#[args(allow_ptr, ptr_read)] Value),
13    Store(StoreOperands),
14    Swap(AtomicBinaryOperands),
15    Add(AtomicBinaryOperands),
16    Sub(AtomicBinaryOperands),
17    Max(AtomicBinaryOperands),
18    Min(AtomicBinaryOperands),
19    And(AtomicBinaryOperands),
20    Or(AtomicBinaryOperands),
21    Xor(AtomicBinaryOperands),
22    CompareAndSwap(CompareAndSwapOperands),
23}
24
25impl Display for AtomicOp {
26    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
27        match self {
28            AtomicOp::Load(val) => write!(f, "atomic_load({val})"),
29            AtomicOp::Store(op) => write!(f, "atomic_store({}, {})", op.ptr, op.value),
30            AtomicOp::Swap(op) => {
31                write!(f, "atomic_swap({}, {})", op.ptr, op.value)
32            }
33            AtomicOp::Add(op) => write!(f, "atomic_add({}, {})", op.ptr, op.value),
34            AtomicOp::Sub(op) => write!(f, "atomic_sub({}, {})", op.ptr, op.value),
35            AtomicOp::Max(op) => write!(f, "atomic_max({}, {})", op.ptr, op.value),
36            AtomicOp::Min(op) => write!(f, "atomic_min({}, {})", op.ptr, op.value),
37            AtomicOp::And(op) => write!(f, "atomic_and({}, {})", op.ptr, op.value),
38            AtomicOp::Or(op) => write!(f, "atomic_or({}, {})", op.ptr, op.value),
39            AtomicOp::Xor(op) => write!(f, "atomic_xor({}, {})", op.ptr, op.value),
40            AtomicOp::CompareAndSwap(op) => {
41                write!(f, "compare_and_swap({}, {}, {})", op.ptr, op.cmp, op.val)
42            }
43        }
44    }
45}
46
47#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
48#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
49#[allow(missing_docs)]
50pub struct CompareAndSwapOperands {
51    #[args(allow_ptr, ptr_read, ptr_write)]
52    pub ptr: Value,
53    pub cmp: Value,
54    pub val: Value,
55}
cubecl_ir/atomic.rs

cubecl_ir/
atomic.rs
