[mlir][Transforms][NFC] Improve builder/listener API of OperationFolder

[springerm]

Summary

The constructor of OperationFolder takes a listener. Therefore, the remaining API should not take any builder/rewriters. This could lead to double notifications in case a listener is attached to the builder/rewriter.

As an internal cleanup, OperationFolder now has an IRRewriter instead of a RewriterBase::Listener. In most cases, OperationFolder no longer has to notify/deal with listeners. This is done by the rewriter.

mlir/include/mlir/IR/PatternMatch.h

#endif
   }
   template <typename ArgT>
   LogicalResult notifyMatchFailure(ArgT &&arg, const char *msg) {
     return notifyMatchFailure(std::forward<ArgT>(arg), Twine(msg));
   }
 
 protected:
   /// Initialize the builder.
-  explicit RewriterBase(MLIRContext *ctx) : OpBuilder(ctx) {}
+  explicit RewriterBase(MLIRContext *ctx,
+                        OpBuilder::Listener *listener = nullptr)
+      : OpBuilder(ctx, listener) {}
   explicit RewriterBase(const OpBuilder &otherBuilder)
       : OpBuilder(otherBuilder) {}
   virtual ~RewriterBase();
 
 private:
   void operator=(const RewriterBase &) = delete;
   RewriterBase(const RewriterBase &) = delete;
 
   /// 'op' and 'newOp' are known to have the same number of results, replace the
   /// uses of op with uses of newOp.
   void replaceOpWithResultsOfAnotherOp(Operation *op, Operation *newOp);
 };
 
 //===----------------------------------------------------------------------===//
 // IRRewriter
 //===----------------------------------------------------------------------===//
 
 /// This class coordinates rewriting a piece of IR outside of a pattern rewrite,
 /// providing a way to keep track of the mutations made to the IR. This class
 /// should only be used in situations where another `RewriterBase` instance,
 /// such as a `PatternRewriter`, is not available.
 class IRRewriter : public RewriterBase {
 public:
-  explicit IRRewriter(MLIRContext *ctx) : RewriterBase(ctx) {}
+  explicit IRRewriter(MLIRContext *ctx, OpBuilder::Listener *listener = nullptr)
+      : RewriterBase(ctx, listener) {}
   explicit IRRewriter(const OpBuilder &builder) : RewriterBase(builder) {}
 };
 
 //===----------------------------------------------------------------------===//
 // PatternRewriter
 //===----------------------------------------------------------------------===//
 
 /// A special type of `RewriterBase` that coordinates the application of a

mlir/include/mlir/Transforms/FoldUtils.h

 //===--------------------------------------------------------------------===//
 // OperationFolder
 //===--------------------------------------------------------------------===//
 
 /// A utility class for folding operations, and unifying duplicated constants
 /// generated along the way.
 class OperationFolder {
 public:
-  OperationFolder(MLIRContext *ctx, RewriterBase::Listener *listener = nullptr)
-      : interfaces(ctx), listener(listener) {}
+  OperationFolder(MLIRContext *ctx, OpBuilder::Listener *listener = nullptr)
+      : interfaces(ctx), rewriter(ctx, 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. 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, bool *inPlaceUpdate = nullptr);
 
   /// Tries to fold a pre-existing constant operation. `constValue` represents
public:
   ///
   /// Note: this method must be called if a constant op is to be deleted
   /// externally to this OperationFolder. `op` must be a constant op.
   void notifyRemoval(Operation *op);
 
   /// Clear out any constants cached inside of the folder.
   void clear();
 
-  /// Get or create a constant using the given builder. On success this returns
-  /// the constant operation, nullptr otherwise.
-  Value getOrCreateConstant(OpBuilder &builder, Dialect *dialect,
-                            Attribute value, Type type, Location loc);
+  /// Get or create a constant for use in the specified block. The constant may
+  /// be created in a parent block. On success this returns the constant
+  /// operation, nullptr otherwise.
+  Value getOrCreateConstant(Block *block, Dialect *dialect, Attribute value,
+                            Type type, Location loc);
 
 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);
