[mlir][IR] RewriterBase: Use composition instead of inheritance for Listener

[springerm]

Summary

This change makes RewriterBase symmetric to OpBuilder.

  OpBuilder           OpBuilder::Listener
      ^                        ^
      |                        |
RewriterBase        RewriterBase::Listener

• Clients can listen to IR modifications with RewriterBase::Listener.
• RewriterBase no longer inherits from OpBuilder::Listener.
• Only a single listener can be registered at the moment (same as OpBuilder).

RFC: https://discourse.llvm.org/t/rfc-listeners-for-rewriterbase/68198

[springerm]

update

[springerm]

update

[rriddle]

mlir/include/mlir/IR/Builders.h

267

Why does this have classof as a base class?

mlir/lib/IR/PatternMatch.cpp

219–220

Generally we hide this complexity in a CRTP base class, can you do that for Listener too? Generally top-level classes shouldn't be mucking around with TypeID directly.

[springerm]

mlir/include/mlir/IR/Builders.h

267

Not needed, deleted.

mlir/lib/IR/PatternMatch.cpp

219–220

Added a new base class to store the "kind".

I don't know what's the best way to use CRTP here. The problem is that we have a deep class hierarchy with user-derived classes. E.g.:

  ListenerBase
       ^
       |
OpBuilder::Listener
   ^    ^
   |    |
   |  MyBuilderListener (user-derived class)
   |
RewriterBase::Listener
        ^
        |
      MyRewriterBaseListener (user-derived class)

If ListenerBase<T> is CRTP-templatized, there must be a way to pass RewriterBase::Listener as a template argument to ListenerBase<T>. So OpBuilder::Listener must also be templatized. That would bloat up the API because users have to pass a template arg when creating their own OpBuilder::Listener subclass.

Another issue is with RewriterBase::Listener::classof. If I understand correctly, you were thinking of something like:

template <typename T>
ListenerBase<T>::classof(ListenerBaseBase *listener) {
  return TypeID::get<T> == listener->getTypeID();
}

This will probably not work for user-derived classes because TypeID(MyRewriterBaseListener) != TypeID(RewriterBase), so we wouldn't be able to cast MyRewriterBaseListener to RewriterBase::Listener. Also, I need to add yet another superclass ListenerBaseBase.

I changed the TypeID field to Kind. That captures it more accurately because user-derived classes have the same "kind" as their superclasses. Overall, this design is a bit like Value/ValueImpl: E.g., BlockArgumentImpl also has a classof function that checks the kind field.

[springerm]

address comments

[springerm]

update

[springerm]

update

[springerm]

Friendly ping. This change would help us clean up some our code.

[mehdi_amini]

LGTM

[mehdi_amini]

River can you confirm you don't have any other issues with this?

[mehdi_amini]

In D143339#4139425, @mehdi_amini wrote:

River can you confirm you don't have any other issues with this?

Did you get feedback on this @springerm ?

[springerm]

In D143339#4143679, @mehdi_amini wrote:

In D143339#4139425, @mehdi_amini wrote:

River can you confirm you don't have any other issues with this?

Did you get feedback on this @springerm ?

I haven't heard back on this change for the last 2 weeks, so I think he's not following this change anymore. But he approved D143380 (which adds stuff to the new RewriterBase::Listener class), so I think this is fine.

[kiranchandramohan]

I believe this causes some errors in certain flang builds and is causing one buildbot CI to fail. https://lab.llvm.org/buildbot/#/builders/160/builds/16365

../llvm-project/flang/include/flang/Optimizer/Builder/FIRBuilder.h: In copy constructor ‘fir::FirOpBuilder::FirOpBuilder(const fir::FirOpBuilder&)’:
../llvm-project/flang/include/flang/Optimizer/Builder/FIRBuilder.h:52:3: error: base class ‘struct mlir::OpBuilder::Listener’ should be explicitly initialized in the copy constructor [-Werror=extra]
   52 |   FirOpBuilder(const FirOpBuilder &other)
      |   ^~~~~~~~~~~~

