[mlir][OpBuilder] Remove the vtable from OpBuilder in favor of using the listener pattern

[rriddle]

Summary

The current OpBuilder has a set of virtual functions required by the fact that the PatternRewriter inherits from it for convenience. The PatternRewriter is required to know about IR mutations for correctness. This revision changes the relationship to be explicit by having users register a listener with the builder instead of using inheritance/vtables. This still requires that users properly transfer the listener when creating new builders, but has several benefits:

• More than one builder can be created during pattern rewrites(assuming that the listener is properly forwarded)
• OpBuilder no longer requires a vtable, and thus does not incur the cost when a listener isn't present.

[rriddle]

Update comment to use ///

[rriddle]

Rebase

[ftynse]

Thanks, River!

I am not sure I completely understand the relation of this with EDSC/ScopedContext (I expect it to be just transparent), but we intend to clean the ScopedContext soon anyway.

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

389 ↗

(On Diff #261354)

Could you please elaborate why this is necessary? There are other uses of GenericLoopNestRangeBuilder in LinalgToLoops that are not modified...

[rriddle]

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

389 ↗

(On Diff #261354)

b is being copied, so wasn't previously notifying the PatterRewriter of any generated operations. One of the tests relied on this behavior and broke, but Nicolas removed that test so this can be removed now. This was the only builder that a test relied on for the old(incorrect) behavior, so it was the only one I updated.

[rriddle]

Remove FIXME that is no longer necessary

mlir/include/mlir/IR/Builders.h

 protected:
   MLIRContext *context;
 };
 
 /// 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) : Builder(ctx) {}
+  explicit OpBuilder(MLIRContext *ctx, Listener *listener = nullptr)
+      : Builder(ctx), listener(listener) {}
 
   /// Create a builder and set the insertion point to the start of the region.
-  explicit OpBuilder(Region *region) : Builder(region->getContext()) {
+  explicit OpBuilder(Region *region, Listener *listener = nullptr)
+      : OpBuilder(region->getContext(), listener) {
     if (!region->empty())
       setInsertionPoint(&region->front(), region->front().begin());
   }
-  explicit OpBuilder(Region &region) : OpBuilder(&region) {}
-
-  virtual ~OpBuilder();
+  explicit OpBuilder(Region &region, Listener *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) : Builder(op->getContext()) {
+  explicit OpBuilder(Operation *op, Listener *listener = nullptr)
+      : OpBuilder(op->getContext(), listener) {
     setInsertionPoint(op);
   }
 
-  OpBuilder(Block *block, Block::iterator insertPoint)
-      : OpBuilder(block->getParent()) {
+  OpBuilder(Block *block, Block::iterator insertPoint,
+            Listener *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) {
-    return OpBuilder(block, block->begin());
+  static OpBuilder atBlockBegin(Block *block, Listener *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) {
-    return OpBuilder(block, block->end());
+  static OpBuilder atBlockEnd(Block *block, Listener *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) {
+  static OpBuilder atBlockTerminator(Block *block,
+                                     Listener *listener = nullptr) {
     auto *terminator = block->getTerminator();
     assert(terminator != nullptr && "the block has no terminator");
-    return OpBuilder(block, terminator->getIterator());
+    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; }
+
+  /// Returns the current listener of this builder, or nullptr if this builder
+  /// doesn't have a listener.
+  Listener *getListener() const { return listener; }
+
+  //===--------------------------------------------------------------------===//
+  // Insertion Point Management
+  //===--------------------------------------------------------------------===//
+
   /// This class represents a saved insertion point.
   class InsertPoint {
   public:
     /// Creates a new insertion point which doesn't point to anything.
     InsertPoint() = default;
 
     /// Creates a new insertion point at the given location.
     InsertPoint(Block *insertBlock, Block::iterator insertPt)
public:
 
   /// Return the block the current insertion point belongs to.  Note that the
   /// the insertion point is not necessarily the end of the block.
   Block *getInsertionBlock() const { return block; }
 
   /// Returns the current insertion point of the builder.
   Block::iterator getInsertionPoint() const { return insertPoint; }
 
-  /// Insert the given operation at the current insertion point and return it.
-  virtual Operation *insert(Operation *op);
+  /// Returns the current block of the builder.
+  Block *getBlock() const { return block; }
+
+  //===--------------------------------------------------------------------===//
+  // Block Creation
+  //===--------------------------------------------------------------------===//
 
   /// Add new block with 'argTypes' arguments and set the insertion point to the
   /// end of it. The block is inserted at the provided insertion point of
   /// 'parent'.
-  virtual Block *createBlock(Region *parent, Region::iterator insertPt = {},
-                             TypeRange argTypes = llvm::None);
+  Block *createBlock(Region *parent, Region::iterator insertPt = {},
 
   /// Add new block with 'argTypes' arguments and set the insertion point to the
   /// end of it. The block is placed before 'insertBefore'.
   Block *createBlock(Block *insertBefore, TypeRange argTypes = llvm::None);
 
-  /// Returns the current block of the builder.
-  Block *getBlock() const { return block; }
+  //===--------------------------------------------------------------------===//
+  // Operation Creation
+  //===--------------------------------------------------------------------===//
+
+  /// Insert the given operation at the current insertion point and return it.
+  Operation *insert(Operation *op);
 
   /// Creates an operation given the fields represented as an OperationState.
   Operation *createOperation(const OperationState &state);
 
   /// Create an operation of specific op type at the current insertion point.
   template <typename OpTy, typename... Args>
   OpTy create(Location location, Args &&... args) {
     OperationState state(location, OpTy::getOperationName());
public:
   Operation *cloneWithoutRegions(Operation &op) {
     return insert(op.cloneWithoutRegions());
   }
   template <typename OpT> OpT cloneWithoutRegions(OpT op) {
     return cast<OpT>(cloneWithoutRegions(*op.getOperation()));
   }
 
 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

 ///
 /// 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 {
+class PatternRewriter : public OpBuilder, public OpBuilder::Listener {
 public:
   /// Create operation of specific op type at the current insertion point
   /// without verifying to see if it is valid.
   template <typename OpTy, typename... Args>
   OpTy create(Location location, Args... args) {
     OperationState state(location, OpTy::getOperationName());
     OpTy::build(*this, state, args...);
     auto *op = createOperation(state);
OpTy createChecked(Location location, Args... args) {
     }
 
     // Otherwise, the error message got emitted.  Just remove the operation
     // we made.
     op->erase();
     return OpTy();
   }
 
-  /// This is implemented to insert the specified operation and serves as a
-  /// notification hook for rewriters that want to know about new operations.
-  virtual Operation *insert(Operation *op) = 0;
-
   /// 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);
 
#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:
-  explicit PatternRewriter(MLIRContext *ctx) : OpBuilder(ctx) {}
-  virtual ~PatternRewriter();
+  /// Initialize the builder with this rewriter as the listener.
+  explicit PatternRewriter(MLIRContext *ctx)
+      : OpBuilder(ctx, /*listener=*/this) {}
+  ~PatternRewriter() override;
 
-  // These are the callback methods that subclasses can choose to implement if
-  // they would like to be notified about certain types of mutations.
+  /// 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.
   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

mlir/include/mlir/Transforms/DialectConversion.h

public:
   /// operation *must* be made dead by the end of the conversion process,
   /// otherwise an assert will be issued.
   void eraseOp(Operation *op) override;
 
   /// PatternRewriter hook for erase all operations in a block. This is not yet
   /// implemented for dialect conversion.
   void eraseBlock(Block *block) override;
 
-  /// PatternRewriter hook for creating a new block with the given arguments.
-  Block *createBlock(Region *parent, Region::iterator insertPt = {},
-                     TypeRange argTypes = llvm::None) override;
+  /// PatternRewriter hook creating a new block.
+  void notifyBlockCreated(Block *block) override;
 
   /// PatternRewriter hook for splitting a block into two parts.
   Block *splitBlock(Block *block, Block::iterator before) override;
 
   /// PatternRewriter hook for merging a block into another.
   void mergeBlocks(Block *source, Block *dest, ValueRange argValues) override;
 
   /// PatternRewriter hook for moving blocks out of a region.
   void inlineRegionBefore(Region &region, Region &parent,
                           Region::iterator before) override;
   using PatternRewriter::inlineRegionBefore;
 
   /// PatternRewriter hook for cloning blocks of one region into another. The
   /// given region to clone *must* not have been modified as part of conversion
   /// yet, i.e. it must be within an operation that is either in the process of
   /// conversion, or has not yet been converted.
   void cloneRegionBefore(Region &region, Region &parent,
                          Region::iterator before,
                          BlockAndValueMapping &mapping) override;
   using PatternRewriter::cloneRegionBefore;
 
   /// PatternRewriter hook for inserting a new operation.
-  Operation *insert(Operation *op) override;
+  void notifyOperationInserted(Operation *op) override;
 
   /// PatternRewriter hook for updating the root operation in-place.
   /// Note: These methods only track updates to the top-level operation itself,
   /// and not nested regions. Updates to regions will still require notification
   /// through other more specific hooks above.
   void startRootUpdate(Operation *op) override;
 
   /// PatternRewriter hook for updating the root operation in-place.

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.
+// OpBuilder
 //===----------------------------------------------------------------------===//
 
-OpBuilder::~OpBuilder() {}
+OpBuilder::Listener::~Listener() {}
 
 /// 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);
   return op;
 }
 
 /// Add new block with 'argTypes' arguments and set the insertion point to the
 /// end of it. The block is inserted at the provided insertion point of
 /// 'parent'.
 Block *OpBuilder::createBlock(Region *parent, Region::iterator insertPt,
                               TypeRange argTypes) {
   assert(parent && "expected valid parent region");
   if (insertPt == Region::iterator())
     insertPt = parent->end();
 
   Block *b = new Block();
   b->addArguments(argTypes);
   parent->getBlocks().insert(insertPt, b);
   setInsertionPointToEnd(b);
+
+  if (listener)
+    listener->notifyBlockCreated(b);
   return b;
 }
 
 /// Add new block with 'argTypes' arguments and set the insertion point to the
 /// end of it.  The block is placed before 'insertBefore'.
 Block *OpBuilder::createBlock(Block *insertBefore, TypeRange argTypes) {
   assert(insertBefore && "expected valid insertion block");
   return createBlock(insertBefore->getParent(), Region::iterator(insertBefore),

mlir/lib/Transforms/DialectConversion.cpp

 
 /// Return the converted value that replaces 'key'. Return 'key' if there is
 /// no such a converted value.
 Value ConversionPatternRewriter::getRemappedValue(Value key) {
   return impl->mapping.lookupOrDefault(key);
 }
 
 /// PatternRewriter hook for creating a new block with the given arguments.
-Block *ConversionPatternRewriter::createBlock(Region *parent,
-                                              Region::iterator insertPtr,
-                                              TypeRange argTypes) {
-  Block *block = PatternRewriter::createBlock(parent, insertPtr, argTypes);
+void ConversionPatternRewriter::notifyBlockCreated(Block *block) {
   impl->notifyCreatedBlock(block);
-  return block;
 }
 
 /// PatternRewriter hook for splitting a block into two parts.
 Block *ConversionPatternRewriter::splitBlock(Block *block,
                                              Block::iterator before) {
   auto *continuation = PatternRewriter::splitBlock(block, before);
   impl->notifySplitBlock(block, continuation);
   return continuation;
void ConversionPatternRewriter::cloneRegionBefore(
 
   // Collect the range of the cloned blocks.
   auto clonedBeginIt = mapping.lookup(&region.front())->getIterator();
   auto clonedBlocks = llvm::make_range(clonedBeginIt, before);
   impl->notifyRegionWasClonedBefore(clonedBlocks, region.getLoc());
 }
 
 /// PatternRewriter hook for creating a new operation.
-Operation *ConversionPatternRewriter::insert(Operation *op) {
+void ConversionPatternRewriter::notifyOperationInserted(Operation *op) {
   LLVM_DEBUG({
     impl->logger.startLine()
         << "** Insert  : '" << op->getName() << "'(" << op << ")\n";
   });
   impl->createdOps.push_back(op);
-  return OpBuilder::insert(op);
 }
 
 /// PatternRewriter hook for updating the root operation in-place.
 void ConversionPatternRewriter::startRootUpdate(Operation *op) {
 #ifndef NDEBUG
   impl->pendingRootUpdates.insert(op);
 #endif
   impl->rootUpdates.emplace_back(op);

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

if (it != worklistMap.end()) {
       worklistMap.erase(it);
     }
   }
 
   // These are hooks implemented for PatternRewriter.
 protected:
   // Implement the hook for inserting operations, and make sure that newly
   // inserted ops are added to the worklist for processing.
-  Operation *insert(Operation *op) override {
-    addToWorklist(op);
-    return OpBuilder::insert(op);
-  }
+  void notifyOperationInserted(Operation *op) override { addToWorklist(op); }
 
   // If an operation is about to be removed, make sure it is not in our
   // worklist anymore because we'd get dangling references to it.
   void notifyOperationRemoved(Operation *op) override {
     addToWorklist(op->getOperands());
     op->walk([this](Operation *operation) {
       removeFromWorklist(operation);
       folder.notifyRemoval(operation);

 class OpPatternRewriteDriver : public PatternRewriter {
 public:
   explicit OpPatternRewriteDriver(MLIRContext *ctx,
                                   const OwningRewritePatternList &patterns)
       : PatternRewriter(ctx), matcher(patterns), folder(ctx) {}
 
   bool simplifyLocally(Operation *op, int maxIterations, bool &erased);
 
-  /// No additional action needed other than inserting the op.
-  Operation *insert(Operation *op) override { return OpBuilder::insert(op); }
-
   // These are hooks implemented for PatternRewriter.
 protected:
   /// If an operation is about to be removed, mark it so that we can let clients
   /// know.
   void notifyOperationRemoved(Operation *op) override {
     opErasedViaPatternRewrites = true;
   }