cubecl_opt/
control_flow.rs

1#![allow(unknown_lints, unnecessary_transmutes)]
2
3use core::mem::transmute;
4
5use crate::{BasicBlock, BlockUse, Function, GlobalState, NodeIndex};
6use alloc::{boxed::Box, vec::Vec};
7use cubecl_ir::{
8    Arithmetic, BinaryOperands, Branch, Comparison, ConstantValue, ElemType, If, IfElse,
9    Instruction, Loop, Marker, Memory, Operation, RangeLoop, StoreOperands, Switch, Type, Value,
10};
11use petgraph::{Direction, graph::EdgeIndex, visit::EdgeRef};
12use stable_vec::StableVec;
13
14/// Control flow that terminates a block
15#[derive(Default, Debug, Clone)]
16pub enum ControlFlow {
17    /// An if or if-else branch that should be structured if applicable.
18    IfElse {
19        cond: Value,
20        then: NodeIndex,
21        or_else: NodeIndex,
22        merge: Option<NodeIndex>,
23    },
24    /// A switch branch that paths based on `value`
25    Switch {
26        value: Value,
27        default: NodeIndex,
28        branches: Vec<(u32, NodeIndex)>,
29        merge: Option<NodeIndex>,
30    },
31    /// A loop with a header (the block that contains this variant), a `body` and a `continue target`.
32    /// `merge` is the block that gets executed as soon as the loop terminates.
33    Loop {
34        body: NodeIndex,
35        continue_target: NodeIndex,
36        merge: NodeIndex,
37    },
38    /// A loop with a header (the block that contains this variant), a `body` and a `continue target`.
39    /// `merge` is the block that gets executed as soon as the loop terminates. The header contains
40    /// the break condition.
41    LoopBreak {
42        break_cond: Value,
43        body: NodeIndex,
44        continue_target: NodeIndex,
45        merge: NodeIndex,
46    },
47    /// A return statement. This should only occur once in the program and all other returns should
48    /// instead branch to this single return block.
49    Return { value: Option<Value> },
50    /// Unreachable control flow
51    Unreachable,
52    /// No special control flow. The block must have exactly one edge that should be followed.
53    #[default]
54    None,
55}
56
57pub(crate) enum ControlFlowAction {
58    None,
59    AbortBlock,
60}
61
62impl Function {
63    pub(crate) fn parse_control_flow(
64        &mut self,
65        state: &GlobalState,
66        branch: Branch,
67    ) -> ControlFlowAction {
68        match branch {
69            Branch::If(if_) => {
70                self.parse_if(state, *if_);
71                ControlFlowAction::None
72            }
73            Branch::IfElse(if_else) => {
74                self.parse_if_else(state, if_else);
75                ControlFlowAction::None
76            }
77            Branch::Switch(switch) => {
78                self.parse_switch(state, *switch);
79                ControlFlowAction::None
80            }
81            Branch::RangeLoop(range_loop) => {
82                self.parse_for_loop(state, *range_loop);
83                ControlFlowAction::None
84            }
85            Branch::Loop(loop_) => {
86                self.parse_loop(state, *loop_);
87                ControlFlowAction::None
88            }
89            Branch::Unreachable => {
90                let current_block = self.current_block.take().unwrap();
91                *self[current_block].control_flow.borrow_mut() = ControlFlow::Unreachable;
92                ControlFlowAction::AbortBlock
93            }
94            Branch::Return => {
95                let current_block = self.current_block.take().unwrap();
96                let ret = self.ret();
97                self.add_edge(current_block, ret, 0);
98                ControlFlowAction::AbortBlock
99            }
100            Branch::Break => {
101                let current_block = self.current_block.take().unwrap();
102                let loop_break = *self.loop_break.back().expect("Can't break outside loop");
103                self.add_edge(current_block, loop_break, 0);
104                ControlFlowAction::AbortBlock
105            }
106        }
107    }
108
109    pub(crate) fn parse_if(&mut self, state: &GlobalState, if_: If) {
110        let current_block = self.current_block.unwrap();
111        let then = self.add_node(BasicBlock::default());
112        let next = self.add_node(BasicBlock::default());
113        let mut merge = next;
114
115        self.add_edge(current_block, then, 0);
116        self.add_edge(current_block, next, 0);
117
118        self.current_block = Some(then);
119        let is_break = self.parse_scope(state, if_.scope);
120
121        if let Some(current_block) = self.current_block {
122            self.add_edge(current_block, next, 0);
123        } else {
124            // Returned
125            merge = self.ret;
126        }
127
128        let merge = if is_break { None } else { Some(merge) };
129
130        *self[current_block].control_flow.borrow_mut() = ControlFlow::IfElse {
131            cond: if_.cond,
132            then,
133            or_else: next,
134            merge,
135        };
136        if let Some(merge) = merge {
137            self[merge].block_use.push(BlockUse::Merge);
138        }
139        self.current_block = Some(next);
140    }
141
142    pub(crate) fn parse_if_else(&mut self, state: &GlobalState, if_else: Box<IfElse>) {
143        let current_block = self.current_block.unwrap();
144        let then = self.add_node(BasicBlock::default());
145        let or_else = self.add_node(BasicBlock::default());
146        let next = self.add_node(BasicBlock::default());
147        let mut merge = next;
148
149        self.add_edge(current_block, then, 0);
150        self.add_edge(current_block, or_else, 0);
151
152        self.current_block = Some(then);
153        let is_break = self.parse_scope(state, if_else.scope_if);
154
155        if let Some(current_block) = self.current_block {
156            self.add_edge(current_block, next, 0);
157        } else {
158            // Returned
159            merge = self.ret;
160        }
161
162        self.current_block = Some(or_else);
163        let is_break = self.parse_scope(state, if_else.scope_else) || is_break;
164
165        if let Some(current_block) = self.current_block {
166            self.add_edge(current_block, next, 0);
167        } else {
168            // Returned
169            merge = self.ret;
170        }
171
172        let merge = if is_break { None } else { Some(merge) };
173        *self[current_block].control_flow.borrow_mut() = ControlFlow::IfElse {
174            cond: if_else.cond,
175            then,
176            or_else,
177            merge,
178        };
179        if let Some(merge) = merge {
180            self[merge].block_use.push(BlockUse::Merge);
181        }
182        self.current_block = Some(next);
183    }
184
185    pub(crate) fn parse_switch(&mut self, state: &GlobalState, switch: Switch) {
186        let current_block = self.current_block.unwrap();
187        let next = self.add_node(BasicBlock::default());
188
189        let branches = switch
190            .cases
191            .into_iter()
192            .map(|(val, case)| {
193                let case_id = self.add_node(BasicBlock::default());
194                self.add_edge(current_block, case_id, 0);
195                self.current_block = Some(case_id);
196                let is_break = self.parse_scope(state, case);
197                let is_ret = if let Some(current_block) = self.current_block {
198                    self.add_edge(current_block, next, 0);
199                    false
200                } else {
201                    !is_break
202                };
203                let val = match val.as_const().expect("Switch value must be constant") {
204                    ConstantValue::Int(val) => unsafe { transmute::<i32, u32>(val as i32) },
205                    ConstantValue::UInt(val) => val as u32,
206                    _ => unreachable!("Switch cases must be integer"),
207                };
208                (val, case_id, is_break, is_ret)
209            })
210            .collect::<Vec<_>>();
211
212        let is_break_branch = branches.iter().any(|it| it.2);
213        let mut is_ret = branches.iter().any(|it| it.3);
214        let branches = branches
215            .into_iter()
216            .map(|it| (it.0, it.1))
217            .collect::<Vec<_>>();
218
219        let default = self.add_node(BasicBlock::default());
220        self.add_edge(current_block, default, 0);
221        self.current_block = Some(default);
222        let is_break_def = self.parse_scope(state, switch.scope_default);
223
224        if let Some(current_block) = self.current_block {
225            self.add_edge(current_block, next, 0);
226        } else {
227            is_ret = !is_break_def;
228        }
229
230        let merge = if is_break_def || is_break_branch {
231            None
232        } else if is_ret {
233            Some(self.ret)
234        } else {
235            self[next].block_use.push(BlockUse::Merge);
236            Some(next)
237        };
238
239        *self[current_block].control_flow.borrow_mut() = ControlFlow::Switch {
240            value: switch.value,
241            default,
242            branches,
243            merge,
244        };
245        if let Some(merge) = merge {
246            self[merge].block_use.push(BlockUse::Merge);
247        }
248
249        self.current_block = Some(next);
250    }
251
252    fn parse_loop(&mut self, state: &GlobalState, loop_: Loop) {
253        let current_block = self.current_block.unwrap();
254        let header = self.add_node(BasicBlock::default());
255        self.add_edge(current_block, header, 0);
256
257        let body = self.add_node(BasicBlock::default());
258        let next = self.add_node(BasicBlock::default());
259
260        self.add_edge(header, body, 0);
261
262        self.loop_break.push_back(next);
263
264        self.current_block = Some(body);
265        self.parse_scope(state, loop_.scope);
266        let continue_target = self.add_node(BasicBlock::default());
267        self[continue_target]
268            .block_use
269            .push(BlockUse::ContinueTarget);
270
271        self.loop_break.pop_back();
272
273        if let Some(current_block) = self.current_block {
274            self.add_edge(current_block, continue_target, 0);
275        }
276
277        self.add_edge(continue_target, header, 0);
278
279        *self[header].control_flow.borrow_mut() = ControlFlow::Loop {
280            body,
281            continue_target,
282            merge: next,
283        };
284        self[next].block_use.push(BlockUse::Merge);
285        self.current_block = Some(next);
286    }
287
288    fn parse_for_loop(&mut self, state: &GlobalState, range_loop: RangeLoop) {
289        let step = range_loop.step.unwrap_or(1.into());
290
291        let i = range_loop.i;
292
293        let load_i = |f: &Self, block: NodeIndex| {
294            let tmp_i = state.allocator.create_value(i.ty.unwrap_ptr());
295            f[block]
296                .ops
297                .borrow_mut()
298                .push(Instruction::new(Memory::Load(i), tmp_i));
299            tmp_i
300        };
301        let store_i = |f: &Self, block: NodeIndex, value: Value| {
302            f[block]
303                .ops
304                .borrow_mut()
305                .push(Instruction::no_out(Memory::Store(StoreOperands {
306                    ptr: i,
307                    value,
308                })));
309        };
310
311        store_i(self, self.current_block.unwrap(), range_loop.start);
312
313        let current_block = self.current_block.unwrap();
314        let header = self.add_node(BasicBlock::default());
315        self.add_edge(current_block, header, 0);
316
317        let body = self.add_node(BasicBlock::default());
318        let next = self.add_node(BasicBlock::default());
319
320        self.add_edge(header, body, 0);
321        self.add_edge(header, next, 0);
322
323        self.loop_break.push_back(next);
324
325        self.current_block = Some(body);
326        self.parse_scope(state, range_loop.scope);
327
328        self.loop_break.pop_back();
329
330        let current_block = self.current_block.expect("For loop has no loopback path");
331
332        let continue_target = if self[current_block].block_use.contains(&BlockUse::Merge) {
333            let target = self.add_node(BasicBlock::default());
334            self.add_edge(current_block, target, 0);
335            target
336        } else {
337            current_block
338        };
339
340        self.add_edge(continue_target, header, 0);
341
342        self[continue_target]
343            .block_use
344            .push(BlockUse::ContinueTarget);
345        self[next].block_use.push(BlockUse::Merge);
346        self.current_block = Some(next);
347
348        // For loop constructs
349        {
350            let op = match range_loop.inclusive {
351                true => Comparison::LowerEqual,
352                false => Comparison::Lower,
353            };
354            let tmp = state.allocator.create_value(Type::scalar(ElemType::Bool));
355            let val_i = load_i(self, header);
356            self[header].ops.borrow_mut().push(Instruction::new(
357                op(BinaryOperands {
358                    lhs: val_i,
359                    rhs: range_loop.end,
360                }),
361                tmp,
362            ));
363
364            *self[header].control_flow.borrow_mut() = ControlFlow::LoopBreak {
365                break_cond: tmp,
366                body,
367                continue_target,
368                merge: next,
369            };
370        }
371        let tmp_i = state.allocator.create_value(i.ty.unwrap_ptr());
372        let val_i = load_i(self, current_block);
373        self[current_block].ops.borrow_mut().push(Instruction::new(
374            Arithmetic::Add(BinaryOperands {
375                lhs: val_i,
376                rhs: step,
377            }),
378            tmp_i,
379        ));
380        store_i(self, current_block, tmp_i);
381    }
382
383    pub(crate) fn split_critical_edges(&mut self) {
384        for block in self.node_ids() {
385            let successors = self.edges(block);
386            let successors = successors.map(|edge| (edge.id(), edge.target()));
387            let successors: Vec<_> = successors.collect();
388
389            if successors.len() > 1 {
390                let crit = successors
391                    .iter()
392                    .filter(|(_, b)| self.predecessors(*b).len() > 1)
393                    .collect::<Vec<_>>();
394                for (edge, successor) in crit {
395                    self.remove_edge(*edge);
396                    let new_block = self.add_node(BasicBlock::default());
397                    self.add_edge(block, new_block, 0);
398                    self.add_edge(new_block, *successor, 0);
399                    self.invalidate_structure();
400                    update_phi(self, *successor, block, new_block);
401                    update_control_flow(self, block, *successor, new_block);
402                }
403            }
404        }
405    }
406
407    /// Split blocks at a `free` call because we only track liveness at a block level
408    /// It's easier than doing liveness per-instruction, and free calls are rare anyways
409    pub(crate) fn split_free(&mut self) {
410        let mut splits = 0;
411        while self.split_free_inner() {
412            splits += 1;
413        }
414        if splits > 0 {
415            self.invalidate_structure();
416        }
417    }
418
419    fn split_free_inner(&mut self) -> bool {
420        let is_free =
421            |inst: &Instruction| matches!(inst.operation, Operation::Marker(Marker::Free(_)));
422
423        for block in self.node_ids() {
424            let ops = self.block(block).ops.clone();
425            let len = ops.borrow().num_elements();
426            let idx = ops.borrow().values().position(is_free);
427            if let Some(idx) = idx {
428                // Separate free into its own block. They can be merged again later.
429                if idx > 0 {
430                    self.split_block_after(block, idx - 1);
431                    return true;
432                }
433                if idx < len - 1 {
434                    self.split_block_after(block, idx);
435                    return true;
436                }
437            }
438        }
439
440        false
441    }
442
443    /// Split block after `idx` and return the new block
444    fn split_block_after(&mut self, block: NodeIndex, idx: usize) -> NodeIndex {
445        let successors = self.successors(block);
446        let edges: Vec<EdgeIndex> = self
447            .edges_directed(block, Direction::Outgoing)
448            .map(|it| it.id())
449            .collect();
450        for edge in edges {
451            self.remove_edge(edge);
452        }
453
454        let ops = self.block(block).ops.take();
455        let before: Vec<_> = ops.values().take(idx + 1).cloned().collect();
456        let after: Vec<_> = ops.values().skip(idx + 1).cloned().collect();
457        *self.block(block).ops.borrow_mut() = StableVec::from_iter(before);
458
459        let new_block = BasicBlock::default();
460        new_block.control_flow.swap(&self.block(block).control_flow);
461        new_block.ops.borrow_mut().extend(after);
462        let new_block = self.add_node(new_block);
463
464        self.add_edge(block, new_block, 0);
465        for successor in successors {
466            self.add_edge(new_block, successor, 0);
467        }
468        new_block
469    }
470}
471
472fn update_control_flow(func: &mut Function, block: NodeIndex, from: NodeIndex, to: NodeIndex) {
473    let update = |id: &mut NodeIndex| {
474        if *id == from {
475            *id = to
476        }
477    };
478
479    match &mut *func[block].control_flow.borrow_mut() {
480        ControlFlow::IfElse { then, or_else, .. } => {
481            update(then);
482            update(or_else);
483        }
484        ControlFlow::Switch {
485            default, branches, ..
486        } => {
487            update(default);
488
489            for branch in branches {
490                update(&mut branch.1);
491            }
492        }
493        ControlFlow::Loop {
494            body,
495            continue_target,
496            merge,
497        } => {
498            update(body);
499            update(continue_target);
500            update(merge);
501        }
502        ControlFlow::LoopBreak {
503            body,
504            continue_target,
505            merge,
506            ..
507        } => {
508            update(body);
509            update(continue_target);
510            update(merge);
511        }
512        _ => {}
513    }
514}
515
516fn update_phi(func: &mut Function, block: NodeIndex, from: NodeIndex, to: NodeIndex) {
517    for phi in func[block].phi_nodes.borrow_mut().iter_mut() {
518        for entry in phi.entries.iter_mut() {
519            if entry.block == from {
520                entry.block = to;
521            }
522        }
523    }
524}
cubecl_opt/control_flow.rs

cubecl_opt/
control_flow.rs
