Cranelift: extend alias analysis to eliminate dead stores

cranelift/codegen/src/context.rs

@@ -19,6 +19,7 @@ use crate::isa::TargetIsa;
 use crate::loop_analysis::LoopAnalysis;
 use crate::machinst::{CompiledCode, CompiledCodeStencil};
 use crate::nan_canonicalization::do_nan_canonicalization;
+use crate::post_dominator_tree::PostDominatorTree;
 use crate::remove_constant_phis::do_remove_constant_phis;
 use crate::result::{CodegenResult, CompileResult};
 use crate::settings::{FlagsOrIsa, OptLevel};
@@ -46,6 +47,9 @@ pub struct Context {
     /// Dominator tree for `func`.
     pub domtree: DominatorTree,
 
+    /// Post-dominator tree for `func`.
+    pub post_dom_tree: PostDominatorTree,
+
     /// Loop analysis of `func`.
     pub loop_analysis: LoopAnalysis,
 
@@ -74,6 +78,7 @@ impl Context {
             func,
             cfg: ControlFlowGraph::new(),
             domtree: DominatorTree::new(),
+            post_dom_tree: PostDominatorTree::new(),
             loop_analysis: LoopAnalysis::new(),
             compiled_code: None,
             want_disasm: false,
@@ -182,6 +187,7 @@ impl Context {
         self.func.dfg.resolve_all_aliases();
 
         if opt_level != OptLevel::None {
+            self.compute_post_dom_tree();
             self.egraph_pass(isa, ctrl_plane)?;
         }
 
@@ -302,6 +308,11 @@ impl Context {
         self.domtree.compute(&self.func, &self.cfg);
     }
 
+    /// Compute the post-dominator tree.
+    pub fn compute_post_dom_tree(&mut self) {
+        self.post_dom_tree.compute(&self.func, &self.cfg);
+    }
+
     /// Compute the loop analysis.
     pub fn compute_loop_analysis(&mut self) {
         self.loop_analysis
@@ -332,7 +343,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, &self.post_dom_tree);
         analysis.compute_and_update_aliases(&mut self.func);
         Ok(())
     }
@@ -364,7 +375,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, &self.post_dom_tree);
         let mut pass = EgraphPass::new(
             &mut self.func,
             &self.domtree,

cranelift/codegen/src/egraph/mod.rs

@@ -529,7 +529,7 @@ where
                         }
                         (_, _) => unreachable!(),
                     }
-                    Some(SkeletonInstSimplification::Remove)
+                    return Some(SkeletonInstSimplification::Remove);
                 }
                 ScopedEntry::Vacant(v) => {
                     // Otherwise, insert it into the value-map.
@@ -539,42 +539,44 @@ where
                     }
                     v.insert(result);
                     trace!(" -> inserts as new (no GVN)");
-                    None
                 }
             }
         }
