[mlir][Pattern] Create a new IRRewriter class to enable sharing code with pattern rewrites

[rriddle]

Summary

This revision adds two new classes, RewriterBase and IRRewriter. RewriterBase is a new shared base class between IRRewriter and PatternRewriter. PatternRewriter will continue to be the base class used to perform rewrites within a rewrite pattern. IRRewriter on the other hand, is a new class that allows for tracking IR rewrites from outside of a rewrite pattern. In this revision all of the old API from PatternRewriter is moved to RewriterBase, but the distinction between IRRewriter and PatternRewriter is kept on the chance that a necessary API divergence happens in the future.

Currently if you want to have some utility that transforms a piece of IR and share it between pattern and non-pattern code, you have to duplicate it. This revision enables the creation of utilities that can be invoked from rewrite patterns and normal transformation code:

c++
void someSharedUtility(RewriterBase &rewriter, ...) {
  // Some interesting IR mutation here.
}

// Some RewritePattern
LogicalResult MyPattern::matchAndRewrite(Operation *op, PatternRewriter &rewriter) {
  ...
  someSharedUtility(rewriter, ...);
  ...
}

// Some Pass
void MyPass::runOnOperation() {
  ...
  IRRewriter rewriter(...);
  someSharedUtility(rewriter, ...);
}

Depends On D94632

[rriddle]

I went with IRRewriterImpl as this seems close to a SmallVectorImpl situation, but happy to rename to something else if there are suggestions.

[bollu]

Out of curiosity: may I please have some context for what this base class is being created for? Ie, what consumers other than the existing PatternRewriter will soon exist? :)

mlir/lib/IR/PatternMatch.cpp

315

Nit: spuriously deleted line?

[rriddle]

In D94638#2503020, @bollu wrote:

Out of curiosity: may I please have some context for what this base class is being created for? Ie, what consumers other than the existing PatternRewriter will soon exist? :)

There are existing use cases, i.e. if you want to have some utility that transforms a piece of IR and share it between pattern and non-pattern code you currently have to duplicate it. What you would be able to do after this revision is share it by using an IRRewriterImpl as the driver for the mutation:

void someSharedUtility(IRRewriterImpl &rewriter, ...) {
  // Some interesting IR mutation here.
}

// Some RewritePattern
LogicalResult MyPattern::matchAndRewrite(Operation *op, PatternRewriter &rewriter) {
  ...
  someSharedUtility(rewriter, ...);
  ...
}

// Some Pass
void MyPass::runOnOperation() {
  ...
  IRRewriter rewriter(...);
  someSharedUtility(rewriter, ...);
}

[jpienaar]

In D94638#2503042, @rriddle wrote:

In D94638#2503020, @bollu wrote:

Out of curiosity: may I please have some context for what this base class is being created for? Ie, what consumers other than the existing PatternRewriter will soon exist? :)

There are existing use cases, i.e. if you want to have some utility that transforms a piece of IR and share it between pattern and non-pattern code you currently have to duplicate it. What you would be able to do after this revision is share it by using an IRRewriterImpl as the driver for the mutation:

I think this would be good to capture in change description more explicitly (I liked the small example and what you wrote here for an extensive commit message :))

mlir/include/mlir/IR/PatternMatch.h

425

Does the naming convention here mirror SmallVector / SmallVectorImpl? "impl" makes me think of an implementation detail vs a common base class. This is a "listening builder" or "observed builder" base class (well technically an observable, it is only observed with listeners registered ...)

597

Context is a bit overloaded. Is this a hard requirement or a nice-to-have? E.g., can I mix and match but a bad idea to do so, or just preferred?

[rriddle]

Address feedback

mlir/include/mlir/IR/PatternMatch.h

425

Yeah, the Impl here mirror'd SmallVectorImpl. I couldn't think of a better name at the time, but switched to using RewriterBase instead.

597

It's a hard requirement, given that the rewriter you aren't using wouldn't get the notifications for rewrites you are doing.

[rriddle]

Rebase

mlir/include/mlir/IR/PatternMatch.h

private:
 
   /// The external functions referenced from within the PDL module.
   llvm::StringMap<PDLConstraintFunction> constraintFunctions;
   llvm::StringMap<PDLCreateFunction> createFunctions;
   llvm::StringMap<PDLRewriteFunction> rewriteFunctions;
 };
 
 //===----------------------------------------------------------------------===//
