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

[Chia-hungDuan]

Summary

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

[Chia-hungDuan]

That assertion helps identifying the places we didn't handle properly. So I didn't add additional tests.

[rriddle]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

542

I'd just drop this comment.

611

[Chia-hungDuan]

Address review comment

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

void addOperandsToWorklist(Operands &&operands) {
     for (Value operand : operands) {
       if (auto *defOp = operand.getDefiningOp()) {
         if (!strictMode || strictModeFilteredOps.contains(defOp))
           addToWorklist(defOp);
       }
     }
   }
 
+  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;
 
bool MultiOpPatternRewriteDriver::simplifyLocally(ArrayRef<Operation *> ops) {
   worklistMap.clear();
   for (Operation *op : ops)
     addToWorklist(op);
 
   // These are scratch vectors used in the folding loop below.
   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.
     if (op == nullptr)
       continue;
 
     // If the operation is trivially dead - remove it.
     if (isOpTriviallyDead(op)) {