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

[kuhar]

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.

[kuhar]

Add level verification and update tests.

[kuhar]

Minor cleanup

[• dberlin]

Can you also change IteratedDominanceFrontier to use this level info instead of computing it on it's own?

[kuhar]

Update the diff.

I would like to teach IDF to use levels in a separate patch.

[kuhar]

Update the diff to apply cleanly.

include/llvm/Support/GenericDomTree.h

 
 /// \brief Base class for the actual dominator tree node.
 template <class NodeT> class DomTreeNodeBase {
   friend struct PostDominatorTree;
   template <class N> friend class DominatorTreeBase;
 
   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 =
       typename std::vector<DomTreeNodeBase *>::const_iterator;
 
   iterator begin() { return Children.begin(); }
   iterator end() { return Children.end(); }
   const_iterator begin() const { return Children.begin(); }
   const_iterator end() const { return Children.end(); }
 
   NodeT *getBlock() const { return TheBB; }
   DomTreeNodeBase *getIDom() const { return IDom; }
+  unsigned getLevel() const { return Level; }
 
   const std::vector<DomTreeNodeBase *> &getChildren() const { return Children; }
 
   std::unique_ptr<DomTreeNodeBase> addChild(
       std::unique_ptr<DomTreeNodeBase> C) {
     Children.push_back(C.get());
     return C;
   }
 
   size_t getNumChildren() const { return Children.size(); }
 
   void clearAllChildren() { Children.clear(); }
 
   bool compare(const DomTreeNodeBase *Other) const {
     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();
       OtherChildren.insert(Nd);
     }
 
     for (const DomTreeNodeBase *I : *this) {
       const NodeT *N = I->getBlock();
       if (OtherChildren.count(N) == 0)
         return true;
     }
     return false;
   }
 
   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;
+
+    auto I = find(IDom->Children, this);
 
-      // Switch to new dominator
-      IDom = NewIDom;
-      IDom->Children.push_back(this);
-    }
+
+    UpdateLevel();
   }
 
   /// getDFSNumIn/getDFSNumOut - These return the DFS visitation order for nodes
   /// in the dominator tree. They are only guaranteed valid if
   /// updateDFSNumbers() has been called.
   unsigned getDFSNumIn() const { return DFSNumIn; }
   unsigned getDFSNumOut() const { return DFSNumOut; }
 
 private:
   // Return true if this node is dominated by other. Use this only if DFS info
   // is valid.
   bool DominatedBy(const DomTreeNodeBase *other) const {
     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>
 raw_ostream &operator<<(raw_ostream &O, const DomTreeNodeBase<NodeT> *Node) {
   if (Node->getBlock())
     Node->getBlock()->printAsOperand(O, false);
   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>
 void PrintDomTree(const DomTreeNodeBase<NodeT> *N, raw_ostream &O,
                   unsigned Lev) {
   O.indent(2 * Lev) << "[" << Lev << "] " << N;
   for (typename DomTreeNodeBase<NodeT>::const_iterator I = N->begin(),
                                                        E = N->end();
bool dominates(const DomTreeNodeBase<NodeT> *A,
     // An unreachable node is dominated by anything.
     if (!isReachableFromEntry(B))
       return true;
 
     // And dominates nothing.
     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
     assert((!DFSInfoValid ||
             (dominatedBySlowTreeWalk(A, B) == B->DominatedBy(A))) &&
            "Tree walk disagrees with dfs numbers!");
 #endif
 
NodeT *findNearestCommonDominator(NodeT *A, NodeT *B) {
     // If either A or B is a entry block then it is nearest common dominator
     // (for forward-dominators).
     if (!this->isPostDominator()) {
       NodeT &Entry = A->getParent()->front();
       if (A == &Entry || B == &Entry)
         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();
-    }
-
-    // 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();
+      NodeA = NodeA->IDom;
     }
 
-    return nullptr;
+    return NodeA ? NodeA->getBlock() : nullptr;
   }
 
   const NodeT *findNearestCommonDominator(const NodeT *A, const NodeT *B) {
     // Cast away the const qualifiers here. This is ok since
     // const is re-introduced on the return type.
     return findNearestCommonDominator(const_cast<NodeT *>(A),
                                       const_cast<NodeT *>(B));
   }
DomTreeNodeBase<NodeT> *setNewRoot(NodeT *BB) {
     DFSInfoValid = false;
     DomTreeNodeBase<NodeT> *NewNode = (DomTreeNodes[BB] =
       llvm::make_unique<DomTreeNodeBase<NodeT>>(BB, nullptr)).get();
     if (Roots.empty()) {
       addRoot(BB);
     } 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;
   }
 
   /// changeImmediateDominator - This method is used to update the dominator
   /// tree information when a node's immediate dominator changes.
   ///

include/llvm/Support/GenericDomTreeConstruction.h

for (auto &NodeToTN : DT.DomTreeNodes) {
 
         return false;
       }
     }
 
     return true;
   }
 
+  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;
+  }
+
   bool verifyParentProperty(const DomTreeT &DT) {
     for (auto &NodeToTN : DT.DomTreeNodes) {
       const TreeNodePtr TN = NodeToTN.second.get();
       const NodePtr BB = TN->getBlock();
       if (!BB || TN->getChildren().empty()) continue;
 
       clear();
       NodeToInfo.insert({BB, {}});

 
 template <class NodeT>
 bool Verify(const DominatorTreeBaseByGraphTraits<GraphTraits<NodeT>> &DT) {
   using NodePtr = typename GraphTraits<NodeT>::NodeRef;
   static_assert(std::is_pointer<NodePtr>::value,
                 "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 detail
 }  // namespace llvm
 
 #endif

unittests/IR/DominatorTreeTest.cpp

runWithDomTree(
         EXPECT_EQ(DT->getNode(BB0)->getDFSNumOut(), 7UL);
         EXPECT_EQ(DT->getNode(BB1)->getDFSNumIn(), 1UL);
         EXPECT_EQ(DT->getNode(BB1)->getDFSNumOut(), 2UL);
         EXPECT_EQ(DT->getNode(BB2)->getDFSNumIn(), 5UL);
         EXPECT_EQ(DT->getNode(BB2)->getDFSNumOut(), 6UL);
         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);
         DT->recalculate(F);
 
         // Check DFS Numbers after
         EXPECT_EQ(DT->getNode(BB0)->getDFSNumIn(), 0UL);
         EXPECT_EQ(DT->getNode(BB0)->getDFSNumOut(), 9UL);
         EXPECT_EQ(DT->getNode(BB1)->getDFSNumIn(), 1UL);
         EXPECT_EQ(DT->getNode(BB1)->getDFSNumOut(), 4UL);
         EXPECT_EQ(DT->getNode(BB2)->getDFSNumIn(), 7UL);
         EXPECT_EQ(DT->getNode(BB2)->getDFSNumOut(), 8UL);
         EXPECT_EQ(DT->getNode(BB3)->getDFSNumIn(), 2UL);
         EXPECT_EQ(DT->getNode(BB3)->getDFSNumOut(), 3UL);
         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 =
             BasicBlock::Create(F.getContext(), "new_entry", &F, BB0);
         BranchInst::Create(BB0, NewEntry);
         EXPECT_EQ(F.begin()->getName(), NewEntry->getName());
         EXPECT_TRUE(&F.getEntryBlock() == NewEntry);
         DT->setNewRoot(NewEntry);