core: Walk regions for rewritings instead of ops (#3506)

Currently, the pattern rewriter walker walks an operation, and tries to
apply patterns on all contained operations, including itself. This is
not what we would want, as rewriting a toplevel operation is dangerous,
if that operation does not have a parent.

This commit changes the walker to walk regions instead of operations, so
it is clear that the toplevel operation will not be modified. This is
similar to what MLIR currently does in applyPatternsAndFoldGreedily.

docs/database_example.ipynb

@@ -310,7 +310,7 @@
     ")\n",
     "\n",
     "\n",
-    "walker.rewrite_op(sel)"
+    "walker.rewrite_region(sel.filter)"
    ]
   },
   {
@@ -387,7 +387,7 @@
     "    walk_reverse=False,\n",
     ")\n",
     "\n",
-    "walker.rewrite_op(sel)"
+    "walker.rewrite_region(sel.filter)"
    ]
   },
   {

docs/xdsl-introduction.ipynb

@@ -762,7 +762,7 @@
     "    apply_recursively=True,\n",
     "    walk_reverse=False,\n",
     ")\n",
-    "walker.rewrite_module(filtered)\n",
+    "walker.rewrite_region(filtered.filter)\n",
     "printer.print_op(filtered)"
    ]
   },
@@ -808,7 +808,7 @@
     "    apply_recursively=True,\n",
     "    walk_reverse=False,\n",
     ")\n",
-    "walker.rewrite_module(filtered)\n",
+    "walker.rewrite_region(filtered.filter)\n",
     "printer.print_op(filtered)"
    ]
   },

tests/pattern_rewriter/test_pattern_rewriter.py

@@ -377,20 +377,24 @@ def test_insert_op_at_start():
     """Test rewrites where operations are inserted with a given position."""
 
     prog = """"builtin.module"() ({
-  %0 = "arith.constant"() <{"value" = 5 : i32}> : () -> i32
+  "test.op"() ({
+    %0 = "arith.constant"() <{"value" = 5 : i32}> : () -> i32
+  }) : () -> ()
 }) : () -> ()"""
 
     expected = """"builtin.module"() ({
-  %0 = "arith.constant"() <{"value" = 42 : i32}> : () -> i32
-  %1 = "arith.constant"() <{"value" = 5 : i32}> : () -> i32
+  "test.op"() ({
+    %0 = "arith.constant"() <{"value" = 42 : i32}> : () -> i32
+    %1 = "arith.constant"() <{"value" = 5 : i32}> : () -> i32
+  }) : () -> ()
 }) : () -> ()"""
 
     class Rewrite(RewritePattern):
         @op_type_rewrite_pattern
-        def match_and_rewrite(self, mod: ModuleOp, rewriter: PatternRewriter):
+        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
             new_cst = Constant.from_int_and_width(42, i32)
 
