[mlir][IR] Trigger nested operation/block insertion notifications for clones

* Trigger "operation inserted" and "block inserted" notifications when cloning nested ops. (No such notifications were sent for nested ops when running `OpBuilder::clone(Operation *)` or `RewriterBase::cloneRegionBefore` so far.)
* `cloneRegionBefore` is moved from `RewriterBase` to `OpBuilder`. (`cloneRegionBefore` just builds IR, it does not modify/erase ops.)
* `cloneRegionBefore` is no longer virtual. It used to be virtual so that the dialect conversion can override it and update the internal state in the correct order (for nested ops). The internal state is updated based on listener notifications. Now that the listener notifications are sent for nested clones (and sent in the right order), this workaround in the dialect conversion is no longer needed.

Details:

Ops/blocks are notified in such an order in which they could have been created one-by-one with an `OpBuilder`. This ensures that when a listener is notified about an op being created, it was already notified about the defining ops of the operands (unless there is a block cycle/graph region).

The implementation first clones an entire op (including nested ops) and then sends all notifications. Ideally, notifications should be interleaved with the cloning process, but that would require duplicating `Region::cloneInto` (with listener support). This commit is an incremental improvement over not sending any notifications.

There is a "fast path" in the `clone` functions in case no listener is attached. No IR traversal is needed in that case.

Imported from Phabricator: https://reviews.llvm.org/D146943

BEGIN_PUBLIC
No public commit message needed for presubmit.
END_PUBLIC

Author: matthias-springer

mlir/include/mlir/IR/Builders.h

@@ -451,6 +451,16 @@ class OpBuilder : public Builder {
   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
   //===--------------------------------------------------------------------===//

mlir/include/mlir/IR/PatternMatch.h

@@ -489,16 +489,6 @@ class RewriterBase : public OpBuilder {
                                   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

mlir/include/mlir/Transforms/DialectConversion.h

@@ -727,14 +727,6 @@ class ConversionPatternRewriter final : public PatternRewriter,
                           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;
 

mlir/lib/IR/Builders.cpp

@@ -14,7 +14,11 @@
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Matchers.h"
+#include "mlir/IR/RegionGraphTraits.h"
 #include "mlir/IR/SymbolTable.h"
+#include "llvm/ADT/PostOrderIterator.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVectorExtras.h"
 #include "llvm/Support/raw_ostream.h"
 
@@ -525,22 +529,95 @@ LogicalResult OpBuilder::tryFold(Operation *op,
   return success();
 }
 
+static void notifyOperationCloned(Operation *op, OpBuilder::Listener *listener);
+
+/// Notify the listener that the given range of regions was cloned.
+///
+/// Blocks and operations are enumerated in an order in which they could have
+/// been created separately (without using the `clone` API): Within a region,
+/// blocks are notified according to their successor relationship. Within a
+/// block, operations are notified in forward mode. This ensures that defining
+/// ops are notified before ops that use their results (unless there are
+/// cycles/graph regions).
+static void notifyRegionsCloned(iterator_range<Region::iterator> range,
+                                OpBuilder::Listener *listener) {
+  // Maintain a set of blocks that were not notified yet. This is needed because
+  // the inverse_post_order iterator does not enumerate dead blocks.
+  llvm::SetVector<Block *> remainingBlocks;
+  // The order in which the set is initialized does not matter for correctness.
+  // For better performance, "leaf" blocks with no successors should be starting
+  // point for the block traversal. (Then there are fewer iterations of the
+  // "while" loop.)
+  for (Block &b : llvm::reverse(range))
+    remainingBlocks.insert(&b);
+  // Set of visited blocks that is shared among all inverse_post_order
+  // iterations. This is to avoid that the same block is enumerated multiple
+  // times.
+  llvm::SmallPtrSet<Block *, 4> visited;
+  while (!remainingBlocks.empty()) {
+    // Enumerate predecessors before successors. I.e., reverse post-order.
+    for (Block *b :
+         llvm::inverse_post_order_ext(remainingBlocks.front(), visited)) {
+      auto it = llvm::find(remainingBlocks, b);
+      assert(it != remainingBlocks.end() &&
+             "expected that only remaining blocks are visited");
+      listener->notifyBlockCreated(b);
+      remainingBlocks.erase(it);
+      for (Operation &op : *b)
+        notifyOperationCloned(&op, listener);
+    }
+  }
+}
+
+/// Notify the listener that the given op was cloned.
+static void notifyOperationCloned(Operation *op,
+                                  OpBuilder::Listener *listener) {
+  listener->notifyOperationInserted(op);
+  for (Region &r : op->getRegions())
+    notifyRegionsCloned(r.getBlocks(), listener);
+}
+
 Operation *OpBuilder::clone(Operation &op, IRMapping &mapper) {
+  // TODO: The listener notifications should be interleaved with `clone`.
   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 that got inserted inside those regions as part of cloning.
-  if (listener) {
-    auto walkFn = [&](Operation *walkedOp) {
-      listener->notifyOperationInserted(walkedOp);
-    };
-    for (Region &region : newOp->getRegions())
-      region.walk(walkFn);
-  }
+
+  // Fast path: If no listener is attached, the op can be inserted directly.
+  if (!listener)
+    return insert(newOp);
+
+  // The `insert` call below handles the notification for inserting `newOp`.
+  // Just notify about nested op/block insertion.
+  for (Region &r : newOp->getRegions())
+    notifyRegionsCloned(r.getBlocks(), listener);
   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) {
+  // TODO: The listener notifications should be interleaved with `clone`.
+  region.cloneInto(&parent, before, mapping);
+
+  // Fast path: If no listener is attached, there is no more work to do.
+  if (!listener)
+    return;
+
+  // Notify about op/block insertion.
+  Region::iterator clonedBeginIt =
+      mapping.lookup(&region.front())->getIterator();
+  notifyRegionsCloned(llvm::make_range(clonedBeginIt, before), listener);
+}
+
+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

@@ -458,21 +458,3 @@ void RewriterBase::inlineRegionBefore(Region &region, Region &parent,
 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);
-}
-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

@@ -1588,23 +1588,6 @@ void ConversionPatternRewriter::inlineRegionBefore(Region &region,
   PatternRewriter::inlineRegionBefore(region, parent, before);
 }
 
-void ConversionPatternRewriter::cloneRegionBefore(Region &region,
-                                                  Region &parent,
-                                                  Region::iterator before,
-                                                  IRMapping &mapping) {
-  if (region.empty())
-    return;
-
-  PatternRewriter::cloneRegionBefore(region, parent, before, mapping);
-
-  for (Block &b : ForwardDominanceIterator<>::makeIterable(region)) {
-    Block *cloned = mapping.lookup(&b);
-    impl->notifyCreatedBlock(cloned);
-    cloned->walk<WalkOrder::PreOrder, ForwardDominanceIterator<>>(
-        [&](Operation *op) { notifyOperationInserted(op); });
-  }
-}
-
 void ConversionPatternRewriter::notifyOperationInserted(Operation *op) {
   LLVM_DEBUG({
     impl->logger.startLine()