[DavidSpickett]

Specifically that is a GCC bot, should you need to reproduce it. Clang is ok with it.

[mehdi_amini]

In D143339#4143805, @springerm wrote:

In D143339#4143679, @mehdi_amini wrote:

In D143339#4139425, @mehdi_amini wrote:

River can you confirm you don't have any other issues with this?

Did you get feedback on this @springerm ?

I haven't heard back on this change for the last 2 weeks,

I know he was out sick..

so I think he's not following this change anymore.

Please get explicit acknowledgement in general and don’t just assume please.

[rriddle]

I am out, but this seems alright to me.

mlir/include/mlir/IR/Builders.h

struct 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; }
 
   /// Returns the current listener of this builder, or nullptr if this builder
   /// doesn't have a listener.
   Listener *getListener() const { return listener; }
 
   //===--------------------------------------------------------------------===//
public:
   Operation *cloneWithoutRegions(Operation &op) {
     return insert(op.cloneWithoutRegions());
   }
   template <typename OpT>
   OpT cloneWithoutRegions(OpT op) {
     return cast<OpT>(cloneWithoutRegions(*op.getOperation()));
   }
 
+protected:
+  /// The optional listener for events of this builder.
+  Listener *listener;
+
 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;
 };
 
 } // namespace mlir
 
 #endif

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:
+  struct Listener : public OpBuilder::Listener {
+    /// Notify the listener that the specified operation is about to be replaced
+    /// with the set of values potentially produced by new operations. This is
+    /// called before the uses of the operation have been changed.
+    virtual void notifyRootReplaced(Operation *op, ValueRange replacement) {}
+
+    /// 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 listener 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
+    /// listeners to optionally hook into the reason why a rewrite failed, and
+    /// display it to users.
+    virtual LogicalResult
+    notifyMatchFailure(Location loc,
+                       function_ref<void(Diagnostic &)> reasonCallback) {
+      return failure();
+    }
+  };
+
   /// 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(Location loc, CallbackT &&reasonCallback) {
 #ifndef NDEBUG
-    return notifyMatchFailure(loc,
-                              function_ref<void(Diagnostic &)>(reasonCallback));
+    if (auto *rewriteListener = dynamic_cast<Listener *>(listener))
+      return rewriteListener->notifyMatchFailure(
+          loc, function_ref<void(Diagnostic &)>(reasonCallback));
+    return failure();
 #else
     return failure();
 #endif
   }
   template <typename CallbackT>
   std::enable_if_t<!std::is_convertible<CallbackT, Twine>::value, LogicalResult>
   notifyMatchFailure(Operation *op, CallbackT &&reasonCallback) {
-    return notifyMatchFailure(op->getLoc(),
-                              function_ref<void(Diagnostic &)>(reasonCallback));
+    if (auto *rewriteListener = dynamic_cast<Listener *>(listener))
+      return rewriteListener->notifyMatchFailure(
+          op->getLoc(), function_ref<void(Diagnostic &)>(reasonCallback));
+    return failure();
   }
   template <typename ArgT>
   LogicalResult notifyMatchFailure(ArgT &&arg, const Twine &msg) {
     return notifyMatchFailure(std::forward<ArgT>(arg),
                               [&](Diagnostic &diag) { diag << msg; });
   }
   template <typename ArgT>
   LogicalResult notifyMatchFailure(ArgT &&arg, const char *msg) {
     return notifyMatchFailure(std::forward<ArgT>(arg), Twine(msg));
   }
 
 protected:
-  /// Initialize the builder with this rewriter as the listener.
-  explicit RewriterBase(MLIRContext *ctx) : OpBuilder(ctx, /*listener=*/this) {}
+  /// Initialize the builder.
+  explicit RewriterBase(MLIRContext *ctx) : OpBuilder(ctx) {}
   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 the set of values potentially produced by new operations. This is
