[mlir:MultiOpDriver] Add operands to worklist should be checked

[Chia-hungDuan]

Summary

Operand's defining op may not be valid for adding to the worklist under
stict mode

[Chia-hungDuan]

@rriddle
I'm not sure if we should do the check in AddToWorklist or just make it as this CL. Besides, I found this bug by the assertion we had before. Do you think we need a test (testing remove/replace/insert operations) or still have the assertion is enough?
What do you think?

[rriddle]

Is there a test we can add here?

[rriddle]

In D127180#3562404, @Chia-hungDuan wrote:

@rriddle
I'm not sure if we should do the check in AddToWorklist or just make it as this CL. Besides, I found this bug by the assertion we had before. Do you think we need a test (testing remove/replace/insert operations) or still have the assertion is enough?
What do you think?

Ahah, missed this comment before I sent the other one. What do you mean by if we should do the check in AddToWorklist? Also, I think it'd be nice to have a test for this if we can.

[Chia-hungDuan]

In D127180#3562470, @rriddle wrote:

In D127180#3562404, @Chia-hungDuan wrote:

@rriddle
I'm not sure if we should do the check in AddToWorklist or just make it as this CL. Besides, I found this bug by the assertion we had before. Do you think we need a test (testing remove/replace/insert operations) or still have the assertion is enough?
What do you think?

Ahah, missed this comment before I sent the other one. What do you mean by if we should do the check in AddToWorklist? Also, I think it'd be nice to have a test for this if we can.

Sorry, didn't say that clearly. I mean things this like

virtual void AddToWorklist(...); // in GreedyPatternRewriteDriver
void AddToWorklist(...) override { // in MultiOpPatternRewriteDriver
  if (strictMode) ...
}

Then we may only need to override notifyOperationRemoved. Not sure which one is better in term of design.

Will add a test in later patch.

[rriddle]

In D127180#3562509, @Chia-hungDuan wrote:

In D127180#3562470, @rriddle wrote:

In D127180#3562404, @Chia-hungDuan wrote:

@rriddle
I'm not sure if we should do the check in AddToWorklist or just make it as this CL. Besides, I found this bug by the assertion we had before. Do you think we need a test (testing remove/replace/insert operations) or still have the assertion is enough?
What do you think?

Ahah, missed this comment before I sent the other one. What do you mean by if we should do the check in AddToWorklist? Also, I think it'd be nice to have a test for this if we can.

Sorry, didn't say that clearly. I mean things this like

virtual void AddToWorklist(...); // in GreedyPatternRewriteDriver
void AddToWorklist(...) override { // in MultiOpPatternRewriteDriver
  if (strictMode) ...
}

Then we may only need to override notifyOperationRemoved. Not sure which one is better in term of design.

Yeah, let's try switching to that. From a glance that feels like it could be cleaner.

Will add a test in later patch.

[Chia-hungDuan]

Address review comment

[Chia-hungDuan]

Run linter

[Chia-hungDuan]

h

In D127180#3562533, @rriddle wrote:

In D127180#3562509, @Chia-hungDuan wrote:

In D127180#3562470, @rriddle wrote:

In D127180#3562404, @Chia-hungDuan wrote:

@rriddle
I'm not sure if we should do the check in AddToWorklist or just make it as this CL. Besides, I found this bug by the assertion we had before. Do you think we need a test (testing remove/replace/insert operations) or still have the assertion is enough?
What do you think?

Ahah, missed this comment before I sent the other one. What do you mean by if we should do the check in AddToWorklist? Also, I think it'd be nice to have a test for this if we can.

Sorry, didn't say that clearly. I mean things this like

virtual void AddToWorklist(...); // in GreedyPatternRewriteDriver
void AddToWorklist(...) override { // in MultiOpPatternRewriteDriver
  if (strictMode) ...
}

Then we may only need to override notifyOperationRemoved. Not sure which one is better in term of design.

Yeah, let's try switching to that. From a glance that feels like it could be cleaner.

Done.

Will add a test in later patch.

Done.

[Chia-hungDuan]

Last commit missed some changes

[Chia-hungDuan]

comment minor update

[Chia-hungDuan]

@rriddle , a gentle ping!

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

public:
   explicit GreedyPatternRewriteDriver(MLIRContext *ctx,
                                       const FrozenRewritePatternSet &patterns,
                                       const GreedyRewriteConfig &config);
 
   /// Simplify the operations within the given regions.
   bool simplify(MutableArrayRef<Region> regions);
 
   /// Add the given operation to the worklist.
