[Support] Use block numbers for DomTree construction (#101706)

Similar to #101705, do the same optimization for dominator tree
construction.

Author: aengelke

llvm/include/llvm/ADT/GraphTraits.h

@@ -19,6 +19,7 @@
 #ifndef LLVM_ADT_GRAPHTRAITS_H
 #define LLVM_ADT_GRAPHTRAITS_H
 
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/iterator_range.h"
 
 namespace llvm {
@@ -94,6 +95,16 @@ struct GraphTraits {
   using NodeRef = typename GraphType::UnknownGraphTypeError;
 };
 
+namespace detail {
+template <typename T>
+using has_number_t = decltype(GraphTraits<T>::getNumber(std::declval<T>()));
+} // namespace detail
+
+/// Indicate whether a GraphTraits<NodeT>::getNumber() is supported.
+template <typename NodeT>
+constexpr bool GraphHasNodeNumbers =
+    is_detected<detail::has_number_t, NodeT>::value;
+
 // Inverse - This class is used as a little marker class to tell the graph
 // iterator to iterate over the graph in a graph defined "Inverse" ordering.
 // Not all graphs define an inverse ordering, and if they do, it depends on

llvm/include/llvm/Support/GenericDomTreeConstruction.h

@@ -73,7 +73,11 @@ struct SemiNCAInfo {
   // Number to node mapping is 1-based. Initialize the mapping to start with
   // a dummy element.
   SmallVector<NodePtr, 64> NumToNode = {nullptr};
-  DenseMap<NodePtr, InfoRec> NodeToInfo;
+  // If blocks have numbers (e.g., BasicBlock, MachineBasicBlock), store node
+  // infos in a vector. Otherwise, store them in a map.
+  std::conditional_t<GraphHasNodeNumbers<NodePtr>, SmallVector<InfoRec, 64>,
+                     DenseMap<NodePtr, InfoRec>>
+      NodeInfos;
 
   using UpdateT = typename DomTreeT::UpdateType;
   using UpdateKind = typename DomTreeT::UpdateKind;
@@ -99,7 +103,7 @@ struct SemiNCAInfo {
 
   void clear() {
     NumToNode = {nullptr}; // Restore to initial state with a dummy start node.
-    NodeToInfo.clear();
+    NodeInfos.clear();
     // Don't reset the pointer to BatchUpdateInfo here -- if there's an update
     // in progress, we need this information to continue it.
   }
@@ -123,13 +127,25 @@ struct SemiNCAInfo {
     return Res;
   }
 
-  NodePtr getIDom(NodePtr BB) const {
-    auto InfoIt = NodeToInfo.find(BB);
-    if (InfoIt == NodeToInfo.end()) return nullptr;
-
-    return InfoIt->second.IDom;
+  InfoRec &getNodeInfo(NodePtr BB) {
+    if constexpr (GraphHasNodeNumbers<NodePtr>) {
+      unsigned Idx = BB ? GraphTraits<NodePtr>::getNumber(BB) + 1 : 0;
+      if (Idx >= NodeInfos.size()) {
+        unsigned Max = 0;
+        if (BB)
+          Max = GraphTraits<decltype(BB->getParent())>::getMaxNumber(
+              BB->getParent());
+        // Max might be zero, graphs might not support getMaxNumber().
+        NodeInfos.resize(Max ? Max + 1 : Idx + 1);
+      }
+      return NodeInfos[Idx];
+    } else {
+      return NodeInfos[BB];
+    }
   }
 
+  NodePtr getIDom(NodePtr BB) { return getNodeInfo(BB).IDom; }
+
   TreeNodePtr getNodeForBlock(NodePtr BB, DomTreeT &DT) {
     if (TreeNodePtr Node = DT.getNode(BB)) return Node;
 
@@ -181,11 +197,11 @@ struct SemiNCAInfo {
                   const NodeOrderMap *SuccOrder = nullptr) {
     assert(V);
     SmallVector<std::pair<NodePtr, unsigned>, 64> WorkList = {{V, AttachToNum}};
-    NodeToInfo[V].Parent = AttachToNum;
+    getNodeInfo(V).Parent = AttachToNum;
 
     while (!WorkList.empty()) {
       const auto [BB, ParentNum] = WorkList.pop_back_val();
-      auto &BBInfo = NodeToInfo[BB];
+      auto &BBInfo = getNodeInfo(BB);
       BBInfo.ReverseChildren.push_back(ParentNum);
 
       // Visited nodes always have positive DFS numbers.
@@ -264,7 +280,7 @@ struct SemiNCAInfo {
     // Initialize IDoms to spanning tree parents.
     for (unsigned i = 1; i < NextDFSNum; ++i) {
       const NodePtr V = NumToNode[i];
-      auto &VInfo = NodeToInfo[V];
+      auto &VInfo = getNodeInfo(V);
       VInfo.IDom = NumToNode[VInfo.Parent];
       NumToInfo.push_back(&VInfo);
     }
@@ -292,7 +308,7 @@ struct SemiNCAInfo {
       const unsigned SDomNum = NumToInfo[WInfo.Semi]->DFSNum;
       NodePtr WIDomCandidate = WInfo.IDom;
       while (true) {
-        auto &WIDomCandidateInfo = NodeToInfo.find(WIDomCandidate)->second;
+        auto &WIDomCandidateInfo = getNodeInfo(WIDomCandidate);
         if (WIDomCandidateInfo.DFSNum <= SDomNum)
           break;
         WIDomCandidate = WIDomCandidateInfo.IDom;
@@ -311,7 +327,7 @@ struct SemiNCAInfo {
     assert(IsPostDom && "Only postdominators have a virtual root");
     assert(NumToNode.size() == 1 && "SNCAInfo must be freshly constructed");
 
-    auto &BBInfo = NodeToInfo[nullptr];
+    auto &BBInfo = getNodeInfo(nullptr);
     BBInfo.DFSNum = BBInfo.Semi = BBInfo.Label = 1;
 
     NumToNode.push_back(nullptr);  // NumToNode[1] = nullptr;
@@ -393,7 +409,7 @@ struct SemiNCAInfo {
       auto InitSuccOrderOnce = [&]() {
         SuccOrder = NodeOrderMap();
         for (const auto Node : nodes(DT.Parent))
-          if (SNCA.NodeToInfo.count(Node) == 0)
+          if (SNCA.getNodeInfo(Node).DFSNum == 0)
             for (const auto Succ : getChildren<false>(Node, SNCA.BatchUpdates))
               SuccOrder->try_emplace(Succ, 0);
 
@@ -417,7 +433,7 @@ struct SemiNCAInfo {
       // unreachable node once, we may just visit it in two directions,
       // depending on how lucky we get.
       for (const NodePtr I : nodes(DT.Parent)) {
-        if (SNCA.NodeToInfo.count(I) == 0) {
+        if (SNCA.getNodeInfo(I).DFSNum == 0) {
           LLVM_DEBUG(dbgs()
                      << "\t\t\tVisiting node " << BlockNamePrinter(I) << "\n");
           // Find the furthest away we can get by following successors, then
@@ -449,7 +465,7 @@ struct SemiNCAInfo {
             const NodePtr N = SNCA.NumToNode[i];
             LLVM_DEBUG(dbgs() << "\t\t\t\tRemoving DFS info for "
                               << BlockNamePrinter(N) << "\n");
-            SNCA.NodeToInfo.erase(N);
+            SNCA.getNodeInfo(N) = {};
             SNCA.NumToNode.pop_back();
           }
           const unsigned PrevNum = Num;
@@ -582,7 +598,7 @@ struct SemiNCAInfo {
 
   void attachNewSubtree(DomTreeT& DT, const TreeNodePtr AttachTo) {
     // Attach the first unreachable block to AttachTo.
-    NodeToInfo[NumToNode[1]].IDom = AttachTo->getBlock();
+    getNodeInfo(NumToNode[1]).IDom = AttachTo->getBlock();
     // Loop over all of the discovered blocks in the function...
     for (NodePtr W : llvm::drop_begin(NumToNode)) {
       if (DT.getNode(W))
@@ -600,11 +616,11 @@ struct SemiNCAInfo {
   }
 
   void reattachExistingSubtree(DomTreeT &DT, const TreeNodePtr AttachTo) {
-    NodeToInfo[NumToNode[1]].IDom = AttachTo->getBlock();
+    getNodeInfo(NumToNode[1]).IDom = AttachTo->getBlock();
     for (const NodePtr N : llvm::drop_begin(NumToNode)) {
       const TreeNodePtr TN = DT.getNode(N);
       assert(TN);
-      const TreeNodePtr NewIDom = DT.getNode(NodeToInfo[N].IDom);
+      const TreeNodePtr NewIDom = DT.getNode(getNodeInfo(N).IDom);
       TN->setIDom(NewIDom);
     }
   }
@@ -1239,7 +1255,7 @@ struct SemiNCAInfo {
       // Virtual root has a corresponding virtual CFG node.
       if (DT.isVirtualRoot(TN)) continue;
 
-      if (NodeToInfo.count(BB) == 0) {
+      if (getNodeInfo(BB).DFSNum == 0) {
         errs() << "DomTree node " << BlockNamePrinter(BB)
                << " not found by DFS walk!\n";
         errs().flush();
@@ -1453,7 +1469,7 @@ struct SemiNCAInfo {
       });
 
       for (TreeNodePtr Child : TN->children())
-        if (NodeToInfo.count(Child->getBlock()) != 0) {
+        if (getNodeInfo(Child->getBlock()).DFSNum != 0) {
           errs() << "Child " << BlockNamePrinter(Child)
                  << " reachable after its parent " << BlockNamePrinter(BB)
                  << " is removed!\n";
@@ -1491,7 +1507,7 @@ struct SemiNCAInfo {
         for (const TreeNodePtr S : TN->children()) {
           if (S == N) continue;
 
-          if (NodeToInfo.count(S->getBlock()) == 0) {
+          if (getNodeInfo(S->getBlock()).DFSNum == 0) {
             errs() << "Node " << BlockNamePrinter(S)
                    << " not reachable when its sibling " << BlockNamePrinter(N)
                    << " is removed!\n";