[BranchProbabilityInfo] Fix eraseBlock

[kazu]

Summary

This patch ensures that BranchProbabilityInfo::eraseBlock(BB) deletes
all entries in Probs associated with with BB.

Without this patch, stale entries for BB may remain in Probs after
eraseBlock(BB), leading to a situation where a newly created basic
block has an edge probability associated with it even before the pass
responsible for creating the basic block adds any edge probability to
it.

Consider the current implementation of eraseBlock(BB):

for (const_succ_iterator I = succ_begin(BB), E = succ_end(BB); I != E; ++I) {
  auto MapI = Probs.find(std::make_pair(BB, I.getSuccessorIndex()));
  if (MapI != Probs.end())
    Probs.erase(MapI);
}

Notice that it uses succ_begin(BB) and succ_end(BB), which are based
on BB->getTerminator(). This means that if the terminator changes
between calls to setEdgeProbability and eraseBlock, then we may not
examine all pairs associated with BB.

This is exactly what happens in MaybeMergeBasicBlockIntoOnlyPred,
which merges basic blocks A into B if A is the sole predecessor of B,
and B is the sole successor of A. It replaces the terminator of A
with UnreachableInst before (indirectly) calling eraseBlock(A).

The patch fixes the problem by keeping track of all edge probablities
entered with setEdgeProbability in a map from BasicBlock* to a
successor index.

[wmi]

LGTM.

llvm/lib/Analysis/BranchProbabilityInfo.cpp

1141–1144

If Src doesn't exist in MaxSuccIdx, MaxSuccIdx[Src] will return 0. Can it be simplified to one statement?
MaxSuccIdx[Src] = std::max(MaxSuccIdx[Src], IndexInSuccessors);

[MaskRay]

Thanks for figuring out the root cause!

It replaces the terminator of A with UnreachableInst before (indirectly) calling eraseBlock(A).

Question: Why is it UnreachableInst?

llvm/lib/Analysis/BranchProbabilityInfo.cpp

1141

This can be simply: int &SuccIdx = MaxSuccIdx[Src]; SuccIdx = std::max(SuccIdx, IndexInSuccessors);

Alternatively,

try_emplace + std::max

1182–1187

Add a comment that we cannot simply use successors?

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)), LastF(Arg.LastF),
+      : Probs(std::move(Arg.Probs)), MaxSuccIdx(std::move(Arg.MaxSuccIdx)),
+        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 &);
 
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>();

 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");
+
+  if (MaxSuccIdx.find(Src) == MaxSuccIdx.end())
+    MaxSuccIdx[Src] = IndexInSuccessors;
+  else
+    MaxSuccIdx[Src] = std::max(MaxSuccIdx[Src], 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());
   if (Probs.size() == 0)
     return; // Nothing to set.
BranchProbabilityInfo::printEdgeProbability(raw_ostream &OS,
   OS << "edge " << Src->getName() << " -> " << Dst->getName()
      << " probability is " << Prob
      << (isEdgeHot(Src, Dst) ? " [HOT edge]\n" : "\n");
 
   return OS;
 }
 
 void BranchProbabilityInfo::eraseBlock(const BasicBlock *BB) {
-  for (const_succ_iterator I = succ_begin(BB), E = succ_end(BB); I != E; ++I) {
-    auto MapI = Probs.find(std::make_pair(BB, I.getSuccessorIndex()));
+  auto It = MaxSuccIdx.find(BB);
+  if (It == MaxSuccIdx.end())
+    return;
+
+  for (unsigned I = 0, E = It->second; I <= E; ++I) {
+    auto MapI = Probs.find(std::make_pair(BB, I));
     if (MapI != Probs.end())
       Probs.erase(MapI);
   }
+  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.