cubecl_core/ir/
plane.rs

1use std::fmt::Display;
2
3use super::{BinaryOperator, UnaryOperator};
4use serde::{Deserialize, Serialize};
5
6/// All plane operations.
7///
8/// Note that not all backends support plane (warp/subgroup) operations. Use the [runtime flag](crate::Feature::Plane).
9#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
10#[allow(dead_code, missing_docs)] // Some variants might not be used with different flags
11pub enum Plane {
12    Elect,
13    All(UnaryOperator),
14    Any(UnaryOperator),
15    Broadcast(BinaryOperator),
16    Sum(UnaryOperator),
17    Prod(UnaryOperator),
18    Min(UnaryOperator),
19    Max(UnaryOperator),
20}
21
22impl Display for Plane {
23    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
24        match self {
25            Plane::Elect => writeln!(f, "plane_elect()"),
26            Plane::All(op) => writeln!(f, "plane_all({})", op.input),
27            Plane::Any(op) => writeln!(f, "plane_any({})", op.input),
28            Plane::Broadcast(op) => {
29                writeln!(f, "plane_broadcast({}, {})", op.lhs, op.rhs)
30            }
31            Plane::Sum(op) => writeln!(f, "plane_sum({})", op.input),
32            Plane::Prod(op) => writeln!(f, "plane_product({})", op.input),
33            Plane::Min(op) => writeln!(f, "plane_min({})", op.input),
34            Plane::Max(op) => writeln!(f, "plane_max({})", op.input),
35        }
36    }
37}