[mlir][IR] Move cloneRegionBefore to OpBuilder

[springerm]

Summary

OpBuilder provides helper functions for cloning ops. MLIR also provides helper functions for cloning regions, but they are currently part of RewriterBase. They can be part of OpBuilder because they only insert new IR.

To be compatible with the dialect conversion, blocks (and their containing ops) must be notified according to their successor relationship. Without dialect conversion in mind, this may seem arbitrary, but it is a valid notification order for any listener. For a general listener, the only requirement when cloning IR is that we send out notifications in way that could that have been the result of creating the operations (and blocks) one-by-one through the builder API.

Dialect conversion is simplified: cloneRegionBefore is no longer overridden. Furthermore, regions that are inserted via OpBuilder::clone are now notified in an order that is suitable for dialect conversion. (Dialect conversion used to override RewriterBase::cloneRegionBefore but not OpBuilder::clone; this was a bug.)

Depends On: D144266

[springerm]

update

[mehdi_amini]

mlir/include/mlir/IR/Builders.h

231

This class is free of a vtable right now and as a default destructor. We're making it not just free anymore to just use freely an OpBuilder anywhere just as a convenient access to the method it provides to get builtin attributes/types.

[springerm]

split revision and rebase

[springerm]

mlir/include/mlir/IR/Builders.h

231

I think I can make cloneRegionBefore non-virtual. It is overridden in the rewriter of the dialect conversion. If we insert the ops in the "right" order in the OpBuilder, it should not be necessary to override the function. This may also make the processing a bit more efficient in the GreedyPatternRewriteDriver because the ops would be enqueued following the linking of blocks. (Not really sure about this, maybe it does nothing...)

[springerm]

OpBuilder no longer virtual

[springerm]

This needs a test case (for the OpBuilder::clone bug that is fixed in the dialect conversion) but I wanted to gather people's opinions first.

mlir/include/mlir/IR/Builders.h

public:
   }
 
   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, 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, Listener *listener = nullptr) {
public:
 
   /// Add new block with 'argTypes' arguments and set the insertion point to the
   /// end of it. The block is placed before 'insertBefore'. `locs` contains the
   /// locations of the inserted arguments, and should match the size of
   /// `argTypes`.
   Block *createBlock(Block *insertBefore, TypeRange argTypes = std::nullopt,
                      ArrayRef<Location> locs = std::nullopt);
 
+  /// 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 cloneRegionBefore(Region &region, Region &parent,
+                         Region::iterator before, IRMapping &mapping);
+  void cloneRegionBefore(Region &region, Region &parent,
+                         Region::iterator before);
+  void cloneRegionBefore(Region &region, Block *before);
+
   //===--------------------------------------------------------------------===//
   // 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.

mlir/include/mlir/IR/PatternMatch.h

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, IRMapping &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 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.

mlir/include/mlir/Transforms/DialectConversion.h

public:
   /// 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, IRMapping &mapping) override;
-  using PatternRewriter::cloneRegionBefore;
-
   /// PatternRewriter hook for inserting a new operation.
   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;

mlir/lib/IR/Builders.cpp

LogicalResult OpBuilder::tryFold(Operation *op,
 
   // If we were successful, insert any generated constants.
   for (Operation *cst : generatedConstants)
     insert(cst);
 
   return success();
 }
 
