[mlir][IR] Trigger notifyBlockCreated when cloning ops

[springerm]

Summary

When cloning an op or a region, we currently trigger the notifyOperationInserted callback but not the notifyBlockCreated callback.

Depends On: D144194

mlir/lib/IR/Builders.cpp

LogicalResult OpBuilder::tryFold(Operation *op,
 
   return success();
 }
 
 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 that got inserted inside those regions as part of cloning.
+  // about any ops/blocks that got inserted inside those regions as part of
+  // cloning.
   if (listener) {
-    auto walkFn = [&](Operation *walkedOp) {
-      // TODO: Notify listener of block creation.
-      listener->notifyOperationInserted(walkedOp);
+    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);
     };
-    for (Region &region : newOp->getRegions())
-      region.walk(walkFn);
+    walkAndNotify(newOp);
   }
   return insert(newOp);
 }
 
 Operation *OpBuilder::clone(Operation &op) {
   IRMapping mapper;
   return clone(op, mapper);
 }

mlir/lib/IR/PatternMatch.cpp

 /// 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) {
-      // TODO: Notify listener of nested blocks.
-      it->walk([&](Operation *op) { listener->notifyOperationInserted(op); });
-      it++;
+      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());
 }