[Mlir-commits] [mlir] [mlir][IR] Trigger nested operation/block insertion notifications for clones (PR #66871)

Matthias Springer llvmlistbot at llvm.org
Wed Sep 20 01:21:19 PDT 2023


https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/66871

* Trigger "operation inserted" and "block inserted" notifications when cloning nested ops. (No such were sent for `RewriterBase::cloneRegionBefore` so far. "Block inserted" notifications were missing for `OpBuilder::clone`.)
* `cloneRegionBefore` is moved from `RewriterBase` to `OpBuilder`. (`cloneRegionBefore` just builds IR, it does not modify/erase ops.)
* `cloneRegionBefore` is no longer virtual. It used to be virtual so that the dialect conversion can override it and update the internal state in the correct order (for nested ops). The internal state is updated based on listener notifications. Now that the listener notifications are sent for nested clones (and sent in the right order), this workaround in the dialect conversion is no longer needed.

Details:

Ops/blocks are notified in such an order in which they could have been created one-by-one with an `OpBuilder`. This ensures that when a listener is notified about an op being created, it was already notified about the defining ops of the operands (unless there is a block cycle/graph region).

The implementation first clones an entire op (including nested ops) and then sends all notifications. Ideally, notifications should be interleaved with the cloning process, but that would require duplicating `Region::cloneInto` (with listener support). This commit is an incremental improvement over not sending any notifications.

It is not possible to use the current implementation of `ConversionPatternRewriter::cloneRegionBefore` to enumerate blocks/ops (using `ForwardDominanceIterator`) because it does not enumerate unreachable (dead) blocks. However, unreachable blocks are cloned. Therefore, notifications should be sent for unreachable blocks. (The dialect conversion does not require notifications for unreachable blocks, so they could get away with this simpler implementation.)

There is a "fast path" in the `clone` functions in case no listener is attached. No IR traversal is needed in that case.

Imported from Phabricator: https://reviews.llvm.org/D146943


>From 26be7a8cc63b2ae40b250a8bf109b97f7a3be6c5 Mon Sep 17 00:00:00 2001
From: Matthias Springer <me at m-sp.org>
Date: Wed, 20 Sep 2023 10:13:33 +0200
Subject: [PATCH] [mlir][IR] Trigger nested operation/block insertion
 notifications for clones

* Trigger "operation inserted" and "block inserted" notifications when cloning nested ops. (No such notifications were sent for nested ops when running `OpBuilder::clone(Operation *)` or `RewriterBase::cloneRegionBefore` so far.)
* `cloneRegionBefore` is moved from `RewriterBase` to `OpBuilder`. (`cloneRegionBefore` just builds IR, it does not modify/erase ops.)
* `cloneRegionBefore` is no longer virtual. It used to be virtual so that the dialect conversion can override it and update the internal state in the correct order (for nested ops). The internal state is updated based on listener notifications. Now that the listener notifications are sent for nested clones (and sent in the right order), this workaround in the dialect conversion is no longer needed.

Details:

Ops/blocks are notified in such an order in which they could have been created one-by-one with an `OpBuilder`. This ensures that when a listener is notified about an op being created, it was already notified about the defining ops of the operands (unless there is a block cycle/graph region).

The implementation first clones an entire op (including nested ops) and then sends all notifications. Ideally, notifications should be interleaved with the cloning process, but that would require duplicating `Region::cloneInto` (with listener support). This commit is an incremental improvement over not sending any notifications.

There is a "fast path" in the `clone` functions in case no listener is attached. No IR traversal is needed in that case.

Imported from Phabricator: https://reviews.llvm.org/D146943

BEGIN_PUBLIC
No public commit message needed for presubmit.
END_PUBLIC
---
 mlir/include/mlir/IR/Builders.h               |  10 +
 mlir/include/mlir/IR/PatternMatch.h           |  10 -
 .../mlir/Transforms/DialectConversion.h       |   8 -
 mlir/lib/IR/Builders.cpp                      |  97 ++++++-
 mlir/lib/IR/PatternMatch.cpp                  |  18 --
 .../Transforms/Utils/DialectConversion.cpp    |  17 --
 .../test-strict-pattern-driver.mlir           | 240 +++++++++++++++++-
 mlir/test/lib/Dialect/Test/TestPatterns.cpp   |  43 +++-
 8 files changed, 366 insertions(+), 77 deletions(-)

diff --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h
index 3eca8bd12f43364..4f5144d821b7e2f 100644
--- a/mlir/include/mlir/IR/Builders.h
+++ b/mlir/include/mlir/IR/Builders.h
@@ -451,6 +451,16 @@ class OpBuilder : public Builder {
   Block *createBlock(Block *insertBefore, TypeRange argTypes = std::nullopt,
                      ArrayRef<Location> locs = std::nullopt);
 
+  /// Clone the blocks that belong to "region" before the given position in
+  /// another region "parent". The two regions must be different. The caller is
+  /// responsible for creating or updating the operation transferring flow of
+  /// control to the region and passing it the correct block arguments.
+  void cloneRegionBefore(Region &region, Region &parent,
+                         Region::iterator before, IRMapping &mapping);
+  void cloneRegionBefore(Region &region, Region &parent,
+                         Region::iterator before);
+  void cloneRegionBefore(Region &region, Block *before);
+
   //===--------------------------------------------------------------------===//
   // Operation Creation
   //===--------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/IR/PatternMatch.h b/mlir/include/mlir/IR/PatternMatch.h
index 6625ef553eba21f..276ef079ad7682f 100644
--- a/mlir/include/mlir/IR/PatternMatch.h
+++ b/mlir/include/mlir/IR/PatternMatch.h
@@ -489,16 +489,6 @@ class RewriterBase : public OpBuilder {
                                   Region::iterator before);
   void inlineRegionBefore(Region &region, Block *before);
 
-  /// Clone the blocks that belong to "region" before the given position in
-  /// another region "parent". The two regions must be different. The caller is
-  /// responsible for creating or updating the operation transferring flow of
-  /// control to the region and passing it the correct block arguments.
-  virtual void cloneRegionBefore(Region &region, Region &parent,
-                                 Region::iterator before, IRMapping &mapping);
-  void cloneRegionBefore(Region &region, Region &parent,
-                         Region::iterator before);
-  void cloneRegionBefore(Region &region, Block *before);
-
   /// This method replaces the uses of the results of `op` with the values in
   /// `newValues` when the provided `functor` returns true for a specific use.
   /// The number of values in `newValues` is required to match the number of
diff --git a/mlir/include/mlir/Transforms/DialectConversion.h b/mlir/include/mlir/Transforms/DialectConversion.h
index 6de981d35c8c3ab..1763344425dbaa7 100644
--- a/mlir/include/mlir/Transforms/DialectConversion.h
+++ b/mlir/include/mlir/Transforms/DialectConversion.h
@@ -727,14 +727,6 @@ class ConversionPatternRewriter final : public PatternRewriter,
                           Region::iterator before) override;
   using PatternRewriter::inlineRegionBefore;
 
-  /// PatternRewriter hook for cloning blocks of one region into another. The
-  /// given region to clone *must* not have been modified as part of conversion
-  /// yet, i.e. it must be within an operation that is either in the process of
-  /// conversion, or has not yet been converted.
-  void cloneRegionBefore(Region &region, Region &parent,
-                         Region::iterator before, IRMapping &mapping) override;
-  using PatternRewriter::cloneRegionBefore;
-
   /// PatternRewriter hook for inserting a new operation.
   void notifyOperationInserted(Operation *op) override;
 
diff --git a/mlir/lib/IR/Builders.cpp b/mlir/lib/IR/Builders.cpp
index ab20f4863e11c23..3a4a6b5b4eaa290 100644
--- a/mlir/lib/IR/Builders.cpp
+++ b/mlir/lib/IR/Builders.cpp
@@ -14,7 +14,11 @@
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Matchers.h"
+#include "mlir/IR/RegionGraphTraits.h"
 #include "mlir/IR/SymbolTable.h"
+#include "llvm/ADT/PostOrderIterator.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVectorExtras.h"
 #include "llvm/Support/raw_ostream.h"
 
@@ -525,18 +529,66 @@ LogicalResult OpBuilder::tryFold(Operation *op,
   return success();
 }
 
+static void notifyOperationCloned(Operation *op, OpBuilder::Listener *listener);
+
+/// Notify the listener that the given range of regions was cloned.
+///
+/// Blocks and operations are enumerated in an order in which they could have
+/// been created separately (without using the `clone` API): Within a region,
+/// blocks are notified according to their successor relationship. Within a
+/// block, operations are notified in forward mode. This ensures that defining
+/// ops are notified before ops that use their results (unless there are
+/// cycles/graph regions).
+static void notifyRegionsCloned(iterator_range<Region::iterator> range,
+                                OpBuilder::Listener *listener) {
+  // Maintain a set of blocks that were not notified yet. This is needed because
+  // the inverse_post_order iterator does not enumerate dead blocks.
+  llvm::SetVector<Block *> remainingBlocks;
+  // The order in which the set is initialized does not matter for correctness.
+  // For better performance, "leaf" blocks with no successors should be starting
+  // point for the block traversal. (Then there are fewer iterations of the
+  // "while" loop.)
+  for (Block &b : llvm::reverse(range))
+    remainingBlocks.insert(&b);
+  // Set of visited blocks that is shared among all inverse_post_order
+  // iterations. This is to avoid that the same block is enumerated multiple
+  // times.
+  llvm::SmallPtrSet<Block *, 4> visited;
+  while (!remainingBlocks.empty()) {
+    // Enumerate predecessors before successors. I.e., reverse post-order.
+    for (Block *b :
+         llvm::inverse_post_order_ext(remainingBlocks.front(), visited)) {
+      auto it = llvm::find(remainingBlocks, b);
+      assert(it != remainingBlocks.end() &&
+             "expected that only remaining blocks are visited");
+      listener->notifyBlockCreated(b);
+      remainingBlocks.erase(it);
+      for (Operation &op : *b)
+        notifyOperationCloned(&op, listener);
+    }
+  }
+}
+
+/// Notify the listener that the given op was cloned.
+static void notifyOperationCloned(Operation *op,
+                                  OpBuilder::Listener *listener) {
+  listener->notifyOperationInserted(op);
+  for (Region &r : op->getRegions())
+    notifyRegionsCloned(r.getBlocks(), listener);
+}
+
 Operation *OpBuilder::clone(Operation &op, IRMapping &mapper) {
+  // TODO: The listener notifications should be interleaved with `clone`.
   Operation *newOp = op.clone(mapper);
-  // The `insert` call below handles the notification for inserting `newOp`
-  // itself. But if `newOp` has any regions, we need to notify the listener
-  // about any ops that got inserted inside those regions as part of cloning.
-  if (listener) {
-    auto walkFn = [&](Operation *walkedOp) {
-      listener->notifyOperationInserted(walkedOp);
-    };
-    for (Region &region : newOp->getRegions())
-      region.walk(walkFn);
-  }
+
+  // Fast path: If no listener is attached, the op can be inserted directly.
+  if (!listener)
+    return insert(newOp);
+
+  // The `insert` call below handles the notification for inserting `newOp`.
+  // Just notify about nested op/block insertion.
+  for (Region &r : newOp->getRegions())
+    notifyRegionsCloned(r.getBlocks(), listener);
   return insert(newOp);
 }
 
@@ -544,3 +596,28 @@ Operation *OpBuilder::clone(Operation &op) {
   IRMapping mapper;
   return clone(op, mapper);
 }
+
+void OpBuilder::cloneRegionBefore(Region &region, Region &parent,
+                                  Region::iterator before, IRMapping &mapping) {
+  // TODO: The listener notifications should be interleaved with `clone`.
+  region.cloneInto(&parent, before, mapping);
+
+  // Fast path: If no listener is attached, there is no more work to do.
+  if (!listener)
+    return;
+
+  // Notify about op/block insertion.
+  Region::iterator clonedBeginIt =
+      mapping.lookup(&region.front())->getIterator();
+  notifyRegionsCloned(llvm::make_range(clonedBeginIt, before), listener);
+}
+
+void OpBuilder::cloneRegionBefore(Region &region, Region &parent,
+                                  Region::iterator before) {
+  IRMapping mapping;
+  cloneRegionBefore(region, parent, before, mapping);
+}
+
+void OpBuilder::cloneRegionBefore(Region &region, Block *before) {
+  cloneRegionBefore(region, *before->getParent(), before->getIterator());
+}
diff --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index 5e9b9b2a810a4c5..a6f09f1609b3bae 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -458,21 +458,3 @@ void RewriterBase::inlineRegionBefore(Region &region, Region &parent,
 void RewriterBase::inlineRegionBefore(Region &region, Block *before) {
   inlineRegionBefore(region, *before->getParent(), before->getIterator());
 }
-
-/// Clone the blocks that belong to "region" before the given position in
-/// another region "parent". The two regions must be different. The caller is
-/// responsible for creating or updating the operation transferring flow of
-/// control to the region and passing it the correct block arguments.
-void RewriterBase::cloneRegionBefore(Region &region, Region &parent,
-                                     Region::iterator before,
-                                     IRMapping &mapping) {
-  region.cloneInto(&parent, before, mapping);
-}
-void RewriterBase::cloneRegionBefore(Region &region, Region &parent,
-                                     Region::iterator before) {
-  IRMapping mapping;
-  cloneRegionBefore(region, parent, before, mapping);
-}
-void RewriterBase::cloneRegionBefore(Region &region, Block *before) {
-  cloneRegionBefore(region, *before->getParent(), before->getIterator());
-}
diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 4d2afe462b9281d..e0b8e0373605462 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -1588,23 +1588,6 @@ void ConversionPatternRewriter::inlineRegionBefore(Region &region,
   PatternRewriter::inlineRegionBefore(region, parent, before);
 }
 
-void ConversionPatternRewriter::cloneRegionBefore(Region &region,
-                                                  Region &parent,
-                                                  Region::iterator before,
-                                                  IRMapping &mapping) {
-  if (region.empty())
-    return;
-
-  PatternRewriter::cloneRegionBefore(region, parent, before, mapping);
-
-  for (Block &b : ForwardDominanceIterator<>::makeIterable(region)) {
-    Block *cloned = mapping.lookup(&b);
-    impl->notifyCreatedBlock(cloned);
-    cloned->walk<WalkOrder::PreOrder, ForwardDominanceIterator<>>(
-        [&](Operation *op) { notifyOperationInserted(op); });
-  }
-}
-
 void ConversionPatternRewriter::notifyOperationInserted(Operation *op) {
   LLVM_DEBUG({
     impl->logger.startLine()
diff --git a/mlir/test/Transforms/test-strict-pattern-driver.mlir b/mlir/test/Transforms/test-strict-pattern-driver.mlir
index a5ab8f97c74ce33..f7c1200ca8ad2a4 100644
--- a/mlir/test/Transforms/test-strict-pattern-driver.mlir
+++ b/mlir/test/Transforms/test-strict-pattern-driver.mlir
@@ -18,7 +18,7 @@
 func.func @test_erase() {
   %0 = "test.arg0"() : () -> (i32)
   %1 = "test.arg1"() : () -> (i32)
-  %erase = "test.erase_op"(%0, %1) : (i32, i32) -> (i32)
+  %erase = "test.erase_op"(%0, %1) {worklist} : (i32, i32) -> (i32)
   return
 }
 
@@ -29,7 +29,7 @@ func.func @test_erase() {
 //       CHECK-EN:   "test.insert_same_op"() {skip = true}
 //       CHECK-EN:   "test.insert_same_op"() {skip = true}
 func.func @test_insert_same_op() {
-  %0 = "test.insert_same_op"() : () -> (i32)
+  %0 = "test.insert_same_op"() {worklist} : () -> (i32)
   return
 }
 
@@ -41,7 +41,7 @@ func.func @test_insert_same_op() {
 //       CHECK-EN:   "test.dummy_user"(%[[n]])
 //       CHECK-EN:   "test.dummy_user"(%[[n]])
 func.func @test_replace_with_new_op() {
-  %0 = "test.replace_with_new_op"() : () -> (i32)
+  %0 = "test.replace_with_new_op"() {worklist} : () -> (i32)
   %1 = "test.dummy_user"(%0) : (i32) -> (i32)
   %2 = "test.dummy_user"(%0) : (i32) -> (i32)
   return
@@ -59,7 +59,7 @@ func.func @test_replace_with_new_op() {
 //   CHECK-EX-NOT:   "test.replace_with_new_op"
 //       CHECK-EX:   "test.erase_op"
 func.func @test_replace_with_erase_op() {
-  "test.replace_with_new_op"() {create_erase_op} : () -> ()
+  "test.replace_with_new_op"() {create_erase_op, worklist} : () -> ()
   return
 }
 
@@ -74,14 +74,14 @@ func.func @test_trigger_rewrite_through_block() {
   return
 ^bb1:
   // Uses bb1. ChangeBlockOp replaces that and all other usages of bb1 with bb2.
-  "test.change_block_op"() [^bb1, ^bb2] : () -> ()
+  "test.change_block_op"() [^bb1, ^bb2] {worklist} : () -> ()
 ^bb2:
   return
 ^bb3:
   // Also uses bb1. ChangeBlockOp replaces that usage with bb2. This triggers
   // this op being put on the worklist, which triggers ImplicitChangeOp, which,
   // in turn, replaces the successor with bb3.
-  "test.implicit_change_op"() [^bb1] : () -> ()
+  "test.implicit_change_op"() [^bb1] {worklist} : () -> ()
 }
 
 // -----
@@ -98,7 +98,7 @@ func.func @test_remove_graph_region() {
       %0 = "test.foo_a"(%1) : (i1) -> (i1)
       %1 = "test.foo_b"(%0) : (i1) -> (i1)
     }
-  }) : () -> ()
+  }) {worklist} : () -> ()
   return
 }
 
@@ -123,7 +123,7 @@ func.func @test_remove_cyclic_blocks() {
   ^bb2(%arg1: i1):
     "test.bar"(%x) : (i1) -> ()
     cf.br ^bb1(%arg1: i1)
-  }) : () -> ()
+  }) {worklist} : () -> ()
   return
 }
 
