[Dominators] Rewrite the dominator implementation for efficiency. NFC.

[lattner]

Summary

The previous impl densely scanned the entire region starting with an op
when dominators were created, creating a DominatorTree for every region.

This is extremely expensive up front -- particularly for clients like
Linalg/Transforms/Fusion.cpp that construct DominanceInfo for a single
query. It is also extremely memory wasteful for IRs that use single
block regions commonly (e.g. affine.for) because it's making a
dominator tree for a region that has trivial dominance. The
implementation also had numerous unnecessary minor efficiencies, e.g.
doing multiple walks of the region tree or tryGetBlocksInSameRegion
building a DenseMap that it didn't need.

This patch switches to an approach where [Post]DominanceInfo is free
to construct, and which lazily constructs DominatorTree's for any
multiblock regions that it needs. This avoids the up-front cost
entirely, making its runtime proportional to the complexity of the
region tree instead of # ops in a region. This also avoids the memory
and time cost of creating DominatorTree's for single block regions.

Finally this rewrites the implementation for simplicity and to avoids
the constant factor problems the old implementation had.

[lattner]

Adding a few others who have touched this code.

[bondhugula]

I'll be able to review this tomorrow, but a few really minor things I noticed.

mlir/include/mlir/IR/Region.h

68

return llvm::hasSingleElement(blocks); will do exactly the same thing.

mlir/lib/IR/Dominance.cpp

87

But a Block is always contained in an op unless it's not linked into the IR tree - should the comment be updated or I'm missing something?

mlir/lib/Transforms/CSE.cpp

216

Nit: You've added a redundant emptiness check with this NFC change FWIW.

[rriddle]

Thanks! This is another thing that's been lingering for a while to be cleaned up. Took a cursory glance and generally seems alright, most comments are stylistic. Will look again when I get back to the office if some else hasn't approved by then.

mlir/include/mlir/IR/Dominance.h

46–49

nit: Can you please remove all of the double space after periods? It doesn't match the rest of the comment style and leads to a weird cognitive mismatch.

82

88

97

Any reason not to use PointerIntPair here?

128

164

Can you document this method?

mlir/include/mlir/IR/Region.h

68

We generally use llvm::hasSingleElement instead of adding methods like this.

mlir/lib/IR/Dominance.cpp

63–64

style nit: If one branch has braces, all branches should have braces.

145

305

mlir/lib/Transforms/BufferOptimizations.cpp

198

nit: I'd prefer not to have comments at the end of lines. It doesn't match the style anywhere else here.

[stephenneuendorffer]

mlir/include/mlir/IR/Dominance.h

109–110

Is "Memoize something that is complex to compute" a common pattern that should be factored out?

139

need trailing tick for a?

mlir/lib/IR/Dominance.cpp

47

Seems like an extra comment about how this code is using the fact that multi-block regions are not allowed in Graph Regions to optimize the fast path and avoid the RegionKind query would be good. I think the logic here isn't really obvious.

220

Change in style? should be a and be with ticks?

249–253

comment style for a and b?

296

missing a close tick

[lattner]

thanks for all the feedback, incorporated!

mlir/include/mlir/IR/Dominance.h

46–49

I'll make an effort, but can make no promises.

This is pretty nit'y, we're not at all consistent on this, even in the mlir codebase:

fgrep -r '. ' mlir/include | wc -l ===> 354 instances

97

Not anymore. That makes things simpler in fact, changed.

109–110

I'm open to suggestions, but that is a very general concept and there are many ways to do it. I don't think it's one general thing.

139

nice catch thx

164

Sure, copied from the .cpp file.

mlir/include/mlir/IR/Region.h

68

I don't want to pull STLExtras.h into Region.h. We shouldn't add dependencies to core IR headers like this particularly, when it is just a one liner.

mlir/lib/IR/Dominance.cpp

47

good point, done. Switching to returning PointerIntPair also allows early exits which makes it easier to read the logic.

87

You're missing something - this is returning the _block_ it is contained in, not the op it is contained in. The top level block (the body of the mlir::ModuleOp) has no enclosing Block, even though it has an enclosing Op (the ModuleOp)

mlir/lib/Transforms/BufferOptimizations.cpp

198

fixed

mlir/lib/Transforms/CSE.cpp

216

That's true. This microoptimization isn't important and is likely picked up by the compiler. Readability is more important here IMO.

[lattner]

Incorporate feedback from review.

[lattner]

This got a lot of valuable feedback already, I'm going to land this to hill climb on top of it (e.g. https://reviews.llvm.org/D103373). I value any other input of course!

