diff --git a/mlir/lib/IR/Builders.cpp b/mlir/lib/IR/Builders.cpp --- a/mlir/lib/IR/Builders.cpp +++ b/mlir/lib/IR/Builders.cpp @@ -519,14 +519,21 @@ 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 walkAndNotify = [&](Operation *op) { + for (Region ®ion : op->getRegions()) { + for (Block &block : region) { + listener->notifyBlockCreated(&block); + for (Operation &op : block) + walkAndNotify(&op); + } + } + if (op != newOp) + listener->notifyOperationInserted(op); }; - for (Region ®ion : newOp->getRegions()) - region.walk(walkFn); + walkAndNotify(newOp); } return insert(newOp); } diff --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp --- a/mlir/lib/IR/PatternMatch.cpp +++ b/mlir/lib/IR/PatternMatch.cpp @@ -383,9 +383,16 @@ if (Listener *listener = getListener()) { Region::iterator it = mapping.lookup(®ion.front())->getIterator(); while (it != before) { - // TODO: Notify listener of nested blocks. - it->walk([&](Operation *op) { listener->notifyOperationInserted(op); }); - it++; + std::function walkAndNotify = [&](Block &block) { + listener->notifyBlockCreated(&block); + for (Operation &op : block) { + for (Region ®ion : op.getRegions()) + for (Block &block : region) + walkAndNotify(block); + listener->notifyOperationInserted(&op); + } + }; + walkAndNotify(*it++); } } }