@@ -155,7 +155,7 @@ func.func @test_remove_dead_blocks() {
       "test.qux_unreachable"() : () -> ()
     }) : () -> ()
     "test.bar"() : () -> ()
-  }) : () -> ()
+  }) {worklist} : () -> ()
   return
 }
 
@@ -201,7 +201,7 @@ func.func @test_remove_nested_ops() {
   ^bb2(%arg1: i1):
     "test.bar"(%x) : (i1) -> ()
     cf.br ^bb1(%arg1: i1)
-  }) : () -> ()
+  }) {worklist} : () -> ()
   return
 }
 
@@ -226,6 +226,224 @@ func.func @test_remove_diamond(%c: i1) {
     cf.br ^bb3
   ^bb3:
     "test.qux"() : () -> ()
-  }) : () -> ()
+  }) {worklist} : () -> ()
   return
 }
+
+// -----
+
+// Make sure that test.erase_op is deleted.
+
+// CHECK-AN: notifyBlockCreated: block ^bb1 from region 0 from operation 'test.dummy_op'
+// CHECK-AN: notifyOperationInserted: test.inner_op
+// CHECK-AN: notifyOperationInserted: test.erase_op
+// CHECK-AN: notifyOperationRemoved: test.erase_op
+// CHECK-AN: notifyOperationRemoved: test.inner_op
+// CHECK-AN: notifyOperationRemoved: test.clone_region_in_parent
+// CHECK-AN: notifyOperationRemoved: test.erase_op
+// CHECK-AN-LABEL: func @test_add_cloned_ops_to_worklist
+
+// CHECK-EN-LABEL: func @test_add_cloned_ops_to_worklist
+//  CHECK-EN-NEXT:   "test.dummy_op"() ({
+//  CHECK-EN-NEXT:     "test.another_op"() : () -> ()
+//  CHECK-EN-NEXT:   ^bb1:  // no predecessors
+//  CHECK-EN-NEXT:     "test.inner_op"() : () -> ()
+//  CHECK-EN-NEXT:   }) : () -> ()
+//  CHECK-EN-NEXT:  }
+
+// When strictness=ExistingOps, test.erase_op is not deleted because it was not
+// on the initial worklist.
+
+// CHECK-EX-LABEL: func @test_add_cloned_ops_to_worklist
+//  CHECK-EX-NEXT:   "test.dummy_op"() ({
+//  CHECK-EX-NEXT:     "test.another_op"() : () -> ()
+//  CHECK-EX-NEXT:   ^bb1:  // no predecessors
+//  CHECK-EX-NEXT:     "test.inner_op"() : () -> ()
+//  CHECK-EX-NEXT:     "test.erase_op"() : () -> ()
+//  CHECK-EX-NEXT:   }) : () -> ()
+//  CHECK-EX-NEXT:  }
+func.func @test_add_cloned_ops_to_worklist() {
+  "test.dummy_op"() ({
+    ^bb0():
+    "test.clone_region_in_parent"() ({
+    ^bb1():
+      "test.inner_op"() : () -> ()
+      "test.erase_op"() {worklist} : () -> ()
+    }) {worklist} : () -> ()
+    "test.another_op"() : () -> ()
+  }) {worklist} : () -> ()
+}
+
+// -----
+
+// CHECK-AN: notifyBlockCreated: block ^bb1 from region 0 from operation 'test.dummy_op'
+// CHECK-AN: notifyOperationInserted: test.inner_op
+// CHECK-AN: notifyBlockCreated: block ^bb0 from region 0 from operation 'test.inner_op'
+// CHECK-AN: notifyOperationInserted: test.graph_region
+// CHECK-AN: notifyBlockCreated: block ^bb0 from region 0 from operation 'test.graph_region'
+// CHECK-AN: notifyOperationInserted: test.foo_a
+// CHECK-AN: notifyOperationInserted: test.foo_b
+// CHECK-AN: notifyOperationInserted: test.foo_c
+// The original op is removed.
+// CHECK-AN: notifyOperationRemoved: test.foo_c
+// CHECK-AN: notifyOperationRemoved: test.foo_b
+// CHECK-AN: notifyOperationRemoved: test.foo_a
+// CHECK-AN: notifyOperationRemoved: test.graph_region
+// CHECK-AN: notifyOperationRemoved: test.inner_op
+// CHECK-AN: notifyOperationRemoved: test.clone_region_in_parent
+// CHECK-AN-LABEL: func @test_clone_nested_graph_region
+//       CHECK-AN:   "test.dummy_op"() ({
+//  CHECK-AN-NEXT:     "test.another_op"() : () -> ()
+//  CHECK-AN-NEXT:     ^bb1:  // no predecessors
+//  CHECK-AN-NEXT:       "test.inner_op"() ({
+//  CHECK-AN-NEXT:         test.graph_region {
+//  CHECK-AN-NEXT:           %{{.*}} = "test.foo_a"(%{{.*}}) : (i1) -> i1
+//  CHECK-AN-NEXT:           %{{.*}} = "test.foo_b"(%{{.*}}) : (i1) -> i1
+//  CHECK-AN-NEXT:         }
+//  CHECK-AN-NEXT:         "test.foo_c"() : () -> ()
+//  CHECK-AN-NEXT:       }) : () -> ()
+//  CHECK-AN-NEXT:     }) : () -> ()
+func.func @test_clone_nested_graph_region() {
+  "test.dummy_op"() ({
+    ^bb0():
+    "test.clone_region_in_parent"() ({
+    ^bb1():
+      "test.inner_op"() ({
+        test.graph_region {
+          %0 = "test.foo_a"(%1) : (i1) -> (i1)
+          %1 = "test.foo_b"(%0) : (i1) -> (i1)
+        }
+        "test.foo_c"() : () -> ()
+      }): () -> ()
+    }) {worklist} : () -> ()
+    "test.another_op"() : () -> ()
+  }) {worklist} : () -> ()
+}
+
+// -----
+
+// CHECK-AN: notifyBlockCreated: block ^bb1 from region 0 from operation 'test.dummy_op'
+// CHECK-AN: notifyOperationInserted: test.dummy_op
+// CHECK-AN: notifyOperationInserted: cf.br
+// CHECK-AN: notifyBlockCreated: block ^bb2 from region 0 from operation 'test.dummy_op'
+// CHECK-AN: notifyOperationInserted: test.foo
+// CHECK-AN: notifyOperationInserted: cf.br
+// CHECK-AN: notifyBlockCreated: block ^bb3 from region 0 from operation 'test.dummy_op'
+// CHECK-AN: notifyOperationInserted: test.bar
+// CHECK-AN: notifyOperationInserted: cf.br
+// CHECK-AN: notifyOperationRemoved: cf.br
+// CHECK-AN: notifyOperationRemoved: test.bar
+// CHECK-AN: notifyOperationRemoved: cf.br
+// CHECK-AN: notifyOperationRemoved: test.foo
+// CHECK-AN: notifyOperationRemoved: cf.br
+// CHECK-AN: notifyOperationRemoved: test.dummy_op
+// CHECK-AN: notifyOperationRemoved: test.clone_region_in_parent
+// CHECK-AN-LABEL: func @test_clone_cyclic_blocks
+func.func @test_clone_cyclic_blocks() {
+  "test.dummy_op"() ({
+    ^bb0():
+    "test.clone_region_in_parent"() ({
+      %x = "test.dummy_op"() : () -> (i1)
+      cf.br ^bb1(%x: i1)
+      ^bb1(%arg0: i1):
+        "test.foo"(%x) : (i1) -> ()
+        cf.br ^bb2(%arg0: i1)
+      ^bb2(%arg1: i1):
+        "test.bar"(%x) : (i1) -> ()
+        cf.br ^bb1(%arg1: i1)
+      }) {worklist} : () -> ()
+      "test.qux"() : () -> ()
+    }) : () -> ()
+    "test.another_op"() : () -> ()
+  return
+}
+
+// -----
+
+// CHECK-AN: notifyBlockCreated: block ^bb3 from region 0 from operation 'test.dummy_op'
+// CHECK-AN: notifyOperationInserted: test.nested_dummy
+// CHECK-AN: notifyBlockCreated: block ^bb1 from region 0 from operation 'test.nested_dummy'
+// CHECK-AN: notifyOperationInserted: test.qux_unreachable
+// CHECK-AN: notifyBlockCreated: block ^bb0 from region 0 from operation 'test.nested_dummy'
+// CHECK-AN: notifyOperationInserted: test.qux_reachable
+// CHECK-AN: notifyOperationInserted: test.bar
+// CHECK-AN: notifyOperationInserted: cf.br
+// CHECK-AN: notifyBlockCreated: block ^bb2 from region 0 from operation 'test.dummy_op'
+// CHECK-AN: notifyOperationInserted: test.foo
+// CHECK-AN: notifyBlockCreated: block ^bb1 from region 0 from operation 'test.dummy_op'
+// CHECK-AN: notifyOperationInserted: test.dummy_op
+// CHECK-AN: notifyOperationRemoved: test.dummy_op
+// CHECK-AN: notifyOperationRemoved: test.foo
+// CHECK-AN: notifyOperationRemoved: cf.br
+// CHECK-AN: notifyOperationRemoved: test.bar
+// CHECK-AN: notifyOperationRemoved: test.qux_reachable
+// CHECK-AN: notifyOperationRemoved: test.qux_unreachable
+// CHECK-AN: notifyOperationRemoved: test.nested_dummy
+// CHECK-AN: notifyOperationRemoved: test.clone_region_in_parent
+// CHECK-AN-LABEL: func @test_clone_dead_blocks
+func.func @test_clone_dead_blocks() {
+  "test.dummy_op"() ({
+    ^bb0():
+    "test.clone_region_in_parent"() ({
+      "test.dummy_op"() : () -> (i1)
+      // The following blocks are not reachable. Still, ^bb2 should be notified
+      // before ^bb1.
+      ^bb1(%arg0: i1):
+        "test.foo"() : () -> ()
+      ^bb2(%arg1: i1):
+        "test.nested_dummy"() ({
+          "test.qux_reachable"() : () -> ()
+        // The following block is unreachable.
+        ^bb3:
+          "test.qux_unreachable"() : () -> ()
+        }) : () -> ()
+        "test.bar"() : () -> ()
+        cf.br ^bb1(%arg0: i1)
+      }) {worklist} : () -> ()
+      "test.qux"() : () -> ()
+    }) : () -> ()
+    "test.another_op"() : () -> ()
+  return
+}
+
+// -----
+
+// CHECK-AN: notifyBlockCreated: block ^bb1 from region 0 from operation 'test.dummy_op'
+// CHECK-AN: notifyOperationInserted: cf.br
+// CHECK-AN: notifyBlockCreated: block ^bb5 from region 0 from operation 'test.dummy_op'
+// CHECK-AN: notifyOperationInserted: cf.cond_br
+// CHECK-AN: notifyBlockCreated: block ^bb4 from region 0 from operation 'test.dummy_op'
+// CHECK-AN: notifyOperationInserted: test.bar
+// CHECK-AN: notifyOperationInserted: cf.br
+// CHECK-AN: notifyBlockCreated: block ^bb3 from region 0 from operation 'test.dummy_op'
+// CHECK-AN: notifyOperationInserted: test.foo
+// CHECK-AN: notifyOperationInserted: cf.br
+// CHECK-AN: notifyBlockCreated: block ^bb2 from region 0 from operation 'test.dummy_op'
+// CHECK-AN: notifyOperationInserted: test.qux
+// CHECK-AN: notifyOperationRemoved: test.qux
+// CHECK-AN: notifyOperationRemoved: cf.br
+// CHECK-AN: notifyOperationRemoved: test.foo
+// CHECK-AN: notifyOperationRemoved: cf.br
+// CHECK-AN: notifyOperationRemoved: test.bar
+// CHECK-AN: notifyOperationRemoved: cf.cond_br
+// CHECK-AN: notifyOperationRemoved: cf.br
+// CHECK-AN: notifyOperationRemoved: test.clone_region_in_parent
+// CHECK-AN-LABEL: func @test_clone_reverse_diamond
+func.func @test_clone_reverse_diamond(%c: i1) {
+  "test.dummy_op"() ({
+    "test.clone_region_in_parent"() ({
+      cf.br ^bb3
+    ^bb0:
+      "test.qux"() : () -> ()
+    ^bb1:
+      "test.foo"() : () -> ()
+      cf.br ^bb0
+    ^bb2:
+      "test.bar"() : () -> ()
+      cf.br ^bb0
+    ^bb3:
+      cf.cond_br %c, ^bb1, ^bb2
+    }) {worklist} : () -> ()
+    "test.another_op"() : () -> ()
+  }) {worklist} : () -> ()
+}
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index 2e3bc76009ca208..7f3daefe941ea2e 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -239,7 +239,27 @@ struct TestPatternDriver
       llvm::cl::init(GreedyRewriteConfig().maxIterations)};
 };
 
