[BranchProbabilityInfo] Get rid of MaxSuccIdx. NFC

[yrouban]

Summary

This refactoring allows to eliminate the MaxSuccIdx map proposed in D90272.
The idea is to remove probabilities for a block BB for all its successors one by one from first, second, ... till N-th until they are defined in Probs. This works because probabilities for the block are set at once for all its successors from number 0 to N-1 and the rest are removed if there were stale probs.
I propose to remove the protected method setEdgeProbability(), which sets probabilities for individual successors, to make it clear that they are set in bulk in the public method with the same name.

[kazu]

Thank you for cleaning this up.

llvm/lib/Analysis/BranchProbabilityInfo.cpp

1176

s/exist/exists/

1176

s/numbers/indexes/

1178

How about removing the public method because we now have only one setEdgeProbability?

[MaskRay]

LGTM.

llvm/lib/Analysis/BranchProbabilityInfo.cpp

1180

Nit: delete true

1183

find -> count == 0

llvm/include/llvm/Analysis/BranchProbabilityInfo.h

public:
 
   BranchProbabilityInfo(const Function &F, const LoopInfo &LI,
                         const TargetLibraryInfo *TLI = nullptr,
                         PostDominatorTree *PDT = nullptr) {
     calculate(F, LI, TLI, PDT);
   }
 
   BranchProbabilityInfo(BranchProbabilityInfo &&Arg)
-      : Probs(std::move(Arg.Probs)), MaxSuccIdx(std::move(Arg.MaxSuccIdx)),
-        LastF(Arg.LastF),
+      : Probs(std::move(Arg.Probs)), LastF(Arg.LastF),
         PostDominatedByUnreachable(std::move(Arg.PostDominatedByUnreachable)),
         PostDominatedByColdCall(std::move(Arg.PostDominatedByColdCall)) {}
 
   BranchProbabilityInfo(const BranchProbabilityInfo &) = delete;
   BranchProbabilityInfo &operator=(const BranchProbabilityInfo &) = delete;
 
   BranchProbabilityInfo &operator=(BranchProbabilityInfo &&RHS) {
     releaseMemory();
     Probs = std::move(RHS.Probs);
-    MaxSuccIdx = std::move(RHS.MaxSuccIdx);
     PostDominatedByColdCall = std::move(RHS.PostDominatedByColdCall);
     PostDominatedByUnreachable = std::move(RHS.PostDominatedByUnreachable);
     return *this;
   }
 
   bool invalidate(Function &, const PreservedAnalyses &PA,
                   FunctionAnalysisManager::Invalidator &);
 
public:
   /// Print an edge's probability.
   ///
   /// Retrieves an edge's probability similarly to \see getEdgeProbability, but
   /// then prints that probability to the provided stream. That stream is then
   /// returned.
   raw_ostream &printEdgeProbability(raw_ostream &OS, const BasicBlock *Src,
                                     const BasicBlock *Dst) const;
 
-protected:
-  /// Set the raw edge probability for the given edge.
-  ///
-  /// This allows a pass to explicitly set the edge probability for an edge. It
-  /// can be used when updating the CFG to update and preserve the branch
-  /// probability information. Read the implementation of how these edge
-  /// probabilities are calculated carefully before using!
-  void setEdgeProbability(const BasicBlock *Src, unsigned IndexInSuccessors,
-                          BranchProbability Prob);
-
 public:
   /// Set the raw probabilities for all edges from the given block.
   ///
   /// This allows a pass to explicitly set edge probabilities for a block. It
   /// can be used when updating the CFG to update the branch probability
   /// information.
   void setEdgeProbability(const BasicBlock *Src,
                           const SmallVectorImpl<BranchProbability> &Probs);
private:
   // the successors have a weight yet. But it doesn't make sense when providing
   // weight to an edge that may have siblings with non-zero weights. This can
   // be handled various ways, but it's probably fine for an edge with unknown
   // weight to just "inherit" the non-zero weight of an adjacent successor.
   static const uint32_t DEFAULT_WEIGHT = 16;
 
   DenseMap<Edge, BranchProbability> Probs;
 
-  // The maximum successor index ever entered for a given basic block.
-  DenseMap<const BasicBlock *, unsigned> MaxSuccIdx;
-
   /// Track the last function we run over for printing.
   const Function *LastF = nullptr;
 
   /// Keeps information about all SCCs in a function.
   std::unique_ptr<const SccInfo> SccI;
 
   /// Track the set of blocks directly succeeded by a returning block.
   SmallPtrSet<const BasicBlock *, 16> PostDominatedByUnreachable;

llvm/lib/Analysis/BranchProbabilityInfo.cpp