+  LogicalResult tryToFold(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,
+  /// Try to process a set of fold results. Populates `results` on success,
+  /// otherwise leaves it unchanged.
+  LogicalResult processFoldResults(Operation *op,
                                    SmallVectorImpl<Value> &results,
                                    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);
+                                    Dialect *dialect, 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;
+  /// A rewriter that performs all IR modifications.
+  IRRewriter rewriter;
 };
 
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_FOLDUTILS_H

mlir/lib/Transforms/SCCP.cpp

static LogicalResult replaceWithConstant(DataFlowSolver &solver,
   if (!lattice || lattice->getValue().isUninitialized())
     return failure();
   const ConstantValue &latticeValue = lattice->getValue();
   if (!latticeValue.getConstantValue())
     return failure();
 
   // Attempt to materialize a constant for the given value.
   Dialect *dialect = latticeValue.getConstantDialect();
-  Value constant = folder.getOrCreateConstant(builder, dialect,
-                                              latticeValue.getConstantValue(),
-                                              value.getType(), value.getLoc());
+  Value constant = folder.getOrCreateConstant(
+      builder.getInsertionBlock(), dialect, latticeValue.getConstantValue(),
   if (!constant)
     return failure();
 
   value.replaceAllUsesWith(constant);
   return success();
 }
 
 /// Rewrite the given regions using the computing analysis. This replaces the

mlir/lib/Transforms/Utils/FoldUtils.cpp

if (isFolderOwnedConstant(op)) {
     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, listener);
-  if (failed(tryToFold(builder, op, results)))
+  if (failed(tryToFold(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);
+    if (auto *rewriteListener = dyn_cast_if_present<RewriterBase::Listener>(
+            rewriter.getListener())) {
+      // Folding API does not notify listeners, so we have to notify manually.
+      rewriteListener->notifyOperationModified(op);
+    }
     return success();
   }
 
   // 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]);
-  eraseOp(op);
+  notifyRemoval(op);
+  rewriter.replaceOp(op, results);
   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);
-    eraseOp(op);
+    notifyRemoval(op);
+    rewriter.replaceOp(op, folderConstOp->getResults());
     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;
 

 /// Clear out any constants cached inside of the folder.
 void OperationFolder::clear() {
   foldScopes.clear();
   referencedDialects.clear();
 }
 
 /// Get or create a constant using the given builder. On success this returns
 /// the constant operation, nullptr otherwise.
-Value OperationFolder::getOrCreateConstant(OpBuilder &builder, Dialect *dialect,
+Value OperationFolder::getOrCreateConstant(Block *block, Dialect *dialect,
                                            Attribute value, Type type,
                                            Location loc) {
-  OpBuilder::InsertionGuard foldGuard(builder);
-
-  // Use the builder insertion block to find an insertion point for the
-  // constant.
-  auto *insertRegion =
-      getInsertionRegion(interfaces, builder.getInsertionBlock());
+  // Find an insertion point for the constant.
+  auto *insertRegion = getInsertionRegion(interfaces, block);
   auto &entry = insertRegion->front();
-  builder.setInsertionPoint(&entry, entry.begin());
+  rewriter.setInsertionPoint(&entry, entry.begin());
 
   // Get the constant map for the insertion region of this operation.
   auto &uniquedConstants = foldScopes[insertRegion];
-  Operation *constOp = tryGetOrCreateConstant(uniquedConstants, dialect,
-                                              builder, value, type, loc);
+  Operation *constOp =
+      tryGetOrCreateConstant(uniquedConstants, dialect, value, type, loc);
   return constOp ? constOp->getResult(0) : Value();
 }
 
 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,
+LogicalResult OperationFolder::tryToFold(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(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)))
+      failed(processFoldResults(op, results, foldResults)))
     return success(updatedOpOperands);
   return success();
 }
 
 LogicalResult
