cubecl_ir/
operator.rs

1use core::fmt::Display;
2
3use alloc::{format, vec::Vec};
4
5use crate::TypeHash;
6
7use crate::{BinaryOperator, OperationArgs, OperationReflect, UnaryOperator, Variable};
8
9/// Operators available on the GPU
10#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
11#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationReflect)]
12#[operation(opcode_name = OperatorOpCode)]
13pub enum Operator {
14    #[operation(pure)]
15    Index(BinaryOperator),
16    CopyMemory(CopyMemoryOperator),
17    CopyMemoryBulk(CopyMemoryBulkOperator),
18    #[operation(pure)]
19    Slice(SliceOperator),
20    #[operation(pure)]
21    ReinterpretSlice(ReinterpretSliceOperator),
22    #[operation(pure)]
23    UncheckedIndex(BinaryOperator),
24    IndexAssign(BinaryOperator),
25    #[operation(pure)]
26    InitLine(LineInitOperator),
27    UncheckedIndexAssign(BinaryOperator),
28    #[operation(commutative, pure)]
29    And(BinaryOperator),
30    #[operation(commutative, pure)]
31    Or(BinaryOperator),
32    #[operation(pure)]
33    Not(UnaryOperator),
34    #[operation(pure)]
35    Cast(UnaryOperator),
36    #[operation(pure)]
37    Reinterpret(UnaryOperator),
38    /// A select statement/ternary
39    #[operation(pure)]
40    Select(Select),
41}
42
43impl Display for Operator {
44    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
45        match self {
46            Operator::Index(op) => write!(f, "{}[{}]", op.lhs, op.rhs),
47            Operator::CopyMemory(op) => {
48                write!(f, "[{}] = {}[{}]", op.out_index, op.input, op.in_index)
49            }
50            Operator::CopyMemoryBulk(op) => write!(
51                f,
52                "memcpy([{}], {}[{}], {})",
53                op.input, op.in_index, op.out_index, op.len
54            ),
55            Operator::Slice(op) => write!(f, "{}[{}..{}]", op.input, op.start, op.end),
56            Operator::ReinterpretSlice(op) => {
57                write!(f, "with_line_size({}, {})", op.input, op.line_size)
58            }
59            Operator::UncheckedIndex(op) => {
60                write!(f, "unchecked {}[{}]", op.lhs, op.rhs)
61            }
62            Operator::IndexAssign(op) => write!(f, "[{}] = {}", op.lhs, op.rhs),
63            Operator::UncheckedIndexAssign(op) => {
64                write!(f, "unchecked [{}] = {}", op.lhs, op.rhs)
65            }
66            Operator::And(op) => write!(f, "{} && {}", op.lhs, op.rhs),
67            Operator::Or(op) => write!(f, "{} || {}", op.lhs, op.rhs),
68            Operator::Not(op) => write!(f, "!{}", op.input),
69            Operator::InitLine(init) => {
70                let inits = init
71                    .inputs
72                    .iter()
73                    .map(|input| format!("{input}"))
74                    .collect::<Vec<_>>();
75                write!(f, "vec({})", inits.join(", "))
76            }
77            Operator::Select(op) => {
78                write!(f, "{} ? {} : {}", op.cond, op.then, op.or_else)
79            }
80            Operator::Cast(op) => write!(f, "cast({})", op.input),
81            Operator::Reinterpret(op) => write!(f, "reinterpret({})", op.input),
82        }
83    }
84}
85
86#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
87#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
88#[allow(missing_docs)]
89pub struct SliceOperator {
90    pub input: Variable,
91    pub start: Variable,
92    pub end: Variable,
93}
94
95#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
96#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
97#[allow(missing_docs)]
98pub struct ReinterpretSliceOperator {
99    pub input: Variable,
100    pub line_size: u32,
101}
102
103#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
104#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
105#[allow(missing_docs)]
106pub struct LineInitOperator {
107    pub inputs: Vec<Variable>,
108}
109
110#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
111#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
112#[allow(missing_docs)]
113pub struct CopyMemoryOperator {
114    pub out_index: Variable,
115    pub input: Variable,
116    pub in_index: Variable,
117}
118
119#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
120#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
121#[allow(missing_docs)]
122pub struct CopyMemoryBulkOperator {
123    pub out_index: Variable,
124    pub input: Variable,
125    pub in_index: Variable,
126    pub len: Variable,
127}
128
129#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
130#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
131#[allow(missing_docs)]
132pub struct Select {
133    pub cond: Variable,
134    pub then: Variable,
135    pub or_else: Variable,
136}
cubecl_ir/operator.rs

cubecl_ir/
operator.rs