BranchProbability TakenProb(IH_TAKEN_WEIGHT,
                               IH_TAKEN_WEIGHT + IH_NONTAKEN_WEIGHT);
   setEdgeProbability(
       BB, SmallVector<BranchProbability, 2>({TakenProb, TakenProb.getCompl()}));
   return true;
 }
 
 void BranchProbabilityInfo::releaseMemory() {
   Probs.clear();
-  MaxSuccIdx.clear();
   Handles.clear();
 }
 
 bool BranchProbabilityInfo::invalidate(Function &, const PreservedAnalyses &PA,
                                        FunctionAnalysisManager::Invalidator &) {
   // Check whether the analysis, all analyses on functions, or the function's
   // CFG have been preserved.
   auto PAC = PA.getChecker<BranchProbabilityAnalysis>();

 /// 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
 /// index to the successors.
 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))};
 }
 
 BranchProbability
if (*I == Dst) {
         FoundProb = true;
         Prob += MapI->second;
       }
     }
   uint32_t succ_num = std::distance(succ_begin(Src), succ_end(Src));
   return FoundProb ? Prob : BranchProbability(EdgeCount, succ_num);
 }
 
-/// Set the edge probability for a given edge specified by PredBlock and an
-/// index to the successors.
-void BranchProbabilityInfo::setEdgeProbability(const BasicBlock *Src,
-                                               unsigned IndexInSuccessors,
-                                               BranchProbability Prob) {
-  Probs[std::make_pair(Src, IndexInSuccessors)] = Prob;
-  Handles.insert(BasicBlockCallbackVH(Src, this));
-  LLVM_DEBUG(dbgs() << "set edge " << Src->getName() << " -> "
-                    << IndexInSuccessors << " successor probability to " << Prob
-                    << "\n");
-
-  unsigned &SuccIdx = MaxSuccIdx[Src];
-  SuccIdx = std::max(SuccIdx, IndexInSuccessors);
-}
-
 /// Set the edge probability for all edges at once.
 void BranchProbabilityInfo::setEdgeProbability(
     const BasicBlock *Src, const SmallVectorImpl<BranchProbability> &Probs) {
   assert(Src->getTerminator()->getNumSuccessors() == Probs.size());
+  eraseBlock(Src); // Erase stale data if any.
   if (Probs.size() == 0)
     return; // Nothing to set.
 
   uint64_t TotalNumerator = 0;
   for (unsigned SuccIdx = 0; SuccIdx < Probs.size(); ++SuccIdx) {
-    setEdgeProbability(Src, SuccIdx, Probs[SuccIdx]);
+    this->Probs[std::make_pair(Src, SuccIdx)] = Probs[SuccIdx];
+    Handles.insert(BasicBlockCallbackVH(Src, this));
+    LLVM_DEBUG(dbgs() << "set edge " << Src->getName() << " -> " << SuccIdx
+                      << " successor probability to " << Probs[SuccIdx]
+                      << "\n");
     TotalNumerator += Probs[SuccIdx].getNumerator();
   }
 
   // Because of rounding errors the total probability cannot be checked to be
   // 1.0 exactly. That is TotalNumerator == BranchProbability::getDenominator.
   // Instead, every single probability in Probs must be as accurate as possible.
   // This results in error 1/denominator at most, thus the total absolute error
   // should be within Probs.size / BranchProbability::getDenominator.
OS << "edge " << Src->getName() << " -> " << Dst->getName()
      << (isEdgeHot(Src, Dst) ? " [HOT edge]\n" : "\n");
 
   return 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.
-  auto It = MaxSuccIdx.find(BB);
-  if (It == MaxSuccIdx.end())
-    return;
-
-  for (unsigned I = 0, E = It->second; I <= E; ++I) {
+  // Instead we remove prob data for the block by iterating successors by their
+  // numbers from 0 till the last which exist. 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 public method
+  // setEdgeProbability().
+  for (unsigned I = 0; true; ++I) {
     auto MapI = Probs.find(std::make_pair(BB, I));
-    if (MapI != Probs.end())
-      Probs.erase(MapI);
+    if (MapI == Probs.end()) {
+      assert(Probs.find(std::make_pair(BB, I + 1)) == Probs.end() &&
+             "Must be no more successors");
+      return;
+    }
   }
-  MaxSuccIdx.erase(BB);
 }
 
 void BranchProbabilityInfo::calculate(const Function &F, const LoopInfo &LI,
                                       const TargetLibraryInfo *TLI,
                                       PostDominatorTree *PDT) {
   LLVM_DEBUG(dbgs() << "---- Branch Probability Info : " << F.getName()
                     << " ----\n\n");
   LastF = &F; // Store the last function we ran on for printing.