+static void printRegion(Region *region) {
+  llvm::outs() << "region " << region->getRegionNumber() << " from operation '"
+               << region->getParentOp()->getName() << "'";
+}
+
+static void printBlock(Block *block) {
+  llvm::outs() << "block ";
+  block->printAsOperand(llvm::outs(), /*printType=*/false);
+  llvm::outs() << " from ";
+  printRegion(block->getParent());
+}
+
 struct DumpNotifications : public RewriterBase::Listener {
+  void notifyOperationInserted(Operation *op) override {
+    llvm::outs() << "notifyOperationInserted: " << op->getName() << "\n";
+  }
+  void notifyBlockCreated(Block *b) override {
+    llvm::outs() << "notifyBlockCreated: ";
+    printBlock(b);
+    llvm::outs() << "\n";
+  }
   void notifyOperationRemoved(Operation *op) override {
     llvm::outs() << "notifyOperationRemoved: " << op->getName() << "\n";
   }
@@ -269,14 +289,14 @@ struct TestStrictPatternDriver
         ReplaceWithNewOp,
         EraseOp,
         ChangeBlockOp,
+        CloneRegionInParentOp,
         ImplicitChangeOp
         // clang-format on
         >(ctx);
     SmallVector<Operation *> ops;
     getOperation()->walk([&](Operation *op) {
-      StringRef opName = op->getName().getStringRef();
-      if (opName == "test.insert_same_op" || opName == "test.change_block_op" ||
-          opName == "test.replace_with_new_op" || opName == "test.erase_op") {
+      if (op->hasAttr("worklist")) {
+        op->removeAttr("worklist");
         ops.push_back(op);
       }
     });
@@ -313,6 +333,23 @@ struct TestStrictPatternDriver
       llvm::cl::init("AnyOp")};
 
 private:
+  // Clone region into parent op.
+  class CloneRegionInParentOp : public RewritePattern {
+  public:
+    CloneRegionInParentOp(MLIRContext *context)
+        : RewritePattern("test.clone_region_in_parent", /*benefit=*/1,
+                         context) {}
+
+    LogicalResult matchAndRewrite(Operation *op,
+                                  PatternRewriter &rewriter) const override {
+      auto &parentRegion = *op->getParentRegion();
+      auto &opRegion = op->getRegion(0);
+      rewriter.cloneRegionBefore(opRegion, parentRegion, parentRegion.end());
+      rewriter.eraseOp(op);
+      return success();
+    }
+  };
+
   // New inserted operation is valid for further transformation.
   class InsertSameOp : public RewritePattern {
   public:



More information about the Mlir-commits mailing list