-        // Otherwise, if a load or store, process it with the alias
-        // analysis to see if we can optimize it (rewrite in terms of
-        // an earlier load or stored value, or remove an idempotent store).
-        else {
-            match self
-                .alias_analysis
-                .process_inst(self.func, self.alias_analysis_state, inst)
-            {
-                OptResult::AliasedLoad(new_result) => {
-                    self.stats.alias_analysis_removed_load += 1;
-                    let result = self.func.dfg.first_result(inst);
-                    trace!(
-                        " -> inst {} has result {} replaced with {}",
-                        inst, result, new_result
-                    );
-                    self.value_to_opt_value[result] = new_result;
-                    self.available_block[result] = self.available_block[new_result];
-                    Some(SkeletonInstSimplification::Remove)
-                }
-                OptResult::IdempotentStore => {
-                    self.stats.alias_analysis_removed_store += 1;
-                    Some(SkeletonInstSimplification::Remove)
-                }
-                OptResult::None => {
-                    // Generic side-effecting op -- always keep it, and
-                    // set its results to identity-map to original values.
-                    for &result in self.func.dfg.inst_results(inst) {
-                        self.value_to_opt_value[result] = result;
-                        self.available_block[result] = block;
-                    }
-                    None
+
+        // Otherwise, if a load or store, process it with the alias analysis to
+        // see if we can optimize it (rewrite in terms of an earlier load or
+        // stored value, or remove an idempotent store).
+        match self
+            .alias_analysis
+            .process_inst(self.func, self.alias_analysis_state, inst)
+        {
+            OptResult::AliasedLoad(new_result) => {
+                self.stats.alias_analysis_removed_load += 1;
+                let result = self.func.dfg.first_result(inst);
+                trace!(
+                    " -> inst {} has result {} replaced with {}",
+                    inst, result, new_result
+                );
+                self.value_to_opt_value[result] = new_result;
+                self.available_block[result] = self.available_block[new_result];
+                Some(SkeletonInstSimplification::Remove)
+            }
+            OptResult::IdempotentStore => {
+                self.stats.alias_analysis_removed_idempotent_store += 1;
+                Some(SkeletonInstSimplification::Remove)
+            }
+            OptResult::DeadStore { dead, killer } => {
+                self.stats.alias_analysis_removed_dead_store += 1;
+                Some(SkeletonInstSimplification::RemoveDeadStore { dead, killer })
+            }
+            OptResult::None => {
+                // Generic side-effecting op -- always keep it, and
+                // set its results to identity-map to original values.
+                for &result in self.func.dfg.inst_results(inst) {
+                    self.value_to_opt_value[result] = result;
+                    self.available_block[result] = block;
                 }
+                None
             }
         }
     }
@@ -670,6 +672,14 @@ where
                     return Some(SkeletonInstSimplification::ReplaceWithTwo { first, second });
                 }
 
+                SkeletonInstSimplification::RemoveDeadStore { dead, killer } => {
+                    log::trace!(
+                        " -> simplify_skeleton: remove other: {dead}: {}",
+                        ctx.func.dfg.display_inst(dead)
+                    );
+                    return Some(SkeletonInstSimplification::RemoveDeadStore { dead, killer });
+                }
+
                 // For instruction replacement simplification, we want to check
                 // that the replacements define the same number and types of
                 // values as the original instruction, and also determine
@@ -1100,6 +1110,25 @@ impl<'a> EgraphPass<'a> {
                 reprocess_from(cursor, first);
                 return;
             }
+            SkeletonInstSimplification::RemoveDeadStore { dead, killer } => {
+                assert!(!matches!(cursor.position(), CursorPosition::At(inst) if inst == dead));
+                // Copy the trap code (if any) from the dead store to its
+                // killer.
+                if let Some(flags) = cursor.func.dfg.insts[dead].memflags_data(&cursor.func.dfg) {
+                    if let Some(code) = flags.trap_code() {
+                        let flags = cursor.func.dfg.insts[killer]
+                            .memflags_data(&cursor.func.dfg)
+                            .unwrap();
+                        let flags = flags.with_trap_code(Some(code));
+                        let flags = cursor.func.dfg.mem_flags.insert(flags).unwrap();
+                        *cursor.func.dfg.insts[killer].memflags_mut().unwrap() = flags;
+                    }
+                }
+                cursor.func.layout.remove_inst(dead);
+                self_map_operands(&cursor.func.dfg, value_to_opt_value, killer);
+                cursor.prev_inst();
+                return;
+            }
 
             SkeletonInstSimplification::Replace { inst } => (inst, None),
             SkeletonInstSimplification::ReplaceWithVal { inst, val } => (inst, Some(val)),
@@ -1237,7 +1266,8 @@ pub(crate) struct Stats {
     pub(crate) skeleton_inst_simplified: u64,
     pub(crate) skeleton_inst_gvn: u64,
     pub(crate) alias_analysis_removed_load: u64,
-    pub(crate) alias_analysis_removed_store: u64,
+    pub(crate) alias_analysis_removed_idempotent_store: u64,
+    pub(crate) alias_analysis_removed_dead_store: u64,
     pub(crate) new_inst: u64,
     pub(crate) union: u64,
     pub(crate) subsume: u64,

cranelift/codegen/src/inst_predicates.rs

@@ -144,23 +144,6 @@ pub fn inst_store_data(func: &Function, inst: Inst) -> Option<Value> {
     }
 }
 
-/// Determine whether this opcode behaves as a memory fence, i.e.,
-/// prohibits any moving of memory accesses across it.
-pub fn has_memory_fence_semantics(op: Opcode) -> bool {
-    match op {
-        Opcode::AtomicRmw
-        | Opcode::AtomicCas
-        | Opcode::AtomicLoad
-        | Opcode::AtomicStore
-        | Opcode::Fence
-        | Opcode::Debugtrap
-        | Opcode::SequencePoint => true,
-        Opcode::Call | Opcode::CallIndirect | Opcode::TryCall | Opcode::TryCallIndirect => true,
-        op if op.can_trap() => true,
-        _ => false,
-    }
-}
-
 /// Visit all successors of a block with a given visitor closure. The closure
 /// arguments are the branch instruction that is used to reach the successor,
 /// the successor block itself, and a flag indicating whether the block is

cranelift/codegen/src/ir/instructions.rs

@@ -567,12 +567,39 @@ impl InstructionData {
         }
     }
 
+    /// If this is a load/store instruction, return its memory flags.
+    pub fn memflags_mut(&mut self) -> Option<&mut MemFlags> {
+        match self {
+            InstructionData::Load { flags, .. }
+            | InstructionData::LoadNoOffset { flags, .. }
+            | InstructionData::Store { flags, .. }
+            | InstructionData::StoreNoOffset { flags, .. }
+            | InstructionData::AtomicCas { flags, .. }
+            | InstructionData::AtomicRmw { flags, .. } => Some(flags),
+            _ => None,
+        }
+    }
+
     /// If this is a load/store instruction, resolve its memory flags to data
     /// through the DFG.
     pub fn memflags_data(&self, dfg: &super::dfg::DataFlowGraph) -> Option<super::MemFlagsData> {
         self.memflags().map(|f| dfg.mem_flags[f])
     }
 
