cubecl_ir/
operator.rs

1use core::fmt::Display;
2
3use alloc::{format, vec::Vec};
4
5use crate::{Builtin, Id, TypeHash};
6
7use crate::{BinaryOperands, OperationArgs, OperationReflect, UnaryOperands, Value};
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    InitVector(InitVectorOperands),
16    #[operation(pure)]
17    ExtractComponent(BinaryOperands),
18    #[operation(pure)]
19    InsertComponent(VectorInsertOperands),
20    #[operation(commutative, pure)]
21    And(BinaryOperands),
22    #[operation(commutative, pure)]
23    Or(BinaryOperands),
24    #[operation(pure)]
25    Not(UnaryOperands),
26    #[operation(pure)]
27    Cast(UnaryOperands),
28    #[operation(pure)]
29    Reinterpret(UnaryOperands),
30    /// A select statement/ternary
31    #[operation(pure)]
32    Select(SelectOperands),
33    #[operation(pure)]
34    ReadBuiltin(Builtin),
35    #[operation(pure)]
36    ReadScalar(Id),
37}
38
39impl Display for Operator {
40    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
41        match self {
42            Operator::And(op) => write!(f, "{} && {}", op.lhs, op.rhs),
43            Operator::Or(op) => write!(f, "{} || {}", op.lhs, op.rhs),
44            Operator::Not(op) => write!(f, "!{}", op.input),
45            Operator::InitVector(init) => {
46                let inits = init
47                    .inputs
48                    .iter()
49                    .map(|input| format!("{input}"))
50                    .collect::<Vec<_>>();
51                write!(f, "vec({})", inits.join(", "))
52            }
53            Operator::ExtractComponent(op) => write!(f, "extract({}, {})", op.lhs, op.rhs),
54            Operator::InsertComponent(op) => {
55                write!(f, "insert({}, {}, {})", op.vector, op.index, op.value)
56            }
57            Operator::Select(op) => {
58                write!(f, "{} ? {} : {}", op.cond, op.then, op.or_else)
59            }
60            Operator::Cast(op) => write!(f, "cast({})", op.input),
61            Operator::Reinterpret(op) => write!(f, "reinterpret({})", op.input),
62            Operator::ReadBuiltin(builtin) => write!(f, "read_builtin({builtin:?})"),
63            Operator::ReadScalar(id) => write!(f, "read_scalar({id})"),
64        }
65    }
66}
67
68#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
69#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
70#[allow(missing_docs)]
71pub struct SliceOperands {
72    pub input: Value,
73    pub start: Value,
74    pub end: Value,
75}
76
77#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
78#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
79#[allow(missing_docs)]
80pub struct ReinterpretSliceOperands {
81    pub input: Value,
82    pub vector_size: u32,
83}
84
85#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
86#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
87#[allow(missing_docs)]
88pub struct InitVectorOperands {
89    pub inputs: Vec<Value>,
90}
91
92#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
93#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
94#[allow(missing_docs)]
95pub struct VectorInsertOperands {
96    pub vector: Value,
97    pub index: Value,
98    pub value: Value,
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 CopyMemoryOperands {
105    #[args(allow_ptr, ptr_read)]
106    pub source: Value,
107    #[args(allow_ptr, ptr_write)]
108    pub target: Value,
109    pub len: usize,
110}
111
112#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
113#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
114#[allow(missing_docs)]
115pub struct SelectOperands {
116    pub cond: Value,
117    pub then: Value,
118    pub or_else: Value,
119}
cubecl_ir/operator.rs

cubecl_ir/
operator.rs
