diff --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h --- a/mlir/include/mlir/IR/Builders.h +++ b/mlir/include/mlir/IR/Builders.h @@ -421,6 +421,16 @@ Block *createBlock(Block *insertBefore, TypeRange argTypes = std::nullopt, ArrayRef 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 ®ion, Region &parent, + Region::iterator before, IRMapping &mapping); + void cloneRegionBefore(Region ®ion, Region &parent, + Region::iterator before); + void cloneRegionBefore(Region ®ion, Block *before); + //===--------------------------------------------------------------------===// // Operation Creation //===--------------------------------------------------------------------===// diff --git a/mlir/include/mlir/IR/PatternMatch.h b/mlir/include/mlir/IR/PatternMatch.h --- a/mlir/include/mlir/IR/PatternMatch.h +++ b/mlir/include/mlir/IR/PatternMatch.h @@ -406,16 +406,6 @@ Region::iterator before); void inlineRegionBefore(Region ®ion, 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 ®ion, Region &parent, - Region::iterator before, IRMapping &mapping); - void cloneRegionBefore(Region ®ion, Region &parent, - Region::iterator before); - void cloneRegionBefore(Region ®ion, 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 diff --git a/mlir/include/mlir/Transforms/DialectConversion.h b/mlir/include/mlir/Transforms/DialectConversion.h --- a/mlir/include/mlir/Transforms/DialectConversion.h +++ b/mlir/include/mlir/Transforms/DialectConversion.h @@ -709,14 +709,6 @@ 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 ®ion, Region &parent, - Region::iterator before, IRMapping &mapping) override; - using PatternRewriter::cloneRegionBefore; - /// PatternRewriter hook for inserting a new operation. void notifyOperationInserted(Operation *op) override; 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 @@ -515,26 +515,62 @@ 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 range) { + SmallVector worklist; + + // Remaining block that have not been notified about yet. + SetVector 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 ®ion : 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 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); - }; - walkAndNotify(newOp); - } + if (listener) + for (Region ®ion : newOp->getRegions()) + notifyCloneInsertion(listener, region.getBlocks()); return insert(newOp); } @@ -542,3 +578,23 @@ IRMapping mapper; return clone(op, mapper); } + +void OpBuilder::cloneRegionBefore(Region ®ion, Region &parent, + Region::iterator before, IRMapping &mapping) { + region.cloneInto(&parent, before, mapping); + if (listener) { + Region::iterator clonedBeginIt = + mapping.lookup(®ion.front())->getIterator(); + notifyCloneInsertion(listener, llvm::make_range(clonedBeginIt, before)); + } +} + +void OpBuilder::cloneRegionBefore(Region ®ion, Region &parent, + Region::iterator before) { + IRMapping mapping; + cloneRegionBefore(region, parent, before, mapping); +} + +void OpBuilder::cloneRegionBefore(Region ®ion, Block *before) { + cloneRegionBefore(region, *before->getParent(), before->getIterator()); +} 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 @@ -371,36 +371,3 @@ void RewriterBase::inlineRegionBefore(Region ®ion, 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 ®ion, Region &parent, - Region::iterator before, - IRMapping &mapping) { - region.cloneInto(&parent, before, mapping); - if (Listener *listener = getListener()) { - Region::iterator it = mapping.lookup(®ion.front())->getIterator(); - while (it != before) { - 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++); - } - } -} -void RewriterBase::cloneRegionBefore(Region ®ion, Region &parent, - Region::iterator before) { - IRMapping mapping; - cloneRegionBefore(region, parent, before, mapping); -} -void RewriterBase::cloneRegionBefore(Region ®ion, Block *before) { - cloneRegionBefore(region, *before->getParent(), before->getIterator()); -} diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp --- a/mlir/lib/Transforms/Utils/DialectConversion.cpp +++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp @@ -952,10 +952,6 @@ void notifyRegionIsBeingInlinedBefore(Region ®ion, Region &parent, Region::iterator before); - /// Notifies that the blocks of a region were cloned into another. - void notifyRegionWasClonedBefore(iterator_range &blocks, - Location origRegionLoc); - /// Notifies that a pattern match failed for the given reason. LogicalResult notifyMatchFailure(Location loc, @@ -1463,20 +1459,6 @@ blockActions.push_back(BlockAction::getMove(laterBlock, {®ion, nullptr})); } -void ConversionPatternRewriterImpl::notifyRegionWasClonedBefore( - iterator_range &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 reasonCallback) { LLVM_DEBUG({ @@ -1621,25 +1603,6 @@ PatternRewriter::inlineRegionBefore(region, parent, before); } -void ConversionPatternRewriter::cloneRegionBefore(Region ®ion, - 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(®ion.front())->getIterator(); - auto clonedBlocks = llvm::make_range(clonedBeginIt, before); - impl->notifyRegionWasClonedBefore(clonedBlocks, region.getLoc()); -} - void ConversionPatternRewriter::notifyOperationInserted(Operation *op) { LLVM_DEBUG({ impl->logger.startLine()