cubecl_ir/
plane.rs

1use core::fmt::Display;
2
3use crate::OperationReflect;
4
5use super::{BinaryOperator, UnaryOperator};
6use crate::TypeHash;
7
8/// All plane operations.
9///
10/// Note that not all backends support plane (warp/subgroup) operations. Use the [runtime flag](crate::Feature::Plane).
11#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
12#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationReflect)]
13#[operation(opcode_name = PlaneOpCode)]
14#[allow(dead_code, missing_docs)] // Some variants might not be used with different flags
15pub enum Plane {
16    Elect,
17    All(UnaryOperator),
18    Any(UnaryOperator),
19    Ballot(UnaryOperator),
20    Broadcast(BinaryOperator),
21    Shuffle(BinaryOperator),
22    ShuffleXor(BinaryOperator),
23    ShuffleUp(BinaryOperator),
24    ShuffleDown(BinaryOperator),
25    Sum(UnaryOperator),
26    InclusiveSum(UnaryOperator),
27    ExclusiveSum(UnaryOperator),
28    Prod(UnaryOperator),
29    InclusiveProd(UnaryOperator),
30    ExclusiveProd(UnaryOperator),
31    Min(UnaryOperator),
32    Max(UnaryOperator),
33}
34
35impl Display for Plane {
36    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
37        match self {
38            Plane::Elect => write!(f, "plane_elect()"),
39            Plane::All(op) => write!(f, "plane_all({})", op.input),
40            Plane::Any(op) => write!(f, "plane_any({})", op.input),
41            Plane::Ballot(op) => write!(f, "plane_ballot({})", op.input),
42            Plane::Broadcast(op) => {
43                write!(f, "plane_broadcast({}, {})", op.lhs, op.rhs)
44            }
45            Plane::Shuffle(op) => {
46                write!(f, "plane_shuffle({}, {})", op.lhs, op.rhs)
47            }
48            Plane::ShuffleXor(op) => {
49                write!(f, "plane_shuffle_xor({}, {})", op.lhs, op.rhs)
50            }
51            Plane::ShuffleUp(op) => {
52                write!(f, "plane_shuffle_up({}, {})", op.lhs, op.rhs)
53            }
54            Plane::ShuffleDown(op) => {
55                write!(f, "plane_shuffle_down({}, {})", op.lhs, op.rhs)
56            }
57            Plane::Sum(op) => write!(f, "plane_sum({})", op.input),
58            Plane::InclusiveSum(op) => write!(f, "plane_inclusive_sum({})", op.input),
59            Plane::ExclusiveSum(op) => write!(f, "plane_exclusive_sum({})", op.input),
60            Plane::Prod(op) => write!(f, "plane_product({})", op.input),
61            Plane::InclusiveProd(op) => write!(f, "plane_inclusive_product({})", op.input),
62            Plane::ExclusiveProd(op) => write!(f, "plane_exclusive_product({})", op.input),
63            Plane::Min(op) => write!(f, "plane_min({})", op.input),
64            Plane::Max(op) => write!(f, "plane_max({})", op.input),
65        }
66    }
67}
cubecl_ir/plane.rs

cubecl_ir/
plane.rs
