diff options
| -rw-r--r-- | llvm/include/llvm/Analysis/BranchProbabilityInfo.h | 21 | ||||
| -rw-r--r-- | llvm/lib/Analysis/BranchProbabilityInfo.cpp | 9 |
2 files changed, 30 insertions, 0 deletions
diff --git a/llvm/include/llvm/Analysis/BranchProbabilityInfo.h b/llvm/include/llvm/Analysis/BranchProbabilityInfo.h index 54514a9f1f3..6434ba962eb 100644 --- a/llvm/include/llvm/Analysis/BranchProbabilityInfo.h +++ b/llvm/include/llvm/Analysis/BranchProbabilityInfo.h @@ -15,9 +15,11 @@ #define LLVM_ANALYSIS_BRANCHPROBABILITYINFO_H #include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.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" @@ -116,10 +118,29 @@ 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; + // We need to store CallbackVH's in order to correctly handle basic block + // removal. + class BasicBlockCallbackVH final : public CallbackVH { + BranchProbabilityInfo *BPI; + void deleted() override { + assert(BPI != nullptr); + BPI->eraseBlock(cast<BasicBlock>(getValPtr())); + BPI->Handles.erase(*this); + } + + public: + BasicBlockCallbackVH(const Value *V, BranchProbabilityInfo *BPI=nullptr) + : CallbackVH(const_cast<Value *>(V)), BPI(BPI) {} + }; + DenseSet<BasicBlockCallbackVH, DenseMapInfo<Value*>> Handles; + // 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; diff --git a/llvm/lib/Analysis/BranchProbabilityInfo.cpp b/llvm/lib/Analysis/BranchProbabilityInfo.cpp index 8c12a100428..d802552d4e2 100644 --- a/llvm/lib/Analysis/BranchProbabilityInfo.cpp +++ b/llvm/lib/Analysis/BranchProbabilityInfo.cpp @@ -624,6 +624,7 @@ void BranchProbabilityInfo::setEdgeProbability(const BasicBlock *Src, unsigned IndexInSuccessors, BranchProbability Prob) { Probs[std::make_pair(Src, IndexInSuccessors)] = Prob; + Handles.insert(BasicBlockCallbackVH(Src, this)); DEBUG(dbgs() << "set edge " << Src->getName() << " -> " << IndexInSuccessors << " successor probability to " << Prob << "\n"); } @@ -641,6 +642,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"); |

