transforms: (canonicalize) Add full dce

tests/filecheck/dialects/cf/canonicalize.mlir

@@ -46,3 +46,15 @@ func.func @br_passthrough(%arg0 : i32, %arg1 : i32) -> (i32, i32) {
 ^2(%arg4 : i32, %arg5 : i32):
   return %arg4, %arg5 : i32, i32
 }
+
+/// Test that dead branches don't affect passthrough
+// CHECK:      func.func @br_dead_passthrough() {
+// CHECK-NEXT:   func.return
+// CHECK-NEXT: }
+func.func @br_dead_passthrough() {
+cf.br ^1
+^0:
+cf.br ^1
+^1:
+func.return
+}

tests/filecheck/dialects/riscv_cf/canonicalize.mlir

@@ -46,15 +46,9 @@ riscv_func.func @never() {
 
 // CHECK-NEXT:  riscv_func.func @never() {
 // CHECK-NEXT:    %one = riscv.li 1 : !riscv.reg
-// CHECK-NEXT:    %three = riscv.li 3 : !riscv.reg
-// CHECK-NEXT:    %{{.*}} = riscv.mv %one : (!riscv.reg) -> !riscv.reg
-// CHECK-NEXT:    riscv_cf.j ^{{\d+}}(%{{.*}} : !riscv.reg) attributes {"comment" = "Constant folded riscv_cf.bge"}
-// CHECK-NEXT:  ^{{\d+}}(%i : !riscv.reg):
-// CHECK-NEXT:    riscv.label "scf_body_0_for"
-// CHECK-NEXT:    "test.op"(%i) : (!riscv.reg) -> ()
-// CHECK-NEXT:    %{{\d+}} = riscv.addi %i, 1 : (!riscv.reg) -> !riscv.reg
-// CHECK-NEXT:    riscv_cf.blt %{{\d+}} : !riscv.reg, %three : !riscv.reg, ^{{\d+}}(%{{\d+}} : !riscv.reg), ^{{\d+}}(%{{\d+}} : !riscv.reg)
-// CHECK-NEXT:  ^{{\d+}}(%{{\d+}} : !riscv.reg):
+// CHECK-NEXT:    %0 = riscv.mv %one : (!riscv.reg) -> !riscv.reg
+// CHECK-NEXT:    riscv_cf.j ^0(%0 : !riscv.reg) attributes {"comment" = "Constant folded riscv_cf.bge"}
+// CHECK-NEXT:  ^0(%1 : !riscv.reg):
 // CHECK-NEXT:    riscv.label "scf_body_end_0_for"
 // CHECK-NEXT:    riscv_func.return
 // CHECK-NEXT:  }

xdsl/pattern_rewriter.py

@@ -668,6 +668,13 @@ class PatternRewriteWalker:
     That way, all uses are replaced before the definitions.
     """
 
+    post_walk_pass: Callable[[Operation, PatternRewriterListener], bool] | None = field(
+        default=None
+    )
+    """
+    Pass to call between each walk of the IR.
+    """
+
     listener: PatternRewriterListener = field(default_factory=PatternRewriterListener)
     """The listener that will be called when an operation or block is modified."""
 
@@ -755,6 +762,8 @@ def rewrite_op(self, op: Operation) -> bool:
 
         self._populate_worklist(op)
         op_was_modified = self._process_worklist(pattern_listener)
+        if self.post_walk_pass is not None:
+            op_was_modified |= self.post_walk_pass(op, pattern_listener)
 
         if not self.apply_recursively:
             return op_was_modified
@@ -764,6 +773,8 @@ def rewrite_op(self, op: Operation) -> bool:
         while op_was_modified:
             self._populate_worklist(op)
             op_was_modified = self._process_worklist(pattern_listener)
+            if self.post_walk_pass is not None:
+                op_was_modified |= self.post_walk_pass(op, pattern_listener)
 
         return result
 

xdsl/transforms/canonicalize.py

@@ -9,7 +9,7 @@
     RewritePattern,
 )
 from xdsl.traits import HasCanonicalizationPatternsTrait
-from xdsl.transforms.dead_code_elimination import RemoveUnusedOperations
+from xdsl.transforms.dead_code_elimination import RemoveUnusedOperations, op_dce
 
 
 class CanonicalizationRewritePattern(RewritePattern):
@@ -37,4 +37,4 @@ def apply(self, ctx: MLContext, op: builtin.ModuleOp) -> None:
         pattern = GreedyRewritePatternApplier(
             [RemoveUnusedOperations(), CanonicalizationRewritePattern()]
         )
-        PatternRewriteWalker(pattern).rewrite_module(op)
+        PatternRewriteWalker(pattern, post_walk_pass=op_dce).rewrite_module(op)

xdsl/transforms/dead_code_elimination.py

@@ -5,7 +5,12 @@
 from xdsl.ir import Operation, Region, SSAValue
 from xdsl.ir.post_order import PostOrderIterator
 from xdsl.passes import ModulePass
-from xdsl.pattern_rewriter import PatternRewriter, PatternRewriteWalker, RewritePattern
+from xdsl.pattern_rewriter import (
+    PatternRewriter,
+    PatternRewriterListener,
+    PatternRewriteWalker,
+    RewritePattern,
+)
 from xdsl.traits import (
     IsTerminator,
     MemoryEffectKind,
@@ -120,7 +125,7 @@ def propagate_region_liveness(self, region: Region):
             for operation in reversed(block.ops):
                 self.propagate_op_liveness(operation)
 
-    def delete_dead(self, region: Region):
+    def delete_dead(self, region: Region, listener: PatternRewriterListener | None):
         first = region.first_block
         if first is None:
             return
@@ -135,26 +140,33 @@ def delete_dead(self, region: Region):
             for operation in reversed(block.ops):
                 if not self.is_live(operation):
                     self.changed = True
+                    if listener is not None:
+                        listener.handle_operation_removal(operation)
                     block.erase_op(operation, safe_erase=False)
                 else:
                     for r in operation.regions:
-                        self.delete_dead(r)
+                        self.delete_dead(r, listener)
 
 
-def region_dce(region: Region) -> bool:
+def region_dce(region: Region, listener: PatternRewriterListener | None = None) -> bool:
     live_set = LiveSet()
 
     while live_set.changed:
         live_set.changed = False
         live_set.propagate_region_liveness(region)
 
-    live_set.delete_dead(region)
+    live_set.delete_dead(region, listener)
     return live_set.changed
 
 
+def op_dce(op: Operation, listener: PatternRewriterListener | None = None):
+    changed = tuple(region_dce(region, listener) for region in op.regions)
+
+    return any(changed)
+
+
 class DeadCodeElimination(ModulePass):
     name = "dce"
 
     def apply(self, ctx: MLContext, op: ModuleOp) -> None:
-        for region in op.regions:
-            region_dce(region)
+        op_dce(op)