cubecl_ir/
branch.rs

1use alloc::{boxed::Box, format, vec::Vec};
2use core::fmt::Display;
3
4use crate::OperationReflect;
5
6use super::{OperationCode, Scope, Variable};
7use crate::TypeHash;
8
9/// All branching types.
10#[allow(clippy::large_enum_variant)]
11#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
12#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationCode)]
13#[operation(opcode_name = BranchOpCode)]
14pub enum Branch {
15    /// An if statement.
16    If(Box<If>),
17    /// An if else statement.
18    IfElse(Box<IfElse>),
19    /// A switch statement
20    Switch(Box<Switch>),
21    /// A range loop.
22    RangeLoop(Box<RangeLoop>),
23    /// A loop.
24    Loop(Box<Loop>),
25    /// A return statement.
26    Return,
27    /// A break statement.
28    Break,
29}
30
31impl OperationReflect for Branch {
32    type OpCode = BranchOpCode;
33
34    fn op_code(&self) -> Self::OpCode {
35        self.__match_opcode()
36    }
37}
38
39impl Display for Branch {
40    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
41        match self {
42            Branch::If(if_) => write!(f, "if({}) {}", if_.cond, if_.scope),
43            Branch::IfElse(if_else) => write!(
44                f,
45                "if({}) {} else {}",
46                if_else.cond, if_else.scope_if, if_else.scope_else
47            ),
48            Branch::Switch(switch) => write!(
49                f,
50                "switch({}) {:?}",
51                switch.value,
52                switch
53                    .cases
54                    .iter()
55                    .map(|case| format!("{}", case.0))
56                    .collect::<Vec<_>>(),
57            ),
58            Branch::RangeLoop(range_loop) => write!(
59                f,
60                "for({} in {}{}{}) {}",
61                range_loop.i,
62                range_loop.start,
63                if range_loop.inclusive { "..=" } else { ".." },
64                range_loop.end,
65                range_loop.scope
66            ),
67            Branch::Loop(loop_) => write!(f, "loop {}", loop_.scope),
68            Branch::Return => write!(f, "return"),
69            Branch::Break => write!(f, "break"),
70        }
71    }
72}
73
74#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
75#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash)]
76#[allow(missing_docs)]
77pub struct If {
78    pub cond: Variable,
79    pub scope: Scope,
80}
81
82#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
83#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash)]
84#[allow(missing_docs)]
85pub struct IfElse {
86    pub cond: Variable,
87    pub scope_if: Scope,
88    pub scope_else: Scope,
89}
90
91#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
92#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash)]
93#[allow(missing_docs)]
94pub struct Switch {
95    pub value: Variable,
96    pub scope_default: Scope,
97    pub cases: Vec<(Variable, Scope)>,
98}
99
100#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
101#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash)]
102#[allow(missing_docs)]
103pub struct RangeLoop {
104    pub i: Variable,
105    pub start: Variable,
106    pub end: Variable,
107    pub step: Option<Variable>,
108    pub inclusive: bool,
109    pub scope: Scope,
110}
111
112#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
113#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash)]
114#[allow(missing_docs)]
115pub struct Loop {
116    pub scope: Scope,
117}
cubecl_ir/branch.rs

cubecl_ir/
branch.rs
