diff --git a/mlir/include/mlir/Analysis/Dominance.h b/mlir/include/mlir/Analysis/Dominance.h --- a/mlir/include/mlir/Analysis/Dominance.h +++ b/mlir/include/mlir/Analysis/Dominance.h @@ -34,12 +34,20 @@ /// Recalculate the dominance info. void recalculate(Operation *op); + /// Finds the nearest common dominator block for the two given blocks a + /// and b. If no common dominator can be found, this function will return + /// nullptr. + Block *findNearestCommonDominator(Block *a, Block *b) const; + /// Get the root dominance node of the given region. DominanceInfoNode *getRootNode(Region *region) { assert(dominanceInfos.count(region) != 0); return dominanceInfos[region]->getRootNode(); } + /// Return the dominance node from the Region containing block A. + DominanceInfoNode *getNode(Block *a); + protected: using super = DominanceInfoBase; @@ -82,9 +90,6 @@ return super::properlyDominates(a, b); } - /// Return the dominance node from the Region containing block A. - DominanceInfoNode *getNode(Block *a); - /// Update the internal DFS numbers for the dominance nodes. void updateDFSNumbers(); }; diff --git a/mlir/lib/Analysis/Dominance.cpp b/mlir/lib/Analysis/Dominance.cpp --- a/mlir/lib/Analysis/Dominance.cpp +++ b/mlir/lib/Analysis/Dominance.cpp @@ -26,6 +26,47 @@ // DominanceInfoBase //===----------------------------------------------------------------------===// +template +Block * +DominanceInfoBase::findNearestCommonDominator(Block *a, + Block *b) const { + // If either a or b are null, then conservatively return nullptr. + if (!a || !b) + return nullptr; + + // Get and verify dominance information of regionA. + Region *regionA = a->getParent(); + auto infoAIt = dominanceInfos.find(regionA); + if (infoAIt == dominanceInfos.end()) + return nullptr; + + // If both block do not live in the same region, we will have to check their + // parent operations. + Region *regionB = b->getParent(); + if (regionA != regionB) { + // Get parent operations and check whether they exist and belong to blocks. + Operation *parentA = regionA->getParentOp(); + Operation *parentB = regionB->getParentOp(); + if (!parentA || !parentA->getBlock() || !parentB || !parentB->getBlock()) + return nullptr; + +// Try to find the nearest common dominator of the two parent blocks. + return findNearestCommonDominator(parentA->getBlock(), parentB->getBlock()); + } + + // If the blocks live in the same region, we can rely on already + // existing dominance functionality. + assert(regionA == regionB); + return infoAIt->second->findNearestCommonDominator(a, b); +} + +template +DominanceInfoNode *DominanceInfoBase::getNode(Block *a) { + auto *region = a->getParent(); + assert(dominanceInfos.count(region) != 0); + return dominanceInfos[region]->getNode(a); +} + template void DominanceInfoBase::recalculate(Operation *op) { dominanceInfos.clear(); @@ -132,12 +173,6 @@ return dominates(a.cast().getOwner(), b->getBlock()); } -DominanceInfoNode *DominanceInfo::getNode(Block *a) { - auto *region = a->getParent(); - assert(dominanceInfos.count(region) != 0); - return dominanceInfos[region]->getNode(a); -} - void DominanceInfo::updateDFSNumbers() { for (auto &iter : dominanceInfos) iter.second->updateDFSNumbers();