cubecl_ir/
operator.rs

1use core::fmt::Display;
2
3use alloc::{format, vec::Vec};
4
5use crate::{IndexAssignOperator, IndexOperator, 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(IndexOperator),
16    CopyMemory(CopyMemoryOperator),
17    CopyMemoryBulk(CopyMemoryBulkOperator),
18    #[operation(pure)]
19    UncheckedIndex(IndexOperator),
20    IndexAssign(IndexAssignOperator),
21    UncheckedIndexAssign(IndexAssignOperator),
22    #[operation(pure)]
23    InitVector(VectorInitOperator),
24    #[operation(commutative, pure)]
25    And(BinaryOperator),
26    #[operation(commutative, pure)]
27    Or(BinaryOperator),
28    #[operation(pure)]
29    Not(UnaryOperator),
30    #[operation(pure)]
31    Cast(UnaryOperator),
32    #[operation(pure)]
33    Reinterpret(UnaryOperator),
34    #[operation(pure)]
35    Real(UnaryOperator),
36    #[operation(pure)]
37    Imag(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.list, op.index),
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::UncheckedIndex(op) => {
56                write!(f, "unchecked {}[{}]", op.list, op.index)
57            }
58            Operator::IndexAssign(op) => write!(f, "[{}] = {}", op.index, op.value),
59            Operator::UncheckedIndexAssign(op) => {
60                write!(f, "unchecked [{}] = {}", op.index, op.value)
61            }
62            Operator::And(op) => write!(f, "{} && {}", op.lhs, op.rhs),
63            Operator::Or(op) => write!(f, "{} || {}", op.lhs, op.rhs),
64            Operator::Not(op) => write!(f, "!{}", op.input),
65            Operator::InitVector(init) => {
66                let inits = init
67                    .inputs
68                    .iter()
69                    .map(|input| format!("{input}"))
70                    .collect::<Vec<_>>();
71                write!(f, "vec({})", inits.join(", "))
72            }
73            Operator::Select(op) => {
74                write!(f, "{} ? {} : {}", op.cond, op.then, op.or_else)
75            }
76            Operator::Cast(op) => write!(f, "cast({})", op.input),
77            Operator::Reinterpret(op) => write!(f, "reinterpret({})", op.input),
78            Operator::Real(op) => write!(f, "{}.real()", op.input),
79            Operator::Imag(op) => write!(f, "{}.imag()", op.input),
80        }
81    }
82}
83
84#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
85#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
86#[allow(missing_docs)]
87pub struct SliceOperator {
88    pub input: Variable,
89    pub start: Variable,
90    pub end: Variable,
91}
92
93#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
94#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
95#[allow(missing_docs)]
96pub struct ReinterpretSliceOperator {
97    pub input: Variable,
98    pub vector_size: u32,
99}
100
101#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
102#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
103#[allow(missing_docs)]
104pub struct VectorInitOperator {
105    pub inputs: Vec<Variable>,
106}
107
108#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
109#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
110#[allow(missing_docs)]
111pub struct CopyMemoryOperator {
112    pub out_index: Variable,
113    pub input: Variable,
114    pub in_index: Variable,
115}
116
117#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
118#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
119#[allow(missing_docs)]
120pub struct CopyMemoryBulkOperator {
121    pub out_index: Variable,
122    pub input: Variable,
123    pub in_index: Variable,
124    pub len: usize,
125    pub offset_input: Variable,
126    pub offset_out: 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