+    /// Get this load/store instruction's trap code, if any.
+    ///
+    /// Returns `None` when this is not a load/store instruction, or if it is a
+    /// load/store instruction but does not have a trap code.
+    pub fn memflags_trap_code(&self, dfg: &super::dfg::DataFlowGraph) -> Option<TrapCode> {
+        self.memflags_data(dfg)?.trap_code()
+    }
+
+    /// If this is a load/store instruction, get its alias region.
+    pub fn alias_region(&self, dfg: &super::dfg::DataFlowGraph) -> Option<super::AliasRegion> {
+        let flags = self.memflags_data(dfg)?;
+        flags.alias_region()
+    }
+
     /// If this instruction references a stack slot, return it
     pub fn stack_slot(&self) -> Option<StackSlot> {
         match self {

cranelift/codegen/src/prelude_opt.isle

@@ -121,7 +121,12 @@
         ;; Neither instruction may define any results, and if the instruction
         ;; being replaced is a block terminator then `second` must also be a
         ;; block terminator.
-        (ReplaceWithTwo (first Inst) (second Inst))))
+        (ReplaceWithTwo (first Inst) (second Inst))
+
+        ;; Remove a dead store.
+        ;;
+        ;; The trap code from `dead` (if any) is copied to `killer`.
+        (RemoveDeadStore (dead Inst) (killer Inst))))
 
 (decl pure inst_to_skeleton_inst_simplification (Inst) SkeletonInstSimplification)
 (rule (inst_to_skeleton_inst_simplification inst)

cranelift/filetests/filetests/alias/check-unset-reset-flag.clif

@@ -0,0 +1,27 @@
+test optimize precise-output
+set opt_level=speed
+target aarch64
+
+function u1:0(i64 vmctx, i32) -> i32 tail {
+    region0 = 0 "flag"
+
+block0(v0: i64, v1: i32):
+    v2 = load.i64 notrap aligned region0 v0
+    trapz v2, user42
+    v3 = iconst.i64 0
+    store notrap aligned region0 v3, v0
+    v4 = iadd v1, v1
+    store notrap aligned region0 v2, v0
+    return v4
+}
+
+; function u1:0(i64 vmctx, i32) -> i32 tail {
+;     region0 = 0 "flag"
+;
+; block0(v0: i64, v1: i32):
+;     v2 = load.i64 notrap aligned region0 v0
+;     trapz v2, user42
+;     store notrap aligned region0 v2, v0
+;     v4 = iadd v1, v1
+;     return v4
+; }

cranelift/filetests/filetests/alias/dead-store-chain.clif

@@ -0,0 +1,23 @@
+test optimize precise-output
+set opt_level=speed
+target aarch64
+
+;; A chain of stores to the same location collapses to only the last store.
+function %chain(i64, i32, i32, i32, i32) {
+    region0 = 0 "R"
+block0(v0: i64, v1: i32, v2: i32, v3: i32, v4: i32):
+    store notrap aligned region0 v1, v0
+    store notrap aligned region0 v2, v0
+    store notrap aligned region0 v3, v0
+    store notrap aligned region0 v4, v0
+    return
+}
+
+; function %chain(i64, i32, i32, i32, i32) fast {
+;     region0 = 0 "R"
+;
+; block0(v0: i64, v1: i32, v2: i32, v3: i32, v4: i32):
+;     store notrap aligned region0 v4, v0
+;     return
+; }
+

cranelift/filetests/filetests/alias/dead-store-cross-block.clif

@@ -0,0 +1,59 @@
+test optimize precise-output
+set opt_level=speed
+target aarch64
+
+function %straight_line(i64, i32, i32) {
+    region0 = 0 "R0"
+block0(v0: i64, v1: i32, v2: i32):
+    store notrap aligned region0 v1, v0
+    jump block1
+
+block1:
+    store notrap aligned region0 v2, v0
+    return
+}
+
+; function %straight_line(i64, i32, i32) fast {
+;     region0 = 0 "R0"
+;
+; block0(v0: i64, v1: i32, v2: i32):
+;     jump block1
+;
+; block1:
+;     store.i32 notrap aligned region0 v2, v0
+;     return
+; }
+
+function %diamond(i64, i32, i32, i32) {
+    region0 = 0 "R0"
+block0(v0: i64, v1: i32, v2: i32, v3: i32):
+    store notrap aligned region0 v1, v0
+    brif v1, block1, block2
+
+block1:
+    jump block3(v2)
+
+block2:
+    jump block3(v3)
+
+block3(v4: i32):
+    store notrap aligned region0 v4, v0
+    return
+}
+
+; function %diamond(i64, i32, i32, i32) fast {
+;     region0 = 0 "R0"
+;
+; block0(v0: i64, v1: i32, v2: i32, v3: i32):
+;     brif v1, block1, block2
+;
+; block1:
+;     jump block3(v2)
+;
+; block2:
+;     jump block3(v3)
+;
+; block3(v4: i32):
+;     store notrap aligned region0 v4, v0
+;     return
+; }