-            rewriter.insert_op(new_cst, InsertPoint.at_start(mod.regions[0].blocks[0]))
+            rewriter.insert_op(new_cst, InsertPoint.at_start(op.regions[0].blocks[0]))
 
     rewrite_and_compare(
         prog,
@@ -404,20 +408,24 @@ def test_insert_op_before():
     """Test rewrites where operations are inserted before a given operation."""
 
     prog = """"builtin.module"() ({
-  %0 = "arith.constant"() <{"value" = 5 : i32}> : () -> i32
+  "test.op"() ({
+    %0 = "arith.constant"() <{"value" = 5 : i32}> : () -> i32
+  }) : () -> ()
 }) : () -> ()"""
 
     expected = """"builtin.module"() ({
-  %0 = "arith.constant"() <{"value" = 42 : i32}> : () -> i32
-  %1 = "arith.constant"() <{"value" = 5 : i32}> : () -> i32
+  "test.op"() ({
+    %0 = "arith.constant"() <{"value" = 42 : i32}> : () -> i32
+    %1 = "arith.constant"() <{"value" = 5 : i32}> : () -> i32
+  }) : () -> ()
 }) : () -> ()"""
 
     class Rewrite(RewritePattern):
         @op_type_rewrite_pattern
-        def match_and_rewrite(self, mod: ModuleOp, rewriter: PatternRewriter):
+        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
             new_cst = Constant.from_int_and_width(42, i32)
 
-            first_op = mod.ops.first
+            first_op = op.regions[0].block.ops.first
             assert first_op is not None
             rewriter.insert_op(new_cst, InsertPoint.before(first_op))
 
@@ -433,20 +441,24 @@ def test_insert_op_after():
     """Test rewrites where operations are inserted after a given operation."""
 
     prog = """"builtin.module"() ({
-  %0 = "arith.constant"() <{"value" = 5 : i32}> : () -> i32
+  "test.op"() ({
+    %0 = "arith.constant"() <{"value" = 5 : i32}> : () -> i32
+  }) : () -> ()
 }) : () -> ()"""
 
     expected = """"builtin.module"() ({
-  %0 = "arith.constant"() <{"value" = 5 : i32}> : () -> i32
-  %1 = "arith.constant"() <{"value" = 42 : i32}> : () -> i32
+  "test.op"() ({
+    %0 = "arith.constant"() <{"value" = 5 : i32}> : () -> i32
+    %1 = "arith.constant"() <{"value" = 42 : i32}> : () -> i32
+  }) : () -> ()
 }) : () -> ()"""
 
     class Rewrite(RewritePattern):
         @op_type_rewrite_pattern
-        def match_and_rewrite(self, mod: ModuleOp, rewriter: PatternRewriter):
+        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
             new_cst = Constant.from_int_and_width(42, i32)
 
-            first_op = mod.ops.first
+            first_op = op.regions[0].block.ops.first
             assert first_op is not None
             rewriter.insert_op(new_cst, InsertPoint.after(first_op))
 
@@ -592,17 +604,21 @@ def test_delete_inner_op():
     """Test rewrites where an operation inside a region of the matched op is deleted."""
 
     prog = """"builtin.module"() ({
-  %0 = "arith.constant"() <{"value" = 5 : i32}> : () -> i32
+  "test.op"() ({
+    %0 = "arith.constant"() <{"value" = 5 : i32}> : () -> i32
+  }) : () -> ()
 }) : () -> ()"""
 
     expected = """"builtin.module"() ({
-^0:
+  "test.op"() ({
+  ^0:
+  }) : () -> ()
 }) : () -> ()"""
 
     class Rewrite(RewritePattern):
         @op_type_rewrite_pattern
-        def match_and_rewrite(self, op: ModuleOp, rewriter: PatternRewriter):
-            first_op = op.ops.first
+        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
+            first_op = op.regions[0].block.ops.first
 
             assert first_op is not None
             rewriter.erase_op(first_op)
@@ -619,17 +635,21 @@ def test_replace_inner_op():
     """Test rewrites where an operation inside a region of the matched op is deleted."""
 
     prog = """"builtin.module"() ({
-  %0 = "arith.constant"() <{"value" = 5 : i32}> : () -> i32
+  "test.op"() ({
+    %0 = "arith.constant"() <{"value" = 5 : i32}> : () -> i32
+  }) : () -> ()
 }) : () -> ()"""
 
     expected = """"builtin.module"() ({
-  %0 = "arith.constant"() <{"value" = 42 : i32}> : () -> i32
+  "test.op"() ({
+    %0 = "arith.constant"() <{"value" = 42 : i32}> : () -> i32
+  }) : () -> ()
 }) : () -> ()"""
 
     class Rewrite(RewritePattern):
         @op_type_rewrite_pattern
-        def match_and_rewrite(self, op: ModuleOp, rewriter: PatternRewriter):
-            first_op = op.ops.first
+        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
+            first_op = op.regions[0].block.ops.first
 
             assert first_op is not None
             rewriter.replace_op(first_op, [Constant.from_int_and_width(42, i32)])
@@ -1048,29 +1068,36 @@ def test_move_region_contents_to_new_regions():
 
     prog = """\
 "builtin.module"() ({
-  %0 = "test.op"() : () -> !test.type<"int">
-  %1 = "test.op"() ({
-  ^0:
-    %2 = "test.op"() : () -> !test.type<"int">
-  }) : () -> !test.type<"int">
+  "test.op"() ({
+    %0 = "test.op"() : () -> !test.type<"int">
+    %1 = "test.op"() ({
+    ^0:
+      %2 = "test.op"() : () -> !test.type<"int">
+    }) : () -> !test.type<"int">
+  }) : () -> ()
 }) : () -> ()
 """
 
     expected = """\
 "builtin.module"() ({
-  %0 = "test.op"() : () -> !test.type<"int">
-  %1 = "test.op"() ({
-  }) : () -> !test.type<"int">
-  %2 = "test.op"() ({
-    %3 = "test.op"() : () -> !test.type<"int">
-  }) : () -> !test.type<"int">
+  "test.op"() ({
+    %0 = "test.op"() : () -> !test.type<"int">
+    %1 = "test.op"() ({
+    }) : () -> !test.type<"int">
+    %2 = "test.op"() ({
+      %3 = "test.op"() : () -> !test.type<"int">
+    }) : () -> !test.type<"int">
+  }) : () -> ()
 }) : () -> ()
 """
 
     class Rewrite(RewritePattern):
         @op_type_rewrite_pattern
-        def match_and_rewrite(self, op: ModuleOp, rewriter: PatternRewriter):
-            ops_iter = iter(op.ops)
+        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
+            # Match the toplevel test.op
+            if not isinstance(op.parent_op(), ModuleOp):
+                return
+            ops_iter = iter(op.regions[0].block.ops)
 
             _ = next(ops_iter)  # skip first op
             old_op = next(ops_iter)

xdsl/pattern_rewriter.py

@@ -690,7 +690,7 @@ class PatternRewriteWalker:
     That way, all uses are replaced before the definitions.
     """
 
-    post_walk_func: Callable[[Operation, PatternRewriterListener], bool] | None = field(
+    post_walk_func: Callable[[Region, PatternRewriterListener], bool] | None = field(
         default=None
     )
     """
@@ -773,34 +773,34 @@ def rewrite_module(self, module: ModuleOp) -> bool:
         Rewrite operations nested in the given operation by repeatedly applying the
         pattern. Returns `True` if the IR was mutated.
         """
-        return self.rewrite_op(module)
+        return self.rewrite_region(module.body)
 
-    def rewrite_op(self, op: Operation) -> bool:
+    def rewrite_region(self, region: Region) -> bool:
         """
         Rewrite operations nested in the given operation by repeatedly applying the
         pattern. Returns `True` if the IR was mutated.
         """
         pattern_listener = self._get_rewriter_listener()
 
-        self._populate_worklist(op)
+        self._populate_worklist(region)
         op_was_modified = self._process_worklist(pattern_listener)
         if self.post_walk_func is not None:
-            op_was_modified |= self.post_walk_func(op, pattern_listener)
+            op_was_modified |= self.post_walk_func(region, pattern_listener)
 
         if not self.apply_recursively:
             return op_was_modified
 
         result = op_was_modified
 
         while op_was_modified:
-            self._populate_worklist(op)
+            self._populate_worklist(region)
             op_was_modified = self._process_worklist(pattern_listener)
             if self.post_walk_func is not None:
-                op_was_modified |= self.post_walk_func(op, pattern_listener)
+                op_was_modified |= self.post_walk_func(region, pattern_listener)
 
         return result
 
-    def _populate_worklist(self, op: Operation) -> None:
+    def _populate_worklist(self, op: Operation | Region | Block) -> None:
         """Populate the worklist with all nested operations."""
         # We walk in reverse order since we use a stack for our worklist.
         for sub_op in op.walk(

xdsl/transforms/apply_pdl.py

@@ -22,19 +22,18 @@ class ApplyPDLPass(ModulePass):
     pdl_file: str | None = None
 
     def apply(self, ctx: MLContext, op: builtin.ModuleOp) -> None:
-        payload_module = op
         if self.pdl_file is not None:
             assert os.path.exists(self.pdl_file)
             with open(self.pdl_file) as f:
                 pdl_module_str = f.read()
                 parser = Parser(ctx, pdl_module_str)
                 pdl_module = parser.parse_module()
         else:
-            pdl_module = payload_module
+            pdl_module = op
         rewrite_patterns: list[RewritePattern] = [
             PDLRewritePattern(op, ctx, None)
             for op in pdl_module.walk()
             if isinstance(op, pdl.RewriteOp)
         ]
         pattern_applier = GreedyRewritePatternApplier(rewrite_patterns)
-        PatternRewriteWalker(pattern_applier).rewrite_op(payload_module)
+        PatternRewriteWalker(pattern_applier).rewrite_module(op)

xdsl/transforms/canonicalize.py

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

xdsl/transforms/convert_qref_to_qssa.py

@@ -48,4 +48,4 @@ class ConvertQRefToQssa(ModulePass):
     def apply(self, ctx: MLContext, op: builtin.ModuleOp) -> None:
         PatternRewriteWalker(
             ConvertQRefToQssaPattern(), apply_recursively=False
-        ).rewrite_op(op)
+        ).rewrite_module(op)

xdsl/transforms/convert_qssa_to_qref.py

@@ -33,4 +33,4 @@ class ConvertQssaToQRef(ModulePass):
     name = "convert-qssa-to-qref"
 
     def apply(self, ctx: MLContext, op: builtin.ModuleOp) -> None:
-        PatternRewriteWalker(ConvertQssaToQRefPattern()).rewrite_op(op)
+        PatternRewriteWalker(ConvertQssaToQRefPattern()).rewrite_module(op)

xdsl/transforms/convert_scf_to_cf.py

@@ -238,4 +238,4 @@ def apply(self, ctx: MLContext, op: builtin.ModuleOp) -> None:
                     ForLowering(),
                 ]
             )
-        ).rewrite_op(op)
+        ).rewrite_module(op)

xdsl/transforms/convert_stencil_to_csl_stencil.py

@@ -252,7 +252,7 @@ def match_and_rewrite(self, op: dmp.SwapOp, rewriter: PatternRewriter, /):
                 ConvertAccessOpFromPrefetchPattern(arg_idx)
             )
 
-            nested_rewriter.rewrite_op(new_apply_op)
+            nested_rewriter.rewrite_region(new_apply_op.region)
 
 
 def split_ops(
@@ -415,7 +415,7 @@ def match_and_rewrite(self, op: stencil.ApplyOp, rewriter: PatternRewriter, /):
         PatternRewriteWalker(
             SplitVarithOpPattern(op.region.block.args[prefetch_idx]),
             apply_recursively=False,
-        ).rewrite_op(op)
+        ).rewrite_region(op.region)
 
         # determine how ops should be split across the two regions
         chunk_region_ops, done_exchange_ops = split_ops(

xdsl/transforms/dead_code_elimination.py

@@ -159,14 +159,8 @@ def region_dce(region: Region, listener: PatternRewriterListener | None = None)
     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:
-        op_dce(op)
+        region_dce(op.body)