+/// Notifies the listener about operation insertions and block creations for
+/// all blocks in the given range. The order in which the listener is notified
+/// about operations/block creation is one that could have been a result of
+/// creating the operations one-by-one.
+///
+/// Operations are notified post-order and top-to-bottom. Furthermore, to
+/// satisfy a dialect conversion requirement, blocks on the same nesting level
+/// are visited according to their "successor" relationship.
+static void notifyCloneInsertion(OpBuilder::Listener *listener,
+                                 iterator_range<Region::iterator> range) {
+  SmallVector<Block *> worklist;
+
+  // Remaining block that have not been notified about yet.
+  SetVector<Block *> remainingBlocks;
+  for (Block &b : range)
+    remainingBlocks.insert(&b);
+
+  while (!remainingBlocks.empty()) {
+    // Populate worklist with the next remaining block.
+    auto it = remainingBlocks.begin();
+    worklist.push_back(*it);
+    remainingBlocks.erase(it);
+
+    // Iteratively visit the block and all of its successors.
+    while (!worklist.empty()) {
+      Block *block = worklist.pop_back_val();
+      listener->notifyBlockCreated(block);
+
+      // Notify insertion for each operation in the block.
+      for (Operation &op : *block) {
+        for (auto &region : op.getRegions())
+          notifyCloneInsertion(listener, region.getBlocks());
+        listener->notifyOperationInserted(&op);
+      }
+
+      // Add unvisited successors to the worklist.
+      for (Block *succ : block->getSuccessors()) {
+        auto it = llvm::find(remainingBlocks, succ);
+        if (it != remainingBlocks.end()) {
+          worklist.push_back(*it);
+          remainingBlocks.erase(it);
+        }
+      }
+    }
+  }
+}
+
 Operation *OpBuilder::clone(Operation &op, IRMapping &mapper) {
   Operation *newOp = op.clone(mapper);
   // The `insert` call below handles the notification for inserting `newOp`
   // itself. But if `newOp` has any regions, we need to notify the listener
   // about any ops/blocks that got inserted inside those regions as part of
   // cloning.
-  if (listener) {
-    std::function<void(Operation *)> walkAndNotify = [&](Operation *op) {
-      for (Region &region : op->getRegions()) {
-        for (Block &block : region) {
-          listener->notifyBlockCreated(&block);
-          for (Operation &op : block)
-            walkAndNotify(&op);
-        }
-      }
-      if (op != newOp)
-        listener->notifyOperationInserted(op);
-    };
-    walkAndNotify(newOp);
-  }
+  if (listener)
+    for (Region &region : newOp->getRegions())
+      notifyCloneInsertion(listener, region.getBlocks());
   return insert(newOp);
 }
 
 Operation *OpBuilder::clone(Operation &op) {
   IRMapping mapper;
   return clone(op, mapper);
 }
+
+void OpBuilder::cloneRegionBefore(Region &region, Region &parent,
+                                  Region::iterator before, IRMapping &mapping) {
+  region.cloneInto(&parent, before, mapping);
+  if (listener) {
+    Region::iterator clonedBeginIt =
+        mapping.lookup(&region.front())->getIterator();
+    notifyCloneInsertion(listener, llvm::make_range(clonedBeginIt, before));
+  }
+}
+
+void OpBuilder::cloneRegionBefore(Region &region, Region &parent,
+                                  Region::iterator before) {
+  IRMapping mapping;
+  cloneRegionBefore(region, parent, before, mapping);
+}
+
+void OpBuilder::cloneRegionBefore(Region &region, Block *before) {
+  cloneRegionBefore(region, *before->getParent(), before->getIterator());
+}

mlir/lib/IR/PatternMatch.cpp

 /// region and pass it the correct block arguments.
 void RewriterBase::inlineRegionBefore(Region &region, Region &parent,
                                       Region::iterator before) {
   parent.getBlocks().splice(before, region.getBlocks());
 }
 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 RewriterBase::cloneRegionBefore(Region &region, Region &parent,
-                                     Region::iterator before,
-                                     IRMapping &mapping) {
-  region.cloneInto(&parent, before, mapping);
-  if (Listener *listener = getListener()) {
-    Region::iterator it = mapping.lookup(&region.front())->getIterator();
-    while (it != before) {
-      std::function<void(Block &)> walkAndNotify = [&](Block &block) {
-        listener->notifyBlockCreated(&block);
-        for (Operation &op : block) {
-          for (Region &region : op.getRegions())
-            for (Block &block : region)
-              walkAndNotify(block);
-          listener->notifyOperationInserted(&op);
-        }
-      };
-      walkAndNotify(*it++);
-    }
-  }
-}
-void RewriterBase::cloneRegionBefore(Region &region, Region &parent,
-                                     Region::iterator before) {
-  IRMapping mapping;
-  cloneRegionBefore(region, parent, before, mapping);
-}
-void RewriterBase::cloneRegionBefore(Region &region, Block *before) {
-  cloneRegionBefore(region, *before->getParent(), before->getIterator());
-}

