[JumpThreading] Prevent dangling pointer problems in BranchProbabilityInfo

We should update results of the BranchProbabilityInfo after removing block in JumpThreading. Otherwise 
we will get dangling pointer inside BranchProbabilityInfo cache.

Differential Revision: http://reviews.llvm.org/D20957

llvm-svn: 272891

include/llvm/Analysis/BranchProbabilityInfo.h

@@ -18,6 +18,7 @@
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/IR/ValueHandle.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/BranchProbability.h"
@@ -116,13 +117,16 @@ public:
 
   void calculate(const Function &F, const LoopInfo &LI);
 
+  /// Forget analysis results for the given basic block.
+  void eraseBlock(const BasicBlock *BB);
+
 private:
   void operator=(const BranchProbabilityInfo &) = delete;
   BranchProbabilityInfo(const BranchProbabilityInfo &) = delete;
 
   // 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;
+  typedef std::pair<AssertingVH<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

include/llvm/Transforms/Scalar/JumpThreading.h

@@ -134,6 +134,8 @@ private:
                               const char *Suffix);
   void UpdateBlockFreqAndEdgeWeight(BasicBlock *PredBB, BasicBlock *BB,
                                     BasicBlock *NewBB, BasicBlock *SuccBB);
+
+  void dropBlockAnalysisResults(BasicBlock *BB);
 };
 
 } // end namespace llvm

include/llvm/Transforms/Utils/Local.h

@@ -44,6 +44,7 @@ class TargetTransformInfo;
 class DIBuilder;
 class DominatorTree;
 class LazyValueInfo;
+class BranchProbabilityInfo;
 
 template<typename T> class SmallVectorImpl;
 
@@ -300,7 +301,11 @@ void removeUnwindEdge(BasicBlock *BB);
 /// Remove all blocks that can not be reached from the function's entry.
 ///
 /// Returns true if any basic block was removed.
-bool removeUnreachableBlocks(Function &F, LazyValueInfo *LVI = nullptr);
+/// \param F Target function
+/// \param LVI Will update this analysis if non null
+/// \param BPI Will update this analysis if non null
+bool removeUnreachableBlocks(Function &F, LazyValueInfo *LVI = nullptr,
+                             BranchProbabilityInfo *BPI = nullptr);
 
 /// Combine the metadata of two instructions so that K can replace J
 ///

lib/Analysis/BranchProbabilityInfo.cpp

@@ -639,6 +639,14 @@ BranchProbabilityInfo::printEdgeProbability(raw_ostream &OS,
   return OS;
 }
 