-  /// called before the uses of the operation have been changed.
-  virtual void notifyRootReplaced(Operation *op, ValueRange replacement) {}
-
-  /// 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(Location loc,
-                     function_ref<void(Diagnostic &)> reasonCallback) {
-    return failure();
-  }
+      : 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);

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 RewriterBase::Listener {
 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:
   notifyMatchFailure(Location loc,
                      function_ref<void(Diagnostic &)> reasonCallback) override;
   using PatternRewriter::notifyMatchFailure;
 
   /// Return a reference to the internal implementation.
   detail::ConversionPatternRewriterImpl &getImpl();
 
 private:
+  using OpBuilder::getListener;
+  using OpBuilder::setListener;
+
   std::unique_ptr<detail::ConversionPatternRewriterImpl> impl;
 };
 
 //===----------------------------------------------------------------------===//
 // ConversionTarget
 //===----------------------------------------------------------------------===//
 
 /// This class describes a specific conversion target.

mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp

matchAndRewrite(AwaitType op, typename AwaitType::Adaptor adaptor,
 
     // Delay lowering to block wait in case await op is inside async.execute
     if (!isInCoroutine && !shouldLowerBlockingWait)
       return failure();
 
     // Inside regular functions we use the blocking wait operation to wait for
     // the async object (token, value or group) to become available.
     if (!isInCoroutine) {
-      ImplicitLocOpBuilder builder(loc, op, rewriter.getListener());
+      ImplicitLocOpBuilder builder(loc, op, &rewriter);
       builder.create<RuntimeAwaitOp>(loc, operand);
 
       // Assert that the awaited operands is not in the error state.
       Value isError = builder.create<RuntimeIsErrorOp>(i1, operand);
       Value notError = builder.create<arith::XOrIOp>(
           isError, builder.create<arith::ConstantOp>(
                        loc, i1, builder.getIntegerAttr(i1, 1)));
 
       builder.create<cf::AssertOp>(notError,
                                    "Awaited async operand is in error state");
     }
 
     // Inside the coroutine we convert await operation into coroutine suspension
     // point, and resume execution asynchronously.
     if (isInCoroutine) {
       CoroMachinery &coro = funcCoro->getSecond();
       Block *suspended = op->getBlock();
 
-      ImplicitLocOpBuilder builder(loc, op, rewriter.getListener());
+      ImplicitLocOpBuilder builder(loc, op, &rewriter);
       MLIRContext *ctx = op->getContext();
 
       // Save the coroutine state and resume on a runtime managed thread when
       // the operand becomes available.
       auto coroSaveOp =
           builder.create<CoroSaveOp>(CoroStateType::get(ctx), coro.coroHandle);
       builder.create<RuntimeAwaitAndResumeOp>(operand, coro.coroHandle);
 

mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

static bool hasTensorSemantics(Operation *op) {
 
   bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
   bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
   return hasTensorResult || hasTensorOperand;
 }
 
 namespace {
 /// A rewriter that keeps track of extra information during bufferization.
-class BufferizationRewriter : public IRRewriter {
+class BufferizationRewriter : public IRRewriter, public RewriterBase::Listener {
 public:
   BufferizationRewriter(MLIRContext *ctx, DenseSet<Operation *> &erasedOps,
                         DenseSet<Operation *> &toMemrefOps,
                         SmallVector<Operation *> &worklist,
                         const BufferizationOptions &options,
                         const OpFilter *opFilter,
                         BufferizationStatistics *statistics)
       : IRRewriter(ctx), erasedOps(erasedOps), toMemrefOps(toMemrefOps),
         worklist(worklist), analysisState(options), opFilter(opFilter),
-        statistics(statistics) {}
+        statistics(statistics) {
+    setListener(this);
+  }
 
 protected:
   void notifyOperationRemoved(Operation *op) override {
-    IRRewriter::notifyOperationRemoved(op);
     erasedOps.insert(op);
     // Erase if present.
     toMemrefOps.erase(op);
   }
 
   void notifyOperationInserted(Operation *op) override {
-    IRRewriter::notifyOperationInserted(op);
     erasedOps.erase(op);
 
     // Gather statistics about allocs and deallocs.
     if (statistics) {
       if (auto sideEffectingOp = dyn_cast<MemoryEffectOpInterface>(op)) {
         statistics->numBufferAlloc += static_cast<int64_t>(
             sideEffectingOp.hasEffect<MemoryEffects::Allocate>());
         statistics->numBufferDealloc += static_cast<int64_t>(

mlir/lib/IR/PatternMatch.cpp

void PDLPatternModule::registerRewriteFunction(StringRef name,
   // Allow existing mappings in the case multiple patterns depend on the same
   // rewrite.
   rewriteFunctions.try_emplace(name, std::move(rewriteFn));
 }
 
 //===----------------------------------------------------------------------===//
 // RewriterBase
 //===----------------------------------------------------------------------===//
 
 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 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, newValues);
+  if (auto *rewriteListener = dynamic_cast<Listener *>(listener))
+    rewriteListener->notifyRootReplaced(op, newValues);
 
   // 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, newValues);
+  if (auto *rewriteListener = dynamic_cast<Listener *>(listener))
+    rewriteListener->notifyRootReplaced(op, newValues);
 
   assert(op->getNumResults() == newValues.size() &&
          "incorrect # of replacement values");
   op->replaceAllUsesWith(newValues);
 
-  notifyOperationRemoved(op);
+  if (auto *rewriteListener = dynamic_cast<Listener *>(listener))
+    rewriteListener->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 (auto *rewriteListener = dynamic_cast<Listener *>(listener))
+    rewriteListener->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

 }
 
 //===----------------------------------------------------------------------===//
 // ConversionPatternRewriter
 //===----------------------------------------------------------------------===//
 
 ConversionPatternRewriter::ConversionPatternRewriter(MLIRContext *ctx)
     : PatternRewriter(ctx),
-      impl(new detail::ConversionPatternRewriterImpl(*this)) {}
+      impl(new detail::ConversionPatternRewriterImpl(*this)) {
+  setListener(this);
+}
+
 ConversionPatternRewriter::~ConversionPatternRewriter() = default;
 
 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
   // operation. The main change is that now an operation can be replaced

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 
 namespace {
 /// This is a worklist-driven driver for the PatternMatcher, which repeatedly
 /// applies the locally optimal patterns.
 ///
 /// This abstract class manages the worklist and contains helper methods for
 /// rewriting ops on the worklist. Derived classes specify how ops are added
 /// to the worklist in the beginning.
-class GreedyPatternRewriteDriver : public PatternRewriter {
+class GreedyPatternRewriteDriver : public PatternRewriter,
+                                   public RewriterBase::Listener {
 protected:
   explicit GreedyPatternRewriteDriver(MLIRContext *ctx,
                                       const FrozenRewritePatternSet &patterns,
                                       const GreedyRewriteConfig &config);
 
   /// Add the given operation to the worklist.
   void addSingleOpToWorklist(Operation *op);
 
GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
     MLIRContext *ctx, const FrozenRewritePatternSet &patterns,
     const GreedyRewriteConfig &config)
     : PatternRewriter(ctx), folder(ctx), config(config), matcher(patterns) {
   assert(config.scope && "scope is not specified");
   worklist.reserve(64);
 
   // Apply a simple cost model based solely on pattern benefit.
   matcher.applyDefaultCostModel();
+
+  // Set up listener.
+  setListener(this);
 }
 
 bool GreedyPatternRewriteDriver::processWorklist() {
 #ifndef NDEBUG
   const char *logLineComment =
       "//===-------------------------------------------===//\n";
 
   /// A utility function to log a process result for the given reason.