cubecl_opt/
block.rs

1use core::cell::RefCell;
2
3use alloc::{rc::Rc, vec::Vec};
4use cubecl_ir::{Instruction, Value};
5use stable_vec::StableVec;
6
7use crate::{ControlFlow, Function, GlobalState, version::PhiInstruction};
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum BlockUse {
11    ContinueTarget,
12    Merge,
13}
14
15/// A basic block of instructions interrupted by control flow. Phi nodes are assumed to come before
16/// any instructions. See <https://en.wikipedia.org/wiki/Basic_block>
17#[derive(Default, Debug, Clone)]
18pub struct BasicBlock {
19    pub(crate) block_use: Vec<BlockUse>,
20    /// The phi nodes that are required to be generated at the start of this block.
21    pub phi_nodes: Rc<RefCell<Vec<PhiInstruction>>>,
22    /// A stable list of operations performed in this block.
23    pub ops: Rc<RefCell<StableVec<Instruction>>>,
24    /// The control flow that terminates this block.
25    pub control_flow: Rc<RefCell<ControlFlow>>,
26}
27
28impl Function {
29    /// Visit all operations in the program with the specified read and write visitors.
30    pub fn visit_all(
31        &mut self,
32        state: &GlobalState,
33        mut visit_read: impl FnMut(&mut Self, &mut Value) + Clone,
34        mut visit_write: impl FnMut(&mut Self, &mut Value) + Clone,
35    ) {
36        for node in self.node_indices().collect::<Vec<_>>() {
37            let phi = self[node].phi_nodes.clone();
38            let ops = self[node].ops.clone();
39            let control_flow = self[node].control_flow.clone();
40
41            for phi in phi.borrow_mut().iter_mut() {
42                for elem in &mut phi.entries {
43                    visit_read(self, &mut elem.value);
44                }
45                visit_write(self, &mut phi.out);
46            }
47            for op in ops.borrow_mut().values_mut() {
48                self.visit_instruction(state, op, visit_read.clone(), visit_write.clone());
49            }
50            self.visit_control_flow(&mut control_flow.borrow_mut(), visit_read.clone());
51        }
52    }
53}
cubecl_opt/block.rs

cubecl_opt/
block.rs