mlir/include/mlir/IR/Dominance.h

 
 namespace mlir {
 using DominanceInfoNode = llvm::DomTreeNodeBase<Block>;
 class Operation;
 
 namespace detail {
 template <bool IsPostDom>
 class DominanceInfoBase {
-  using base = llvm::DominatorTreeBase<Block, IsPostDom>;
+  using DomTree = llvm::DominatorTreeBase<Block, IsPostDom>;
 
 public:
-  DominanceInfoBase(Operation *op) { recalculate(op); }
+  DominanceInfoBase(Operation *op) {}
   DominanceInfoBase(DominanceInfoBase &&) = default;
   DominanceInfoBase &operator=(DominanceInfoBase &&) = default;
+  ~DominanceInfoBase();
 
   DominanceInfoBase(const DominanceInfoBase &) = delete;
   DominanceInfoBase &operator=(const DominanceInfoBase &) = delete;
 
-  /// Recalculate the dominance info.
-  void recalculate(Operation *op);
+  /// Invalidate dominance info. This can be used by clients that make major
+  /// changes to the CFG and don't have a good way to update it.
+  void invalidate() { dominanceInfos.clear(); }
+  void invalidate(Region *region) { dominanceInfos.erase(region); }
 
   /// 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.
+  /// Get the root dominance node of the given region. Note that this operation
+  /// is only defined for multi-block regions!
   DominanceInfoNode *getRootNode(Region *region) {
-    assert(dominanceInfos.count(region) != 0);
-    return dominanceInfos[region]->getRootNode();
+    auto domInfo = getDominanceInfo(region, /*needsDomTree*/ true).getPointer();
+    assert(domInfo && "Region isn't multiblock");
+    return domInfo->getRootNode();
   }
 
-  /// Return the dominance node from the Region containing block A.
-  DominanceInfoNode *getNode(Block *a);
+  /// Return the dominance node from the Region containing block A. This only
+  /// works for multi-block regions.
+  DominanceInfoNode *getNode(Block *a) {
+    return getDomTree(a->getParent()).getNode(a);
+  }
+
+  /// Return true if the specified block is reachable from the entry
+  /// block of its region.
+  bool isReachableFromEntry(Block *a) const;
+
+  /// Return true if operations in the specified block are known to obey SSA
+  /// dominance requirements. False if the block is a graph region or unknown.
+  bool hasSSADominance(Block *block) const {
+    return hasSSADominance(block->getParent());
+  }
+  /// Return true if operations in the specified block are known to obey SSA
+  /// dominance requirements. False if the block is a graph region or unknown.
+  bool hasSSADominance(Region *region) const {
+    return getDominanceInfo(region, /*needsDomTree=*/false).getInt();
+  }
+
+  DomTree &getDomTree(Region *region) const {
+    assert(!region->hasOneBlock() &&
+           "Can't get DomTree for single block regions");
+    return *getDominanceInfo(region, /*needsDomTree=*/true).getPointer();
+  }
 
 protected:
   using super = DominanceInfoBase<IsPostDom>;
 
+  /// Return the dom tree and "hasSSADominance" bit for the given region. The
+  /// DomTree will be null for single-block regions. This lazily constructs the
+  /// DomTree on demand when needsDomTree=true.
+  llvm::PointerIntPair<DomTree *, 1, bool>
+  getDominanceInfo(Region *region, bool needsDomTree) const;
+
   /// Return true if the specified block A properly dominates block B.
   bool properlyDominates(Block *a, Block *b) const;
 
-  /// Return true if the specified block is reachable from the entry
-  /// block of its region.
-  bool isReachableFromEntry(Block *a) const;
-
-  /// A mapping of regions to their base dominator tree.
-  DenseMap<Region *, std::unique_ptr<base>> dominanceInfos;
+  /// A mapping of regions to their base dominator tree and a cached
+  /// "hasSSADominance" bit. This map does not contain dominator trees for
+  /// single block CFG regions, but we do want to cache the "hasSSADominance"
+  /// bit for them. We may also not have computed the DomTree yet. In either
+  /// case, the DomTree is just null.
+  ///
+  mutable DenseMap<Region *, llvm::PointerIntPair<DomTree *, 1, bool>>
+      dominanceInfos;
 };
 } // end namespace detail
 
 /// A class for computing basic dominance information. Note that this
 /// class is aware of different types of regions and returns a
 /// region-kind specific concept of dominance. See RegionKindInterface.
 class DominanceInfo : public detail::DominanceInfoBase</*IsPostDom=*/false> {
 public:
   using super::super;
 
-  /// Return true if the specified block is reachable from the entry block of
-  /// its region. In an SSACFG region, a block is reachable from the entry block
-  /// if it is the successor of the entry block or another reachable block. In a
-  /// Graph region, all blocks are reachable.
-  bool isReachableFromEntry(Block *a) const {
-    return super::isReachableFromEntry(a);
-  }
-
   /// Return true if operation A properly dominates operation B, i.e. if A and B
   /// are in the same block and A properly dominates B within the block, or if
   /// the block that contains A properly dominates the block that contains B. In
   /// an SSACFG region, Operation A dominates Operation B in the same block if A
   /// preceeds B. In a Graph region, all operations in a block dominate all
   /// other operations in the same block.
-  bool properlyDominates(Operation *a, Operation *b) const;
+  bool properlyDominates(Operation *a, Operation *b) const {
+    return properlyDominatesImpl(a, b, /*enclosingOpOk=*/true);
+  }
 
   /// Return true if operation A dominates operation B, i.e. if A and B are the
   /// same operation or A properly dominates B.
   bool dominates(Operation *a, Operation *b) const {
     return a == b || properlyDominates(a, b);
   }
 
-  /// Return true if value A properly dominates operation B, i.e if the
-  /// operation that defines A properlyDominates B and the operation that
-  /// defines A does not contain B.
+  /// Return true if the `a` value properly dominates operation `b`, i.e if the
+  /// operation that defines `a` properlyDominates `b` and the operation that
+  /// defines `a` does not contain `b`.
   bool properlyDominates(Value a, Operation *b) const;
 
-  /// Return true if operation A dominates operation B, i.e if the operation
-  /// that defines A dominates B.
+  /// Return true if the `a` value dominates operation `b`.
   bool dominates(Value a, Operation *b) const {
     return (Operation *)a.getDefiningOp() == b || properlyDominates(a, b);
   }
 
   /// Return true if the specified block A dominates block B, i.e. if block A
   /// and block B are the same block or block A properly dominates block B.
   bool dominates(Block *a, Block *b) const {
     return a == b || properlyDominates(a, b);
   }
 
   /// Return true if the specified block A properly dominates block B, i.e.: if
   /// block A contains block B, or if the region which contains block A also
   /// contains block B or some parent of block B and block A dominates that
-  /// block in that kind of region.  In an SSACFG region, block A dominates
+  /// block in that kind of region. In an SSACFG region, block A dominates
   /// block B if all control flow paths from the entry block to block B flow
   /// through block A. In a Graph region, all blocks dominate all other blocks.
   bool properlyDominates(Block *a, Block *b) const {
     return super::properlyDominates(a, b);
   }
 
-  /// Update the internal DFS numbers for the dominance nodes.
-  void updateDFSNumbers();
+private:
+  // Return true if operation A properly dominates operation B.  The
+  /// 'enclosingOpOk' flag says whether we should return true if the b op is
+  /// enclosed by a region on 'A'.
+  bool properlyDominatesImpl(Operation *a, Operation *b,
+                             bool enclosingOpOk) const;
 };
 
 /// A class for computing basic postdominance information.
 class PostDominanceInfo : public detail::DominanceInfoBase</*IsPostDom=*/true> {
 public:
   using super::super;
 
-  /// Return true if the specified block is reachable from the entry
-  /// block of its region.
-  bool isReachableFromEntry(Block *a) const {
-    return super::isReachableFromEntry(a);
-  }
-
   /// Return true if operation A properly postdominates operation B.
   bool properlyPostDominates(Operation *a, Operation *b);
 
   /// Return true if operation A postdominates operation B.
   bool postDominates(Operation *a, Operation *b) {
     return a == b || properlyPostDominates(a, b);
   }
 

mlir/include/mlir/IR/Region.h

public:
 
   bool empty() { return blocks.empty(); }
   void push_back(Block *block) { blocks.push_back(block); }
   void push_front(Block *block) { blocks.push_front(block); }
 
   Block &back() { return blocks.back(); }
   Block &front() { return blocks.front(); }
 
+  /// Return true if this region has exactly one block.
+  bool hasOneBlock() { return !empty() && std::next(begin()) == end(); }
+
   /// getSublistAccess() - Returns pointer to member of region.
   static BlockListType Region::*getSublistAccess(Block *) {
     return &Region::blocks;
   }
 
   //===--------------------------------------------------------------------===//
   // Argument Handling
   //===--------------------------------------------------------------------===//
void takeBody(Region &other) {
     blocks.splice(blocks.end(), other.getBlocks());
   }
 
   /// Returns 'block' if 'block' lies in this region, or otherwise finds the
   /// ancestor of 'block' that lies in this region. Returns nullptr if the
   /// latter fails.
   Block *findAncestorBlockInRegion(Block &block);
 
+  /// Returns 'op' if 'op' lies in this region, or otherwise finds the
+  /// ancestor of 'op' that lies in this region. Returns nullptr if the
+  /// latter fails.
+  Operation *findAncestorOpInRegion(Operation &op);
+
   /// Drop all operand uses from operations within this region, which is
   /// an essential step in breaking cyclic dependences between references when
   /// they are to be deleted.
   void dropAllReferences();
 
   //===--------------------------------------------------------------------===//
   // Operation Walkers
   //===--------------------------------------------------------------------===//

mlir/lib/IR/Dominance.cpp

 
 using namespace mlir;
 using namespace mlir::detail;
 
 template class llvm::DominatorTreeBase<Block, /*IsPostDom=*/false>;
 template class llvm::DominatorTreeBase<Block, /*IsPostDom=*/true>;
 template class llvm::DomTreeNodeBase<Block>;
 
-/// Return true if the region with the given index inside the operation
-/// has SSA dominance.
-static bool hasSSADominance(Operation *op, unsigned index) {
-  if (!op->isRegistered())
-    return false;
-
-  auto kindInterface = dyn_cast<RegionKindInterface>(op);
-  return !kindInterface || kindInterface.hasSSADominance(index);
-}
-
 //===----------------------------------------------------------------------===//
 // DominanceInfoBase
 //===----------------------------------------------------------------------===//
 
 template <bool IsPostDom>
-void DominanceInfoBase<IsPostDom>::recalculate(Operation *op) {
-  dominanceInfos.clear();
+DominanceInfoBase<IsPostDom>::~DominanceInfoBase() {
+  for (auto entry : dominanceInfos)
+    delete entry.second.getPointer();
+}
 
-  // Build the dominance for each of the operation regions.
-  op->walk([&](Operation *op) {
-    unsigned numRegions = op->getNumRegions();
-    if (numRegions == 0)
-      return;
-
-    auto kindInterface = dyn_cast<RegionKindInterface>(op);
-    for (unsigned i = 0; i < numRegions; i++) {
-      Region &region = op->getRegion(i);
-      // Don't compute dominance if the region is empty.
-      if (region.empty())
-        continue;
-
-      // Dominance changes based on the region type. Avoid the helper
-      // function here so we don't do the region cast repeatedly.
-      bool hasSSADominance =
-          op->isRegistered() &&
-          (!kindInterface || kindInterface.hasSSADominance(i));
-      // If a region has SSADominance, then compute detailed dominance
-      // info.  Otherwise, all values in the region are live anywhere
-      // in the region, which is represented as an empty entry in the
-      // dominanceInfos map.
-      if (hasSSADominance) {
-        auto opDominance = std::make_unique<base>();
-        opDominance->recalculate(region);
-        dominanceInfos.try_emplace(&region, std::move(opDominance));
-      }
+/// Return the dom tree and "hasSSADominance" bit for the given region.  The
+/// DomTree will be null for single-block regions.  This lazily constructs the
+/// DomTree on demand when needsDomTree=true.
+template <bool IsPostDom>
+auto DominanceInfoBase<IsPostDom>::getDominanceInfo(Region *region,
+                                                    bool needsDomTree) const
+    -> llvm::PointerIntPair<DomTree *, 1, bool> {
+  // Check to see if we already have this information.
+  auto itAndInserted = dominanceInfos.insert({region, {nullptr, true}});
+  auto &entry = itAndInserted.first->second;
+
+  // This method builds on knowledge that multi-block regions always have
+  // SSADominance.  Graph regions are only allowed to be single-block regions,
+  // but of course single-block regions may also have SSA dominance.
+  if (!itAndInserted.second) {
+    // We do have it, so we know the 'hasSSADominance' bit is correct, but we
+    // may not have constructed a DominatorTree yet.  If we need it, build it.
+    if (needsDomTree && !entry.getPointer() && !region->hasOneBlock()) {
+      auto *domTree = new DomTree();
+      domTree->recalculate(*region);
+      entry.setPointer(domTree);
+    return entry;
+  }
+
+  // Nope, lazily construct it.  Create a DomTree if this is a multi-block
+  // region.
+  if (!region->hasOneBlock()) {
+    auto *domTree = new DomTree();
+    domTree->recalculate(*region);
+    entry.setPointer(domTree);
+    // Multiblock regions always have SSA dominance, leave `second` set to true.
+    return entry;
+  }
+
+  // Single block regions have a more complicated predicate.
+  if (Operation *parentOp = region->getParentOp()) {
+    if (!parentOp->isRegistered()) { // We don't know about unregistered ops.
+      entry.setInt(false);
+    } else if (auto regionKindItf = dyn_cast<RegionKindInterface>(parentOp)) {
+      // Registered ops can opt-out of SSA dominance with
+      // RegionKindInterface.
+      entry.setInt(regionKindItf.hasSSADominance(region->getRegionNumber()));
     }
-  });
+  }
+
+  return entry;
+}
+
+/// Return the ancestor block enclosing the specified block.  This returns null
+/// if we reach the top of the hierarchy.
+static Block *getAncestorBlock(Block *block) {
+  if (Operation *ancestorOp = block->getParentOp())
+    return ancestorOp->getBlock();
+  return nullptr;
 }
 
 /// Walks up the list of containers of the given block and calls the
 /// user-defined traversal function for every pair of a region and block that
 /// could be found during traversal. If the user-defined function returns true
 /// for a given pair, traverseAncestors will return the current block. Nullptr
 /// otherwise.
 template <typename FuncT>
-Block *traverseAncestors(Block *block, const FuncT &func) {
-  // Invoke the user-defined traversal function in the beginning for the current
-  // block.
+static Block *traverseAncestors(Block *block, const FuncT &func) {
+  do {
+    // Invoke the user-defined traversal function for each block.
     if (func(block))
       return block;
-
-  Region *region = block->getParent();
-  while (region) {
-    Operation *ancestor = region->getParentOp();
-    // If we have reached to top... return.
-    if (!ancestor || !(block = ancestor->getBlock()))
-      break;
-
-    // Update the nested region using the new ancestor block.
-    region = block->getParent();
-
-    // Invoke the user-defined traversal function and check whether we can
-    // already return.
-    if (func(block))
-      return block;
-  }
+  } while ((block = getAncestorBlock(block)));
   return nullptr;
 }
 
 /// Tries to update the given block references to live in the same region by
 /// exploring the relationship of both blocks with respect to their regions.
 static bool tryGetBlocksInSameRegion(Block *&a, Block *&b) {
   // If both block do not live in the same region, we will have to check their
   // parent operations.
-  if (a->getParent() == b->getParent())
+  Region *aRegion = a->getParent();
+  Region *bRegion = b->getParent();
+  if (aRegion == bRegion)
     return true;
 
-  // Iterate over all ancestors of a and insert them into the map. This allows
-  // for efficient lookups to find a commonly shared region.
-  llvm::SmallDenseMap<Region *, Block *, 4> ancestors;
-  traverseAncestors(a, [&](Block *block) {
-    ancestors[block->getParent()] = block;
-    return false;
-  });
+  // Iterate over all ancestors of `a`, counting the depth of `a`. If one of
+  // `a`s ancestors are in the same region as `b`, then we stop early because we
+  // found our NCA.
+  size_t aRegionDepth = 0;
+  if (Block *aResult = traverseAncestors(a, [&](Block *block) {
+        ++aRegionDepth;
+        return block->getParent() == bRegion;
+      })) {
+    a = aResult;
+    return true;
+  }
 
-  // Try to find a common ancestor starting with regionB.
-  b = traverseAncestors(
-      b, [&](Block *block) { return ancestors.count(block->getParent()) > 0; });
+  // Iterate over all ancestors of `b`, counting the depth of `b`. If one of
+  // `b`s ancestors are in the same region as `a`, then we stop early because
+  // we found our NCA.
+  size_t bRegionDepth = 0;
+  if (Block *bResult = traverseAncestors(b, [&](Block *block) {
+        ++bRegionDepth;
+        return block->getParent() == aRegion;
+      })) {
+    b = bResult;
+    return true;
+  }
 
-  // If there is no match, we will not be able to find a common dominator since
-  // both regions do not share a common parent region.
-  if (!b)
-    return false;
+  // Otherwise we found two blocks that are siblings at some level.  Walk the
+  // deepest one up until we reach the top or find an NCA.
+  while (true) {
+    if (aRegionDepth > bRegionDepth) {
+      a = getAncestorBlock(a);
+      --aRegionDepth;
+    } else if (aRegionDepth < bRegionDepth) {
+      b = getAncestorBlock(b);
+      --bRegionDepth;
+    } else {
+      break;
+    }
+  }
 
-  // We have found a common parent region. Update block a to refer to this
-  // region.
-  auto it = ancestors.find(b->getParent());
-  assert(it != ancestors.end());
-  a = it->second;
+  // If we found something with the same level, then we can march both up at the
+  // same time from here on out.
+  while (a) {
+    // If they are at the same level, and have the same parent region then we
+    // succeeded.
+    if (a->getParent() == b->getParent())
       return true;
+
+    a = getAncestorBlock(a);
+    b = getAncestorBlock(b);
+  }
+
+  // They don't share an NCA, perhaps they are in different modules or
+  // something.
+  return false;
 }
 
 template <bool IsPostDom>
 Block *
 DominanceInfoBase<IsPostDom>::findNearestCommonDominator(Block *a,
                                                          Block *b) const {
   // If either a or b are null, then conservatively return nullptr.
   if (!a || !b)
     return nullptr;
 
+  // If they are the same block, then we are done.
+  if (a == b)
+    return a;
+
   // Try to find blocks that are in the same region.
   if (!tryGetBlocksInSameRegion(a, b))
     return nullptr;
 
-  // Get and verify dominance information of the common parent region.
-  Region *parentRegion = a->getParent();
-  auto infoAIt = dominanceInfos.find(parentRegion);
-  if (infoAIt == dominanceInfos.end())
-    return nullptr;
+  // If the common ancestor in a common region is the same block, then return
+  // it.
+  if (a == b)
+    return a;
 
-  // Since the blocks live in the same region, we can rely on already
-  // existing dominance functionality.
-  return infoAIt->second->findNearestCommonDominator(a, b);
-}
-
-template <bool IsPostDom>
-DominanceInfoNode *DominanceInfoBase<IsPostDom>::getNode(Block *a) {
-  Region *region = a->getParent();
-  assert(dominanceInfos.count(region) != 0);
-  return dominanceInfos[region]->getNode(a);
+  // Otherwise, there must be multiple blocks in the region, check the
+  // DomTree.
+  return getDomTree(a->getParent()).findNearestCommonDominator(a, b);
 }
 
 /// Return true if the specified block A properly dominates block B.
 template <bool IsPostDom>
 bool DominanceInfoBase<IsPostDom>::properlyDominates(Block *a, Block *b) const {
+  assert(a && b && "null blocks not allowed");
+
   // A block dominates itself but does not properly dominate itself.
   if (a == b)
     return false;
 
-  // If either a or b are null, then conservatively return false.
-  if (!a || !b)
-    return false;
-
-  // If both blocks are not in the same region, 'a' properly dominates 'b' if
-  // 'b' is defined in an operation region that (recursively) ends up being
-  // dominated by 'a'. Walk up the list of containers enclosing B.
+  // If both blocks are not in the same region, `a` properly dominates `b` if
+  // `b` is defined in an operation region that (recursively) ends up being
+  // dominated by `a`. Walk up the list of containers enclosing B.
   Region *regionA = a->getParent();
   if (regionA != b->getParent()) {
-    b = traverseAncestors(
-        b, [&](Block *block) { return block->getParent() == regionA; });
-
+    b = regionA ? regionA->findAncestorBlockInRegion(*b) : nullptr;
     // If we could not find a valid block b then it is a not a dominator.
-    if (!b)
+    if (b == nullptr)
       return false;
 
-    // Check to see if the ancestor of 'b' is the same block as 'a'.
+    // Check to see if the ancestor of `b` is the same block as `a`.  A properly
+    // dominates B if it contains an op that contains the B block.
     if (a == b)
       return true;
   }
 
-  // Otherwise, use the standard dominance functionality.
-
-  // If we don't have a dominance information for this region, assume that b is
-  // dominated by anything.
-  auto baseInfoIt = dominanceInfos.find(regionA);
-  if (baseInfoIt == dominanceInfos.end())
-    return true;
-  return baseInfoIt->second->properlyDominates(a, b);
+  // Otherwise, they are two different blocks in the same region, use DomTree.
+  return getDomTree(regionA).properlyDominates(a, b);
 }
 
-/// Return true if the specified block is reachable from the entry block of its
-/// region.
+/// Return true if the specified block is reachable from the entry block of
+/// its region.
 template <bool IsPostDom>
 bool DominanceInfoBase<IsPostDom>::isReachableFromEntry(Block *a) const {
-  Region *regionA = a->getParent();
-  auto baseInfoIt = dominanceInfos.find(regionA);
-  if (baseInfoIt == dominanceInfos.end())
+  // If this is the first block in its region, then it is obviously reachable.
+  Region *region = a->getParent();
+  if (&region->front() == a)
     return true;
-  return baseInfoIt->second->isReachableFromEntry(a);
+
+  // Otherwise this is some block in a multi-block region.  Check DomTree.
+  return getDomTree(region).isReachableFromEntry(a);
 }
 
 template class detail::DominanceInfoBase</*IsPostDom=*/true>;
 template class detail::DominanceInfoBase</*IsPostDom=*/false>;
 
 //===----------------------------------------------------------------------===//
 // DominanceInfo
 //===----------------------------------------------------------------------===//
 
-/// Return true if operation A properly dominates operation B.
-bool DominanceInfo::properlyDominates(Operation *a, Operation *b) const {
+/// Return true if operation `a` properly dominates operation `b`.  The
+/// 'enclosingOpOk' flag says whether we should return true if the `b` op is
+/// enclosed by a region on 'a'.
+bool DominanceInfo::properlyDominatesImpl(Operation *a, Operation *b,
+                                          bool enclosingOpOk) const {
   Block *aBlock = a->getBlock(), *bBlock = b->getBlock();
-  Region *aRegion = a->getParentRegion();
-  unsigned aRegionNum = aRegion->getRegionNumber();
-  Operation *ancestor = aRegion->getParentOp();
+  assert(aBlock && bBlock && "operations must be in a block");
 
-  // If a or b are not within a block, then a does not dominate b.
-  if (!aBlock || !bBlock)
+  // An instruction dominates, but does not properlyDominate, itself unless this
+  // is a graph region.
+  if (a == b)
+    return !hasSSADominance(aBlock);
+
+  // If these ops are in different regions, then normalize one into the other.
+  Region *aRegion = aBlock->getParent();
+  if (aRegion != bBlock->getParent()) {
+    // Scoot up b's region tree until we find an operation in A's region that
+    // encloses it.  If this fails, then we know there is no post-dom relation.
+    b = aRegion ? aRegion->findAncestorOpInRegion(*b) : nullptr;
+    if (!b)
       return false;
+    bBlock = b->getBlock();
+    assert(bBlock->getParent() == aRegion);
 
+    // If 'a' encloses 'b', then we consider it to dominate.
+    if (a == b && enclosingOpOk)
+      return true;
+  }
+
+  // Ok, they are in the same region now.
   if (aBlock == bBlock) {
     // Dominance changes based on the region type. In a region with SSA
     // dominance, uses inside the same block must follow defs. In other
     // regions kinds, uses and defs can come in any order inside a block.
-    if (hasSSADominance(ancestor, aRegionNum)) {
+    if (hasSSADominance(aBlock)) {
       // If the blocks are the same, then check if b is before a in the block.
       return a->isBeforeInBlock(b);
     }
     return true;
   }
 
-  // Traverse up b's hierarchy to check if b's block is contained in a's.
-  if (auto *bAncestor = aBlock->findAncestorOpInBlock(*b)) {
-    // Since we already know that aBlock != bBlock, here bAncestor != b.
-    // a and bAncestor are in the same block; check if 'a' dominates
-    // bAncestor.
-    return dominates(a, bAncestor);
-  }
-
-  // If the blocks are different, check if a's block dominates b's.
-  return properlyDominates(aBlock, bBlock);
+  // If the blocks are different, use DomTree to resolve the query.
+  return getDomTree(aRegion).properlyDominates(aBlock, bBlock);
 }
 
-/// Return true if value A properly dominates operation B.
+/// Return true if the `a` value properly dominates operation `b`, i.e if the
+/// operation that defines `a` properlyDominates `b` and the operation that
+/// defines `a` does not contain `b`.
 bool DominanceInfo::properlyDominates(Value a, Operation *b) const {
-  if (auto *aOp = a.getDefiningOp()) {
-    // Dominance changes based on the region type.
-    auto *aRegion = aOp->getParentRegion();
-    unsigned aRegionNum = aRegion->getRegionNumber();
-    Operation *ancestor = aRegion->getParentOp();
-    // Dominance changes based on the region type. In a region with SSA
-    // dominance, values defined by an operation cannot be used by the
-    // operation. In other regions kinds they can be used the operation.
-    if (hasSSADominance(ancestor, aRegionNum)) {
-      // The values defined by an operation do *not* dominate any nested
-      // operations.
-      if (aOp->getParentRegion() != b->getParentRegion() && aOp->isAncestor(b))
-        return false;
-    }
-    return properlyDominates(aOp, b);
-  }
-
   // block arguments properly dominate all operations in their own block, so
   // we use a dominates check here, not a properlyDominates check.
-  return dominates(a.cast<BlockArgument>().getOwner(), b->getBlock());
-}
+  if (auto blockArg = a.dyn_cast<BlockArgument>())
+    return dominates(blockArg.getOwner(), b->getBlock());
 
-void DominanceInfo::updateDFSNumbers() {
-  for (auto &iter : dominanceInfos)
-    iter.second->updateDFSNumbers();
+  // `a` properlyDominates `b` if the operation defining `a` properlyDominates
+  // `b`, but `a` does not itself enclose `b` in one of its regions.
+  return properlyDominatesImpl(a.getDefiningOp(), b, /*enclosingOpOk=*/false);
 }
 
 //===----------------------------------------------------------------------===//
 // PostDominanceInfo
 //===----------------------------------------------------------------------===//
 
 /// Returns true if statement 'a' properly postdominates statement b.
 bool PostDominanceInfo::properlyPostDominates(Operation *a, Operation *b) {
   auto *aBlock = a->getBlock(), *bBlock = b->getBlock();
-  auto *aRegion = a->getParentRegion();
-  unsigned aRegionNum = aRegion->getRegionNumber();
-  Operation *ancestor = aRegion->getParentOp();
+  assert(aBlock && bBlock && "operations must be in a block");
 
-  // If a or b are not within a block, then a does not post dominate b.
-  if (!aBlock || !bBlock)
+  // An instruction postDominates, but does not properlyPostDominate, itself
+  // unless this is a graph region.
+  if (a == b)
+    return !hasSSADominance(aBlock);
+
+  // If these ops are in different regions, then normalize one into the other.
+  Region *aRegion = aBlock->getParent();
+  if (aRegion != bBlock->getParent()) {
+    // Scoot up b's region tree until we find an operation in A's region that
+    // encloses it.  If this fails, then we know there is no post-dom relation.
+    b = aRegion ? aRegion->findAncestorOpInRegion(*b) : nullptr;
+    if (!b)
       return false;
+    bBlock = b->getBlock();
+    assert(bBlock->getParent() == aRegion);
 
-  // If the blocks are the same, check if b is before a in the block.
+    // If 'a' encloses 'b', then we consider it to postdominate.
+    if (a == b)
+      return true;
+  }
+
+  // Ok, they are in the same region.  If they are in the same block, check if b
+  // is before a in the block.
   if (aBlock == bBlock) {
     // Dominance changes based on the region type.
-    if (hasSSADominance(ancestor, aRegionNum)) {
+    if (hasSSADominance(aBlock)) {
       // If the blocks are the same, then check if b is before a in the block.
       return b->isBeforeInBlock(a);
     }
     return true;
   }
 
-  // Traverse up b's hierarchy to check if b's block is contained in a's.
-  if (auto *bAncestor = a->getBlock()->findAncestorOpInBlock(*b))
-    // Since we already know that aBlock != bBlock, here bAncestor != b.
-    // a and bAncestor are in the same block; check if 'a' postdominates
-    // bAncestor.
-    return postDominates(a, bAncestor);
-
   // If the blocks are different, check if a's block post dominates b's.
-  return properlyDominates(aBlock, bBlock);
+  return getDomTree(aRegion).properlyDominates(aBlock, bBlock);
 }

mlir/lib/IR/Region.cpp

void Region::cloneInto(Region *dest, Region::iterator destPos,
   for (iterator it(mapper.lookup(&front())); it != destPos; ++it)
     it->walk(remapOperands);
 }
 
 /// Returns 'block' if 'block' lies in this region, or otherwise finds the
 /// ancestor of 'block' that lies in this region. Returns nullptr if the latter
 /// fails.
 Block *Region::findAncestorBlockInRegion(Block &block) {
-  auto currBlock = █
+  Block *currBlock = █
   while (currBlock->getParent() != this) {
     Operation *parentOp = currBlock->getParentOp();
     if (!parentOp || !parentOp->getBlock())
       return nullptr;
     currBlock = parentOp->getBlock();
   }
   return currBlock;
 }
 
+/// Returns 'op' if 'op' lies in this region, or otherwise finds the
+/// ancestor of 'op' that lies in this region. Returns nullptr if the
+/// latter fails.
+Operation *Region::findAncestorOpInRegion(Operation &op) {
+  Operation *curOp = &op;
+  while (Region *opRegion = curOp->getParentRegion()) {
+    if (opRegion == this)
+      return curOp;
+
+    curOp = opRegion->getParentOp();
+    if (!curOp)
+      return nullptr;
+  }
+  return nullptr;
+}
+
 void Region::dropAllReferences() {
   for (Block &b : *this)
     b.dropAllReferences();
 }
 
 Region *llvm::ilist_traits<::mlir::Block>::getParentRegion() {
   size_t Offset(
       size_t(&((Region *)nullptr->*Region::getSublistAccess(nullptr))));

mlir/lib/Transforms/BufferOptimizations.cpp

Block *findPlacementBlock(StateT &state, Block *upperBound) {
     Operation *parentOp;
     Block *parentBlock;
     while ((parentOp = currentBlock->getParentOp()) &&
            (parentBlock = parentOp->getBlock()) &&
            (!upperBound ||
             dominators.properlyDominates(upperBound, currentBlock))) {
       // Try to find an immediate dominator and check whether the parent block
       // is above the immediate dominator (if any).
-      DominanceInfoNode *idom = dominators.getNode(currentBlock)->getIDom();
+      DominanceInfoNode *idom = nullptr;
+
+      // DominanceInfo doesn't support getNode queries for single-block regions.
+      if (!currentBlock->isEntryBlock())
+        idom = dominators.getNode(currentBlock)->getIDom();
+
       if (idom && dominators.properlyDominates(parentBlock, idom->getBlock())) {
         // If the current immediate dominator is below the placement block, move
         // to the immediate dominator block.
         currentBlock = idom->getBlock();
         state.recordMoveToDominator(currentBlock);
       } else {
         // We have to move to our parent block since an immediate dominator does
         // either not exist or is above our parent block. If we cannot move to

mlir/lib/Transforms/CSE.cpp

 
 void CSE::simplifyRegion(ScopedMapTy &knownValues, Region &region,
                          bool hasSSADominance) {
   // If the region is empty there is nothing to do.
   if (region.empty())
     return;
 
   // If the region only contains one block, then simplify it directly.
-  if (std::next(region.begin()) == region.end()) {
+  if (region.hasOneBlock()) {
     ScopedMapTy::ScopeTy scope(knownValues);
     simplifyBlock(knownValues, &region.front(), hasSSADominance);
     return;
   }
 
   // If the region does not have dominanceInfo, then skip it.
   // TODO: Regions without SSA dominance should define a different
   // traversal order which is appropriate and can be used here.