core: Walk regions for rewritings instead of ops

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)