[JumpThreading] Prevent dangling pointer problems in BranchProbabilityInfo

[igor-laevsky]

Summary

We should update results of the BranchProbabilityInfo after removing block in JumpThreading. Otherwise we will get dangling pointer inside BranchProbabilityInfo cache. Which can lead to various problems. Note that there is similar issue for the BlockFrequencyInfo but it's a subject for separate change.

[sanjoy]

include/llvm/Analysis/BranchProbabilityInfo.h

120

New code should not use \brief, since we have autobrief enabled for our doxygen.

include/llvm/Transforms/Utils/Local.h

302

Nit: document the new parameter (and maybe the \p LVI as well)?

lib/Analysis/BranchProbabilityInfo.cpp

643

Why do you need to do a linear search here? Isn't Probs a DenseMap<Edge, BranchProbability>? IOW, why not:

for (int i = 0; i < BB->numSuccessors(); i++)
  Probs.erase({BB, i});

unittests/Analysis/BlockFrequencyInfoTest.cpp

58

Should this be called releaseBPI?

59

Nit: spelling. Also add a period after the sentence.

[igor-laevsky]

lib/Analysis/BranchProbabilityInfo.cpp

643

Problem with this approach is that contents of the BB could have been changed since it was first recorded. If conditional terminator was replaced with unconditional jump, we may leave some BB references hanging in the map.

unittests/Analysis/BlockFrequencyInfoTest.cpp

58

This is to be symmetric with buildBFI. Besides I suspect that BFI has the same issue and we will need to add "BFI->releaseMemory()" here.

[sanjoy]

lib/Analysis/BranchProbabilityInfo.cpp

644

I see, but I'm not sure that this won't be a compile time regression. Maybe @dnovillo or @congh (names that show up on git-blame) will have an idea?

unittests/Analysis/BlockFrequencyInfoTest.cpp

58

In that case, let's call it destroyBFI?

[igor-laevsky]

lib/Analysis/BranchProbabilityInfo.cpp

644

Yes, compile time might be an issue here. It would be possible to avoid linear search by changing data type of the branch probability storage from:

DenseMap<pair<BasicBlock*, unsigned>, BranchProbability> Probs;

into:

DenseMap<BasicBlock*, IndexedMap<unsigned>> Probs;

However it would be fairly large change and I don't wont to mix it with this one. Does it sound reasonable if I will submit this version first and then follow up with a data type change?

[sanjoy]

lib/Analysis/BranchProbabilityInfo.cpp

644

Right, other than that issue, this lgtm; but I'd prefer waiting till someone more familiar with this code has had time to chime in about the potential compile time problems.

[dnovillo]

I agree with Igor's analysis. Let's first fix the correctness issue. Any compile-time problems that show up, can be fixed with a data structure change later. Igor, do you have any test cases that show the problem? Without a test case, it's harder to justify this change.

Thanks. Diego.

[igor-laevsky]

In D20957#454115, @dnovillo wrote:

I agree with Igor's analysis. Let's first fix the correctness issue. Any compile-time problems that show up, can be fixed with a data structure change later. Igor, do you have any test cases that show the problem? Without a test case, it's harder to justify this change.

Thanks. Diego.

Hi,

Noticeable effect of this failure appears only if new BasicBlock was allocated at the same place as the removed one. It will depend on many things which are impossible to guarantee for a single test case. I've added test which will fail AssertingVH in case if BasicBlock was not removed. Should be enough to guarantee that no one forgets to call eraseBlock in the future.

[davidxl]

lib/Transforms/Scalar/JumpThreading.cpp

208

maybe just check nullness of BPI

236

Same here.

718

This pattern repeats many times, it is better to create a small wrapper member function to erase BB.

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

2

What does this test do? THere is no check string?

[davidxl]

lib/Transforms/Scalar/JumpThreading.cpp

1955

nit: ForgetBlockAnalysisResults --> dropBlockAnalysisResults(..)

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

2

should 2>&1 be needed ?

[davidxl]

lgtm

[dnovillo]

LGTM

Thanks for the test case.

[igor-laevsky]

