cubecl_ir/
arithmetic.rs

1use core::fmt::Display;
2
3use crate::TypeHash;
4
5use crate::{BinaryOperands, OperationArgs, OperationReflect, UnaryOperands, Value};
6
7/// Arithmetic operations
8#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
9#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationReflect)]
10#[operation(opcode_name = ArithmeticOpCode, pure)]
11pub enum Arithmetic {
12    #[operation(commutative)]
13    Add(BinaryOperands),
14    #[operation(commutative)]
15    SaturatingAdd(BinaryOperands),
16    Fma(FmaOperands),
17    Sub(BinaryOperands),
18    SaturatingSub(BinaryOperands),
19    #[operation(commutative)]
20    Mul(BinaryOperands),
21    Div(BinaryOperands),
22    Abs(UnaryOperands),
23    Exp(UnaryOperands),
24    Log(UnaryOperands),
25    Log1p(UnaryOperands),
26    Expm1(UnaryOperands),
27    Cos(UnaryOperands),
28    Sin(UnaryOperands),
29    Tan(UnaryOperands),
30    Tanh(UnaryOperands),
31    Sinh(UnaryOperands),
32    Cosh(UnaryOperands),
33    ArcCos(UnaryOperands),
34    ArcSin(UnaryOperands),
35    ArcTan(UnaryOperands),
36    ArcSinh(UnaryOperands),
37    ArcCosh(UnaryOperands),
38    ArcTanh(UnaryOperands),
39    Degrees(UnaryOperands),
40    Radians(UnaryOperands),
41    ArcTan2(BinaryOperands),
42    Powf(BinaryOperands),
43    Powi(BinaryOperands),
44    Hypot(BinaryOperands),
45    Rhypot(BinaryOperands),
46    Sqrt(UnaryOperands),
47    InverseSqrt(UnaryOperands),
48    Round(UnaryOperands),
49    Floor(UnaryOperands),
50    Ceil(UnaryOperands),
51    Trunc(UnaryOperands),
52    Erf(UnaryOperands),
53    Recip(UnaryOperands),
54    Clamp(ClampOperands),
55    Neg(UnaryOperands),
56    #[operation(commutative)]
57    Max(BinaryOperands),
58    #[operation(commutative)]
59    Min(BinaryOperands),
60    /// Rust `Rem::rem`
61    Rem(BinaryOperands),
62    /// Pytorch %, or mod in SPIR-V
63    ModFloor(BinaryOperands),
64    Magnitude(UnaryOperands),
65    Normalize(UnaryOperands),
66    #[operation(commutative)]
67    Dot(BinaryOperands),
68    #[operation(commutative)]
69    MulHi(BinaryOperands),
70    VectorSum(UnaryOperands),
71}
72
73impl Display for Arithmetic {
74    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
75        match self {
76            Arithmetic::Add(op) => write!(f, "{} + {}", op.lhs, op.rhs),
77            Arithmetic::SaturatingAdd(op) => write!(f, "saturating_add({}, {})", op.lhs, op.rhs),
78            Arithmetic::Fma(op) => write!(f, "{} * {} + {}", op.a, op.b, op.c),
79            Arithmetic::Sub(op) => write!(f, "{} - {}", op.lhs, op.rhs),
80            Arithmetic::SaturatingSub(op) => write!(f, "saturating_sub({}, {})", op.lhs, op.rhs),
81            Arithmetic::Mul(op) => write!(f, "{} * {}", op.lhs, op.rhs),
82            Arithmetic::Div(op) => write!(f, "{} / {}", op.lhs, op.rhs),
83            Arithmetic::Abs(op) => write!(f, "{}.abs()", op.input),
84            Arithmetic::Exp(op) => write!(f, "{}.exp()", op.input),
85            Arithmetic::Log(op) => write!(f, "{}.log()", op.input),
86            Arithmetic::Log1p(op) => write!(f, "{}.log_1p()", op.input),
87            Arithmetic::Expm1(op) => write!(f, "{}.exp_m1()", op.input),
88            Arithmetic::Cos(op) => write!(f, "{}.cos()", op.input),
89            Arithmetic::Sin(op) => write!(f, "{}.sin()", op.input),
90            Arithmetic::Tan(op) => write!(f, "{}.tan()", op.input),
91            Arithmetic::Tanh(op) => write!(f, "{}.tanh()", op.input),
92            Arithmetic::Sinh(op) => write!(f, "{}.sinh()", op.input),
93            Arithmetic::Cosh(op) => write!(f, "{}.cosh()", op.input),
94            Arithmetic::ArcCos(op) => write!(f, "{}.acos()", op.input),
95            Arithmetic::ArcSin(op) => write!(f, "{}.asin()", op.input),
96            Arithmetic::ArcTan(op) => write!(f, "{}.atan()", op.input),
97            Arithmetic::ArcSinh(op) => write!(f, "{}.asinh()", op.input),
98            Arithmetic::ArcCosh(op) => write!(f, "{}.acosh()", op.input),
99            Arithmetic::ArcTanh(op) => write!(f, "{}.atanh()", op.input),
100            Arithmetic::Degrees(op) => write!(f, "{}.degrees()", op.input),
101            Arithmetic::Radians(op) => write!(f, "{}.radians()", op.input),
102            Arithmetic::ArcTan2(op) => write!(f, "{}.atan2({})", op.lhs, op.rhs),
103            Arithmetic::Powf(op) => write!(f, "{}.pow({})", op.lhs, op.rhs),
104            Arithmetic::Powi(op) => write!(f, "{}.powi({})", op.lhs, op.rhs),
105            Arithmetic::Hypot(op) => write!(f, "{}.hypot({})", op.lhs, op.rhs),
106            Arithmetic::Rhypot(op) => write!(f, "{}.rhypot({})", op.lhs, op.rhs),
107            Arithmetic::Sqrt(op) => write!(f, "{}.sqrt()", op.input),
108            Arithmetic::InverseSqrt(op) => write!(f, "{}.inverse_sqrt()", op.input),
109            Arithmetic::Round(op) => write!(f, "{}.round()", op.input),
110            Arithmetic::Floor(op) => write!(f, "{}.floor()", op.input),
111            Arithmetic::Ceil(op) => write!(f, "{}.ceil()", op.input),
112            Arithmetic::Trunc(op) => write!(f, "{}.trunc()", op.input),
113            Arithmetic::Erf(op) => write!(f, "{}.erf()", op.input),
114            Arithmetic::Recip(op) => write!(f, "{}.recip()", op.input),
115            Arithmetic::Clamp(op) => {
116                write!(f, "{}.clamp({}, {})", op.input, op.min_value, op.max_value)
117            }
118            Arithmetic::Neg(op) => write!(f, "-{}", op.input),
119            Arithmetic::Max(op) => write!(f, "{}.max({})", op.lhs, op.rhs),
120            Arithmetic::Min(op) => write!(f, "{}.min({})", op.lhs, op.rhs),
121            Arithmetic::Rem(op) => write!(f, "{} % {}", op.lhs, op.rhs),
122            Arithmetic::ModFloor(op) => write!(f, "{}.mod_floor({})", op.lhs, op.rhs),
123            Arithmetic::Magnitude(op) => write!(f, "{}.length()", op.input),
124            Arithmetic::Normalize(op) => write!(f, "{}.normalize()", op.input),
125            Arithmetic::Dot(op) => write!(f, "{}.dot({})", op.lhs, op.rhs),
126            Arithmetic::MulHi(op) => write!(f, "mul_hi({}, {})", op.lhs, op.rhs),
127            Arithmetic::VectorSum(op) => write!(f, "{}.vector_sum()", op.input),
128        }
129    }
130}
131
132#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
133#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
134#[allow(missing_docs)]
135pub struct ClampOperands {
136    pub input: Value,
137    pub min_value: Value,
138    pub max_value: Value,
139}
140
141#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
142#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationArgs)]
143#[allow(missing_docs)]
144pub struct FmaOperands {
145    pub a: Value,
146    pub b: Value,
147    pub c: Value,
148}
cubecl_ir/arithmetic.rs

cubecl_ir/
arithmetic.rs
