cubecl-ir 0.11.0-pre.1

Intermediate representation for CubeCL
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

use core::fmt::Display;

use crate::{AtomicBinaryOperands, StoreOperands, TypeHash};

use crate::{OperationArgs, OperationReflect, Value};

/// Operations that operate on atomics
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationReflect)]
#[operation(opcode_name = AtomicOpCode)]
pub enum AtomicOp {
    Load(#[args(allow_ptr, ptr_read)] Value),
    Store(StoreOperands),
    Swap(AtomicBinaryOperands),
    Add(AtomicBinaryOperands),
    Sub(AtomicBinaryOperands),
    Max(AtomicBinaryOperands),
    Min(AtomicBinaryOperands),
    And(AtomicBinaryOperands),
    Or(AtomicBinaryOperands),
    Xor(AtomicBinaryOperands),
    CompareAndSwap(CompareAndSwapOperands),
}

impl Display for AtomicOp {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            AtomicOp::Load(val) => write!(f, "atomic_load({val})"),
            AtomicOp::Store(op) => write!(f, "atomic_store({}, {})", op.ptr, op.value),
            AtomicOp::Swap(op) => {
                write!(f, "atomic_swap({}, {})", op.ptr, op.value)
            }
            AtomicOp::Add(op) => write!(f, "atomic_add({}, {})", op.ptr, op.value),
            AtomicOp::Sub(op) => write!(f, "atomic_sub({}, {})", op.ptr, op.value),
            AtomicOp::Max(op) => write!(f, "atomic_max({}, {})", op.ptr, op.value),
            AtomicOp::Min(op) => write!(f, "atomic_min({}, {})", op.ptr, op.value),
            AtomicOp::And(op) => write!(f, "atomic_and({}, {})", op.ptr, op.value),
            AtomicOp::Or(op) => write!(f, "atomic_or({}, {})", op.ptr, op.value),
            AtomicOp::Xor(op) => write!(f, "atomic_xor({}, {})", op.ptr, op.value),
            AtomicOp::CompareAndSwap(op) => {
                write!(f, "compare_and_swap({}, {}, {})", op.ptr, op.cmp, op.val)
            }
        }
    }
}

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
#[allow(missing_docs)]
pub struct CompareAndSwapOperands {
    #[args(allow_ptr, ptr_read, ptr_write)]
    pub ptr: Value,
    pub cmp: Value,
    pub val: Value,
}