[mlir] Fix worklist bug in MultiOpPatternRewriteDriver

[springerm]

Summary

When strict = true, only pre-existing and newly-created ops are rewritten and/or folded. Such ops are stored in strictModeFilteredOps.

/// Applies the specified rewrite patterns on `ops` while also trying to fold
/// these ops as well as any other ops that were in turn created due to such
/// rewrites. [...]
bool applyOpPatternsAndFold(...);

Newly-created ops were previously added to strictModeFilteredOps after calling addToWorklist (via GreedyPatternRewriteDriver::notifyOperationInserted). Therefore, newly-created ops were never added to the worklist.

Also fix a test case that should have gone into an infinite loop (test.replace_with_new_op was replaced with itself, which should have caused the op to be rewritten over and over), but did not due to this bug.

Related Revision: https://reviews.llvm.org/D127180

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

public:
 
   void addToWorklist(Operation *op) override {
     if (!strictMode || strictModeFilteredOps.contains(op))
       GreedyPatternRewriteDriver::addToWorklist(op);
   }
 
 private:
   void notifyOperationInserted(Operation *op) override {
-    GreedyPatternRewriteDriver::notifyOperationInserted(op);
     if (strictMode)
       strictModeFilteredOps.insert(op);
+    GreedyPatternRewriteDriver::notifyOperationInserted(op);
   }
 
   void notifyOperationRemoved(Operation *op) override {
     GreedyPatternRewriteDriver::notifyOperationRemoved(op);
     if (strictMode)
       strictModeFilteredOps.erase(op);
   }
 

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

 // RUN: mlir-opt -allow-unregistered-dialect -test-strict-pattern-driver %s | FileCheck %s
 
-// CHECK-LABEL: @test_erase
+// CHECK-LABEL: func @test_erase
+//       CHECK:   test.arg0
+//       CHECK:   test.arg1
+//   CHECK-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)
   return
 }
 
-// CHECK-LABEL: @test_insert_same_op
+// CHECK-LABEL: func @test_insert_same_op
+//       CHECK:   "test.insert_same_op"() {skip = true}
+//       CHECK:   "test.insert_same_op"() {skip = true}
 func.func @test_insert_same_op() {
   %0 = "test.insert_same_op"() : () -> (i32)
   return
 }
 
-// CHECK-LABEL: @test_replace_with_same_op
-func.func @test_replace_with_same_op() {
-  %0 = "test.replace_with_same_op"() : () -> (i32)
+// CHECK-LABEL: func @test_replace_with_new_op
+//       CHECK:   %[[n:.*]] = "test.new_op"
+//       CHECK:   "test.dummy_user"(%[[n]])
+//       CHECK:   "test.dummy_user"(%[[n]])
+func.func @test_replace_with_new_op() {
+  %0 = "test.replace_with_new_op"() : () -> (i32)
   %1 = "test.dummy_user"(%0) : (i32) -> (i32)
   %2 = "test.dummy_user"(%0) : (i32) -> (i32)
   return
 }
+
+// CHECK-LABEL: func @test_replace_with_erase_op
+//   CHECK-NOT:   test.replace_with_new_op
+//   CHECK-NOT:   test.erase_op
+func.func @test_replace_with_erase_op() {
+  "test.replace_with_new_op"() {create_erase_op} : () -> ()
+  return
+}

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

public:
 
   StringRef getArgument() const final { return "test-strict-pattern-driver"; }
   StringRef getDescription() const final {
     return "Run strict mode of pattern driver";
   }
 
   void runOnOperation() override {
     mlir::RewritePatternSet patterns(&getContext());
-    patterns.add<InsertSameOp, ReplaceWithSameOp, EraseOp>(&getContext());
+    patterns.add<InsertSameOp, ReplaceWithNewOp, EraseOp>(&getContext());
     SmallVector<Operation *> ops;
     getOperation()->walk([&](Operation *op) {
       StringRef opName = op->getName().getStringRef();
       if (opName == "test.insert_same_op" ||
-          opName == "test.replace_with_same_op" || opName == "test.erase_op") {
+          opName == "test.replace_with_new_op" || opName == "test.erase_op") {
         ops.push_back(op);
       }
     });
 
     // Check if these transformations introduce visiting of operations that
     // are not in the `ops` set (The new created ops are valid). An invalid
     // operation will trigger the assertion while processing.
     (void)applyOpPatternsAndFold(ArrayRef(ops), std::move(patterns),
LogicalResult matchAndRewrite(Operation *op,
       op->setAttr("skip", rewriter.getBoolAttr(true));
       newOp->setAttr("skip", rewriter.getBoolAttr(true));
 
       return success();
     }
   };
 
   // Replace an operation may introduce the re-visiting of its users.
-  class ReplaceWithSameOp : public RewritePattern {
+  class ReplaceWithNewOp : public RewritePattern {
   public:
-    ReplaceWithSameOp(MLIRContext *context)
-        : RewritePattern("test.replace_with_same_op", /*benefit=*/1, context) {}
+    ReplaceWithNewOp(MLIRContext *context)
+        : RewritePattern("test.replace_with_new_op", /*benefit=*/1, context) {}
 
     LogicalResult matchAndRewrite(Operation *op,
                                   PatternRewriter &rewriter) const override {
-      Operation *newOp =
-          rewriter.create(op->getLoc(), op->getName().getIdentifier(),
-                          op->getOperands(), op->getResultTypes());
+      Operation *newOp;
+      if (op->hasAttr("create_erase_op")) {
+        newOp = rewriter.create(
+            op->getLoc(),
+            OperationName("test.erase_op", op->getContext()).getIdentifier(),
+            ValueRange(), TypeRange());
+      } else {
+        newOp = rewriter.create(
+            op->getLoc(),
+            OperationName("test.new_op", op->getContext()).getIdentifier(),
+      }
       rewriter.replaceOp(op, newOp->getResults());
       return success();
     }
   };
 
   // Remove an operation may introduce the re-visiting of its opreands.
   class EraseOp : public RewritePattern {
   public: