Skip to main content

cubecl_opt/
phi_frontiers.rs

1use alloc::vec::Vec;
2use cubecl_ir::{AddressSpace, Id, Type, Value};
3use hashbrown::HashMap;
4use petgraph::graph::NodeIndex;
5
6use crate::{
7    Function, GlobalState, MemoryBlock,
8    analyses::{dominance::DomFrontiers, liveness::Liveness, writes::LocalStores},
9};
10
11use super::version::{PhiEntry, PhiInstruction};
12
13impl Function {
14    /// Places a phi node for each live variable at each frontier
15    pub fn place_phi_nodes(&mut self, state: &GlobalState) {
16        let locals = self.destructurable_local_memories();
17        let writes = self.analysis::<LocalStores>(state);
18        let liveness = self.analysis::<Liveness>(state);
19        let dom_frontiers = self.analysis::<DomFrontiers>(state);
20
21        for (local_id, mem) in locals {
22            let mut workset: Vec<_> = self
23                .node_ids()
24                .iter()
25                .filter(|index| writes[*index].contains(&local_id))
26                .copied()
27                .collect();
28            let mut considered = workset.clone();
29            let mut already_inserted = Vec::new();
30
31            while let Some(node) = workset.pop() {
32                for frontier in dom_frontiers[&node].clone() {
33                    if already_inserted.contains(&frontier) || liveness.is_dead(frontier, local_id)
34                    {
35                        continue;
36                    }
37                    self.insert_phi(frontier, local_id, mem.value_ty);
38                    already_inserted.push(frontier);
39                    if !considered.contains(&frontier) {
40                        workset.push(frontier);
41                        considered.push(frontier);
42                    }
43                }
44            }
45        }
46    }
47
48    /// Insert a phi node for variable `id` at `block`
49    pub fn insert_phi(&mut self, block: NodeIndex, id: Id, item: Type) {
50        let val = Value::new(id, item);
51        let entries = self.predecessors(block).into_iter().map(|pred| PhiEntry {
52            block: pred,
53            value: val,
54        });
55        let phi = PhiInstruction {
56            out: val,
57            entries: entries.collect(),
58        };
59        self[block].phi_nodes.borrow_mut().push(phi);
60    }
61
62    /// Returns all pointers to local stack space that are [destructurable](Type::is_destructurable)
63    pub fn destructurable_local_memories(&self) -> HashMap<Id, MemoryBlock> {
64        let locals = self.memories.iter().filter(|(_, mem)| {
65            matches!(mem.address_space, AddressSpace::Local) && mem.value_ty.is_destructurable()
66        });
67        locals.map(|(k, v)| (*k, *v)).collect()
68    }
69}