[Mlir-commits] [mlir] 2aeffc6 - [mlir:MultiOpDriver] Don't add ops which are not in the allowed list

[Chia-hung Duan]

Author: Chia-hung Duan
Date: 2022-06-02T18:27:37Z
New Revision: 2aeffc6d8ddf3a6eab3b0481869b73678d91856c

URL: https://github.com/llvm/llvm-project/commit/2aeffc6d8ddf3a6eab3b0481869b73678d91856c
DIFF: https://github.com/llvm/llvm-project/commit/2aeffc6d8ddf3a6eab3b0481869b73678d91856c.diff

LOG: [mlir:MultiOpDriver] Don't add ops which are not in the allowed list

In strict mode, only the new inserted operation is allowed to add to the
worklist. Before this change, it would add the users of a replaced op
and it didn't check if the users are allowed to be pushed into the
worklist

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D126899

Added: 
    

Modified: 
    mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index a08657cb7f6eb..be7924e153088 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -539,12 +539,27 @@ class MultiOpPatternRewriteDriver : public GreedyPatternRewriteDriver {
     }
   }
 
+  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
@@ -592,6 +607,8 @@ bool MultiOpPatternRewriteDriver::simplifyLocally(ArrayRef<Operation *> ops) {
   SmallVector<Value, 8> originalOperands, resultValues;
   while (!worklist.empty()) {
     Operation *op = popFromWorklist();
+    assert((!strictMode || strictModeFilteredOps.contains(op)) &&
+           "unexpected op was inserted under strict mode");
 
     // Nulls get added to the worklist when operations are removed, ignore
     // them.