[mlir][IR] Trigger notifyOperationInserted for nested ops during cloneRegionBefore

[springerm]

Summary

Fixes an inconsistency where cloneRegionBefore did not notify the listener of "op insertion" (whereas OpBuilder::clone does). Various parts of the core infrastructure such as the GreedyPatternRewriteDriver depend on such notifications.

mlir/lib/IR/Builders.cpp

 
 Operation *OpBuilder::clone(Operation &op, IRMapping &mapper) {
   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) {
+      // TODO: Notify listener of block creation.
       listener->notifyOperationInserted(walkedOp);
     };
     for (Region &region : newOp->getRegions())
       region.walk(walkFn);
   }
   return insert(newOp);
 }
 
 Operation *OpBuilder::clone(Operation &op) {
   IRMapping mapper;
   return clone(op, mapper);
 }

mlir/lib/IR/PatternMatch.cpp

 /// 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);
+  if (Listener *listener = getListener()) {
+    Region::iterator it = mapping.lookup(&region.front())->getIterator();
+    while (it != before) {
+      // TODO: Notify listener of nested blocks.
+      it->walk([&](Operation *op) { listener->notifyOperationInserted(op); });
+      it++;
+    }
+  }
 }
 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());
 }

mlir/lib/Transforms/Utils/DialectConversion.cpp

void ConversionPatternRewriter::inlineRegionBefore(Region &region,
   impl->notifyRegionIsBeingInlinedBefore(region, parent, before);
   PatternRewriter::inlineRegionBefore(region, parent, before);
 }
 
 void ConversionPatternRewriter::cloneRegionBefore(Region &region,
                                                   Region &parent,
                                                   Region::iterator before,
                                                   IRMapping &mapping) {
+  // Note: Do not call PatternRewriter::cloneRegionBefore. We cannot rely on its
+  // "op created" notifications because we need to populate the list of created
+  // ops in a certain order. This is done in `notifyRegionWasClonedBefore`.
+
   if (region.empty())
     return;
-  PatternRewriter::cloneRegionBefore(region, parent, before, mapping);
+
+  region.cloneInto(&parent, before, mapping);
 
   // Collect the range of the cloned blocks.
   auto clonedBeginIt = mapping.lookup(&region.front())->getIterator();
   auto clonedBlocks = llvm::make_range(clonedBeginIt, before);
   impl->notifyRegionWasClonedBefore(clonedBlocks, region.getLoc());
 }
 
 void ConversionPatternRewriter::notifyOperationInserted(Operation *op) {

mlir/test/Transforms/test-strict-pattern-driver.mlir

 // CHECK-EN-LABEL: func @test_erase
 //  CHECK-EN-SAME:     pattern_driver_all_erased = true, pattern_driver_changed = true}
 //       CHECK-EN:   test.arg0
 //       CHECK-EN:   test.arg1
 //   CHECK-EN-NOT:   test.erase_op
 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
 }
 
 // -----
 
 // CHECK-EN-LABEL: func @test_insert_same_op
 //  CHECK-EN-SAME:     {pattern_driver_all_erased = false, pattern_driver_changed = true}
 //       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
 }
 
 // -----
 
 // CHECK-EN-LABEL: func @test_replace_with_new_op
 //  CHECK-EN-SAME:     {pattern_driver_all_erased = true, pattern_driver_changed = true}
 //       CHECK-EN:   %[[n:.*]] = "test.new_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
 }
 
 // -----
 
 // CHECK-EN-LABEL: func @test_replace_with_erase_op
 //  CHECK-EN-SAME:     {pattern_driver_all_erased = true, pattern_driver_changed = true}
 //   CHECK-EN-NOT:   test.replace_with_new_op
 //   CHECK-EN-NOT:   test.erase_op
 
 // CHECK-EX-LABEL: func @test_replace_with_erase_op
 //  CHECK-EX-SAME:     {pattern_driver_all_erased = true, pattern_driver_changed = true}
 //   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
 }
 
 // -----
 
 // CHECK-AN-LABEL: func @test_trigger_rewrite_through_block
 //       CHECK-AN: "test.change_block_op"()[^[[BB0:.*]], ^[[BB0]]]
 //       CHECK-AN: return
 //       CHECK-AN: ^[[BB1:[^:]*]]:
 //       CHECK-AN: "test.implicit_change_op"()[^[[BB1]]]
 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] : () -> ()
 }
+
+// -----
+
+// Make sure that test.erase_op is deleted.
+
+// 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} : () -> ()
+    "test.another_op"() : () -> ()
+  }) : () -> ()
+}

mlir/test/lib/Dialect/Test/TestPatterns.cpp

StringRef getDescription() const final {
     return "Test strict mode of pattern driver";
   }
 
   void runOnOperation() override {
     MLIRContext *ctx = &getContext();
     mlir::RewritePatternSet patterns(ctx);
     patterns.add<
         // clang-format off
-        InsertSameOp,
-        ReplaceWithNewOp,
         EraseOp,
         ChangeBlockOp,
-        ImplicitChangeOp
+        CloneRegionInParentOp,
+        ImplicitChangeOp,
+        InsertSameOp,
+        ReplaceWithNewOp
         // 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);
       }
     });
 
     GreedyRewriteConfig config;
     if (strictMode == "AnyOp") {
       config.strictMode = GreedyRewriteStrictness::AnyOp;
     } else if (strictMode == "ExistingAndNewOps") {
public:
   }
 
   Option<std::string> strictMode{
       *this, "strictness",
       llvm::cl::desc("Can be {AnyOp, ExistingAndNewOps, ExistingOps}"),
       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:
     InsertSameOp(MLIRContext *context)
         : RewritePattern("test.insert_same_op", /*benefit=*/1, context) {}
 
     LogicalResult matchAndRewrite(Operation *op,
                                   PatternRewriter &rewriter) const override {