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

[springerm]

Summary

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.)

Depends On: D143380

[springerm]

update

mlir/include/mlir/Transforms/FoldUtils.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TRANSFORMS_FOLDUTILS_H
 #define MLIR_TRANSFORMS_FOLDUTILS_H
 
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectInterface.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/FoldInterfaces.h"
 
 namespace mlir {
 class Operation;
 class Value;
 
 //===--------------------------------------------------------------------===//
 // OperationFolder
 //===--------------------------------------------------------------------===//
 
 /// A utility class for folding operations, and unifying duplicated constants
 /// generated along the way.
 class OperationFolder {
 public:
-  OperationFolder(MLIRContext *ctx) : interfaces(ctx) {}
+  OperationFolder(MLIRContext *ctx, RewriterBase::Listener *listener = nullptr)
+      : interfaces(ctx), listener(listener) {}
 
   /// Tries to perform folding on the given `op`, including unifying
   /// deduplicated constants. If successful, replaces `op`'s uses with
-  /// folded results, and returns success. `preReplaceAction` is invoked on `op`
-  /// before it is replaced. 'processGeneratedConstants' is invoked for any new
-  /// operations generated when folding. If the op was completely folded it is
+  /// folded results, and returns success. If the op was completely folded it is
   /// erased. If it is just updated in place, `inPlaceUpdate` is set to true.
-  LogicalResult
-  tryToFold(Operation *op,
-            function_ref<void(Operation *)> processGeneratedConstants = nullptr,
-            function_ref<void(Operation *)> preReplaceAction = nullptr,
-            bool *inPlaceUpdate = nullptr);
+  LogicalResult tryToFold(Operation *op, bool *inPlaceUpdate = nullptr);
 
   /// Tries to fold a pre-existing constant operation. `constValue` represents
   /// the value of the constant, and can be optionally passed if the value is
   /// already known (e.g. if the constant was discovered by m_Constant). This is
   /// purely an optimization opportunity for callers that already know the value
   /// of the constant. Returns false if an existing constant for `op` already
   /// exists in the folder, in which case `op` is replaced and erased.
   /// Otherwise, returns true and `op` is inserted into the folder (and

 private:
   /// This map keeps track of uniqued constants by dialect, attribute, and type.
   /// A constant operation materializes an attribute with a type. Dialects may
   /// generate different constants with the same input attribute and type, so we
   /// also need to track per-dialect.
   using ConstantMap =
       DenseMap<std::tuple<Dialect *, Attribute, Type>, Operation *>;
 
+  /// Erase the given operation and notify the listener.
+  void eraseOp(Operation *op);
+
   /// Returns true if the given operation is an already folded constant that is
   /// owned by this folder.
   bool isFolderOwnedConstant(Operation *op) const;
 
   /// Tries to perform folding on the given `op`. If successful, populates
   /// `results` with the results of the folding.
-  LogicalResult tryToFold(
-      OpBuilder &builder, Operation *op, SmallVectorImpl<Value> &results,
-      function_ref<void(Operation *)> processGeneratedConstants = nullptr);
+  LogicalResult tryToFold(OpBuilder &builder, Operation *op,
+                          SmallVectorImpl<Value> &results);
 
   /// Try to process a set of fold results, generating constants as necessary.
   /// Populates `results` on success, otherwise leaves it unchanged.
-  LogicalResult
-  processFoldResults(OpBuilder &builder, Operation *op,
+  LogicalResult processFoldResults(OpBuilder &builder, Operation *op,
                                    SmallVectorImpl<Value> &results,
-                     ArrayRef<OpFoldResult> foldResults,
-                     function_ref<void(Operation *)> processGeneratedConstants);
+                                   ArrayRef<OpFoldResult> foldResults);
 
   /// Try to get or create a new constant entry. On success this returns the
   /// constant operation, nullptr otherwise.
   Operation *tryGetOrCreateConstant(ConstantMap &uniquedConstants,
                                     Dialect *dialect, OpBuilder &builder,
                                     Attribute value, Type type, Location loc);
 
   /// A mapping between an insertion region and the constants that have been
   /// created within it.
   DenseMap<Region *, ConstantMap> foldScopes;
 
   /// This map tracks all of the dialects that an operation is referenced by;
   /// given that many dialects may generate the same constant.
   DenseMap<Operation *, SmallVector<Dialect *, 2>> referencedDialects;
 
   /// A collection of dialect folder interfaces.
   DialectInterfaceCollection<DialectFoldInterface> interfaces;
+
+  /// An optional listener that is notified of all IR changes.
+  RewriterBase::Listener *listener = nullptr;
 };
 
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_FOLDUTILS_H

mlir/lib/Transforms/Utils/FoldUtils.cpp

static Operation *materializeConstant(Dialect *dialect, OpBuilder &builder,
 
   return nullptr;
 }
 
 //===----------------------------------------------------------------------===//
 // OperationFolder
 //===----------------------------------------------------------------------===//
 
-LogicalResult OperationFolder::tryToFold(
-    Operation *op, function_ref<void(Operation *)> processGeneratedConstants,
-    function_ref<void(Operation *)> preReplaceAction, bool *inPlaceUpdate) {
+LogicalResult OperationFolder::tryToFold(Operation *op, bool *inPlaceUpdate) {
   if (inPlaceUpdate)
     *inPlaceUpdate = false;
 
   // If this is a unique'd constant, return failure as we know that it has
   // already been folded.
   if (isFolderOwnedConstant(op)) {
     // Check to see if we should rehoist, i.e. if a non-constant operation was
     // inserted before this one.
     Block *opBlock = op->getBlock();
     if (&opBlock->front() != op && !isFolderOwnedConstant(op->getPrevNode()))
       op->moveBefore(&opBlock->front());
     return failure();
   }
 
   // Try to fold the operation.
   SmallVector<Value, 8> results;
-  OpBuilder builder(op);
-  if (failed(tryToFold(builder, op, results, processGeneratedConstants)))
+  OpBuilder builder(op, listener);
+  if (failed(tryToFold(builder, op, results)))
     return failure();
 
   // Check to see if the operation was just updated in place.
   if (results.empty()) {
     if (inPlaceUpdate)
       *inPlaceUpdate = true;
+    if (listener)
+      listener->notifyOperationModified(op);
     return success();
   }
 
-  // Constant folding succeeded. We will start replacing this op's uses and
-  // erase this op. Invoke the callback provided by the caller to perform any
-  // pre-replacement action.
-  if (preReplaceAction)
-    preReplaceAction(op);
-
-  // Replace all of the result values and erase the operation.
+  // Constant folding succeeded. Replace all of the result values and erase the
+  // operation.
+  if (listener)
+    listener->notifyOperationReplaced(op, results);
   for (unsigned i = 0, e = results.size(); i != e; ++i)
     op->getResult(i).replaceAllUsesWith(results[i]);
-  op->erase();
+  eraseOp(op);
   return success();
 }
 
 bool OperationFolder::insertKnownConstant(Operation *op, Attribute constValue) {
   Block *opBlock = op->getBlock();
 
   // If this is a constant we unique'd, we don't need to insert, but we can
   // check to see if we should rehoist it.
#endif
   // Check for an existing constant operation for the attribute value.
   Region *insertRegion = getInsertionRegion(interfaces, opBlock);
   auto &uniquedConstants = foldScopes[insertRegion];
   Operation *&folderConstOp = uniquedConstants[std::make_tuple(
       op->getDialect(), constValue, *op->result_type_begin())];
 
   // If there is an existing constant, replace `op`.
   if (folderConstOp) {
+    if (listener)
+      listener->notifyOperationReplaced(op, folderConstOp->getResults());
     op->replaceAllUsesWith(folderConstOp);
-    op->erase();
+    eraseOp(op);
     return false;
   }
 
   // Otherwise, we insert `op`. If `op` is in the insertion block and is either
   // already at the front of the block, or the previous operation is already a
   // constant we unique'd (i.e. one we inserted), then we don't need to do
   // anything. Otherwise, we move the constant to the insertion block.
   Block *insertBlock = &insertRegion->front();
   if (opBlock != insertBlock || (&insertBlock->front() != op &&
                                  !isFolderOwnedConstant(op->getPrevNode())))
     op->moveBefore(&insertBlock->front());
 
   folderConstOp = op;
   referencedDialects[op].push_back(op->getDialect());
   return true;
 }
 
+void OperationFolder::eraseOp(Operation *op) {
+  notifyRemoval(op);
+  if (listener)
+    listener->notifyOperationRemoved(op);
+  op->erase();
+}
+
 /// Notifies that the given constant `op` should be remove from this
 /// OperationFolder's internal bookkeeping.
 void OperationFolder::notifyRemoval(Operation *op) {
   // Check to see if this operation is uniqued within the folder.
   auto it = referencedDialects.find(op);
   if (it == referencedDialects.end())
     return;
 

 }
 
 bool OperationFolder::isFolderOwnedConstant(Operation *op) const {
   return referencedDialects.count(op);
 }
 
 /// Tries to perform folding on the given `op`. If successful, populates
 /// `results` with the results of the folding.
-LogicalResult OperationFolder::tryToFold(
-    OpBuilder &builder, Operation *op, SmallVectorImpl<Value> &results,
-    function_ref<void(Operation *)> processGeneratedConstants) {
+LogicalResult OperationFolder::tryToFold(OpBuilder &builder, Operation *op,
+                                         SmallVectorImpl<Value> &results) {
   SmallVector<Attribute, 8> operandConstants;
 
   // If this is a commutative operation, move constants to be trailing operands.
   bool updatedOpOperands = false;
   if (op->getNumOperands() >= 2 && op->hasTrait<OpTrait::IsCommutative>()) {
     auto isNonConstant = [&](OpOperand &o) {
       return !matchPattern(o.get(), m_Constant());
     };
LogicalResult OperationFolder::tryToFold(OpBuilder &builder, Operation *op,
   for (unsigned i = 0, e = op->getNumOperands(); i != e; ++i)
     matchPattern(op->getOperand(i), m_Constant(&operandConstants[i]));
 
   // Attempt to constant fold the operation. If we failed, check to see if we at
   // least updated the operands of the operation. We treat this as an in-place
   // fold.
   SmallVector<OpFoldResult, 8> foldResults;
   if (failed(op->fold(operandConstants, foldResults)) ||
-      failed(processFoldResults(builder, op, results, foldResults,
-                                processGeneratedConstants)))
+      failed(processFoldResults(builder, op, results, foldResults)))
     return success(updatedOpOperands);
   return success();
 }
 
-LogicalResult OperationFolder::processFoldResults(
-    OpBuilder &builder, Operation *op, SmallVectorImpl<Value> &results,
-    ArrayRef<OpFoldResult> foldResults,
-    function_ref<void(Operation *)> processGeneratedConstants) {
+LogicalResult
+OperationFolder::processFoldResults(OpBuilder &builder, Operation *op,
+                                    SmallVectorImpl<Value> &results,
+                                    ArrayRef<OpFoldResult> foldResults) {
   // Check to see if the operation was just updated in place.
   if (foldResults.empty())
     return success();
   assert(foldResults.size() == op->getNumResults());
 
   // Create a builder to insert new operations into the entry block of the
   // insertion region.
   auto *insertRegion =
if (auto *constOp =
         constOp->moveBefore(&opBlock->front());
 
       results.push_back(constOp->getResult(0));
       continue;
     }
     // If materialization fails, cleanup any operations generated for the
     // previous results and return failure.
     for (Operation &op : llvm::make_early_inc_range(
-             llvm::make_range(entry.begin(), builder.getInsertionPoint()))) {
-      notifyRemoval(&op);
-      op.erase();
-    }
+             llvm::make_range(entry.begin(), builder.getInsertionPoint())))
+      eraseOp(&op);
+
     results.clear();
     return failure();
   }
 
-  // Process any newly generated operations.
-  if (processGeneratedConstants) {
-    for (auto i = entry.begin(), e = builder.getInsertionPoint(); i != e; ++i)
-      processGeneratedConstants(&*i);
-  }
-
   return success();
 }
 
 /// Try to get or create a new constant entry. On success this returns the
 /// constant operation value, nullptr otherwise.
 Operation *OperationFolder::tryGetOrCreateConstant(
     ConstantMap &uniquedConstants, Dialect *dialect, OpBuilder &builder,
     Attribute value, Type type, Location loc) {
Operation *OperationFolder::tryGetOrCreateConstant(
 
   // If it isn't, then we also need to make sure that the mapping for the new
   // dialect is valid.
   auto newKey = std::make_tuple(newDialect, value, type);
 
   // If an existing operation in the new dialect already exists, delete the
   // materialized operation in favor of the existing one.
   if (auto *existingOp = uniquedConstants.lookup(newKey)) {
-    constOp->erase();
+    eraseOp(constOp);
     referencedDialects[existingOp].push_back(dialect);
     return constOp = existingOp;
   }
 
   // Otherwise, update the new dialect to the materialized operation.
   referencedDialects[constOp].assign({dialect, newDialect});
   auto newIt = uniquedConstants.insert({newKey, constOp});
   return newIt.first->second;
 }

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

#endif
   /// The low-level pattern applicator.
   PatternApplicator matcher;
 };
 } // namespace
 
 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.
   matcher.applyDefaultCostModel();
 
   // Set up listener.
   setListener(this);
 }