-  void addToWorklist(Operation *op);
+  virtual void addToWorklist(Operation *op);
 
   /// Pop the next operation from the worklist.
   Operation *popFromWorklist();
 
   /// If the specified operation is in the worklist, remove it.
   void removeFromWorklist(Operation *op);
 
 protected:
   // Implement the hook for inserting operations, and make sure that newly
   // inserted ops are added to the worklist for processing.
   void notifyOperationInserted(Operation *op) override;
 
   // Look over the provided operands for any defining operations that should
   // be re-added to the worklist. This function should be called when an
   // operation is modified or removed, as it may trigger further
   // simplifications.
-  template <typename Operands>
-  void addToWorklist(Operands &&operands);
+  void addOperandsToWorklist(ValueRange operands);
 
   // If an operation is about to be removed, make sure it is not in our
   // worklist anymore because we'd get dangling references to it.
   void notifyOperationRemoved(Operation *op) override;
 
   // When the root of a pattern is about to be replaced, it can trigger
   // simplifications to its users - make sure to add them to the worklist
   // before the root is changed.
while (!worklist.empty()) {
         continue;
       }
 
       // Collects all the operands and result uses of the given `op` into work
       // list. Also remove `op` and nested ops from worklist.
       originalOperands.assign(op->operand_begin(), op->operand_end());
       auto preReplaceAction = [&](Operation *op) {
         // Add the operands to the worklist for visitation.
-        addToWorklist(originalOperands);
+        addOperandsToWorklist(originalOperands);
 
         // Add all the users of the result to the worklist so we make sure
         // to revisit them.
         for (auto result : op->getResults())
           for (auto *userOp : result.getUsers())
             addToWorklist(userOp);
 
         notifyOperationRemoved(op);

 void GreedyPatternRewriteDriver::notifyOperationInserted(Operation *op) {
   LLVM_DEBUG({
     logger.startLine() << "** Insert  : '" << op->getName() << "'(" << op
                        << ")\n";
   });
   addToWorklist(op);
 }
 
-template <typename Operands>
-void GreedyPatternRewriteDriver::addToWorklist(Operands &&operands) {
+void GreedyPatternRewriteDriver::addOperandsToWorklist(ValueRange operands) {
   for (Value operand : operands) {
     // If the use count of this operand is now < 2, we re-add the defining
     // operation to the worklist.
     // TODO: This is based on the fact that zero use operations
     // may be deleted, and that single use values often have more
     // canonicalization opportunities.
     if (!operand || (!operand.use_empty() && !operand.hasOneUse()))
       continue;
     if (auto *defOp = operand.getDefiningOp())
       addToWorklist(defOp);
   }
 }
 
 void GreedyPatternRewriteDriver::notifyOperationRemoved(Operation *op) {
-  addToWorklist(op->getOperands());
+  addOperandsToWorklist(op->getOperands());
   op->walk([this](Operation *operation) {
     removeFromWorklist(operation);
     folder.notifyRemoval(operation);
   });
 }
 
 void GreedyPatternRewriteDriver::notifyRootReplaced(Operation *op) {
   LLVM_DEBUG({
public:
   explicit MultiOpPatternRewriteDriver(MLIRContext *ctx,
                                        const FrozenRewritePatternSet &patterns,
                                        bool strict)
       : GreedyPatternRewriteDriver(ctx, patterns, GreedyRewriteConfig()),
         strictMode(strict) {}
 
   bool simplifyLocally(ArrayRef<Operation *> op);
 
-private:
-  // Look over the provided operands for any defining operations that should
-  // be re-added to the worklist. This function should be called when an
-  // operation is modified or removed, as it may trigger further
-  // simplifications. If `strict` is set to true, only ops in
-  // `strictModeFilteredOps` are considered.
-  template <typename Operands>
-  void addOperandsToWorklist(Operands &&operands) {
-    for (Value operand : operands) {
-      if (auto *defOp = operand.getDefiningOp()) {
-        if (!strictMode || strictModeFilteredOps.contains(defOp))
-          addToWorklist(defOp);
-      }
-    }
+  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);
   }
 
   void notifyOperationRemoved(Operation *op) override {
     GreedyPatternRewriteDriver::notifyOperationRemoved(op);
     if (strictMode)
       strictModeFilteredOps.erase(op);
   }
 
-  void notifyRootReplaced(Operation *op) override {
-    for (auto result : op->getResults()) {
-      for (auto *user : result.getUsers()) {
-        if (!strictMode || strictModeFilteredOps.contains(user))
-          addToWorklist(user);
-      }
-    }
-  }
-
   /// If `strictMode` is true, any pre-existing ops outside of
   /// `strictModeFilteredOps` remain completely untouched by the rewrite driver.
   /// If `strictMode` is false, operations that use results of (or supply
   /// operands to) any rewritten ops stemming from the simplification of the
   /// provided ops are in turn simplified; any other ops still remain untouched
   /// (i.e., regardless of `strictMode`).
   bool strictMode = false;
 
while (!worklist.empty()) {
     // list. Also remove `op` and nested ops from worklist.
     originalOperands.assign(op->operand_begin(), op->operand_end());
     auto preReplaceAction = [&](Operation *op) {
       // Add the operands to the worklist for visitation.
       addOperandsToWorklist(originalOperands);
 
       // Add all the users of the result to the worklist so we make sure
       // to revisit them.
-      for (Value result : op->getResults())
-        for (Operation *userOp : result.getUsers()) {
-          if (!strictMode || strictModeFilteredOps.contains(userOp))
-            addToWorklist(userOp);
-        }
+      for (Value result : op->getResults()) {
+        for (Operation *userOp : result.getUsers())
+
       notifyOperationRemoved(op);
     };
 
     // Add the given operation generated by the folder to the worklist.
     auto processGeneratedConstants = [this](Operation *op) {
-      // Newly created ops are also simplified -- these are also "local".
-      addToWorklist(op);
-      // When strict mode is off, we don't need to maintain
-      // strictModeFilteredOps.
-      if (strictMode)
-        strictModeFilteredOps.insert(op);
+      notifyOperationInserted(op);
     };
 
     // Try to fold this op.
     bool inPlaceUpdate;
     if (succeeded(folder.tryToFold(op, processGeneratedConstants,
                                    preReplaceAction, &inPlaceUpdate))) {
       changed = true;
       if (!inPlaceUpdate) {

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
+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
+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)
+  %1 = "test.dummy_user"(%0) : (i32) -> (i32)
+  %2 = "test.dummy_user"(%0) : (i32) -> (i32)
+  return
+}

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

(void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns),
                                        config);
   }
 
   Option<bool> useTopDownTraversal{
       *this, "top-down",
       llvm::cl::desc("Seed the worklist in general top-down order"),
       llvm::cl::init(GreedyRewriteConfig().useTopDownTraversal)};
 };
+
+struct TestStrictPatternDriver
+    : public PassWrapper<TestStrictPatternDriver, OperationPass<func::FuncOp>> {
+public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestStrictPatternDriver)
+
+  TestStrictPatternDriver() = default;
+  TestStrictPatternDriver(const TestStrictPatternDriver &other)
+      : PassWrapper(other) {}
+
+  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());
+    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") {
+        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(makeArrayRef(ops), std::move(patterns),
+                                 /*strict=*/true);
+  }
+
+private:
+  // 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 {
+      if (op->hasAttr("skip"))
+        return failure();
+
+      Operation *newOp =
+          rewriter.create(op->getLoc(), op->getName().getIdentifier(),
+                          op->getOperands(), op->getResultTypes());
+      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 {
+  public:
+    ReplaceWithSameOp(MLIRContext *context)
+        : RewritePattern("test.replace_with_same_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());
+      rewriter.replaceOp(op, newOp->getResults());
+      return success();
+    }
+  };
+
+  // Remove an operation may introduce the re-visiting of its opreands.
+  class EraseOp : public RewritePattern {
+  public:
+    EraseOp(MLIRContext *context)
+        : RewritePattern("test.erase_op", /*benefit=*/1, context) {}
+    LogicalResult matchAndRewrite(Operation *op,
+                                  PatternRewriter &rewriter) const override {
+      rewriter.eraseOp(op);
+      return success();
+    }
+  };
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // ReturnType Driver.
 //===----------------------------------------------------------------------===//
 
 namespace {
 // Generate ops for each instance where the type can be successfully inferred.

 namespace mlir {
 namespace test {
 void registerPatternsTestPass() {
   PassRegistration<TestReturnTypeDriver>();
 
   PassRegistration<TestDerivedAttributeDriver>();
 
   PassRegistration<TestPatternDriver>();
+  PassRegistration<TestStrictPatternDriver>();
 
   PassRegistration<TestLegalizePatternDriver>([] {
     return std::make_unique<TestLegalizePatternDriver>(legalizerConversionMode);
   });
 
   PassRegistration<TestRemappedValue>();
 
   PassRegistration<TestUnknownRootOpDriver>();