[mlir] Separate Listener from PatternRewriter

[Mogball]

Summary

Introduce a new RewriteListener class that replaces the previous OpBuilder::Listener and takes some of the methods that were in PatternRewriter. Moving the listening logic out of PatternRewriter and consolidating it in a new class gives users and developers more control over listening to rewrite events and makes listeners more composable.

For example, this patch also adds a ListenerList which is just a RewriteListener that forwards events to multiple listeners. Rewrite drivers that were previously implemented by subclassing PatternRewriter can be implemented by subclassing RewriteListener instead, and then they can accept a user-provided listener that will also be notified of any rewrite events.

Also adds an optional listener to the greedy rewrite driver. Previously, there was no way to be made aware of IR changes made by the greedy rewriter, which makes it difficult to pass information about the IR through multiple rewrite passes, e.g. tracking whether a certain operation was deleted.

For example, with listeners I can

1. Apply pattern set A, and keep track of every foo.op that is created.
2. Apply canonicalization patterns, and track which foo.op are deleted.
3. Do something with the foo.op that remain.

[Mogball]

FYI I might change this to accept a more limited RewriteListener class that only has the notify operations.

[rriddle]

This is giving off a really bad design smell. It would be good to articulate the exact use cases you are trying to support. Intermixing multiple pattern rewriters does not sound like a good idea.

[Mogball]

Yeah I agree. I reworked the design to split out listener from pattern rewriter and hopefully got rid of some of the smell.

[Mogball]

Removed unused constructor arg

[rriddle]

I'd prefer to have proper design discussions about this (preferably when I'm not OOO) the direction here isn't obvious that it's something to support.

mlir/include/mlir/Dialect/SCF/SCFOps.td

