[Dominators] Keep tree level in DomTreeNode and use it to find NCD and answer dominance queries

Summary:
This patch makes DomTreeNodes keep their level (depth) in the DomTree. By having this information always available, it is possible to speedup and simplify findNearestCommonDominator and certain dominance queries.

In the future, level information will be also needed to perform incremental updates.

My testing doesn't show any noticeable performance differences after applying this patch. There may be some improvements when other passes are thought to use the level information.

Reviewers: dberlin, sanjoy, chandlerc, grosser

Reviewed By: dberlin

Subscribers: llvm-commits

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

llvm-svn: 306892

3 files changed, 103 insertions, 58 deletions

diff --git a/llvm/include/llvm/Support/GenericDomTree.h b/llvm/include/llvm/Support/GenericDomTree.h
index 66a04502d11..290515f7776 100644
--- a/llvm/include/llvm/Support/GenericDomTree.h
+++ b/llvm/include/llvm/Support/GenericDomTree.h
@@ -65,12 +65,14 @@ template <class NodeT> class DomTreeNodeBase {
 
   NodeT *TheBB;
   DomTreeNodeBase *IDom;
+  unsigned Level;
   std::vector<DomTreeNodeBase *> Children;
   mutable unsigned DFSNumIn = ~0;
   mutable unsigned DFSNumOut = ~0;
 
  public:
-  DomTreeNodeBase(NodeT *BB, DomTreeNodeBase *iDom) : TheBB(BB), IDom(iDom) {}
+  DomTreeNodeBase(NodeT *BB, DomTreeNodeBase *iDom)
+      : TheBB(BB), IDom(iDom), Level(IDom ? IDom->Level + 1 : 0) {}
 
   using iterator = typename std::vector<DomTreeNodeBase *>::iterator;
   using const_iterator =
@@ -83,6 +85,7 @@ template <class NodeT> class DomTreeNodeBase {
 
   NodeT *getBlock() const { return TheBB; }
   DomTreeNodeBase *getIDom() const { return IDom; }
+  unsigned getLevel() const { return Level; }
 
   const std::vector<DomTreeNodeBase *> &getChildren() const { return Children; }
 
@@ -100,6 +103,8 @@ template <class NodeT> class DomTreeNodeBase {
     if (getNumChildren() != Other->getNumChildren())
       return true;
 
+    if (Level != Other->Level) return true;
+
     SmallPtrSet<const NodeT *, 4> OtherChildren;
     for (const DomTreeNodeBase *I : *Other) {
       const NodeT *Nd = I->getBlock();
@@ -116,18 +121,19 @@ template <class NodeT> class DomTreeNodeBase {
 
   void setIDom(DomTreeNodeBase *NewIDom) {
     assert(IDom && "No immediate dominator?");
-    if (IDom != NewIDom) {
-      typename std::vector<DomTreeNodeBase *>::iterator I =
-          find(IDom->Children, this);
-      assert(I != IDom->Children.end() &&
-             "Not in immediate dominator children set!");
-      // I am no longer your child...
-      IDom->Children.erase(I);
+    if (IDom == NewIDom) return;
 
-      // Switch to new dominator
-      IDom = NewIDom;
-      IDom->Children.push_back(this);
-    }
+    auto I = find(IDom->Children, this);
+    assert(I != IDom->Children.end() &&
+           "Not in immediate dominator children set!");
+    // I am no longer your child...
+    IDom->Children.erase(I);
+
+    // Switch to new dominator
+    IDom = NewIDom;
+    IDom->Children.push_back(this);
+
+    UpdateLevel();
   }
 
   /// getDFSNumIn/getDFSNumOut - These return the DFS visitation order for nodes
@@ -143,6 +149,23 @@ private:
     return this->DFSNumIn >= other->DFSNumIn &&
            this->DFSNumOut <= other->DFSNumOut;
   }
+
+  void UpdateLevel() {
+    assert(IDom);
+    if (Level == IDom->Level + 1) return;
+
+    SmallVector<DomTreeNodeBase *, 64> WorkStack = {this};
+
+    while (!WorkStack.empty()) {
+      DomTreeNodeBase *Current = WorkStack.pop_back_val();
+      Current->Level = Current->IDom->Level + 1;
+
+      for (DomTreeNodeBase *C : *Current) {
+        assert(C->IDom);
+        if (C->Level != C->IDom->Level + 1) WorkStack.push_back(C);
+      }
+    }
+  }
 };
 
 template <class NodeT>
@@ -152,9 +175,10 @@ raw_ostream &operator<<(raw_ostream &O, const DomTreeNodeBase<NodeT> *Node) {
   else
     O << " <<exit node>>";
 
-  O << " {" << Node->getDFSNumIn() << "," << Node->getDFSNumOut() << "}";
+  O << " {" << Node->getDFSNumIn() << "," << Node->getDFSNumOut() << "} ["
+    << Node->getLevel() << "]\n";
 
-  return O << "\n";
+  return O;
 }
 
 template <class NodeT>
@@ -345,6 +369,13 @@ template <class NodeT> class DominatorTreeBase {
     if (!isReachableFromEntry(A))
       return false;
 
+    if (B->getIDom() == A) return true;
+
+    if (A->getIDom() == B) return false;
+
+    // A can only dominate B if it is higher in the tree.
+    if (A->getLevel() >= B->getLevel()) return false;
+
     // Compare the result of the tree walk and the dfs numbers, if expensive
     // checks are enabled.
 #ifdef EXPENSIVE_CHECKS
@@ -388,51 +419,18 @@ template <class NodeT> class DominatorTreeBase {
         return &Entry;
     }
 
-    // If B dominates A then B is nearest common dominator.
-    if (dominates(B, A))
-      return B;
-
-    // If A dominates B then A is nearest common dominator.
-    if (dominates(A, B))
-      return A;
-
     DomTreeNodeBase<NodeT> *NodeA = getNode(A);
     DomTreeNodeBase<NodeT> *NodeB = getNode(B);
 
-    // If we have DFS info, then we can avoid all allocations by just querying
-    // it from each IDom. Note that because we call 'dominates' twice above, we
-    // expect to call through this code at most 16 times in a row without
-    // building valid DFS information. This is important as below is a *very*
-    // slow tree walk.
-    if (DFSInfoValid) {
-      DomTreeNodeBase<NodeT> *IDomA = NodeA->getIDom();
-      while (IDomA) {
-        if (NodeB->DominatedBy(IDomA))
-          return IDomA->getBlock();
-        IDomA = IDomA->getIDom();
-      }
-      return nullptr;
-    }
+    // Use level information to go up the tree until the levels match. Then
+    // continue going up til we arrive at the same node.
+    while (NodeA && NodeA != NodeB) {
+      if (NodeA->getLevel() < NodeB->getLevel()) std::swap(NodeA, NodeB);
 
-    // Collect NodeA dominators set.
-    SmallPtrSet<DomTreeNodeBase<NodeT> *, 16> NodeADoms;
-    NodeADoms.insert(NodeA);
-    DomTreeNodeBase<NodeT> *IDomA = NodeA->getIDom();
-    while (IDomA) {
-      NodeADoms.insert(IDomA);
-      IDomA = IDomA->getIDom();
+      NodeA = NodeA->IDom;
     }
 
-    // Walk NodeB immediate dominators chain and find common dominator node.
-    DomTreeNodeBase<NodeT> *IDomB = NodeB->getIDom();
-    while (IDomB) {
-      if (NodeADoms.count(IDomB) != 0)
-        return IDomB->getBlock();
-
-      IDomB = IDomB->getIDom();
-    }
-
-    return nullptr;
+    return NodeA ? NodeA->getBlock() : nullptr;
   }
 
   const NodeT *findNearestCommonDominator(const NodeT *A, const NodeT *B) {
@@ -481,8 +479,10 @@ template <class NodeT> class DominatorTreeBase {
     } else {
       assert(Roots.size() == 1);
       NodeT *OldRoot = Roots.front();
-      DomTreeNodes[OldRoot] =
-        NewNode->addChild(std::move(DomTreeNodes[OldRoot]));
+      auto &OldNode = DomTreeNodes[OldRoot];
+      OldNode = NewNode->addChild(std::move(DomTreeNodes[OldRoot]));
+      OldNode->IDom = NewNode;
+      OldNode->UpdateLevel();
       Roots[0] = BB;
     }
     return RootNode = NewNode;
diff --git a/llvm/include/llvm/Support/GenericDomTreeConstruction.h b/llvm/include/llvm/Support/GenericDomTreeConstruction.h
index 038d34fd5c6..679469ec42b 100644
--- a/llvm/include/llvm/Support/GenericDomTreeConstruction.h
+++ b/llvm/include/llvm/Support/GenericDomTreeConstruction.h
@@ -318,6 +318,39 @@ struct SemiNCAInfo {
     return true;
   }
 
+  // Check if for every parent with a level L in the tree all of its children
+  // have level L + 1.
+  static bool VerifyLevels(const DomTreeT &DT) {
+    for (auto &NodeToTN : DT.DomTreeNodes) {
+      const TreeNodePtr TN = NodeToTN.second.get();
+      const NodePtr BB = TN->getBlock();
+      if (!BB) continue;
+
+      const TreeNodePtr IDom = TN->getIDom();
+      if (!IDom && TN->getLevel() != 0) {
+        errs() << "Node without an IDom ";
+        PrintBlockOrNullptr(errs(), BB);
+        errs() << " has a nonzero level " << TN->getLevel() << "!\n";
+        errs().flush();
+
+        return false;
+      }
+
+      if (IDom && TN->getLevel() != IDom->getLevel() + 1) {
+        errs() << "Node ";
+        PrintBlockOrNullptr(errs(), BB);
+        errs() << " has level " << TN->getLevel() << " while it's IDom ";
+        PrintBlockOrNullptr(errs(), IDom->getBlock());
+        errs() << " has level " << IDom->getLevel() << "!\n";
+        errs().flush();
+
+        return false;
+      }
+    }
+
+    return true;
+  }
+
   // Checks if the tree has the parent property: if for all edges from V to W in
   // the input graph, such that V is reachable, the parent of W in the tree is
   // an ancestor of V in the tree.
@@ -405,9 +438,8 @@ bool Verify(const DominatorTreeBaseByGraphTraits<GraphTraits<NodeT>> &DT) {
                 "NodePtr should be a pointer type");
   SemiNCAInfo<typename std::remove_pointer<NodePtr>::type> SNCA;
 
-  return SNCA.verifyReachability(DT) && SNCA.verifyParentProperty(DT) &&
-         SNCA.verifySiblingProperty(DT);
-
+  return SNCA.verifyReachability(DT) && SNCA.VerifyLevels(DT) &&
+         SNCA.verifyParentProperty(DT) && SNCA.verifySiblingProperty(DT);
 }
 
 }  // namespace DomTreeBuilder
diff --git a/llvm/unittests/IR/DominatorTreeTest.cpp b/llvm/unittests/IR/DominatorTreeTest.cpp
index 5d3ef30177a..fa3dad8a2ab 100644
--- a/llvm/unittests/IR/DominatorTreeTest.cpp
+++ b/llvm/unittests/IR/DominatorTreeTest.cpp
@@ -230,6 +230,12 @@ TEST(DominatorTree, Unreachable) {
         EXPECT_EQ(DT->getNode(BB4)->getDFSNumIn(), 3UL);
         EXPECT_EQ(DT->getNode(BB4)->getDFSNumOut(), 4UL);
 
+        // Check levels before
+        EXPECT_EQ(DT->getNode(BB0)->getLevel(), 0U);
+        EXPECT_EQ(DT->getNode(BB1)->getLevel(), 1U);
+        EXPECT_EQ(DT->getNode(BB2)->getLevel(), 1U);
+        EXPECT_EQ(DT->getNode(BB4)->getLevel(), 1U);
+
         // Reattach block 3 to block 1 and recalculate
         BB1->getTerminator()->eraseFromParent();
         BranchInst::Create(BB4, BB3, ConstantInt::getTrue(F.getContext()), BB1);
@@ -248,6 +254,13 @@ TEST(DominatorTree, Unreachable) {
         EXPECT_EQ(DT->getNode(BB4)->getDFSNumIn(), 5UL);
         EXPECT_EQ(DT->getNode(BB4)->getDFSNumOut(), 6UL);
 
+        // Check levels after
+        EXPECT_EQ(DT->getNode(BB0)->getLevel(), 0U);
+        EXPECT_EQ(DT->getNode(BB1)->getLevel(), 1U);
+        EXPECT_EQ(DT->getNode(BB2)->getLevel(), 1U);
+        EXPECT_EQ(DT->getNode(BB3)->getLevel(), 2U);
+        EXPECT_EQ(DT->getNode(BB4)->getLevel(), 1U);
+
         // Change root node
         DT->verifyDomTree();
         BasicBlock *NewEntry =