feat(cranelift): Use DominatorTreePreorder in more places (#11098)

Author: PhantomInTheWire

cranelift/codegen/src/alias_analysis.rs

@@ -63,7 +63,7 @@
 
 use crate::{
     cursor::{Cursor, FuncCursor},
-    dominator_tree::DominatorTree,
+    dominator_tree::DominatorTreePreorder,
     inst_predicates::{
         has_memory_fence_semantics, inst_addr_offset_type, inst_store_data, visit_block_succs,
     },
@@ -176,7 +176,7 @@ struct MemoryLoc {
 /// An alias-analysis pass.
 pub struct AliasAnalysis<'a> {
     /// The domtree for the function.
-    domtree: &'a DominatorTree,
+    domtree: &'a DominatorTreePreorder,
 
     /// Input state to a basic block.
     block_input: FxHashMap<Block, LastStores>,
@@ -191,7 +191,7 @@ pub struct AliasAnalysis<'a> {
 
 impl<'a> AliasAnalysis<'a> {
     /// Perform an alias analysis pass.
-    pub fn new(func: &Function, domtree: &'a DominatorTree) -> AliasAnalysis<'a> {
+    pub fn new(func: &Function, domtree: &'a DominatorTreePreorder) -> AliasAnalysis<'a> {
         trace!("alias analysis: input is:\n{:?}", func);
         let mut analysis = AliasAnalysis {
             domtree,
@@ -333,7 +333,7 @@ impl<'a> AliasAnalysis<'a> {
                             value.index(),
                             def_inst.index()
                         );
-                        if self.domtree.dominates(def_inst, inst, &func.layout) {
+                        if self.domtree.dominates_inst(def_inst, inst, &func.layout) {
                             trace!(
                                 " -> dominates; value equiv from v{} to v{} inserted",
                                 load_result.index(),

cranelift/codegen/src/context.rs

@@ -11,6 +11,7 @@
 
 use crate::alias_analysis::AliasAnalysis;
 use crate::dominator_tree::DominatorTree;
+use crate::dominator_tree::DominatorTreePreorder;
 use crate::egraph::EgraphPass;
 use crate::flowgraph::ControlFlowGraph;
 use crate::ir::Function;
@@ -46,6 +47,9 @@ pub struct Context {
     /// Dominator tree for `func`.
     pub domtree: DominatorTree,
 
+    /// Dominator tree with dominance stored implicitly via visit-order indices for `func`
+    domtree_preorder: DominatorTreePreorder,
+
     /// Loop analysis of `func`.
     pub loop_analysis: LoopAnalysis,
 
@@ -74,6 +78,7 @@ impl Context {
             func,
             cfg: ControlFlowGraph::new(),
             domtree: DominatorTree::new(),
+            domtree_preorder: DominatorTreePreorder::new(),
             loop_analysis: LoopAnalysis::new(),
             compiled_code: None,
             want_disasm: false,
@@ -305,7 +310,8 @@ impl Context {
 
     /// Compute dominator tree.
     pub fn compute_domtree(&mut self) {
-        self.domtree.compute(&self.func, &self.cfg)
+        self.domtree.compute(&self.func, &self.cfg);
+        self.domtree_preorder.compute(&self.domtree);
     }
 
     /// Compute the loop analysis.
@@ -335,7 +341,7 @@ impl Context {
     /// by a store instruction to the same instruction (so-called
     /// "store-to-load forwarding").
     pub fn replace_redundant_loads(&mut self) -> CodegenResult<()> {
-        let mut analysis = AliasAnalysis::new(&self.func, &self.domtree);
+        let mut analysis = AliasAnalysis::new(&self.func, &self.domtree_preorder);
         analysis.compute_and_update_aliases(&mut self.func);
         Ok(())
     }
@@ -367,7 +373,7 @@ impl Context {
         );
         let fisa = fisa.into();
         self.compute_loop_analysis();
-        let mut alias_analysis = AliasAnalysis::new(&self.func, &self.domtree);
+        let mut alias_analysis = AliasAnalysis::new(&self.func, &self.domtree_preorder);
         let mut pass = EgraphPass::new(
             &mut self.func,
             &self.domtree,

cranelift/codegen/src/dominator_tree.rs

@@ -2,7 +2,7 @@
 
 use crate::entity::SecondaryMap;
 use crate::flowgraph::{BlockPredecessor, ControlFlowGraph};
-use crate::ir::{Block, Function, Layout, ProgramPoint};
+use crate::ir::{Block, Function, Inst, Layout, ProgramPoint};
 use crate::packed_option::PackedOption;
 use crate::timing;
 use alloc::vec::Vec;
@@ -583,6 +583,20 @@ impl DominatorTreePreorder {
         na.pre_number <= nb.pre_number && na.pre_max >= nb.pre_max
     }
 
+    /// Checks if one instruction dominates another.
+    pub fn dominates_inst(&self, a: Inst, b: Inst, layout: &Layout) -> bool {
+        match (layout.inst_block(a), layout.inst_block(b)) {
+            (Some(block_a), Some(block_b)) => {
+                if block_a == block_b {
+                    layout.pp_cmp(a, b) != Ordering::Greater
+                } else {
+                    self.dominates(block_a, block_b)
+                }
+            }
+            _ => false,
+        }
+    }
+
     /// Compare two blocks according to the dominator pre-order.
     pub fn pre_cmp_block(&self, a: Block, b: Block) -> Ordering {
         self.nodes[a].pre_number.cmp(&self.nodes[b].pre_number)

cranelift/codegen/src/verifier/mod.rs

@@ -65,6 +65,7 @@
 
 use crate::dbg::DisplayList;
 use crate::dominator_tree::DominatorTree;
+use crate::dominator_tree::DominatorTreePreorder;
 use crate::entity::SparseSet;
 use crate::flowgraph::{BlockPredecessor, ControlFlowGraph};
 use crate::ir::entities::AnyEntity;
@@ -301,6 +302,7 @@ enum BlockCallTargetType {
 struct Verifier<'a> {
     func: &'a Function,
     expected_cfg: ControlFlowGraph,
+    expected_domtree_preorder: DominatorTreePreorder,
     expected_domtree: DominatorTree,
     isa: Option<&'a dyn TargetIsa>,
 }
@@ -309,10 +311,13 @@ impl<'a> Verifier<'a> {
     pub fn new(func: &'a Function, fisa: FlagsOrIsa<'a>) -> Self {
         let expected_cfg = ControlFlowGraph::with_function(func);
         let expected_domtree = DominatorTree::with_function(func, &expected_cfg);
+        let mut expected_domtree_preorder = DominatorTreePreorder::new();
+        expected_domtree_preorder.compute(&expected_domtree);
         Self {
             func,
             expected_cfg,
             expected_domtree,
+            expected_domtree_preorder,
             isa: fisa.isa,
         }
     }
@@ -1000,10 +1005,11 @@ impl<'a> Verifier<'a> {
                 }
                 // Defining instruction dominates the instruction that uses the value.
                 if is_reachable {
-                    if !self
-                        .expected_domtree
-                        .dominates(def_inst, loc_inst, &self.func.layout)
-                    {
+                    if !self.expected_domtree_preorder.dominates_inst(
+                        def_inst,
+                        loc_inst,
+                        &self.func.layout,
+                    ) {
                         return errors.fatal((
                             loc_inst,
                             self.context(loc_inst),
@@ -1036,12 +1042,9 @@ impl<'a> Verifier<'a> {
                         format!("{v} is defined by {block} which is not in the layout"),
                     ));
                 }
+                let user_block = self.func.layout.inst_block(loc_inst).expect("Expected instruction to be in a block as we're traversing code already in layout");
                 // The defining block dominates the instruction using this value.
-                if is_reachable
-                    && !self
-                        .expected_domtree
-                        .dominates(block, loc_inst, &self.func.layout)
-                {
+                if is_reachable && !self.expected_domtree_preorder.dominates(block, user_block) {
                     return errors.fatal((
                         loc_inst,
                         self.context(loc_inst),