[mlir][IR] Use Listener for IR callbacks in OperationFolder

Remove the IR modification callbacks from `OperationFolder`. Instead, an optional `RewriterBase::Listener` can be specified.
* `processGeneratedConstants` => `notifyOperationCreated`
* `preReplaceAction` => `notifyOperationReplaced`

This simplifies the GreedyPatternRewriterDriver because we no longer need special handling for IR modifications due to op folding.

A folded operation is now enqueued on the GreedyPatternRewriteDriver's worklist if it was modified in-place. (There may be new patterns that apply after folding.)

Also fixes a bug in `TestOpInPlaceFold::fold`. The folder could previously be applied over and over and did not return a "null" OpFoldResult if the IR was not modified. (This is similar to a pattern that returns `success` without modifying IR; it can trigger an infinite loop in the GreedyPatternRewriteDriver.)

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

1 files changed, 5 insertions, 30 deletions

diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 15495ef..a5977cc 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -127,7 +127,8 @@ private:
 GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
     MLIRContext *ctx, const FrozenRewritePatternSet &patterns,
     const GreedyRewriteConfig &config)
-    : PatternRewriter(ctx), folder(ctx), config(config), matcher(patterns) {
+    : PatternRewriter(ctx), folder(ctx, this), config(config),
+      matcher(patterns) {
   worklist.reserve(64);
 
   // Apply a simple cost model based solely on pattern benefit.
@@ -156,9 +157,6 @@ bool GreedyPatternRewriteDriver::processWorklist() {
   };
 #endif
 
-  // These are scratch vectors used in the folding loop below.
-  SmallVector<Value, 8> originalOperands;
-
   bool changed = false;
   int64_t numRewrites = 0;
   while (!worklist.empty() &&
@@ -197,34 +195,11 @@ bool GreedyPatternRewriteDriver::processWorklist() {
       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.
-      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);
-    };
-
-    // Add the given operation to the worklist.
-    auto collectOps = [this](Operation *op) { addToWorklist(op); };
-
     // Try to fold this op.
-    bool inPlaceUpdate;
-    if ((succeeded(folder.tryToFold(op, collectOps, preReplaceAction,
-                                    &inPlaceUpdate)))) {
+    if (succeeded(folder.tryToFold(op))) {
       LLVM_DEBUG(logResultWithLine("success", "operation was folded"));
-
       changed = true;
-      if (!inPlaceUpdate)
-        continue;
+      continue;
     }
 
     // Try to match one of the patterns. The rewriter is automatically
@@ -465,7 +440,7 @@ LogicalResult RegionPatternRewriteDriver::simplify() && {
       // Add all nested operations to the worklist in preorder.
       region.walk<WalkOrder::PreOrder>([&](Operation *op) {
         if (!insertKnownConstant(op)) {
-          worklist.push_back(op);
+          addToWorklist(op);
           return WalkResult::advance();
         }
         return WalkResult::skip();