[Mlir-commits] [mlir] 412ae15 - [Dominators] Rewrite the dominator implementation for efficiency. NFC.

Chris Lattner llvmlistbot at llvm.org
Tue Jun 1 14:58:24 PDT 2021


Author: Chris Lattner
Date: 2021-06-01T14:46:37-07:00
New Revision: 412ae15de49a227de25a695735451f8908ebf999

URL: https://github.com/llvm/llvm-project/commit/412ae15de49a227de25a695735451f8908ebf999
DIFF: https://github.com/llvm/llvm-project/commit/412ae15de49a227de25a695735451f8908ebf999.diff

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

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.

Differential Revision: https://reviews.llvm.org/D103384

Added: 
    

Modified: 
    mlir/include/mlir/IR/Dominance.h
    mlir/include/mlir/IR/Region.h
    mlir/lib/IR/Dominance.cpp
    mlir/lib/IR/Region.cpp
    mlir/lib/Transforms/BufferOptimizations.cpp
    mlir/lib/Transforms/CSE.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/Dominance.h b/mlir/include/mlir/IR/Dominance.h
index c2df20255ae88..68513c27aa0ea 100644
--- a/mlir/include/mlir/IR/Dominance.h
+++ b/mlir/include/mlir/IR/Dominance.h
@@ -32,45 +32,82 @@ 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
 
@@ -81,21 +118,15 @@ 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.
@@ -103,13 +134,12 @@ class DominanceInfo : public detail::DominanceInfoBase</*IsPostDom=*/false> {
     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);
   }
@@ -123,15 +153,19 @@ class DominanceInfo : public detail::DominanceInfoBase</*IsPostDom=*/false> {
   /// 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.
@@ -139,12 +173,6 @@ 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);
 

diff  --git a/mlir/include/mlir/IR/Region.h b/mlir/include/mlir/IR/Region.h
index 17a5072faa11c..720a5ae09eefe 100644
--- a/mlir/include/mlir/IR/Region.h
+++ b/mlir/include/mlir/IR/Region.h
@@ -64,6 +64,9 @@ class Region {
   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;
@@ -235,6 +238,11 @@ class Region {
   /// 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.

diff  --git a/mlir/lib/IR/Dominance.cpp b/mlir/lib/IR/Dominance.cpp
index c13a57dbf3004..68cc2039ab646 100644
--- a/mlir/lib/IR/Dominance.cpp
+++ b/mlir/lib/IR/Dominance.cpp
@@ -24,53 +24,71 @@ 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();
-
-  // 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));
-      }
+DominanceInfoBase<IsPostDom>::~DominanceInfoBase() {
+  for (auto entry : dominanceInfos)
+    delete entry.second.getPointer();
+}
+
+/// 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
@@ -79,27 +97,12 @@ void DominanceInfoBase<IsPostDom>::recalculate(Operation *op) {
 /// 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.
-  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.
+static Block *traverseAncestors(Block *block, const FuncT &func) {
+  do {
+    // Invoke the user-defined traversal function for each block.
     if (func(block))
       return block;
-  }
+  } while ((block = getAncestorBlock(block)));
   return nullptr;
 }
 
@@ -108,32 +111,64 @@ Block *traverseAncestors(Block *block, const FuncT &func) {
 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;
-  return true;
+  // 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 
diff erent modules or
+  // something.
+  return false;
 }
 
 template <bool IsPostDom>
@@ -144,75 +179,64 @@ DominanceInfoBase<IsPostDom>::findNearestCommonDominator(Block *a,
   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;
-
-  // Since the blocks live in the same region, we can rely on already
-  // existing dominance functionality.
-  return infoAIt->second->findNearestCommonDominator(a, b);
-}
+  // If the common ancestor in a common region is the same block, then return
+  // it.
+  if (a == b)
+    return a;
 
-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 
diff erent 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>;
@@ -222,67 +246,63 @@ 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)
-    return false;
+  // 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 
diff erent 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 
diff erent, check if a's block dominates b's.
-  return properlyDominates(aBlock, bBlock);
+  // If the blocks are 
diff erent, 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);
 }
 
 //===----------------------------------------------------------------------===//
@@ -292,31 +312,40 @@ void DominanceInfo::updateDFSNumbers() {
 /// 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)
-    return false;
+  // 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 
diff erent 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 
diff erent, check if a's block post dominates b's.
-  return properlyDominates(aBlock, bBlock);
+  return getDomTree(aRegion).properlyDominates(aBlock, bBlock);
 }

diff  --git a/mlir/lib/IR/Region.cpp b/mlir/lib/IR/Region.cpp
index a06b89fc3ffa6..b2600b58e9422 100644
--- a/mlir/lib/IR/Region.cpp
+++ b/mlir/lib/IR/Region.cpp
@@ -122,7 +122,7 @@ void Region::cloneInto(Region *dest, Region::iterator destPos,
 /// 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())
@@ -132,6 +132,22 @@ Block *Region::findAncestorBlockInRegion(Block &block) {
   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();

diff  --git a/mlir/lib/Transforms/BufferOptimizations.cpp b/mlir/lib/Transforms/BufferOptimizations.cpp
index b5dad2b5d393a..1b87b8c32b11f 100644
--- a/mlir/lib/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Transforms/BufferOptimizations.cpp
@@ -194,7 +194,12 @@ class BufferAllocationHoisting : public BufferPlacementTransformationBase {
             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.

diff  --git a/mlir/lib/Transforms/CSE.cpp b/mlir/lib/Transforms/CSE.cpp
index 4ca76f4425725..541b54ae2a6b8 100644
--- a/mlir/lib/Transforms/CSE.cpp
+++ b/mlir/lib/Transforms/CSE.cpp
@@ -213,7 +213,7 @@ void CSE::simplifyRegion(ScopedMapTy &knownValues, Region &region,
     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;


        


More information about the Mlir-commits mailing list