transforms: (canonicalize) improve dead code elimination (#3252)

Moves the dead code elimination into the main pattern rewriting loop of canonicalization, allowing more canonicalization patterns to be done after dce.
xdslproject · Oct 7, 2024 · 5243d81 · 5243d81
1 parent e6d51d8
commit 5243d81
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/tests/filecheck/dialects/stencil/canonicalize.mlir b/tests/filecheck/dialects/stencil/canonicalize.mlir
@@ -39,8 +39,7 @@ func.func @unused_res(%f1 : !stencil.field<[0,64]xf64>, %f2 : !stencil.field<[0,
 
 // CHECK:         func.func @unused_res(%f1 : !stencil.field<[0,64]xf64>, %f2 : !stencil.field<[0,64]xf64>, %of : !stencil.field<[0,64]xf64>) {
 // CHECK-NEXT:      %t1 = stencil.load %f1 : !stencil.field<[0,64]xf64> -> !stencil.temp<?xf64>
-// CHECK-NEXT:      %t2 = stencil.load %f2 : !stencil.field<[0,64]xf64> -> !stencil.temp<?xf64>
-// CHECK-NEXT:      %o1 = stencil.apply(%one = %t1 : !stencil.temp<?xf64>, %two = %t2 : !stencil.temp<?xf64>) -> (!stencil.temp<?xf64>) {
+// CHECK-NEXT:      %o1 = stencil.apply(%one = %t1 : !stencil.temp<?xf64>) -> (!stencil.temp<?xf64>) {
 // CHECK-NEXT:        %0 = stencil.access %one[0] : !stencil.temp<?xf64>
 // CHECK-NEXT:        stencil.return %0 : f64
 // CHECK-NEXT:      }

diff --git a/xdsl/transforms/canonicalize.py b/xdsl/transforms/canonicalize.py
@@ -9,7 +9,7 @@
     RewritePattern,
 )
 from xdsl.traits import HasCanonicalizationPatternsTrait
-from xdsl.transforms.dead_code_elimination import dce
+from xdsl.transforms.dead_code_elimination import RemoveUnusedOperations
 
 
 class CanonicalizationRewritePattern(RewritePattern):
@@ -34,5 +34,7 @@ class CanonicalizePass(ModulePass):
     name = "canonicalize"
 
     def apply(self, ctx: MLContext, op: builtin.ModuleOp) -> None:
-        PatternRewriteWalker(CanonicalizationRewritePattern()).rewrite_module(op)
-        dce(op)
+        pattern = GreedyRewritePatternApplier(
+            [RemoveUnusedOperations(), CanonicalizationRewritePattern()]
+        )
+        PatternRewriteWalker(pattern).rewrite_module(op)