-OperationFolder::processFoldResults(OpBuilder &builder, Operation *op,
+OperationFolder::processFoldResults(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 =
-      getInsertionRegion(interfaces, builder.getInsertionBlock());
+  auto *insertRegion = getInsertionRegion(interfaces, op->getBlock());
   auto &entry = insertRegion->front();
-  OpBuilder::InsertionGuard foldGuard(builder);
-  builder.setInsertionPoint(&entry, entry.begin());
+  rewriter.setInsertionPoint(&entry, entry.begin());
 
   // Get the constant map for the insertion region of this operation.
   auto &uniquedConstants = foldScopes[insertRegion];
 
   // Create the result constants and replace the results.
   auto *dialect = op->getDialect();
   for (unsigned i = 0, e = op->getNumResults(); i != e; ++i) {
     assert(!foldResults[i].isNull() && "expected valid OpFoldResult");
 
     // Check if the result was an SSA value.
     if (auto repl = foldResults[i].dyn_cast<Value>()) {
       if (repl.getType() != op->getResult(i).getType()) {
         results.clear();
         return failure();
       }
       results.emplace_back(repl);
       continue;
     }
 
     // Check to see if there is a canonicalized version of this constant.
     auto res = op->getResult(i);
     Attribute attrRepl = foldResults[i].get<Attribute>();
-    if (auto *constOp =
-            tryGetOrCreateConstant(uniquedConstants, dialect, builder, attrRepl,
-                                   res.getType(), op->getLoc())) {
+    if (auto *constOp = tryGetOrCreateConstant(
+            uniquedConstants, dialect, attrRepl, res.getType(), op->getLoc())) {
       // Ensure that this constant dominates the operation we are replacing it
       // with. This may not automatically happen if the operation being folded
       // was inserted before the constant within the insertion block.
       Block *opBlock = op->getBlock();
       if (opBlock == constOp->getBlock() && &opBlock->front() != 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())))
-      eraseOp(&op);
+             llvm::make_range(entry.begin(), rewriter.getInsertionPoint()))) {
+      notifyRemoval(&op);
+      rewriter.eraseOp(&op);
+    }
 
     results.clear();
     return failure();
   }
 
   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(ConstantMap &uniquedConstants,
+                                        Dialect *dialect, Attribute value,
+                                        Type type, Location loc) {
   // Check if an existing mapping already exists.
   auto constKey = std::make_tuple(dialect, value, type);
   Operation *&constOp = uniquedConstants[constKey];
   if (constOp)
     return constOp;
 
   // If one doesn't exist, try to materialize one.
-  if (!(constOp = materializeConstant(dialect, builder, value, type, loc)))
+  if (!(constOp = materializeConstant(dialect, rewriter, value, type, loc)))
     return nullptr;
 
   // Check to see if the generated constant is in the expected dialect.
   auto *newDialect = constOp->getDialect();
   if (newDialect == dialect) {
     referencedDialects[constOp].push_back(dialect);
     return constOp;
   }
 
   // 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)) {
-    eraseOp(constOp);
+    notifyRemoval(constOp);
+    rewriter.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/test/lib/Transforms/TestIntRangeInference.cpp

static LogicalResult replaceWithConstant(DataFlowSolver &solver, OpBuilder &b,
   if (!maybeConstValue.has_value())
     return failure();
 
   Operation *maybeDefiningOp = value.getDefiningOp();
   Dialect *valueDialect =
       maybeDefiningOp ? maybeDefiningOp->getDialect()
                       : value.getParentRegion()->getParentOp()->getDialect();
   Attribute constAttr = b.getIntegerAttr(value.getType(), *maybeConstValue);
-  Value constant = folder.getOrCreateConstant(b, valueDialect, constAttr,
-                                              value.getType(), value.getLoc());
+  Value constant =
+      folder.getOrCreateConstant(b.getInsertionBlock(), valueDialect, constAttr,
   if (!constant)
     return failure();
 
   value.replaceAllUsesWith(constant);
   return success();
 }
 
 static void rewrite(DataFlowSolver &solver, MLIRContext *context,