let builders = [
     OpBuilder<(ins "Value":$cond,
       CArg<"function_ref<void(OpBuilder &, Location)>",
            "buildTerminatedBody">:$thenBuilder,
       CArg<"function_ref<void(OpBuilder &, Location)>",
            "nullptr">:$elseBuilder)>
   ];
 
   let extraClassDeclaration = [{
-    OpBuilder getThenBodyBuilder(OpBuilder::Listener *listener = nullptr) {
+    OpBuilder getThenBodyBuilder(RewriteListener *listener = nullptr) {
       Block* body = getBody(0);
       return getResults().empty() ? OpBuilder::atBlockTerminator(body, listener)
                                   : OpBuilder::atBlockEnd(body, listener);
     }
-    OpBuilder getElseBodyBuilder(OpBuilder::Listener *listener = nullptr) {
+    OpBuilder getElseBodyBuilder(RewriteListener *listener = nullptr) {
       Block* body = getBody(1);
       return getResults().empty() ? OpBuilder::atBlockTerminator(body, listener)
                                   : OpBuilder::atBlockEnd(body, listener);
     }
     Block* thenBlock();
     YieldOp thenYield();
     Block* elseBlock();
     YieldOp elseYield();

mlir/include/mlir/IR/Builders.h

 class SymbolRefAttr;
 class ElementsAttr;
 class DenseElementsAttr;
 class DenseIntElementsAttr;
 class AffineMapAttr;
 class AffineMap;
 class UnitAttr;
 
+//===----------------------------------------------------------------------===//
+// Builder
+//===----------------------------------------------------------------------===//
+
 /// This class is a general helper class for creating context-global objects
 /// like types, attributes, and affine expressions.
 class Builder {
 public:
   explicit Builder(MLIRContext *context) : context(context) {}
   explicit Builder(Operation *op) : Builder(op->getContext()) {}
 
   MLIRContext *getContext() const { return context; }
public:
   /// Eg: input: (d0, d1)[s0] -> (d0, d1 + s0), shift = 2
   ///   returns:    (d0, d1)[s0] -> (d0 + 2, d1 + s0 + 2)
   AffineMap getShiftedAffineMap(AffineMap map, int64_t shift);
 
 protected:
   MLIRContext *context;
 };
 
+//===----------------------------------------------------------------------===//
+// RewriteListener
+//===----------------------------------------------------------------------===//
+
+/// This class represents a listener that can be used to hook on to various
+/// rewrite events in an `OpBuilder` or `PatternRewriter`. The class is notified
+/// by when:
+///
+/// - an operation is removed
+/// - an operation is inserted
+/// - an operation is replaced
+/// - a block is created
+/// - a pattern match failed
+///
+/// Listeners can be used to track IR mutations throughout pattern rewrites.
+struct RewriteListener {
+  virtual ~RewriteListener();
+
+  /// These are the callback methods that subclasses can choose to implement if
+  /// they would like to be notified about certain types of mutations.
+
+  /// Notification handler for when an operation is inserted into the builder.
+  /// op` is the operation that was inserted.
+  virtual void notifyOperationInserted(Operation *op) {}
+
+  /// Notification handler for when a block is created using the builder.
+  /// `block` is the block that was created.
+  virtual void notifyBlockCreated(Block *block) {}
+
+  /// Notification handler for when 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) {}
+
+  /// Notification handler for when an the specified operation is about to be
+  /// deleted. At this point, the operation has zero uses.
+  virtual void notifyOperationRemoved(Operation *op) {}
+
+  /// Notify the listener that a 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 listeners to optionally
+  /// hook into the reason why a rewrite failed, and display it to users.
+  virtual void
+  notifyMatchFailure(Operation *op,
+                     function_ref<void(Diagnostic &)> reasonCallback) {}
+};
+
+//===----------------------------------------------------------------------===//
+// ListenerList
+//===----------------------------------------------------------------------===//
+
+/// This class contains multiple listeners to which rewrite events can be sent.
+class ListenerList : public RewriteListener {
+public:
+  /// Add a listener to the list.
+  void addListener(RewriteListener *listener) { listeners.push_back(listener); }
+
+  /// Send notification of an operation being inserted to all listeners.
+  void notifyOperationInserted(Operation *op) override;
+
+  /// Send notification of a block being created to all listeners.
+  void notifyBlockCreated(Block *block) override;
+
+  /// Send notification that an operation has been replaced to all listeners.
+  void notifyRootReplaced(Operation *op) override;
+
+  /// Send notification that an operation is about to be deleted to all
+  /// listeners.
+  void notifyOperationRemoved(Operation *op) override;
+
+  /// Notify all listeners that a pattern match failed.
+  void
+  notifyMatchFailure(Operation *op,
+                     function_ref<void(Diagnostic &)> reasonCallback) override;
+
+private:
+  /// The list of listeners to send events to.
+  SmallVector<RewriteListener *, 1> listeners;
+};
+
+//===----------------------------------------------------------------------===//
+// OpBuilder
+//===----------------------------------------------------------------------===//
+
 /// This class helps build Operations. Operations that are created are
 /// automatically inserted at an insertion point. The builder is copyable.
 class OpBuilder : public Builder {
 public:
-  struct Listener;
-
   /// Create a builder with the given context.
-  explicit OpBuilder(MLIRContext *ctx, Listener *listener = nullptr)
+  explicit OpBuilder(MLIRContext *ctx, RewriteListener *listener = nullptr)
       : Builder(ctx), listener(listener) {}
 
   /// Create a builder and set the insertion point to the start of the region.
-  explicit OpBuilder(Region *region, Listener *listener = nullptr)
+  explicit OpBuilder(Region *region, RewriteListener *listener = nullptr)
       : OpBuilder(region->getContext(), listener) {
     if (!region->empty())
       setInsertionPoint(&region->front(), region->front().begin());
   }
-  explicit OpBuilder(Region &region, Listener *listener = nullptr)
+  explicit OpBuilder(Region &region, RewriteListener *listener = nullptr)
       : OpBuilder(&region, listener) {}
 
   /// Create a builder and set insertion point to the given operation, which
   /// will cause subsequent insertions to go right before it.
-  explicit OpBuilder(Operation *op, Listener *listener = nullptr)
+  explicit OpBuilder(Operation *op, RewriteListener *listener = nullptr)
       : OpBuilder(op->getContext(), listener) {
     setInsertionPoint(op);
   }
 
   OpBuilder(Block *block, Block::iterator insertPoint,
-            Listener *listener = nullptr)
+            RewriteListener *listener = nullptr)
       : OpBuilder(block->getParent()->getContext(), listener) {
     setInsertionPoint(block, insertPoint);
   }
 
   /// Create a builder and set the insertion point to before the first operation
   /// in the block but still inside the block.
-  static OpBuilder atBlockBegin(Block *block, Listener *listener = nullptr) {
+  static OpBuilder atBlockBegin(Block *block,
+                                RewriteListener *listener = nullptr) {
     return OpBuilder(block, block->begin(), listener);
   }
 
   /// Create a builder and set the insertion point to after the last operation
   /// in the block but still inside the block.
-  static OpBuilder atBlockEnd(Block *block, Listener *listener = nullptr) {
+  static OpBuilder atBlockEnd(Block *block,
+                              RewriteListener *listener = nullptr) {
     return OpBuilder(block, block->end(), listener);
   }
 
   /// Create a builder and set the insertion point to before the block
   /// terminator.
   static OpBuilder atBlockTerminator(Block *block,
-                                     Listener *listener = nullptr) {
+                                     RewriteListener *listener = nullptr) {
     auto *terminator = block->getTerminator();
     assert(terminator != nullptr && "the block has no terminator");
     return OpBuilder(block, Block::iterator(terminator), listener);
   }
 
   //===--------------------------------------------------------------------===//
   // Listeners
   //===--------------------------------------------------------------------===//
 
-  /// This class represents a listener that may be used to hook into various
-  /// actions within an OpBuilder.
-  struct Listener {
-    virtual ~Listener();
-
-    /// Notification handler for when an operation is inserted into the builder.
-    /// `op` is the operation that was inserted.
-    virtual void notifyOperationInserted(Operation *op) {}
-
-    /// Notification handler for when a block is created using the builder.
-    /// `block` is the block that was created.
-    virtual void notifyBlockCreated(Block *block) {}
-  };
-
   /// Sets the listener of this builder to the one provided.
-  void setListener(Listener *newListener) { listener = newListener; }
+  void setListener(RewriteListener *newListener) { listener = newListener; }
 
   /// Returns the current listener of this builder, or nullptr if this builder
   /// doesn't have a listener.
-  Listener *getListener() const { return listener; }
+  RewriteListener *getListener() const { return listener; }
 
   //===--------------------------------------------------------------------===//
   // Insertion Point Management
   //===--------------------------------------------------------------------===//
 
   /// This class represents a saved insertion point.
   class InsertPoint {
   public:

 
 private:
   /// The current block this builder is inserting into.
   Block *block = nullptr;
   /// The insertion point within the block that this builder is inserting
   /// before.
   Block::iterator insertPoint;
   /// The optional listener for events of this builder.
-  Listener *listener;
+  RewriteListener *listener;
 };
 
 } // namespace mlir
 
 #endif

mlir/include/mlir/IR/ImplicitLocOpBuilder.h

public:
   /// OpBuilder has a bunch of convenience constructors - we support them all
   /// with the additional Location.
   template <typename... T>
   ImplicitLocOpBuilder(Location loc, T &&...operands)
       : OpBuilder(std::forward<T>(operands)...), curLoc(loc) {}
 
   /// Create a builder and set the insertion point to before the first operation
   /// in the block but still inside the block.
-  static ImplicitLocOpBuilder atBlockBegin(Location loc, Block *block,
-                                           Listener *listener = nullptr) {
+  static ImplicitLocOpBuilder
+  atBlockBegin(Location loc, Block *block,
+               RewriteListener *listener = nullptr) {
     return ImplicitLocOpBuilder(loc, block, block->begin(), listener);
   }
 
   /// Create a builder and set the insertion point to after the last operation
   /// in the block but still inside the block.
   static ImplicitLocOpBuilder atBlockEnd(Location loc, Block *block,
-                                         Listener *listener = nullptr) {
+                                         RewriteListener *listener = nullptr) {
     return ImplicitLocOpBuilder(loc, block, block->end(), listener);
   }
 
   /// Create a builder and set the insertion point to before the block
   /// terminator.
-  static ImplicitLocOpBuilder atBlockTerminator(Location loc, Block *block,
-                                                Listener *listener = nullptr) {
+  static ImplicitLocOpBuilder
+  atBlockTerminator(Location loc, Block *block,
+                    RewriteListener *listener = nullptr) {
     auto *terminator = block->getTerminator();
     assert(terminator != nullptr && "the block has no terminator");
     return ImplicitLocOpBuilder(loc, block, Block::iterator(terminator),
                                 listener);
   }
 
   /// Accessors for the implied location.
   Location getLoc() const { return curLoc; }

mlir/include/mlir/IR/PatternMatch.h

 // RewriterBase
 //===----------------------------------------------------------------------===//
 
 /// 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 {
+class RewriterBase : public OpBuilder {
 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);
public:
   /// 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();
+    if (RewriteListener *listener = getListener())
+      listener->notifyMatchFailure(
+          op, function_ref<void(Diagnostic &)>(reasonCallback));
 #endif
+    return failure();
   }
   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 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 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 rewrite is removing, right before
-  /// the operation is deleted. At this point, the operation has zero uses.
-  virtual void notifyOperationRemoved(Operation *op) {}
-
-  /// 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();
-  }
+  /// Inherit constructors.
+  using OpBuilder::OpBuilder;
+
+  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(const OpBuilder &builder) : RewriterBase(builder) {}
+  using RewriterBase::RewriterBase;
 };
 
 //===----------------------------------------------------------------------===//
 // 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

mlir/include/mlir/Transforms/DialectConversion.h

 
 namespace detail {
 struct ConversionPatternRewriterImpl;
 } // namespace detail
 
 /// This class implements a pattern rewriter for use with ConversionPatterns. It
 /// extends the base PatternRewriter and provides special conversion specific
 /// hooks.
-class ConversionPatternRewriter final : public PatternRewriter {
+class ConversionPatternRewriter final : public PatternRewriter,
+                                        public RewriteListener {
 public:
   explicit ConversionPatternRewriter(MLIRContext *ctx);
   ~ConversionPatternRewriter() override;
 
   /// Apply a signature conversion to the entry block of the given region. This
   /// replaces the entry block with a new block containing the updated
   /// signature. The new entry block to the region is returned for convenience.
   ///
public:
 
   /// PatternRewriter hook for updating the root operation in-place.
   void finalizeRootUpdate(Operation *op) override;
 
   /// PatternRewriter hook for updating the root operation in-place.
   void cancelRootUpdate(Operation *op) override;
 
   /// PatternRewriter hook for notifying match failure reasons.
-  LogicalResult
+  void
   notifyMatchFailure(Operation *op,
                      function_ref<void(Diagnostic &)> reasonCallback) override;
   using PatternRewriter::notifyMatchFailure;
 
   /// Return a reference to the internal implementation.
   detail::ConversionPatternRewriterImpl &getImpl();
 
 private:

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

 ///
 /// Return success if the iterative process converged and no more patterns can
 /// be matched in the result operation regions.
 ///
 /// Note: This does not apply patterns to the top-level operation itself.
 ///       These methods also perform folding and simple dead-code elimination
 ///       before attempting to match any of the provided patterns.
 ///
-/// You may configure several aspects of this with GreedyRewriteConfig.
-LogicalResult applyPatternsAndFoldGreedily(
-    MutableArrayRef<Region> regions, const FrozenRewritePatternSet &patterns,
-    GreedyRewriteConfig config = GreedyRewriteConfig());
-
-/// Rewrite the given regions, which must be isolated from above.
-inline LogicalResult applyPatternsAndFoldGreedily(
-    Operation *op, const FrozenRewritePatternSet &patterns,
-    GreedyRewriteConfig config = GreedyRewriteConfig()) {
+/// You may configure several aspects of this with GreedyRewriteConfig. A
+/// rewrite listener can be supplied to hook on to rewrite events, such as
+/// operations being removed or replaced.
+LogicalResult
+applyPatternsAndFoldGreedily(MutableArrayRef<Region> regions,
+                             const FrozenRewritePatternSet &patterns,
+                             GreedyRewriteConfig config = GreedyRewriteConfig(),
+                             RewriteListener *listener = nullptr);
+
+/// Rewrite the regions of the given operation, which must be isolated from
+/// above.
+inline LogicalResult
+applyPatternsAndFoldGreedily(Operation *op,
+                             const FrozenRewritePatternSet &patterns,
+                             GreedyRewriteConfig config = GreedyRewriteConfig(),
+                             RewriteListener *listener = nullptr) {
   return applyPatternsAndFoldGreedily(op->getRegions(), patterns, config);
 }
 
 /// Applies the specified patterns on `op` alone while also trying to fold it,
 /// by selecting the highest benefits patterns in a greedy manner. Returns
 /// success if no more patterns can be matched. `erased` is set to true if `op`
 /// was folded away or erased as a result of becoming dead. Note: This does not
 /// apply any patterns recursively to the regions of `op`.

mlir/lib/IR/Builders.cpp

AffineMap Builder::getShiftedAffineMap(AffineMap map, int64_t shift) {
   shiftedResults.reserve(map.getNumResults());
   for (auto resultExpr : map.getResults())
     shiftedResults.push_back(resultExpr + shift);
   return AffineMap::get(map.getNumDims(), map.getNumSymbols(), shiftedResults,
                         context);
 }
 
 //===----------------------------------------------------------------------===//
-// OpBuilder
+// RewriteListener
+//===----------------------------------------------------------------------===//
+
+RewriteListener::~RewriteListener() {}
+
+//===----------------------------------------------------------------------===//
+// ListenerList
 //===----------------------------------------------------------------------===//
 
-OpBuilder::Listener::~Listener() {}
+void ListenerList::notifyOperationInserted(Operation *op) {
+  for (RewriteListener *listener : listeners)
+    listener->notifyOperationInserted(op);
+}
+
+void ListenerList::notifyBlockCreated(Block *block) {
+  for (RewriteListener *listener : listeners)
+    listener->notifyBlockCreated(block);
+}
+
+void ListenerList::notifyRootReplaced(Operation *op) {
+  for (RewriteListener *listener : listeners)
+    listener->notifyRootReplaced(op);
+}
+
+void ListenerList::notifyOperationRemoved(Operation *op) {
+  for (RewriteListener *listener : listeners)
+    listener->notifyOperationRemoved(op);
+}
+
+void ListenerList::notifyMatchFailure(
+    Operation *op, function_ref<void(Diagnostic &)> reasonCallback) {
+  for (RewriteListener *listener : listeners)
+    listener->notifyMatchFailure(op, reasonCallback);
+}
+
+//===----------------------------------------------------------------------===//
+// OpBuilder
+//===----------------------------------------------------------------------===//
 
 /// Insert the given operation at the current insertion point and return it.
 Operation *OpBuilder::insert(Operation *op) {
   if (block)
     block->getOperations().insert(insertPoint, op);
 
   if (listener)
     listener->notifyOperationInserted(op);

mlir/lib/IR/PatternMatch.cpp

 /// results of `op`.
 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);
+  if (RewriteListener *listener = getListener())
+    listener->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)
void RewriterBase::replaceOpWithinBlock(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);
+  if (RewriteListener *listener = getListener())
+    listener->notifyRootReplaced(op);
 
   assert(op->getNumResults() == newValues.size() &&
          "incorrect # of replacement values");
   op->replaceAllUsesWith(newValues);
 
-  notifyOperationRemoved(op);
+  if (RewriteListener *listener = getListener())
+    listener->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 RewriterBase::eraseOp(Operation *op) {
   assert(op->use_empty() && "expected 'op' to have no uses");
-  notifyOperationRemoved(op);
+  if (RewriteListener *listener = getListener())
+    listener->notifyOperationRemoved(op);
   op->erase();
 }
 
 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);
   }

mlir/lib/Transforms/Utils/DialectConversion.cpp

LogicalResult ConversionPatternRewriterImpl::notifyMatchFailure(
   return failure();
 }
 
 //===----------------------------------------------------------------------===//
 // ConversionPatternRewriter
 //===----------------------------------------------------------------------===//
 
 ConversionPatternRewriter::ConversionPatternRewriter(MLIRContext *ctx)
-    : PatternRewriter(ctx),
+    : PatternRewriter(ctx, /*listener=*/this),
       impl(new detail::ConversionPatternRewriterImpl(*this)) {}
 ConversionPatternRewriter::~ConversionPatternRewriter() {}
 
 void ConversionPatternRewriter::replaceOpWithIf(
     Operation *op, ValueRange newValues, bool *allUsesReplaced,
     llvm::unique_function<bool(OpOperand &) const> functor) {
   // TODO: To support this we will need to rework a bit of how replacements are
   // tracked, given that this isn't guranteed to replace all of the uses of an
#endif
   auto &rootUpdates = impl->rootUpdates;
   auto it = llvm::find_if(llvm::reverse(rootUpdates), stateHasOp);
   assert(it != rootUpdates.rend() && "no root update started on op");
   (*it).resetOperation();
   int updateIdx = std::prev(rootUpdates.rend()) - it;
   rootUpdates.erase(rootUpdates.begin() + updateIdx);
 }
 
-LogicalResult ConversionPatternRewriter::notifyMatchFailure(
+void ConversionPatternRewriter::notifyMatchFailure(
     Operation *op, function_ref<void(Diagnostic &)> reasonCallback) {
-  return impl->notifyMatchFailure(op->getLoc(), reasonCallback);
+  (void)impl->notifyMatchFailure(op->getLoc(), reasonCallback);
 }
 
 detail::ConversionPatternRewriterImpl &ConversionPatternRewriter::getImpl() {
   return *impl;
 }
 
 //===----------------------------------------------------------------------===//
 // ConversionPattern

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 
 //===----------------------------------------------------------------------===//
 // GreedyPatternRewriteDriver
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This is a worklist-driven driver for the PatternMatcher, which repeatedly
 /// applies the locally optimal patterns in a roughly "bottom up" way.
-class GreedyPatternRewriteDriver : public PatternRewriter {
+class GreedyPatternRewriteDriver : public RewriteListener {
 public:
   explicit GreedyPatternRewriteDriver(MLIRContext *ctx,
                                       const FrozenRewritePatternSet &patterns,
-                                      const GreedyRewriteConfig &config);
+                                      const GreedyRewriteConfig &config,
+                                      RewriteListener *listener);
 
   /// Simplify the operations within the given regions.
   bool simplify(MutableArrayRef<Region> regions);
 
   /// Add the given operation to the worklist.
   void addToWorklist(Operation *op);
 
   /// Pop the next operation from the worklist.
protected:
   // worklist anymore because we'd get dangling references to it.
   void notifyOperationRemoved(Operation *op) override;
 
   // When the root of a pattern is about to be replaced, it can trigger
   // simplifications to its users - make sure to add them to the worklist
   // before the root is changed.
   void notifyRootReplaced(Operation *op) override;
 
-  /// PatternRewriter hook for erasing a dead operation.
-  void eraseOp(Operation *op) override;
-
   /// PatternRewriter hook for notifying match failure reasons.
-  LogicalResult
+  void
   notifyMatchFailure(Operation *op,
                      function_ref<void(Diagnostic &)> reasonCallback) override;
 
   /// The low-level pattern applicator.
   PatternApplicator matcher;
 
   /// The worklist for this transformation keeps track of the operations that
   /// need to be revisited, plus their index in the worklist.  This allows us to
   /// efficiently remove operations from the worklist when they are erased, even
   /// if they aren't the root of a pattern.
   std::vector<Operation *> worklist;
   DenseMap<Operation *, unsigned> worklistMap;
 
   /// Non-pattern based folder for operations.
   OperationFolder folder;
 
+  /// The pattern rewriter instance to use to perform rewrite operations.
+  PatternRewriter rewriter;
+
 private:
   /// Configuration information for how to simplify.
   GreedyRewriteConfig config;
 
+  /// The listeners attached to the pattern rewriter, including this and an
+  /// optional user-provided listener to notify of rewrite events.
+  ListenerList listeners;
+
 #ifndef NDEBUG
   /// A logger used to emit information during the application process.
   llvm::ScopedPrinter logger{llvm::dbgs()};
 #endif
 };
 } // namespace
 
 GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
     MLIRContext *ctx, const FrozenRewritePatternSet &patterns,
-    const GreedyRewriteConfig &config)
-    : PatternRewriter(ctx), matcher(patterns), folder(ctx), config(config) {
+    const GreedyRewriteConfig &config, RewriteListener *listener)
+    : matcher(patterns), folder(ctx), rewriter(ctx, &listeners),
+      config(config) {
+  // Attach ourselves as a listener.
+  listeners.addListener(this);
+  // Attach the user-provided listener, if there is one.
+  if (listener)
+    listeners.addListener(listener);
+
   worklist.reserve(64);
 
   // Apply a simple cost model based solely on pattern benefit.
   matcher.applyDefaultCostModel();
 }
 
 bool GreedyPatternRewriteDriver::simplify(MutableArrayRef<Region> regions) {
 #ifndef NDEBUG
#ifndef NDEBUG
         LLVM_DEBUG(logResult("failure", "pattern failed to match"));
       };
       auto onSuccess = [&](const Pattern &pattern) {
         LLVM_DEBUG(logResult("success", "pattern applied successfully"));
         return success();
       };
 
       LogicalResult matchResult =
-          matcher.matchAndRewrite(op, *this, canApply, onFailure, onSuccess);
+          matcher.matchAndRewrite(op, rewriter, canApply, onFailure, onSuccess);
       if (succeeded(matchResult))
         LLVM_DEBUG(logResultWithLine("success", "pattern matched"));
       else
         LLVM_DEBUG(logResultWithLine("failure", "pattern failed to match"));
 #else
-      LogicalResult matchResult = matcher.matchAndRewrite(op, *this);
+      LogicalResult matchResult = matcher.matchAndRewrite(op, rewriter);
 #endif
 
-
 #ifndef NDEBUG
 #endif
 
       changed |= succeeded(matchResult);
     }
 
     // After applying patterns, make sure that the CFG of each of the regions
     // is kept up to date.
     if (config.enableRegionSimplification)
-      changed |= succeeded(simplifyRegions(*this, regions));
+      changed |= succeeded(simplifyRegions(rewriter, regions));
   } while (changed &&
            (++iteration < config.maxIterations ||
             config.maxIterations == GreedyRewriteConfig::kNoIterationLimit));
 
   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
   return !changed;
 }
 
