[llvm] a7b662d - [BranchProbabilityInfo] Fix eraseBlock

[Kazu Hirata via llvm-commits]

Author: Kazu Hirata
Date: 2020-10-27T16:14:25-07:00
New Revision: a7b662d0f4098371b96ce4446fb0eba79b0b649f

URL: https://github.com/llvm/llvm-project/commit/a7b662d0f4098371b96ce4446fb0eba79b0b649f
DIFF: https://github.com/llvm/llvm-project/commit/a7b662d0f4098371b96ce4446fb0eba79b0b649f.diff

LOG: [BranchProbabilityInfo] Fix eraseBlock

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.

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

Added: 
    

Modified: 
    llvm/include/llvm/Analysis/BranchProbabilityInfo.h
    llvm/lib/Analysis/BranchProbabilityInfo.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Analysis/BranchProbabilityInfo.h b/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
index 447f14501cb6..e17716eb8d94 100644
--- a/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
+++ b/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
@@ -62,7 +62,8 @@ class BranchProbabilityInfo {
   }
 
   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)) {}
 
@@ -72,6 +73,7 @@ class BranchProbabilityInfo {
   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;
@@ -273,6 +275,9 @@ class BranchProbabilityInfo {
 
   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;
 

diff  --git a/llvm/lib/Analysis/BranchProbabilityInfo.cpp b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
index eae2c4ea9da8..7d9760e3619c 100644
--- a/llvm/lib/Analysis/BranchProbabilityInfo.cpp
+++ b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
@@ -1031,6 +1031,7 @@ bool BranchProbabilityInfo::calcInvokeHeuristics(const BasicBlock *BB) {
 
 void BranchProbabilityInfo::releaseMemory() {
   Probs.clear();
+  MaxSuccIdx.clear();
   Handles.clear();
 }
 
@@ -1136,6 +1137,11 @@ void BranchProbabilityInfo::setEdgeProbability(const BasicBlock *Src,
   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.
@@ -1173,11 +1179,16 @@ BranchProbabilityInfo::printEdgeProbability(raw_ostream &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,