[BranchProbabilityInfo] Change analysis result storage to simplify basic block removal

[igor-laevsky]

Summary

This is follow up after https://reviews.llvm.org/D20957

There were concerns that in it's current form EraseBlock operation is not efficient enough. Considering that same problem was fixed for LVI (https://reviews.llvm.org/D11651) we should do the same for BranchProbabilityInfo.

This change replaces existing results cache: "(BB, SuccessorIdx) -> BranchProbability" with "BB -> (SuccessorIdx -> BranchProbability)". This allows us to efficiently perform basic block removal in a cost of slightly more expensive lookups.

[igor-laevsky]

ping?

[dnovillo]

Any changes in build time? I guess I'm surprised that removal operations are more frequent than lookups.

[davidxl]

Same question as Diego -- what is the build time improvement? The original version looks cleaner and easier to maintain

lib/Analysis/BranchProbabilityInfo.cpp

627

I don't quite like this EdgeProbablilityRecorded gets checked in multiple different places. It is better to isolate in one private helper function getEdgeProbability_(Src,index) that returns unknown when the index is out of range.

[igor-laevsky]

Hi,

I don't have specific build times. This is mostly motivated by the review comments in the https://reviews.llvm.org/D20957. Evidence for importance of this is that LVI also had linear time EraseBlock and it caused problems there. BPI will call EraseBlock exactly as frequently as LVI did, so it's natural to assume that performance bottleneck will be in the same place.

include/llvm/Analysis/BranchProbabilityInfo.h

 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/IndexedMap.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/ValueHandle.h"

   };
   DenseSet<BasicBlockCallbackVH, DenseMapInfo<Value*>> Handles;
 
-  // Since we allow duplicate edges from one basic block to another, we use
-  // a pair (PredBlock and an index in the successors) to specify an edge.
-  typedef std::pair<const BasicBlock *, unsigned> Edge;
-
   // Default weight value. Used when we don't have information about the edge.
   // TODO: DEFAULT_WEIGHT makes sense during static predication, when none of
   // the successors have a weight yet. But it doesn't make sense when providing

   // weight to just "inherit" the non-zero weight of an adjacent successor.
   static const uint32_t DEFAULT_WEIGHT = 16;
 
-  DenseMap<Edge, BranchProbability> Probs;
+  // Since we allow duplicate edges from one basic block to another, we need to
+  // use successor index to determine edge destination. This results in a
+  // following data structure:
+  //   BB -> (SuccessorIdx -> BranchProbability)
+  DenseMap<const BasicBlock *, IndexedMap<BranchProbability>> Probs;
 
   /// \brief Track the last function we run over for printing.
   const Function *LastF;

lib/Analysis/BranchProbabilityInfo.cpp

   return nullptr;
 }
 
+static bool
+EdgeProbabilityRecorded(const IndexedMap<BranchProbability> &SuccMap,
+                        unsigned Index) {
+  return SuccMap.inBounds(Index) && !SuccMap[Index].isUnknown();
+}
+
 /// Get the raw edge probability for the edge. If can't find it, return a
 /// default probability 1/N where N is the number of successors. Here an edge is
 /// specified using PredBlock and an

 BranchProbability
 BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
                                           unsigned IndexInSuccessors) const {
-  auto I = Probs.find(std::make_pair(Src, IndexInSuccessors));
+  const auto I = Probs.find(Src);
 
-  if (I != Probs.end())
-    return I->second;
+  if (I != Probs.end() &&
+      EdgeProbabilityRecorded(I->second, IndexInSuccessors))
+    return I->second[IndexInSuccessors];
 
   return {1,
           static_cast<uint32_t>(std::distance(succ_begin(Src), succ_end(Src)))};

 BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
                                           const BasicBlock *Dst) const {
   auto Prob = BranchProbability::getZero();
+  uint32_t succ_num = std::distance(succ_begin(Src), succ_end(Src));
+
+  const auto BBProbsIt = Probs.find(Src);
+  if (BBProbsIt == Probs.end())
+    return {1, succ_num};
+  const IndexedMap<BranchProbability> &BBSuccProbs = BBProbsIt->second;
+
   bool FoundProb = false;
-  for (succ_const_iterator I = succ_begin(Src), E = succ_end(Src); I != E; ++I)
-    if (*I == Dst) {
-      auto MapI = Probs.find(std::make_pair(Src, I.getSuccessorIndex()));
-      if (MapI != Probs.end()) {
-        FoundProb = true;
-        Prob += MapI->second;
-      }
+  for (succ_const_iterator I = succ_begin(Src), E = succ_end(Src); I != E;
+       ++I) {
+    unsigned SuccIdx = I.getSuccessorIndex();
+
+    if (*I == Dst && EdgeProbabilityRecorded(BBSuccProbs, SuccIdx)) {
+      FoundProb = true;
+      Prob += BBSuccProbs[SuccIdx];
     }
-  uint32_t succ_num = std::distance(succ_begin(Src), succ_end(Src));
+  }
+
   return FoundProb ? Prob : BranchProbability(1, succ_num);
 }
 

 void BranchProbabilityInfo::setEdgeProbability(const BasicBlock *Src,
                                                unsigned IndexInSuccessors,
                                                BranchProbability Prob) {
-  Probs[std::make_pair(Src, IndexInSuccessors)] = Prob;
+  IndexedMap<BranchProbability> &BBSuccProbs = Probs[Src];
+
+  // This is nop if IndexInSuccessors is already present in the map.
+  // If map was resized new elements will be filled with the 'Unknown' branch
+  // probability.
+  BBSuccProbs.grow(IndexInSuccessors);
+  BBSuccProbs[IndexInSuccessors] = Prob;
+
   Handles.insert(BasicBlockCallbackVH(Src, this));
+
   DEBUG(dbgs() << "set edge " << Src->getName() << " -> " << IndexInSuccessors
                << " successor probability to " << Prob << "\n");
 }

 }
 
 void BranchProbabilityInfo::eraseBlock(const BasicBlock *BB) {
-  for (auto I = Probs.begin(), E = Probs.end(); I != E; ++I) {
-    auto Key = I->first;
-    if (Key.first == BB)
-      Probs.erase(Key);
-  }
+  Probs.erase(BB);
 }
 
 void BranchProbabilityInfo::calculate(const Function &F, const LoopInfo &LI) {