This change was reverted due to the failing AssertingVH on some of the buildbots. Looked very much like a problem described in llvm-dev thread "Should analyses be able to hold AssertingVH to IR? (related to PR28400)". I bypassed this issue by using CallbackVH instead of AssertingVH and resubmitted (r275563).

include/llvm/Analysis/BranchProbabilityInfo.h

 
 #ifndef LLVM_ANALYSIS_BRANCHPROBABILITYINFO_H
 #define LLVM_ANALYSIS_BRANCHPROBABILITYINFO_H
 
 #include "llvm/ADT/DenseMap.h"
 #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"
 
 namespace llvm {
 class LoopInfo;
 class raw_ostream;
 
public:
 
   static BranchProbability getBranchProbStackProtector(bool IsLikely) {
     static const BranchProbability LikelyProb((1u << 20) - 1, 1u << 20);
     return IsLikely ? LikelyProb : LikelyProb.getCompl();
   }
 
   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
   // 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;

include/llvm/Transforms/Scalar/JumpThreading.h

public:
   bool TryToUnfoldSelect(CmpInst *CondCmp, BasicBlock *BB);
   bool TryToUnfoldSelectInCurrBB(BasicBlock *BB);
 
 private:
   BasicBlock *SplitBlockPreds(BasicBlock *BB, ArrayRef<BasicBlock *> Preds,
                               const char *Suffix);
   void UpdateBlockFreqAndEdgeWeight(BasicBlock *PredBB, BasicBlock *BB,
                                     BasicBlock *NewBB, BasicBlock *SuccBB);
+
+  void dropBlockAnalysisResults(BasicBlock *BB);
 };
 
 } // end namespace llvm
 
 #endif

include/llvm/Transforms/Utils/Local.h

 class AssumptionCache;
 class ConstantExpr;
 class DataLayout;
 class TargetLibraryInfo;
 class TargetTransformInfo;
 class DIBuilder;
 class DominatorTree;
 class LazyValueInfo;
+class BranchProbabilityInfo;
 
 template<typename T> class SmallVectorImpl;
 
 //===----------------------------------------------------------------------===//
 //  Local constant propagation.
 //
 
 /// If a terminator instruction is predicated on a constant value, convert it

 /// block. Rewrites `invoke` to `call`, etc. Updates any PHIs in unwind
 /// successor.
 ///
 /// \param BB  Block whose terminator will be replaced.  Its terminator must
 ///            have an unwind successor.
 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
 ///
 /// Metadata not listed as known via KnownIDs is removed
 void combineMetadata(Instruction *K, const Instruction *J, ArrayRef<unsigned> KnownIDs);
 
 /// Replace each use of 'From' with 'To' if that use is dominated by
 /// the given edge.  Returns the number of replacements made.

lib/Analysis/BranchProbabilityInfo.cpp