+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);
+  }
+}
+
 void BranchProbabilityInfo::calculate(const Function &F, const LoopInfo &LI) {
   DEBUG(dbgs() << "---- Branch Probability Info : " << F.getName()
                << " ----\n\n");

lib/Transforms/Scalar/JumpThreading.cpp

@@ -182,7 +182,7 @@ bool JumpThreadingPass::runImpl(Function &F, TargetLibraryInfo *TLI_,
   // back edges. This works for normal cases but not for unreachable blocks as
   // they may have cycle with no back edge.
   bool EverChanged = false;
-  EverChanged |= removeUnreachableBlocks(F, LVI);
+  EverChanged |= removeUnreachableBlocks(F, LVI, BPI.get());
 
   FindLoopHeaders(F);
 
@@ -204,7 +204,7 @@ bool JumpThreadingPass::runImpl(Function &F, TargetLibraryInfo *TLI_,
         DEBUG(dbgs() << "  JT: Deleting dead block '" << BB->getName()
               << "' with terminator: " << *BB->getTerminator() << '\n');
         LoopHeaders.erase(BB);
-        LVI->eraseBlock(BB);
+        dropBlockAnalysisResults(BB);
         DeleteDeadBlock(BB);
         Changed = true;
         continue;
@@ -232,7 +232,7 @@ bool JumpThreadingPass::runImpl(Function &F, TargetLibraryInfo *TLI_,
         // for a block even if it doesn't get erased.  This isn't totally
         // awesome, but it allows us to use AssertingVH to prevent nasty
         // dangling pointer issues within LazyValueInfo.
-        LVI->eraseBlock(BB);
+        dropBlockAnalysisResults(BB);
         if (TryToSimplifyUncondBranchFromEmptyBlock(BB)) {
           Changed = true;
           // If we deleted BB and BB was the header of a loop, then the
@@ -715,7 +715,7 @@ bool JumpThreadingPass::ProcessBlock(BasicBlock *BB) {
       if (LoopHeaders.erase(SinglePred))
         LoopHeaders.insert(BB);
 
-      LVI->eraseBlock(SinglePred);
+      dropBlockAnalysisResults(SinglePred);
       MergeBasicBlockIntoOnlyPred(BB);
 
       return true;
@@ -1949,3 +1949,11 @@ bool JumpThreadingPass::TryToUnfoldSelectInCurrBB(BasicBlock *BB) {
   
   return false;
 }
+
+/// dropBlockAnalysisResults - Inform relevant analyzes that BB is going to
+/// be removed. This is important in order to prevent dangling pointer problems.
+void JumpThreadingPass::dropBlockAnalysisResults(BasicBlock *BB) {
+  LVI->eraseBlock(BB);
+  if (BPI)
+    BPI->eraseBlock(BB);
+}

lib/Transforms/Utils/Local.cpp

@@ -20,6 +20,7 @@
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/EHPersonalities.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
@@ -1486,7 +1487,8 @@ void llvm::removeUnwindEdge(BasicBlock *BB) {
 /// removeUnreachableBlocksFromFn - Remove blocks that are not reachable, even
 /// if they are in a dead cycle.  Return true if a change was made, false
 /// otherwise.
-bool llvm::removeUnreachableBlocks(Function &F, LazyValueInfo *LVI) {
+bool llvm::removeUnreachableBlocks(Function &F, LazyValueInfo *LVI, 
+                                   BranchProbabilityInfo *BPI) {
   SmallPtrSet<BasicBlock*, 16> Reachable;
   bool Changed = markAliveBlocks(F, Reachable);
 
@@ -1509,6 +1511,8 @@ bool llvm::removeUnreachableBlocks(Function &F, LazyValueInfo *LVI) {
         (*SI)->removePredecessor(&*BB);
     if (LVI)
       LVI->eraseBlock(&*BB);
+    if (BPI)
+      BPI->eraseBlock(&*BB);
     BB->dropAllReferences();
   }
 

test/Transforms/JumpThreading/dangling-pointers-in-bpi.ll

@@ -0,0 +1,19 @@
+; RUN: opt -S -jump-threading %s 2>&1 | FileCheck %s
+
+; Test that after removing basic block we had also removed it from BPI. If we
+; didn't there will be dangling pointer inside BPI. It can lead to a
+; miscalculation in the edge weight and possibly other problems.
+; Checks that we are not failing assertion inside AssettingVH
+
+; CHECK-NOT: An asserting value handle still pointed to this value
+
+define void @foo(i32 %n, i1 %cond) !prof !0 {
+; Record this block into BPI cache
+single-pred:
+  br i1 %cond, label %entry, label %entry, !prof !{!"branch_weights", i32 1, i32 1}
+
+entry:
+  ret void
+}
+
+!0 = !{!"function_entry_count", i64 1}

unittests/Analysis/BlockFrequencyInfoTest.cpp

@@ -54,6 +54,11 @@ protected:
     SMDiagnostic Err;
     return parseAssemblyString(ModuleStrig, Err, C);
   }
+
+  void destroyBFI() {
+    // Prevent AssertingVH from failing.
+    BPI->releaseMemory();
+  }
 };
 
 TEST_F(BlockFrequencyInfoTest, Basic) {
@@ -80,6 +85,8 @@ TEST_F(BlockFrequencyInfoTest, Basic) {
   EXPECT_EQ(BFI.getBlockProfileCount(BB3).getValue(), UINT64_C(100));
   EXPECT_EQ(BFI.getBlockProfileCount(BB1).getValue(), 100 * BB1Freq / BB0Freq);
   EXPECT_EQ(BFI.getBlockProfileCount(BB2).getValue(), 100 * BB2Freq / BB0Freq);
+
+  destroyBFI();
 }
 
 } // end anonymous namespace