auto logResult = [&](StringRef result, const llvm::Twine &msg = {}) {
     logger.getOStream() << "\n";
   };
   auto logResultWithLine = [&](StringRef result, const llvm::Twine &msg = {}) {
     logResult(result, msg);
     logger.startLine() << logLineComment;
   };
 #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() &&
          (numRewrites < config.maxNumRewrites ||
           config.maxNumRewrites == GreedyRewriteConfig::kNoLimit)) {
     auto *op = popFromWorklist();
 
     // Nulls get added to the worklist when operations are removed, ignore
if (isOpTriviallyDead(op)) {
       notifyOperationRemoved(op);
       op->erase();
       changed = true;
 
       LLVM_DEBUG(logResultWithLine("success", "operation is trivially dead"));
       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;
     }
 
     // Try to match one of the patterns. The rewriter is automatically
     // notified of any necessary changes, so there is nothing else to do
     // here.
 #ifndef NDEBUG
       auto canApply = [&](const Pattern &pattern) {
         LLVM_DEBUG({
if (!config.useTopDownTraversal) {
       region.walk([&](Operation *op) {
         if (!insertKnownConstant(op))
           addToWorklist(op);
       });
     } else {
       // 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();
       });
 
       // Reverse the list so our pop-back loop processes them in-order.
       std::reverse(worklist.begin(), worklist.end());
       // Remember the reverse index.

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

LogicalResult TestOpWithVariadicResultsAndFolder::fold(
     FoldAdaptor adaptor, SmallVectorImpl<OpFoldResult> &results) {
   for (Value input : this->getOperands()) {
     results.push_back(input);
   }
   return success();
 }
 
 OpFoldResult TestOpInPlaceFold::fold(FoldAdaptor adaptor) {
-  if (adaptor.getOp()) {
+  if (adaptor.getOp() && !(*this)->hasAttr("attr")) {
+    // The folder adds "attr" if not present.
     (*this)->setAttr("attr", adaptor.getOp());
     return getResult();
   }
   return {};
 }
 
 OpFoldResult TestPassthroughFold::fold(FoldAdaptor adaptor) {
   return getOperand();

mlir/test/lib/Dialect/Test/TestOps.td

 }
 
 def TestOpInPlaceFoldAnchor : TEST_Op<"op_in_place_fold_anchor"> {
   let arguments = (ins I32);
   let results = (outs I32);
 }
 
 def TestOpInPlaceFold : TEST_Op<"op_in_place_fold"> {
-  let arguments = (ins I32:$op, I32Attr:$attr);
+  let arguments = (ins I32:$op, OptionalAttr<I32Attr>:$attr);
   let results = (outs I32);
   let hasFolder = 1;
 }
 
 // Test op that simply returns success.
 def TestOpInPlaceFoldSuccess : TEST_Op<"op_in_place_fold_success"> {
   let results = (outs Variadic<I1>);
   let hasFolder = 1;

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

LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
     // Exercise OperationFolder API for a single-result operation that is folded
     // upon construction. The operation being created through the folder has an
     // in-place folder, and it should be still present in the output.
     // Furthermore, the folder should not crash when attempting to recover the
     // (unchanged) operation result.
     OperationFolder folder(op->getContext());
     Value result = folder.create<TestOpInPlaceFold>(
-        rewriter, op->getLoc(), rewriter.getIntegerType(32), op->getOperand(0),
-        rewriter.getI32IntegerAttr(0));
+        rewriter, op->getLoc(), rewriter.getIntegerType(32), op->getOperand(0));
     assert(result);
     rewriter.replaceOp(op, result);
     return success();
   }
 };
 
 /// This pattern creates a foldable operation at the entry point of the block.
 /// This tests the situation where the operation folder will need to replace an

mlir/test/lib/Transforms/TestConstantFold.cpp

 //===- TestConstantFold.cpp - Pass to test constant folding ---------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/FoldUtils.h"
 
 using namespace mlir;
 
 namespace {
 /// Simple constant folding pass.
-struct TestConstantFold
-    : public PassWrapper<TestConstantFold, OperationPass<>> {
+struct TestConstantFold : public PassWrapper<TestConstantFold, OperationPass<>>,
+                          public RewriterBase::Listener {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestConstantFold)
 
   StringRef getArgument() const final { return "test-constant-fold"; }
   StringRef getDescription() const final {
     return "Test operation constant folding";
   }
   // All constants in the operation post folding.
   SmallVector<Operation *> existingConstants;
 
   void foldOperation(Operation *op, OperationFolder &helper);
   void runOnOperation() override;
-};
-} // namespace
 
-void TestConstantFold::foldOperation(Operation *op, OperationFolder &helper) {
-  auto processGeneratedConstants = [this](Operation *op) {
+  void notifyOperationInserted(Operation *op) override {
     existingConstants.push_back(op);
-  };
+  }
+  void notifyOperationRemoved(Operation *op) override {
+    auto it = llvm::find(existingConstants, op);
+    if (it != existingConstants.end())
+      existingConstants.erase(it);
+  }
+} // namespace
 
+void TestConstantFold::foldOperation(Operation *op, OperationFolder &helper) {
   // Attempt to fold the specified operation, including handling unused or
   // duplicated constants.
-  (void)helper.tryToFold(op, processGeneratedConstants);
+  (void)helper.tryToFold(op);
 }
 
 void TestConstantFold::runOnOperation() {
   existingConstants.clear();
 
   // Collect and fold the operations within the operation.
   SmallVector<Operation *, 8> ops;
   getOperation()->walk([&](Operation *op) { ops.push_back(op); });
 
   // Fold the constants in reverse so that the last generated constants from
   // folding are at the beginning. This creates somewhat of a linear ordering to
   // the newly generated constants that matches the operation order and improves
   // the readability of test cases.
-  OperationFolder helper(&getContext());
+  OperationFolder helper(&getContext(), /*listener=*/this);
   for (Operation *op : llvm::reverse(ops))
     foldOperation(op, helper);
 
   // By the time we are done, we may have simplified a bunch of code, leaving
   // around dead constants. Check for them now and remove them.
   for (auto *cst : existingConstants) {
     if (cst->use_empty())
       cst->erase();
   }
 }
 
 namespace mlir {
 namespace test {
 void registerTestConstantFold() { PassRegistration<TestConstantFold>(); }
 } // namespace test
 } // namespace mlir