rustc_codegen_spirv/linker/
ipo.rs

1//! Tools for interprocedural optimizations (aka "IPO"s).
2
3// FIXME(eddyb) perhaps make all IPOs sub-modules of this module?
4
5use indexmap::IndexSet;
6use rspirv::dr::Module;
7use rspirv::spirv::{Op, Word};
8use rustc_data_structures::fx::FxHashMap;
9
10// FIXME(eddyb) use newtyped indices and `IndexVec`.
11type FuncIdx = usize;
12
13pub struct CallGraph {
14    pub entry_points: IndexSet<FuncIdx>,
15
16    /// `callees[i].contains(j)` implies `functions[i]` calls `functions[j]`.
17    pub callees: Vec<IndexSet<FuncIdx>>,
18}
19
20impl CallGraph {
21    pub fn collect(module: &Module) -> Self {
22        Self::collect_with_func_id_to_idx(module).0
23    }
24    pub fn collect_with_func_id_to_idx(module: &Module) -> (Self, FxHashMap<Word, FuncIdx>) {
25        let func_id_to_idx: FxHashMap<_, _> = module
26            .functions
27            .iter()
28            .enumerate()
29            .map(|(i, func)| (func.def_id().unwrap(), i))
30            .collect();
31        let entry_points = module
32            .entry_points
33            .iter()
34            .map(|entry| {
35                assert_eq!(entry.class.opcode, Op::EntryPoint);
36                func_id_to_idx[&entry.operands[1].unwrap_id_ref()]
37            })
38            .collect();
39        let callees = module
40            .functions
41            .iter()
42            .map(|func| {
43                func.all_inst_iter()
44                    .filter(|inst| inst.class.opcode == Op::FunctionCall)
45                    .filter_map(|inst| {
46                        // FIXME(eddyb) `func_id_to_idx` should always have an
47                        // entry for a callee ID, but when ran early enough
48                        // (before zombie removal), the callee ID might not
49                        // point to an `OpFunction` (unsure what, `OpUndef`?).
50                        func_id_to_idx
51                            .get(&inst.operands[0].unwrap_id_ref())
52                            .copied()
53                    })
54                    .collect()
55            })
56            .collect();
57        (
58            Self {
59                entry_points,
60                callees,
61            },
62            func_id_to_idx,
63        )
64    }
65
66    /// Order functions using a post-order traversal, i.e. callees before callers.
67    // FIXME(eddyb) replace this with `rustc_data_structures::graph::iterate`
68    // (or similar).
69    pub fn post_order(&self) -> Vec<FuncIdx> {
70        let num_funcs = self.callees.len();
71
72        // FIXME(eddyb) use a proper bitset.
73        let mut visited = vec![false; num_funcs];
74        let mut post_order = Vec::with_capacity(num_funcs);
75
76        // Visit the call graph with entry points as roots.
77        for &entry in &self.entry_points {
78            self.post_order_step(entry, &mut visited, &mut post_order);
79        }
80
81        // Also visit any functions that were not reached from entry points
82        // (they might be dead but they should be processed nonetheless).
83        for func in 0..num_funcs {
84            if !visited[func] {
85                self.post_order_step(func, &mut visited, &mut post_order);
86            }
87        }
88
89        post_order
90    }
91
92    fn post_order_step(&self, func: FuncIdx, visited: &mut [bool], post_order: &mut Vec<FuncIdx>) {
93        if visited[func] {
94            return;
95        }
96        visited[func] = true;
97
98        for &callee in &self.callees[func] {
99            self.post_order_step(callee, visited, post_order);
100        }
101
102        post_order.push(func);
103    }
104}