[RDF] Lower the sorting complexity in RDFLiveness::getAllReachingDefs

The sorting is needed, because reaching defs are (logically) ordered,
but are not collected in that order. This change will break up the
single call to std::sort into a series of smaller sorts, each of which
should use a cheaper comparison function than the original.

Author: Krzysztof Parzyszek

llvm/lib/CodeGen/RDFLiveness.cpp

@@ -23,8 +23,10 @@
 // <10.1145/2086696.2086706>. <hal-00647369>
 //
 #include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineDominanceFrontier.h"
 #include "llvm/CodeGen/MachineDominators.h"
@@ -108,7 +110,7 @@ NodeList Liveness::getAllReachingDefs(RegisterRef RefRR,
       const RegisterAggr &DefRRs) {
   NodeList RDefs; // Return value.
   SetVector<NodeId> DefQ;
-  SetVector<NodeId> Owners;
+  DenseMap<MachineInstr*, uint32_t> OrdMap;
 
   // Dead defs will be treated as if they were live, since they are actually
   // on the data-flow path. They cannot be ignored because even though they
@@ -151,18 +153,9 @@ NodeList Liveness::getAllReachingDefs(RegisterRef RefRR,
     for (auto S : DFG.getRelatedRefs(TA.Addr->getOwner(DFG), TA))
       if (NodeId RD = NodeAddr<RefNode*>(S).Addr->getReachingDef())
         DefQ.insert(RD);
-  }
-
-  // Remove all non-phi defs that are not aliased to RefRR, and collect
-  // the owners of the remaining defs.
-  SetVector<NodeId> Defs;
-  for (NodeId N : DefQ) {
-    auto TA = DFG.addr<DefNode*>(N);
-    bool IsPhi = TA.Addr->getFlags() & NodeAttrs::PhiRef;
-    if (!IsPhi && !PRI.alias(RefRR, TA.Addr->getRegRef(DFG)))
-      continue;
-    Defs.insert(TA.Id);
-    Owners.insert(TA.Addr->getOwner(DFG).Id);
+    // Don't visit sibling defs. They share the same reaching def (which
+    // will be visited anyway), but they define something not aliased to
+    // this ref.
   }
 
   // Return the MachineBasicBlock containing a given instruction.
@@ -174,38 +167,81 @@ NodeList Liveness::getAllReachingDefs(RegisterRef RefRR,
     NodeAddr<BlockNode*> BA = PA.Addr->getOwner(DFG);
     return BA.Addr->getCode();
   };
-  // Less(A,B) iff instruction A is further down in the dominator tree than B.
-  auto Less = [&Block,this] (NodeId A, NodeId B) -> bool {
+
+  SmallSet<NodeId,32> Defs;
+
+  // Remove all non-phi defs that are not aliased to RefRR, and segregate
+  // the the remaining defs into buckets for containing blocks.
+  std::map<NodeId, NodeAddr<InstrNode*>> Owners;
+  std::map<MachineBasicBlock*, SmallVector<NodeId,32>> Blocks;
+  for (NodeId N : DefQ) {
+    auto TA = DFG.addr<DefNode*>(N);
+    bool IsPhi = TA.Addr->getFlags() & NodeAttrs::PhiRef;
+    if (!IsPhi && !PRI.alias(RefRR, TA.Addr->getRegRef(DFG)))
+      continue;
+    Defs.insert(TA.Id);
+    NodeAddr<InstrNode*> IA = TA.Addr->getOwner(DFG);
+    Owners[TA.Id] = IA;
+    Blocks[Block(IA)].push_back(IA.Id);
+  }
+
+  auto Precedes = [this,&OrdMap] (NodeId A, NodeId B) {
     if (A == B)
       return false;
-    auto OA = DFG.addr<InstrNode*>(A), OB = DFG.addr<InstrNode*>(B);
-    MachineBasicBlock *BA = Block(OA), *BB = Block(OB);
-    if (BA != BB)
-      return MDT.dominates(BB, BA);
-    // They are in the same block.
+    NodeAddr<InstrNode*> OA = DFG.addr<InstrNode*>(A);
+    NodeAddr<InstrNode*> OB = DFG.addr<InstrNode*>(B);
     bool StmtA = OA.Addr->getKind() == NodeAttrs::Stmt;
     bool StmtB = OB.Addr->getKind() == NodeAttrs::Stmt;
-    if (StmtA) {
-      if (!StmtB)   // OB is a phi and phis dominate statements.
-        return true;
-      MachineInstr *CA = NodeAddr<StmtNode*>(OA).Addr->getCode();
-      MachineInstr *CB = NodeAddr<StmtNode*>(OB).Addr->getCode();
-      // The order must be linear, so tie-break such equalities.
-      if (CA == CB)
-        return A < B;
-      return MDT.dominates(CB, CA);
-    } else {
-      // OA is a phi.
-      if (StmtB)
-        return false;
-      // Both are phis. There is no ordering between phis (in terms of
-      // the data-flow), so tie-break this via node id comparison.
+    if (StmtA && StmtB) {
+      const MachineInstr *InA = NodeAddr<StmtNode*>(OA).Addr->getCode();
+      const MachineInstr *InB = NodeAddr<StmtNode*>(OB).Addr->getCode();
+      assert(InA->getParent() == InB->getParent());
+      auto FA = OrdMap.find(InA);
+      if (FA != OrdMap.end())
+        return FA->second < OrdMap.find(InB)->second;
+      const MachineBasicBlock *BB = InA->getParent();
+      for (auto It = BB->begin(), E = BB->end(); It != E; ++It) {
+        if (It == InA->getIterator())
+          return true;
+        if (It == InB->getIterator())
+          return false;
+      }
+      llvm_unreachable("InA and InB should be in the same block");
+    }
+    // One of them is a phi node.
+    if (!StmtA && !StmtB) {
+      // Both are phis, which are unordered. Break the tie by id numbers.
       return A < B;
     }
+    // Only one of them is a phi. Phis always precede statements.
+    return !StmtA;
   };
 
-  std::vector<NodeId> Tmp(Owners.begin(), Owners.end());
-  llvm::sort(Tmp, Less);
+  auto GetOrder = [&OrdMap] (MachineBasicBlock &B) {
+    uint32_t Pos = 0;
+    for (MachineInstr &In : B)
+      OrdMap.insert({&In, ++Pos});
+  };
+
+  // For each block, sort the nodes in it.
+  std::vector<MachineBasicBlock*> TmpBB;
+  for (auto &Bucket : Blocks) {
+    TmpBB.push_back(Bucket.first);
+    if (Bucket.second.size() > 2)
+      GetOrder(*Bucket.first);
+    std::sort(Bucket.second.begin(), Bucket.second.end(), Precedes);
+  }
+
+  // Sort the blocks with respect to dominance.
+  std::sort(TmpBB.begin(), TmpBB.end(), [this](auto A, auto B) {
+    return MDT.dominates(A, B);
+  });
+
+  std::vector<NodeId> TmpInst;
+  for (auto I = TmpBB.rbegin(), E = TmpBB.rend(); I != E; ++I) {
+    auto &Bucket = Blocks[*I];
+    TmpInst.insert(TmpInst.end(), Bucket.rbegin(), Bucket.rend());
+  }
 
   // The vector is a list of instructions, so that defs coming from
   // the same instruction don't need to be artificially ordered.
@@ -220,14 +256,18 @@ NodeList Liveness::getAllReachingDefs(RegisterRef RefRR,
   // *d3<C>              If A \incl BuC, and B \incl AuC, then *d2 would be
   //                     covered if we added A first, and A would be covered
   //                     if we added B first.
+  // In this example we want both A and B, because we don't want to give
+  // either one priority over the other, since they belong to the same
+  // statement.
 
   RegisterAggr RRs(DefRRs);
 
   auto DefInSet = [&Defs] (NodeAddr<RefNode*> TA) -> bool {
     return TA.Addr->getKind() == NodeAttrs::Def &&
            Defs.count(TA.Id);
   };
-  for (NodeId T : Tmp) {
+
+  for (NodeId T : TmpInst) {
     if (!FullChain && RRs.hasCoverOf(RefRR))
       break;
     auto TA = DFG.addr<InstrNode*>(T);