mlir/lib/Transforms/Utils/DialectConversion.cpp

struct ConversionPatternRewriterImpl {
 
   /// Notifies that `block` is being merged with `srcBlock`.
   void notifyBlocksBeingMerged(Block *block, Block *srcBlock);
 
   /// Notifies that the blocks of a region are about to be moved.
   void notifyRegionIsBeingInlinedBefore(Region &region, Region &parent,
                                         Region::iterator before);
 
-  /// Notifies that the blocks of a region were cloned into another.
-  void notifyRegionWasClonedBefore(iterator_range<Region::iterator> &blocks,
-                                   Location origRegionLoc);
-
   /// Notifies that a pattern match failed for the given reason.
   LogicalResult
   notifyMatchFailure(Location loc,
                      function_ref<void(Diagnostic &)> reasonCallback);
 
   //===--------------------------------------------------------------------===//
   // State
   //===--------------------------------------------------------------------===//
void ConversionPatternRewriterImpl::notifyRegionIsBeingInlinedBefore(
   for (auto &earlierBlock : llvm::drop_begin(llvm::reverse(region), 1)) {
     blockActions.push_back(
         BlockAction::getMove(laterBlock, {&region, &earlierBlock}));
     laterBlock = &earlierBlock;
   }
   blockActions.push_back(BlockAction::getMove(laterBlock, {&region, nullptr}));
 }
 
-void ConversionPatternRewriterImpl::notifyRegionWasClonedBefore(
-    iterator_range<Region::iterator> &blocks, Location origRegionLoc) {
-  for (Block &block : blocks)
-    blockActions.push_back(BlockAction::getCreate(&block));
-
-  // Compute the conversion set for the inlined region.
-  auto result = computeConversionSet(blocks, origRegionLoc, createdOps);
-
-  // This original region has already had its conversion set computed, so there
-  // shouldn't be any new failures.
-  (void)result;
-  assert(succeeded(result) && "expected region to have no unreachable blocks");
-}
-
 LogicalResult ConversionPatternRewriterImpl::notifyMatchFailure(
     Location loc, function_ref<void(Diagnostic &)> reasonCallback) {
   LLVM_DEBUG({
     Diagnostic diag(loc, DiagnosticSeverity::Remark);
     reasonCallback(diag);
     logger.startLine() << "** Failure : " << diag.str() << "\n";
     if (notifyCallback)
       notifyCallback(diag);

 
 void ConversionPatternRewriter::inlineRegionBefore(Region &region,
                                                    Region &parent,
                                                    Region::iterator before) {
   impl->notifyRegionIsBeingInlinedBefore(region, parent, before);
   PatternRewriter::inlineRegionBefore(region, parent, before);
 }
 
-void ConversionPatternRewriter::cloneRegionBefore(Region &region,
-                                                  Region &parent,
-                                                  Region::iterator before,
-                                                  IRMapping &mapping) {
-  // Note: Do not call PatternRewriter::cloneRegionBefore. We cannot rely on its
-  // "op created" notifications because we need to populate the list of created
-  // ops in a certain order. This is done in `notifyRegionWasClonedBefore`.
-
-  if (region.empty())
-    return;
-
-  region.cloneInto(&parent, before, mapping);
-
-  // 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());
-}
-
 void ConversionPatternRewriter::notifyOperationInserted(Operation *op) {
   LLVM_DEBUG({
     impl->logger.startLine()
         << "** Insert  : '" << op->getName() << "'(" << op << ")\n";
   });
   impl->createdOps.push_back(op);
 }