cubecl_ir/
atomic.rs

1use core::fmt::Display;
2
3use crate::TypeHash;
4
5use crate::{BinaryOperator, OperationArgs, OperationReflect, UnaryOperator, Variable};
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(UnaryOperator),
13    Store(UnaryOperator),
14    Swap(BinaryOperator),
15    Add(BinaryOperator),
16    Sub(BinaryOperator),
17    Max(BinaryOperator),
18    Min(BinaryOperator),
19    And(BinaryOperator),
20    Or(BinaryOperator),
21    Xor(BinaryOperator),
22    CompareAndSwap(CompareAndSwapOperator),
23}
24
25impl Display for AtomicOp {
26    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
27        match self {
28            AtomicOp::Load(op) => write!(f, "atomic_load({})", op.input),
29            AtomicOp::Store(op) => write!(f, "atomic_store({})", op.input),
30            AtomicOp::Swap(op) => {
31                write!(f, "atomic_swap({}, {})", op.lhs, op.rhs)
32            }
33            AtomicOp::Add(op) => write!(f, "atomic_add({}, {})", op.lhs, op.rhs),
34            AtomicOp::Sub(op) => write!(f, "atomic_sub({}, {})", op.lhs, op.rhs),
35            AtomicOp::Max(op) => write!(f, "atomic_max({}, {})", op.lhs, op.rhs),
36            AtomicOp::Min(op) => write!(f, "atomic_min({}, {})", op.lhs, op.rhs),
37            AtomicOp::And(op) => write!(f, "atomic_and({}, {})", op.lhs, op.rhs),
38            AtomicOp::Or(op) => write!(f, "atomic_or({}, {})", op.lhs, op.rhs),
39            AtomicOp::Xor(op) => write!(f, "atomic_xor({}, {})", op.lhs, op.rhs),
40            AtomicOp::CompareAndSwap(op) => {
41                write!(f, "compare_and_swap({}, {}, {})", op.input, 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 CompareAndSwapOperator {
51    pub input: Variable,
52    pub cmp: Variable,
53    pub val: Variable,
54}
cubecl_ir/atomic.rs

cubecl_ir/
atomic.rs
