core: Walk regions for rewritings instead of ops

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,17 +377,21 @@ 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, mod: 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]))
@@ -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, mod: test.TestOp, rewriter: PatternRewriter):
             new_cst = Constant.from_int_and_width(42, i32)
 
-            first_op = mod.ops.first
+            first_op = mod.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, mod: test.TestOp, rewriter: PatternRewriter):
             new_cst = Constant.from_int_and_width(42, i32)
 
-            first_op = mod.ops.first
+            first_op = mod.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/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,12 @@ 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)
+def op_dce(region: Region, listener: PatternRewriterListener | None = None):
+    return region_dce(region, listener)
 
 
 class DeadCodeElimination(ModulePass):
     name = "dce"
 
     def apply(self, ctx: MLContext, op: ModuleOp) -> None:
-        op_dce(op)
+        op_dce(op.body)

xdsl/transforms/varith_transformations.py

@@ -252,7 +252,7 @@ def apply(self, ctx: MLContext, op: builtin.ModuleOp) -> None:
                 ]
             ),
             walk_reverse=True,
-        ).rewrite_op(op)
+        ).rewrite_module(op)
 
 
 class ConvertVarithToArithPass(ModulePass):
@@ -268,7 +268,7 @@ def apply(self, ctx: MLContext, op: builtin.ModuleOp) -> None:
         PatternRewriteWalker(
             VarithToArithPattern(),
             apply_recursively=False,
-        ).rewrite_op(op)
+        ).rewrite_module(op)
 
 
 class VarithFuseRepeatedOperandsPass(ModulePass):
@@ -285,4 +285,4 @@ def apply(self, ctx: MLContext, op: builtin.ModuleOp) -> None:
         PatternRewriteWalker(
             FuseRepeatedAddArgsPattern(self.min_reps),
             apply_recursively=False,
-        ).rewrite_op(op)
+        ).rewrite_module(op)