for (Value operand : operands) {
     if (!operand || (!operand.use_empty() && !operand.hasOneUse()))
       continue;
     if (auto *defOp = operand.getDefiningOp())
       addToWorklist(defOp);
   }
 }
 
 void GreedyPatternRewriteDriver::notifyOperationRemoved(Operation *op) {
+  LLVM_DEBUG({
+    logger.startLine() << "** Erase   : '" << op->getName() << "'(" << op
+                       << ")\n";
+  });
   addToWorklist(op->getOperands());
   op->walk([this](Operation *operation) {
     removeFromWorklist(operation);
     folder.notifyRemoval(operation);
   });
 }
 
 void GreedyPatternRewriteDriver::notifyRootReplaced(Operation *op) {
   LLVM_DEBUG({
     logger.startLine() << "** Replace : '" << op->getName() << "'(" << op
                        << ")\n";
   });
   for (auto result : op->getResults())
     for (auto *user : result.getUsers())
       addToWorklist(user);
 }
 
-void GreedyPatternRewriteDriver::eraseOp(Operation *op) {
-  LLVM_DEBUG({
-    logger.startLine() << "** Erase   : '" << op->getName() << "'(" << op
-                       << ")\n";
-  });
-  PatternRewriter::eraseOp(op);
-}
-
-LogicalResult GreedyPatternRewriteDriver::notifyMatchFailure(
+void GreedyPatternRewriteDriver::notifyMatchFailure(
     Operation *op, function_ref<void(Diagnostic &)> reasonCallback) {
   LLVM_DEBUG({
     Diagnostic diag(op->getLoc(), DiagnosticSeverity::Remark);
     reasonCallback(diag);
     logger.startLine() << "** Failure : " << diag.str() << "\n";
   });
-  return failure();
 }
 
 /// Rewrite the regions of the specified operation, which must be isolated from
 /// above, by repeatedly applying the highest benefit patterns in a greedy
 /// work-list driven manner. Return success if no more patterns can be matched
 /// in the result operation regions. Note: This does not apply patterns to the
 /// top-level operation itself.
 ///
-LogicalResult
-mlir::applyPatternsAndFoldGreedily(MutableArrayRef<Region> regions,
-                                   const FrozenRewritePatternSet &patterns,
-                                   GreedyRewriteConfig config) {
+LogicalResult mlir::applyPatternsAndFoldGreedily(
+    MutableArrayRef<Region> regions, const FrozenRewritePatternSet &patterns,
+    GreedyRewriteConfig config, RewriteListener *listener) {
   if (regions.empty())
     return success();
 
   // The top-level operation must be known to be isolated from above to
   // prevent performing canonicalizations on operations defined at or above
   // the region containing 'op'.
   auto regionIsIsolated = [](Region &region) {
     return region.getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>();
   };
   (void)regionIsIsolated;
   assert(llvm::all_of(regions, regionIsIsolated) &&
          "patterns can only be applied to operations IsolatedFromAbove");
 
   // Start the pattern driver.
-  GreedyPatternRewriteDriver driver(regions[0].getContext(), patterns, config);
+  GreedyPatternRewriteDriver driver(regions[0].getContext(), patterns, config,
+                                    listener);
   bool converged = driver.simplify(regions);
   LLVM_DEBUG(if (!converged) {
     llvm::dbgs() << "The pattern rewrite doesn't converge after scanning "
                  << config.maxIterations << " times\n";
   });
   return success(converged);
 }
 
 //===----------------------------------------------------------------------===//
 // OpPatternRewriteDriver
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This is a simple driver for the PatternMatcher to apply patterns and perform
 /// folding on a single op. It repeatedly applies locally optimal patterns.
-class OpPatternRewriteDriver : public PatternRewriter {
+class OpPatternRewriteDriver : public RewriteListener {
 public:
   explicit OpPatternRewriteDriver(MLIRContext *ctx,
                                   const FrozenRewritePatternSet &patterns)
-      : PatternRewriter(ctx), matcher(patterns), folder(ctx) {
+      : matcher(patterns), folder(ctx), rewriter(ctx, this) {
     // Apply a simple cost model based solely on pattern benefit.
     matcher.applyDefaultCostModel();
   }
 
   LogicalResult simplifyLocally(Operation *op, int maxIterations, bool &erased);
 
   // These are hooks implemented for PatternRewriter.
 protected:

 
 private:
   /// The low-level pattern applicator.
   PatternApplicator matcher;
 
   /// Non-pattern based folder for operations.
   OperationFolder folder;
 
+  /// The pattern rewriter instance to use to perform rewrite operations.
+  PatternRewriter rewriter;
+
   /// Set to true if the operation has been erased via pattern rewrites.
   bool opErasedViaPatternRewrites = false;
 };
 
 } // namespace
 
 /// Performs the rewrites and folding only on `op`. The simplification
 /// converges if the op is erased as a result of being folded, replaced, or
if (succeeded(folder.tryToFold(op, /*processGeneratedConstants=*/nullptr,
       if (!inPlaceUpdate) {
         erased = true;
         return success();
       }
     }
 
     // Try to match one of the patterns. The rewriter is automatically
     // notified of any necessary changes, so there is nothing else to do here.
-    changed |= succeeded(matcher.matchAndRewrite(op, *this));
+    changed |= succeeded(matcher.matchAndRewrite(op, rewriter));
     if ((erased = opErasedViaPatternRewrites))
       return success();
   } while (changed &&
            (++iterations < maxIterations ||
             maxIterations == GreedyRewriteConfig::kNoIterationLimit));
 
   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
   return failure(changed);

 /// perform folding for a supplied set of ops. It repeatedly simplifies while
 /// restricting the rewrites to only the provided set of ops or optionally
 /// to those directly affected by it (result users or operand providers).
 class MultiOpPatternRewriteDriver : public GreedyPatternRewriteDriver {
 public:
   explicit MultiOpPatternRewriteDriver(MLIRContext *ctx,
                                        const FrozenRewritePatternSet &patterns,
                                        bool strict)
-      : GreedyPatternRewriteDriver(ctx, patterns, GreedyRewriteConfig()),
+      : GreedyPatternRewriteDriver(ctx, patterns, GreedyRewriteConfig(),
+                                   /*listener=*/nullptr),
         strictMode(strict) {}
 
   bool simplifyLocally(ArrayRef<Operation *> op);
 
 private:
   // Look over the provided operands for any defining operations that should
   // be re-added to the worklist. This function should be called when an
   // operation is modified or removed, as it may trigger further
if (succeeded(folder.tryToFold(op, processGeneratedConstants,
         // Op has been erased.
         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.
-    changed |= succeeded(matcher.matchAndRewrite(op, *this));
+    changed |= succeeded(matcher.matchAndRewrite(op, rewriter));
   }
 
   return changed;
 }
 
 /// Rewrites only `op` using the supplied canonicalization patterns and
 /// folding. `erased` is set to true if the op is erased as a result of being
 /// folded, replaced, or dead.