Cranelift: Prevent DSE across divergent paths

Author: subotac

cranelift/codegen/src/alias_analysis.rs

@@ -540,6 +540,30 @@ pub struct AliasAnalysis<'a> {
     mem_values: FxHashMap<MemoryLoc, (Inst, Value)>,
 }
 
+/// Can a path starting at `start` reach a divergent block without first
+/// passing through `stop`?
+fn can_reach_divergent_block(
+    post_dom_tree: &PostDominatorTree,
+    cfg: &ControlFlowGraph,
+    start: Block,
+    stop: Block,
+) -> bool {
+    let mut worklist = vec![start];
+    let mut visited = FxHashSet::default();
+
+    while let Some(block) = worklist.pop() {
+        if block == stop || !visited.insert(block) {
+            continue;
+        }
+        if post_dom_tree.diverges(block) {
+            return true;
+        }
+        worklist.extend(cfg.succ_iter(block));
+    }
+
+    false
+}
+
 impl<'a> AliasAnalysis<'a> {
     /// Perform an alias analysis pass.
     pub fn new(func: &Function, domtree: &'a DominatorTree) -> AliasAnalysis<'a> {
@@ -584,9 +608,16 @@ impl<'a> AliasAnalysis<'a> {
             return func.layout.pp_cmp(overwriter, maybe_dead) != Ordering::Less;
         }
 
-        self.post_dom_tree
-            .get_or_insert_with(|| PostDominatorTree::with_cfg(cfg))
-            .post_dominates(overwriter, maybe_dead, &func.layout)
+        let post_dom_tree = self
+            .post_dom_tree
+            .get_or_insert_with(|| PostDominatorTree::with_cfg(cfg));
+
+        // The post-dominator tree only considers paths that reach an explicit
+        // CFG exit. A path into a divergent region may still observe the store
+        // by trapping, so it must also pass through the overwriter before we
+        // can remove the original store.
+        post_dom_tree.post_dominates(overwriter, maybe_dead, &func.layout)
+            && !can_reach_divergent_block(post_dom_tree, cfg, maybe_dead_block, overwriter_block)
     }
 
     fn compute_block_input_states(&mut self, func: &Function) {

cranelift/filetests/filetests/alias/not-dead-divergent-path.clif

@@ -0,0 +1,23 @@
+test alias-analysis
+set opt_level=speed
+target aarch64
+
+function %divergent_path(i64, i32) {
+    region0 = 0 "R0"
+
+block0(v0: i64, v1: i32):
+    v2 = iconst.i32 111
+    store notrap aligned region0 v2, v0
+    brif v1, block1, block2
+
+block1:
+    v3 = iconst.i32 222
+    store notrap aligned region0 v3, v0
+    return
+
+block2:
+    jump block2
+}
+
+; check: store notrap aligned region0 v2, v0
+; check: store notrap aligned region0 v3, v0