-// PatternRewriter
+// RewriterBase
 //===----------------------------------------------------------------------===//
 
-/// This class coordinates the application of a pattern to the current function,
-/// providing a way to create operations and keep track of what gets deleted.
-///
-/// These class serves two purposes:
-///  1) it is the interface that patterns interact with to make mutations to the
-///     IR they are being applied to.
-///  2) It is a base class that clients of the PatternMatcher use when they want
-///     to apply patterns and observe their effects (e.g. to keep worklists or
-///     other data structures up to date).
-///
-class PatternRewriter : public OpBuilder, public OpBuilder::Listener {
+/// This class coordinates the application of a rewrite on a set of IR,
+/// providing a way for clients to track mutations and create new operations.
+/// This class serves as a common API for IR mutation between pattern rewrites
+/// and non-pattern rewrites, and facilitates the development of shared
+/// IR transformation utilities.
+class RewriterBase : public OpBuilder, public OpBuilder::Listener {
 public:
   /// Move the blocks that belong to "region" before the given position in
   /// another region "parent". The two regions must be different. The caller
   /// is responsible for creating or updating the operation transferring flow
   /// of control to the region and passing it the correct block arguments.
   virtual void inlineRegionBefore(Region &region, Region &parent,
                                   Region::iterator before);
   void inlineRegionBefore(Region &region, Block *before);
 
   /// Clone the blocks that belong to "region" before the given position in
   /// another region "parent". The two regions must be different. The caller is
   /// responsible for creating or updating the operation transferring flow of
   /// control to the region and passing it the correct block arguments.
   virtual void cloneRegionBefore(Region &region, Region &parent,
                                  Region::iterator before,
                                  BlockAndValueMapping &mapping);
   void cloneRegionBefore(Region &region, Region &parent,
                          Region::iterator before);
   void cloneRegionBefore(Region &region, Block *before);
 
   /// This method replaces the uses of the results of `op` with the values in
   /// `newValues` when the provided `functor` returns true for a specific use.
   /// The number of values in `newValues` is required to match the number of
   /// results of `op`. `allUsesReplaced`, if non-null, is set to true if all of
-  /// the uses of `op` were replaced. Note that in some pattern rewriters, the
-  /// given 'functor' may be stored beyond the lifetime of the pattern being
-  /// applied. As such, the function should not capture by reference and instead
-  /// use value capture as necessary.
+  /// the uses of `op` were replaced. Note that in some rewriters, the given
+  /// 'functor' may be stored beyond the lifetime of the rewrite being applied.
+  /// As such, the function should not capture by reference and instead use
+  /// value capture as necessary.
   virtual void
   replaceOpWithIf(Operation *op, ValueRange newValues, bool *allUsesReplaced,
                   llvm::unique_function<bool(OpOperand &) const> functor);
   void replaceOpWithIf(Operation *op, ValueRange newValues,
                        llvm::unique_function<bool(OpOperand &) const> functor) {
     replaceOpWithIf(op, newValues, /*allUsesReplaced=*/nullptr,
                     std::move(functor));
   }
 
   /// This method replaces the uses of the results of `op` with the values in
   /// `newValues` when a use is nested within the given `block`. The number of
   /// values in `newValues` is required to match the number of results of `op`.
   /// If all uses of this operation are replaced, the operation is erased.
   void replaceOpWithinBlock(Operation *op, ValueRange newValues, Block *block,
                             bool *allUsesReplaced = nullptr);
 
-  /// This method performs the final replacement for a pattern, where the
-  /// results of the operation are updated to use the specified list of SSA
-  /// values.
+  /// This method replaces the results of the operation with the specified list
+  /// of values. The number of provided values must match the number of results
+  /// of the operation.
   virtual void replaceOp(Operation *op, ValueRange newValues);
 
   /// Replaces the result op with a new op that is created without verification.
   /// The result values of the two ops must be the same types.
   template <typename OpTy, typename... Args>
   OpTy replaceOpWithNewOp(Operation *op, Args &&... args) {
     auto newOp = create<OpTy>(op->getLoc(), std::forward<Args>(args)...);
     replaceOpWithResultsOfAnotherOp(op, newOp.getOperation());
public:
   /// callable.
   template <typename CallableT>
   void updateRootInPlace(Operation *root, CallableT &&callable) {
     startRootUpdate(root);
     callable();
     finalizeRootUpdate(root);
   }
 
-  /// Notify the pattern rewriter that the pattern is failing to match the given
-  /// operation, and provide a callback to populate a diagnostic with the reason
-  /// why the failure occurred. This method allows for derived rewriters to
-  /// optionally hook into the reason why a pattern failed, and display it to
+  /// Used to notify the rewriter that the IR failed to be rewritten because of
+  /// a match failure, and provide a callback to populate a diagnostic with the
+  /// reason why the failure occurred. This method allows for derived rewriters
+  /// to optionally hook into the reason why a rewrite failed, and display it to
   /// users.
   template <typename CallbackT>
   std::enable_if_t<!std::is_convertible<CallbackT, Twine>::value, LogicalResult>
   notifyMatchFailure(Operation *op, CallbackT &&reasonCallback) {
 #ifndef NDEBUG
     return notifyMatchFailure(op,
                               function_ref<void(Diagnostic &)>(reasonCallback));
 #else
     return failure();
 #endif
   }
   LogicalResult notifyMatchFailure(Operation *op, const Twine &msg) {
     return notifyMatchFailure(op, [&](Diagnostic &diag) { diag << msg; });
   }
   LogicalResult notifyMatchFailure(Operation *op, const char *msg) {
     return notifyMatchFailure(op, Twine(msg));
   }
 
 protected:
   /// Initialize the builder with this rewriter as the listener.
-  explicit PatternRewriter(MLIRContext *ctx)
-      : OpBuilder(ctx, /*listener=*/this) {}
-  ~PatternRewriter() override;
+  explicit RewriterBase(MLIRContext *ctx) : OpBuilder(ctx, /*listener=*/this) {}
+  explicit RewriterBase(const OpBuilder &otherBuilder)
+      : OpBuilder(otherBuilder) {
+    setListener(this);
+  }
+  ~RewriterBase() override;
 
   /// These are the callback methods that subclasses can choose to implement if
   /// they would like to be notified about certain types of mutations.
 
-  /// Notify the pattern rewriter that the specified operation is about to be
-  /// replaced with another set of operations.  This is called before the uses
-  /// of the operation have been changed.
+  /// Notify the rewriter that the specified operation is about to be replaced
+  /// with another set of operations. This is called before the uses of the
+  /// operation have been changed.
   virtual void notifyRootReplaced(Operation *op) {}
 
-  /// This is called on an operation that a pattern match is removing, right
-  /// before the operation is deleted.  At this point, the operation has zero
-  /// uses.
+  /// This is called on an operation that a rewrite is removing, right before
+  /// the operation is deleted. At this point, the operation has zero uses.
   virtual void notifyOperationRemoved(Operation *op) {}
 
-  /// Notify the pattern rewriter that the pattern is failing to match the given
-  /// operation, and provide a callback to populate a diagnostic with the reason
-  /// why the failure occurred. This method allows for derived rewriters to
-  /// optionally hook into the reason why a pattern failed, and display it to
-  /// users.
+  /// Notify the rewriter that the pattern failed to match the given operation,
+  /// and provide a callback to populate a diagnostic with the reason why the
+  /// failure occurred. This method allows for derived rewriters to optionally
+  /// hook into the reason why a rewrite failed, and display it to users.
   virtual LogicalResult
   notifyMatchFailure(Operation *op,
                      function_ref<void(Diagnostic &)> reasonCallback) {
     return failure();
   }
 
 private:
   /// '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(const OpBuilder &builder) : RewriterBase(builder) {}
+};
+
+//===----------------------------------------------------------------------===//
+// PatternRewriter
+//===----------------------------------------------------------------------===//
+
+/// A special type of `RewriterBase` that coordinates the application of a
+/// rewrite pattern on the current IR being matched, providing a way to keep
+/// track of any mutations made. This class should be used to perform all
+/// necessary IR mutations within a rewrite pattern, as the pattern driver may
+/// be tracking various state that would be invalidated when a mutation takes
+/// place.
+class PatternRewriter : public RewriterBase {
+public:
+  using RewriterBase::RewriterBase;
+};
+
+//===----------------------------------------------------------------------===//
 // OwningRewritePatternList
 //===----------------------------------------------------------------------===//
 
 class OwningRewritePatternList {
   using NativePatternListT = std::vector<std::unique_ptr<RewritePattern>>;
 
 public:
   OwningRewritePatternList() = default;

mlir/lib/IR/PatternMatch.cpp

void PDLPatternModule::registerRewriteFunction(StringRef name,
                                                PDLRewriteFunction rewriteFn) {
   auto it = rewriteFunctions.try_emplace(name, std::move(rewriteFn));
   (void)it;
   assert(it.second && "native rewrite function with the given name has "
                       "already been registered");
 }
 
 //===----------------------------------------------------------------------===//
-// PatternRewriter
+// RewriterBase
 //===----------------------------------------------------------------------===//
 
-PatternRewriter::~PatternRewriter() {
+RewriterBase::~RewriterBase() {
   // Out of line to provide a vtable anchor for the class.
 }
 
 /// This method replaces the uses of the results of `op` with the values in
 /// `newValues` when the provided `functor` returns true for a specific use.
 /// The number of values in `newValues` is required to match the number of
 /// results of `op`.
-void PatternRewriter::replaceOpWithIf(
+void RewriterBase::replaceOpWithIf(
     Operation *op, ValueRange newValues, bool *allUsesReplaced,
     llvm::unique_function<bool(OpOperand &) const> functor) {
   assert(op->getNumResults() == newValues.size() &&
          "incorrect number of values to replace operation");
 
   // Notify the rewriter subclass that we're about to replace this root.
   notifyRootReplaced(op);
 
   // Replace each use of the results when the functor is true.
   bool replacedAllUses = true;
   for (auto it : llvm::zip(op->getResults(), newValues)) {
     std::get<0>(it).replaceUsesWithIf(std::get<1>(it), functor);
     replacedAllUses &= std::get<0>(it).use_empty();
   }
   if (allUsesReplaced)
     *allUsesReplaced = replacedAllUses;
 }
 
 /// This method replaces the uses of the results of `op` with the values in
 /// `newValues` when a use is nested within the given `block`. The number of
 /// values in `newValues` is required to match the number of results of `op`.
 /// If all uses of this operation are replaced, the operation is erased.
-void PatternRewriter::replaceOpWithinBlock(Operation *op, ValueRange newValues,
-                                           Block *block,
-                                           bool *allUsesReplaced) {
+void RewriterBase::replaceOpWithinBlock(Operation *op, ValueRange newValues,
+                                        Block *block, bool *allUsesReplaced) {
   replaceOpWithIf(op, newValues, allUsesReplaced, [block](OpOperand &use) {
     return block->getParentOp()->isProperAncestor(use.getOwner());
   });
 }
 
-/// This method performs the final replacement for a pattern, where the
-/// results of the operation are updated to use the specified list of SSA
-/// values.
-void PatternRewriter::replaceOp(Operation *op, ValueRange newValues) {
+/// This method replaces the results of the operation with the specified list of
+/// values. The number of provided values must match the number of results of
+/// the operation.
+void RewriterBase::replaceOp(Operation *op, ValueRange newValues) {
   // Notify the rewriter subclass that we're about to replace this root.
   notifyRootReplaced(op);
 
   assert(op->getNumResults() == newValues.size() &&
          "incorrect # of replacement values");
   op->replaceAllUsesWith(newValues);
 
   notifyOperationRemoved(op);
   op->erase();
 }
 
 /// This method erases an operation that is known to have no uses. The uses of
 /// the given operation *must* be known to be dead.
-void PatternRewriter::eraseOp(Operation *op) {
+void RewriterBase::eraseOp(Operation *op) {
   assert(op->use_empty() && "expected 'op' to have no uses");
   notifyOperationRemoved(op);
   op->erase();
 }
 
-void PatternRewriter::eraseBlock(Block *block) {
+void RewriterBase::eraseBlock(Block *block) {
   for (auto &op : llvm::make_early_inc_range(llvm::reverse(*block))) {
     assert(op.use_empty() && "expected 'op' to have no uses");
     eraseOp(&op);
   }
   block->erase();
 }
 
 /// Merge the operations of block 'source' into the end of block 'dest'.
 /// 'source's predecessors must be empty or only contain 'dest`.
 /// 'argValues' is used to replace the block arguments of 'source' after
 /// merging.
-void PatternRewriter::mergeBlocks(Block *source, Block *dest,
-                                  ValueRange argValues) {
+void RewriterBase::mergeBlocks(Block *source, Block *dest,
   assert(llvm::all_of(source->getPredecessors(),
                       [dest](Block *succ) { return succ == dest; }) &&
          "expected 'source' to have no predecessors or only 'dest'");
   assert(argValues.size() == source->getNumArguments() &&
          "incorrect # of argument replacement values");
 
   // Replace all of the successor arguments with the provided values.
   for (auto it : llvm::zip(source->getArguments(), argValues))
     std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
 
   // Splice the operations of the 'source' block into the 'dest' block and erase
   // it.
   dest->getOperations().splice(dest->end(), source->getOperations());
   source->dropAllUses();
   source->erase();
 }
 
 // Merge the operations of block 'source' before the operation 'op'. Source
 // block should not have existing predecessors or successors.
-void PatternRewriter::mergeBlockBefore(Block *source, Operation *op,
-                                       ValueRange argValues) {
+void RewriterBase::mergeBlockBefore(Block *source, Operation *op,
   assert(source->hasNoPredecessors() &&
          "expected 'source' to have no predecessors");
   assert(source->hasNoSuccessors() &&
          "expected 'source' to have no successors");
 
   // Split the block containing 'op' into two, one containing all operations
   // before 'op' (prologue) and another (epilogue) containing 'op' and all
   // operations after it.
   Block *prologue = op->getBlock();
   Block *epilogue = splitBlock(prologue, op->getIterator());
 
   // Merge the source block at the end of the prologue.
   mergeBlocks(source, prologue, argValues);
 
   // Merge the epilogue at the end the prologue.
   mergeBlocks(epilogue, prologue);
 }
 
 /// Split the operations starting at "before" (inclusive) out of the given
 /// block into a new block, and return it.
-Block *PatternRewriter::splitBlock(Block *block, Block::iterator before) {
+Block *RewriterBase::splitBlock(Block *block, Block::iterator before) {
   return block->splitBlock(before);
 }
 
 /// 'op' and 'newOp' are known to have the same number of results, replace the
 /// uses of op with uses of newOp
-void PatternRewriter::replaceOpWithResultsOfAnotherOp(Operation *op,
-                                                      Operation *newOp) {
+void RewriterBase::replaceOpWithResultsOfAnotherOp(Operation *op,
   assert(op->getNumResults() == newOp->getNumResults() &&
          "replacement op doesn't match results of original op");
   if (op->getNumResults() == 1)
     return replaceOp(op, newOp->getResult(0));
   return replaceOp(op, newOp->getResults());
 }
 
 /// Move the blocks that belong to "region" before the given position in
 /// another region.  The two regions must be different.  The caller is in
 /// charge to update create the operation transferring the control flow to the
 /// region and pass it the correct block arguments.
-void PatternRewriter::inlineRegionBefore(Region &region, Region &parent,
-                                         Region::iterator before) {
+void RewriterBase::inlineRegionBefore(Region &region, Region &parent,
   parent.getBlocks().splice(before, region.getBlocks());
 }
-void PatternRewriter::inlineRegionBefore(Region &region, Block *before) {
+void RewriterBase::inlineRegionBefore(Region &region, Block *before) {
   inlineRegionBefore(region, *before->getParent(), before->getIterator());
 }
 
 /// Clone the blocks that belong to "region" before the given position in
 /// another region "parent". The two regions must be different. The caller is
 /// responsible for creating or updating the operation transferring flow of
 /// control to the region and passing it the correct block arguments.
-void PatternRewriter::cloneRegionBefore(Region &region, Region &parent,
-                                        Region::iterator before,
-                                        BlockAndValueMapping &mapping) {
+void RewriterBase::cloneRegionBefore(Region &region, Region &parent,
   region.cloneInto(&parent, before, mapping);
 }
-void PatternRewriter::cloneRegionBefore(Region &region, Region &parent,
-                                        Region::iterator before) {
+void RewriterBase::cloneRegionBefore(Region &region, Region &parent,
   BlockAndValueMapping mapping;
   cloneRegionBefore(region, parent, before, mapping);
 }
-void PatternRewriter::cloneRegionBefore(Region &region, Block *before) {
+void RewriterBase::cloneRegionBefore(Region &region, Block *before) {
   cloneRegionBefore(region, *before->getParent(), before->getIterator());
 }
-