[BranchProbabilityInfo] Use SmallVector (NFC)

This patch simplifies BranchProbabilityInfo by changing the type of
Probs.

Without this patch:

  DenseMap<Edge, BranchProbability> Probs

maps an ordered pair of a BasicBlock* and a successor index to an edge
probability.

With this patch:

  DenseMap<const BasicBlock *, SmallVector<BranchProbability, 2>> Probs

maps a BasicBlock* to a vector of edge probabilities.

BranchProbabilityInfo has a property that for a given basic block, we
either have edge probabilities for all successors or do not have any
edge probability at all.  This property combined with the current map
type leads to a somewhat complicated algorithm in eraseBlock to erase
map entries one by one while increasing the successor index.

The new map type allows us to remove the all edge probabilities for a
given basic block in a more intuitive manner, namely:

  Probs.erase(BB);

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

1 files changed, 27 insertions, 35 deletions

diff --git a/llvm/lib/Analysis/BranchProbabilityInfo.cpp b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
index f235e9f..37504cf 100644
--- a/llvm/lib/Analysis/BranchProbabilityInfo.cpp
+++ b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
@@ -1091,16 +1091,15 @@ BranchProbabilityInfo::getHotSucc(const BasicBlock *BB) const {
 BranchProbability
 BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
                                           unsigned IndexInSuccessors) const {
-  auto I = Probs.find(std::make_pair(Src, IndexInSuccessors));
-  assert((Probs.end() == Probs.find(std::make_pair(Src, 0))) ==
-             (Probs.end() == I) &&
-         "Probability for I-th successor must always be defined along with the "
-         "probability for the first successor");
-
-  if (I != Probs.end())
-    return I->second;
-
-  return {1, static_cast<uint32_t>(succ_size(Src))};
+  auto I = Probs.find(Src);
+  if (I == Probs.end())
+    return {1, static_cast<uint32_t>(succ_size(Src))};
+
+  const SmallVector<BranchProbability, 2> &SrcProbs = I->second;
+  assert(SrcProbs.size() == Src->getTerminator()->getNumSuccessors() &&
+         "The number of edge probabilities must match the number of "
+         "successors.");
+  return SrcProbs[IndexInSuccessors];
 }
 
 BranchProbability
@@ -1114,13 +1113,18 @@ BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
 BranchProbability
 BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
                                           const BasicBlock *Dst) const {
-  if (!Probs.count(std::make_pair(Src, 0)))
+  auto SrcProbsIt = Probs.find(Src);
+  if (SrcProbsIt == Probs.end())
     return BranchProbability(llvm::count(successors(Src), Dst), succ_size(Src));
 
+  const SmallVector<BranchProbability, 2> &SrcProbs = SrcProbsIt->second;
+  assert(SrcProbs.size() == Src->getTerminator()->getNumSuccessors() &&
+         "The number of edge probabilities must match the number of "
+         "successors.");
   auto Prob = BranchProbability::getZero();
   for (const_succ_iterator I = succ_begin(Src), E = succ_end(Src); I != E; ++I)
     if (*I == Dst)
-      Prob += Probs.find(std::make_pair(Src, I.getSuccessorIndex()))->second;
+      Prob += SrcProbs[I.getSuccessorIndex()];
 
   return Prob;
 }
@@ -1135,8 +1139,10 @@ void BranchProbabilityInfo::setEdgeProbability(
 
   Handles.insert(BasicBlockCallbackVH(Src, this));
   uint64_t TotalNumerator = 0;
+  auto &SrcProbs = this->Probs[Src];
+  SrcProbs.reserve(Probs.size());
   for (unsigned SuccIdx = 0; SuccIdx < Probs.size(); ++SuccIdx) {
-    this->Probs[std::make_pair(Src, SuccIdx)] = Probs[SuccIdx];
+    SrcProbs.push_back(Probs[SuccIdx]);
     LLVM_DEBUG(dbgs() << "set edge " << Src->getName() << " -> " << SuccIdx
                       << " successor probability to " << Probs[SuccIdx]
                       << "\n");
@@ -1159,16 +1165,17 @@ void BranchProbabilityInfo::copyEdgeProbabilities(BasicBlock *Src,
   assert(NumSuccessors == Dst->getTerminator()->getNumSuccessors());
   if (NumSuccessors == 0)
     return; // Nothing to set.
-  if (this->Probs.find(std::make_pair(Src, 0)) == this->Probs.end())
+  auto SrcProbsIt = this->Probs.find(Src);
+  if (SrcProbsIt == this->Probs.end())
     return; // No probability is set for edges from Src. Keep the same for Dst.
 
   Handles.insert(BasicBlockCallbackVH(Dst, this));
-  for (unsigned SuccIdx = 0; SuccIdx < NumSuccessors; ++SuccIdx) {
-    auto Prob = this->Probs[std::make_pair(Src, SuccIdx)];
-    this->Probs[std::make_pair(Dst, SuccIdx)] = Prob;
+  auto &DstProbs = this->Probs[Dst];
+  DstProbs = SrcProbsIt->second;
+  for (unsigned SuccIdx = 0; SuccIdx < NumSuccessors; ++SuccIdx)
     LLVM_DEBUG(dbgs() << "set edge " << Dst->getName() << " -> " << SuccIdx
-                      << " successor probability to " << Prob << "\n");
-  }
+                      << " successor probability to " << DstProbs[SuccIdx]
+                      << "\n");
 }
 
 raw_ostream &
@@ -1184,23 +1191,8 @@ BranchProbabilityInfo::printEdgeProbability(raw_ostream &OS,
 }
 
 void BranchProbabilityInfo::eraseBlock(const BasicBlock *BB) {
-  // Note that we cannot use successors of BB because the terminator of BB may
-  // have changed when eraseBlock is called as a BasicBlockCallbackVH callback.
-  // Instead we remove prob data for the block by iterating successors by their
-  // indices from 0 till the last which exists. There could not be prob data for
-  // a pair (BB, N) if there is no data for (BB, N-1) because the data is always
-  // set for all successors from 0 to M at once by the method
-  // setEdgeProbability().
   Handles.erase(BasicBlockCallbackVH(BB, this));
-  for (unsigned I = 0;; ++I) {
-    auto MapI = Probs.find(std::make_pair(BB, I));
-    if (MapI == Probs.end()) {
-      assert(Probs.count(std::make_pair(BB, I + 1)) == 0 &&
-             "Must be no more successors");
-      return;
-    }
-    Probs.erase(MapI);
-  }
+  Probs.erase(BB);
 }
 
 void BranchProbabilityInfo::calculate(const Function &F, const LoopInfo &LI,