BranchProbabilityInfo::printEdgeProbability(raw_ostream &OS,
   const BranchProbability Prob = getEdgeProbability(Src, Dst);
   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 (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");
   LastF = &F; // Store the last function we ran on for printing.
   assert(PostDominatedByUnreachable.empty());
   assert(PostDominatedByColdCall.empty());
 
   // Walk the basic blocks in post-order so that we can build up state about

lib/Transforms/Scalar/JumpThreading.cpp

bool JumpThreadingPass::runImpl(Function &F, TargetLibraryInfo *TLI_,
   // Remove unreachable blocks from function as they may result in infinite
   // loop. We do threading if we found something profitable. Jump threading a
   // branch can create other opportunities. If these opportunities form a cycle
   // i.e. if any jump threading is undoing previous threading in the path, then
   // we will loop forever. We take care of this issue by not jump threading for
   // 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);
 
   bool Changed;
   do {
     Changed = false;
     for (Function::iterator I = F.begin(), E = F.end(); I != E;) {
       BasicBlock *BB = &*I;
       // Thread all of the branches we can over this block.
       while (ProcessBlock(BB))
         Changed = true;
 
       ++I;
 
       // If the block is trivially dead, zap it.  This eliminates the successor
       // edges which simplifies the CFG.
       if (pred_empty(BB) &&
           BB != &BB->getParent()->getEntryBlock()) {
         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;
       }
 
       BranchInst *BI = dyn_cast<BranchInst>(BB->getTerminator());
 
       // Can't thread an unconditional jump, but if the block is "almost
       // empty", we can replace uses of it with uses of the successor and make
for (Function::iterator I = F.begin(), E = F.end(); I != E;) {
         // reinsert afterward if needed.
         bool ErasedFromLoopHeaders = LoopHeaders.erase(BB);
         BasicBlock *Succ = BI->getSuccessor(0);
 
         // FIXME: It is always conservatively correct to drop the info
         // 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
           // successor is now the header of the loop.
           BB = Succ;
         }
 
         if (ErasedFromLoopHeaders)
           LoopHeaders.insert(BB);
bool JumpThreadingPass::ProcessBlock(BasicBlock *BB) {
   if (BasicBlock *SinglePred = BB->getSinglePredecessor()) {
     const TerminatorInst *TI = SinglePred->getTerminator();
     if (!TI->isExceptional() && TI->getNumSuccessors() == 1 &&
         SinglePred != BB && !hasAddressTakenAndUsed(BB)) {
       // If SinglePred was a loop header, BB becomes one.
       if (LoopHeaders.erase(SinglePred))
         LoopHeaders.insert(BB);
 
-      LVI->eraseBlock(SinglePred);
+      dropBlockAnalysisResults(SinglePred);
       MergeBasicBlockIntoOnlyPred(BB);
 
       return true;
     }
   }
 
   if (TryToUnfoldSelectInCurrBB(BB))
     return true;
if (HasConst) {
       SI->replaceAllUsesWith(NewPN);
       SI->eraseFromParent();
       return true;
     }
   }
   
   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

 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/STLExtras.h"
 #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"
 #include "llvm/Analysis/LazyValueInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DIBuilder.h"
void llvm::removeUnwindEdge(BasicBlock *BB) {
   UnwindDest->removePredecessor(BB);
   TI->replaceAllUsesWith(NewTI);
   TI->eraseFromParent();
 }
 
 /// 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);
 
   // If there are unreachable blocks in the CFG...
   if (Reachable.size() == F.size())
     return Changed;
 
   assert(Reachable.size() < F.size());
   NumRemoved += F.size()-Reachable.size();
 
   // Loop over all of the basic blocks that are not reachable, dropping all of
   // their internal references...
   for (Function::iterator BB = ++F.begin(), E = F.end(); BB != E; ++BB) {
     if (Reachable.count(&*BB))
       continue;
 
     for (succ_iterator SI = succ_begin(&*BB), SE = succ_end(&*BB); SI != SE;
          ++SI)
       if (Reachable.count(*SI))
         (*SI)->removePredecessor(&*BB);
     if (LVI)
       LVI->eraseBlock(&*BB);
+    if (BPI)
+      BPI->eraseBlock(&*BB);
     BB->dropAllReferences();
   }
 
   for (Function::iterator I = ++F.begin(); I != F.end();)
     if (!Reachable.count(&*I))
       I = F.getBasicBlockList().erase(I);
     else
       ++I;

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

+; 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

const char *ModuleStrig = "define i32 @f(i32 %x) {\n"
                               "  br label %bb3\n"
                               "bb3:\n"
                               "  %y2 = phi i32 [0, %bb1], [1, %bb2] \n"
                               "  ret i32 %y2\n"
                               "}\n";
     SMDiagnostic Err;
     return parseAssemblyString(ModuleStrig, Err, C);
   }
+
+  void destroyBFI() {
+    // Prevent AssertingVH from failing.
+    BPI->releaseMemory();
+  }
 };
 
 TEST_F(BlockFrequencyInfoTest, Basic) {
   auto M = makeLLVMModule();
   Function *F = M->getFunction("f");
   F->setEntryCount(100);
 
   BlockFrequencyInfo BFI = buildBFI(*F);
TEST_F(BlockFrequencyInfoTest, Basic) {
   EXPECT_EQ(BB0Freq, BB3Freq);
   EXPECT_EQ(BB0Freq, BB1Freq + BB2Freq);
   EXPECT_EQ(BB0Freq, BB3Freq);
 
   EXPECT_EQ(BFI.getBlockProfileCount(&BB0).getValue(), UINT64_C(100));
   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
 } // end namespace llvm