cranelift_codegen/
licm.rs

1//! A Loop Invariant Code Motion optimization pass
2
3use crate::cursor::{Cursor, FuncCursor};
4use crate::dominator_tree::DominatorTree;
5use crate::entity::{EntityList, ListPool};
6use crate::flowgraph::{BlockPredecessor, ControlFlowGraph};
7use crate::fx::FxHashSet;
8use crate::ir::{
9    Block, DataFlowGraph, Function, Inst, InstBuilder, InstructionData, Layout, Opcode, Type, Value,
10};
11use crate::loop_analysis::{Loop, LoopAnalysis};
12use crate::timing;
13use alloc::vec::Vec;
14
15/// Performs the LICM pass by detecting loops within the CFG and moving
16/// loop-invariant instructions out of them.
17/// Changes the CFG and domtree in-place during the operation.
18pub fn do_licm(
19    func: &mut Function,
20    cfg: &mut ControlFlowGraph,
21    domtree: &mut DominatorTree,
22    loop_analysis: &mut LoopAnalysis,
23) {
24    let _tt = timing::licm();
25    debug_assert!(cfg.is_valid());
26    debug_assert!(domtree.is_valid());
27    debug_assert!(loop_analysis.is_valid());
28
29    for lp in loop_analysis.loops() {
30        // For each loop that we want to optimize we determine the set of loop-invariant
31        // instructions
32        let invariant_insts = remove_loop_invariant_instructions(lp, func, cfg, loop_analysis);
33        // Then we create the loop's pre-header and fill it with the invariant instructions
34        // Then we remove the invariant instructions from the loop body
35        if !invariant_insts.is_empty() {
36            // If the loop has a natural pre-header we use it, otherwise we create it.
37            let mut pos;
38            match has_pre_header(&func.layout, cfg, domtree, loop_analysis.loop_header(lp)) {
39                None => {
40                    let pre_header =
41                        create_pre_header(loop_analysis.loop_header(lp), func, cfg, domtree);
42                    pos = FuncCursor::new(func).at_last_inst(pre_header);
43                }
44                // If there is a natural pre-header we insert new instructions just before the
45                // related jumping instruction (which is not necessarily at the end).
46                Some((_, last_inst)) => {
47                    pos = FuncCursor::new(func).at_inst(last_inst);
48                }
49            };
50            // The last instruction of the pre-header is the termination instruction (usually
51            // a jump) so we need to insert just before this.
52            for inst in invariant_insts {
53                pos.insert_inst(inst);
54            }
55        }
56    }
57    // We have to recompute the domtree to account for the changes
58    cfg.compute(func);
59    domtree.compute(func, cfg);
60}
61
62/// Insert a pre-header before the header, modifying the function layout and CFG to reflect it.
63/// A jump instruction to the header is placed at the end of the pre-header.
64fn create_pre_header(
65    header: Block,
66    func: &mut Function,
67    cfg: &mut ControlFlowGraph,
68    domtree: &DominatorTree,
69) -> Block {
70    let pool = &mut ListPool::<Value>::new();
71    let header_args_values = func.dfg.block_params(header).to_vec();
72    let header_args_types: Vec<Type> = header_args_values
73        .into_iter()
74        .map(|val| func.dfg.value_type(val))
75        .collect();
76    let pre_header = func.dfg.make_block();
77    let mut pre_header_args_value: EntityList<Value> = EntityList::new();
78    for typ in header_args_types {
79        pre_header_args_value.push(func.dfg.append_block_param(pre_header, typ), pool);
80    }
81
82    for BlockPredecessor {
83        inst: last_inst, ..
84    } in cfg.pred_iter(header)
85    {
86        // We only follow normal edges (not the back edges)
87        if !domtree.dominates(header, last_inst, &func.layout) {
88            func.rewrite_branch_destination(last_inst, header, pre_header);
89        }
90    }
91
92    // Inserts the pre-header at the right place in the layout.
93    let mut pos = FuncCursor::new(func).at_top(header);
94    pos.insert_block(pre_header);
95    pos.next_inst();
96    pos.ins().jump(header, pre_header_args_value.as_slice(pool));
97
98    pre_header
99}
100
101/// Detects if a loop header has a natural pre-header.
102///
103/// A loop header has a pre-header if there is only one predecessor that the header doesn't
104/// dominate.
105/// Returns the pre-header Block and the instruction jumping to the header.
106fn has_pre_header(
107    layout: &Layout,
108    cfg: &ControlFlowGraph,
109    domtree: &DominatorTree,
110    header: Block,
111) -> Option<(Block, Inst)> {
112    let mut result = None;
113    for BlockPredecessor {
114        block: pred_block,
115        inst: branch_inst,
116    } in cfg.pred_iter(header)
117    {
118        // We only count normal edges (not the back edges)
119        if !domtree.dominates(header, branch_inst, layout) {
120            if result.is_some() {
121                // We have already found one, there are more than one
122                return None;
123            }
124            if branch_inst != layout.last_inst(pred_block).unwrap()
125                || cfg.succ_iter(pred_block).nth(1).is_some()
126            {
127                // It's along a critical edge, so don't use it.
128                return None;
129            }
130            result = Some((pred_block, branch_inst));
131        }
132    }
133    result
134}
135
136/// Test whether the given opcode is unsafe to even consider for LICM.
137fn trivially_unsafe_for_licm(opcode: Opcode) -> bool {
138    opcode.can_store()
139        || opcode.is_call()
140        || opcode.is_branch()
141        || opcode.is_terminator()
142        || opcode.is_return()
143        || opcode.can_trap()
144        || opcode.other_side_effects()
145}
146
147fn is_unsafe_load(inst_data: &InstructionData) -> bool {
148    match *inst_data {
149        InstructionData::Load { flags, .. } => !flags.readonly() || !flags.notrap(),
150        _ => inst_data.opcode().can_load(),
151    }
152}
153
154/// Test whether the given instruction is loop-invariant.
155fn is_loop_invariant(inst: Inst, dfg: &DataFlowGraph, loop_values: &FxHashSet<Value>) -> bool {
156    if trivially_unsafe_for_licm(dfg.insts[inst].opcode()) {
157        return false;
158    }
159
160    if is_unsafe_load(&dfg.insts[inst]) {
161        return false;
162    }
163
164    for arg in dfg.inst_values(inst) {
165        let arg = dfg.resolve_aliases(arg);
166        if loop_values.contains(&arg) {
167            return false;
168        }
169    }
170    true
171}
172
173/// Traverses a loop in reverse post-order from a header block and identify loop-invariant
174/// instructions. These loop-invariant instructions are then removed from the code and returned
175/// (in reverse post-order) for later use.
176fn remove_loop_invariant_instructions(
177    lp: Loop,
178    func: &mut Function,
179    cfg: &ControlFlowGraph,
180    loop_analysis: &LoopAnalysis,
181) -> Vec<Inst> {
182    let mut loop_values: FxHashSet<Value> = FxHashSet();
183    let mut invariant_insts: Vec<Inst> = Vec::new();
184    let mut pos = FuncCursor::new(func);
185    // We traverse the loop block in reverse post-order.
186    for block in postorder_blocks_loop(loop_analysis, cfg, lp).iter().rev() {
187        // Arguments of the block are loop values
188        for val in pos.func.dfg.block_params(*block) {
189            loop_values.insert(*val);
190        }
191        pos.goto_top(*block);
192        #[cfg_attr(feature = "cargo-clippy", allow(clippy::block_in_if_condition_stmt))]
193        while let Some(inst) = pos.next_inst() {
194            if is_loop_invariant(inst, &pos.func.dfg, &loop_values) {
195                // If all the instruction's argument are defined outside the loop
196                // then this instruction is loop-invariant
197                invariant_insts.push(inst);
198                // We remove it from the loop
199                pos.remove_inst_and_step_back();
200            } else {
201                // If the instruction is not loop-invariant we push its results in the set of
202                // loop values
203                for out in pos.func.dfg.inst_results(inst) {
204                    loop_values.insert(*out);
205                }
206            }
207        }
208    }
209    invariant_insts
210}
211
212/// Return blocks from a loop in post-order, starting from an entry point in the block.
213fn postorder_blocks_loop(
214    loop_analysis: &LoopAnalysis,
215    cfg: &ControlFlowGraph,
216    lp: Loop,
217) -> Vec<Block> {
218    let mut grey = FxHashSet();
219    let mut black = FxHashSet();
220    let mut stack = vec![loop_analysis.loop_header(lp)];
221    let mut postorder = Vec::new();
222
223    while !stack.is_empty() {
224        let node = stack.pop().unwrap();
225        if !grey.contains(&node) {
226            // This is a white node. Mark it as gray.
227            grey.insert(node);
228            stack.push(node);
229            // Get any children we've never seen before.
230            for child in cfg.succ_iter(node) {
231                if loop_analysis.is_in_loop(child, lp) && !grey.contains(&child) {
232                    stack.push(child);
233                }
234            }
235        } else if !black.contains(&node) {
236            postorder.push(node);
237            black.insert(node);
238        }
239    }
240    postorder
241}