[llvm] [Analysis][NFC] Use block numbers for BranchProbabilityInfo (PR #186658)
via llvm-commits
llvm-commits at lists.llvm.org
Sun Mar 15 04:48:39 PDT 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-transforms
@llvm/pr-subscribers-llvm-analysis
Author: Alexis Engelke (aengelke)
<details>
<summary>Changes</summary>
Instead of a hash map mapping pairs of blocks and successor index to the probability, store the probabilities as flat array and start indices into this array in a per-block information vector.
Also drop value handles: no stored pointers => no stale pointers. If a block is removed, the block number is not reused unless the function is renumbered, and BPI doesn't support renumbering.
[stage2-O3 -0.12%](http://llvm-compile-time-tracker.com/compare.php?from=a7aebd809d49fd1c8205e8f0fcdc8d97777e5103&to=53669f7fa569ec2f1f1e1d179d7238c770b8a20b&stat=instructions:u)
---
Full diff: https://github.com/llvm/llvm-project/pull/186658.diff
3 Files Affected:
- (modified) llvm/include/llvm/Analysis/BranchProbabilityInfo.h (+7-49)
- (modified) llvm/lib/Analysis/BranchProbabilityInfo.cpp (+58-59)
- (modified) llvm/test/Transforms/JumpThreading/thread-prob-8.ll (-1)
``````````diff
diff --git a/llvm/include/llvm/Analysis/BranchProbabilityInfo.h b/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
index 7827808fd1c02..7dfb6a0ca200e 100644
--- a/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
+++ b/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
@@ -13,13 +13,9 @@
#ifndef LLVM_ANALYSIS_BRANCHPROBABILITYINFO_H
#define LLVM_ANALYSIS_BRANCHPROBABILITYINFO_H
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/DenseMapInfo.h"
-#include "llvm/ADT/DenseSet.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/CFG.h"
#include "llvm/IR/PassManager.h"
-#include "llvm/IR/ValueHandle.h"
#include "llvm/Pass.h"
#include "llvm/Support/BranchProbability.h"
#include "llvm/Support/Compiler.h"
@@ -121,30 +117,9 @@ class BranchProbabilityInfo {
calculate(F, LI, TLI, DT, PDT);
}
- BranchProbabilityInfo(BranchProbabilityInfo &&Arg)
- : Handles(std::move(Arg.Handles)), Probs(std::move(Arg.Probs)),
- LastF(Arg.LastF) {
- for (auto &Handle : Handles)
- Handle.setBPI(this);
- }
-
- BranchProbabilityInfo(const BranchProbabilityInfo &) = delete;
- BranchProbabilityInfo &operator=(const BranchProbabilityInfo &) = delete;
-
- BranchProbabilityInfo &operator=(BranchProbabilityInfo &&RHS) {
- releaseMemory();
- Handles = std::move(RHS.Handles);
- Probs = std::move(RHS.Probs);
- for (auto &Handle : Handles)
- Handle.setBPI(this);
- return *this;
- }
-
LLVM_ABI bool invalidate(Function &, const PreservedAnalyses &PA,
FunctionAnalysisManager::Invalidator &);
- LLVM_ABI void releaseMemory();
-
LLVM_ABI void print(raw_ostream &OS) const;
/// Get an edge's probability, relative to other out-edges of the Src.
@@ -211,33 +186,17 @@ class BranchProbabilityInfo {
LLVM_ABI void eraseBlock(const BasicBlock *BB);
private:
- // 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()));
- }
-
- public:
- void setBPI(BranchProbabilityInfo *BPI) { this->BPI = BPI; }
-
- 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.
- using Edge = std::pair<const BasicBlock *, unsigned>;
+ MutableArrayRef<BranchProbability> allocEdges(const BasicBlock *BB);
+ const BranchProbability *getEdges(const BasicBlock *BB) const;
- DenseMap<Edge, BranchProbability> Probs;
+ // Storage for branch probabilities.
+ SmallVector<BranchProbability> Probs;
+ // Map from block number to first edge.
+ SmallVector<unsigned> EdgeStarts;
/// Track the last function we run over for printing.
const Function *LastF = nullptr;
+ unsigned BlockNumberEpoch;
};
/// Analysis pass which computes \c BranchProbabilityInfo.
@@ -282,7 +241,6 @@ class LLVM_ABI BranchProbabilityInfoWrapperPass : public FunctionPass {
void getAnalysisUsage(AnalysisUsage &AU) const override;
bool runOnFunction(Function &F) override;
- void releaseMemory() override;
void print(raw_ostream &OS, const Module *M = nullptr) const override;
};
diff --git a/llvm/lib/Analysis/BranchProbabilityInfo.cpp b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
index 75ab3ba9d54fa..53c149cd135b6 100644
--- a/llvm/lib/Analysis/BranchProbabilityInfo.cpp
+++ b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
@@ -1284,9 +1284,32 @@ void BPIConstruction::calculate(const Function &F, const LoopInfo &LoopI,
} // end anonymous namespace
-void BranchProbabilityInfo::releaseMemory() {
- Probs.clear();
- Handles.clear();
+MutableArrayRef<BranchProbability>
+BranchProbabilityInfo::allocEdges(const BasicBlock *BB) {
+ assert(BB->getParent() == LastF);
+ assert(BlockNumberEpoch == LastF->getBlockNumberEpoch());
+ unsigned NumSuccs = succ_size(BB);
+ if (NumSuccs == 0) {
+ eraseBlock(BB);
+ return {};
+ }
+ if (EdgeStarts.size() <= BB->getNumber())
+ EdgeStarts.resize(LastF->getMaxBlockNumber(), 0);
+ unsigned EdgeStart = Probs.size();
+ EdgeStarts[BB->getNumber()] = EdgeStart + 1; // 0 = no edges.
+ Probs.append(NumSuccs, {});
+ return MutableArrayRef(&Probs[EdgeStart], NumSuccs);
+}
+
+const BranchProbability *
+BranchProbabilityInfo::getEdges(const BasicBlock *BB) const {
+ assert(BB->getParent() == LastF);
+ assert(BlockNumberEpoch == LastF->getBlockNumberEpoch());
+ if (EdgeStarts.size() <= BB->getNumber())
+ return nullptr;
+ if (unsigned EdgeStart = EdgeStarts[BB->getNumber()])
+ return &Probs[EdgeStart - 1]; // 0 = no edges.
+ return nullptr;
}
bool BranchProbabilityInfo::invalidate(Function &, const PreservedAnalyses &PA,
@@ -1323,15 +1346,8 @@ isEdgeHot(const BasicBlock *Src, const BasicBlock *Dst) const {
BranchProbability
BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
unsigned IndexInSuccessors) const {
- auto I = Probs.find(std::make_pair(Src, IndexInSuccessors));
- assert((Probs.end() == Probs.find(std::make_pair(Src, 0))) ==
- (Probs.end() == I) &&
- "Probability for I-th successor must always be defined along with the "
- "probability for the first successor");
-
- if (I != Probs.end())
- return I->second;
-
+ if (const BranchProbability *P = getEdges(Src))
+ return P[IndexInSuccessors];
return {1, static_cast<uint32_t>(succ_size(Src))};
}
@@ -1346,13 +1362,14 @@ BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
BranchProbability
BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
const BasicBlock *Dst) const {
- if (!Probs.count(std::make_pair(Src, 0)))
+ const BranchProbability *P = getEdges(Src);
+ if (!P)
return BranchProbability(llvm::count(successors(Src), Dst), succ_size(Src));
auto Prob = BranchProbability::getZero();
for (auto It : enumerate(successors(Src)))
if (It.value() == Dst)
- Prob += Probs.find(std::make_pair(Src, It.index()))->second;
+ Prob += P[It.index()];
return Prob;
}
@@ -1361,14 +1378,15 @@ BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
void BranchProbabilityInfo::setEdgeProbability(
const BasicBlock *Src, const SmallVectorImpl<BranchProbability> &Probs) {
assert(Src->getTerminator()->getNumSuccessors() == Probs.size());
- eraseBlock(Src); // Erase stale data if any.
- if (Probs.size() == 0)
- return; // Nothing to set.
+ if (Probs.empty()) {
+ eraseBlock(Src);
+ return;
+ }
- Handles.insert(BasicBlockCallbackVH(Src, this));
+ MutableArrayRef<BranchProbability> P = allocEdges(Src);
uint64_t TotalNumerator = 0;
for (unsigned SuccIdx = 0; SuccIdx < Probs.size(); ++SuccIdx) {
- this->Probs[std::make_pair(Src, SuccIdx)] = Probs[SuccIdx];
+ P[SuccIdx] = Probs[SuccIdx];
LLVM_DEBUG(dbgs() << "set edge " << Src->getName() << " -> " << SuccIdx
<< " successor probability to " << Probs[SuccIdx]
<< "\n");
@@ -1387,31 +1405,25 @@ void BranchProbabilityInfo::setEdgeProbability(
void BranchProbabilityInfo::copyEdgeProbabilities(BasicBlock *Src,
BasicBlock *Dst) {
- eraseBlock(Dst); // Erase stale data if any.
- unsigned NumSuccessors = Src->getTerminator()->getNumSuccessors();
- assert(NumSuccessors == Dst->getTerminator()->getNumSuccessors());
- if (NumSuccessors == 0)
- return; // Nothing to set.
- if (!this->Probs.contains(std::make_pair(Src, 0)))
- return; // No probability is set for edges from Src. Keep the same for Dst.
-
- Handles.insert(BasicBlockCallbackVH(Dst, this));
- for (unsigned SuccIdx = 0; SuccIdx < NumSuccessors; ++SuccIdx) {
- auto Prob = this->Probs[std::make_pair(Src, SuccIdx)];
- this->Probs[std::make_pair(Dst, SuccIdx)] = Prob;
- LLVM_DEBUG(dbgs() << "set edge " << Dst->getName() << " -> " << SuccIdx
- << " successor probability to " << Prob << "\n");
+ assert(succ_size(Src) == succ_size(Dst));
+ // allocEdges can reallocate and must be called first.
+ MutableArrayRef<BranchProbability> DstP = allocEdges(Dst);
+ const BranchProbability *SrcP = getEdges(Src);
+ if (!SrcP) {
+ eraseBlock(Dst);
+ return;
+ }
+ for (unsigned i = 0; i != DstP.size(); ++i) {
+ DstP[i] = SrcP[i];
+ LLVM_DEBUG(dbgs() << "set edge " << Dst->getName() << " -> " << i
+ << " successor probability to " << SrcP[i] << "\n");
}
}
void BranchProbabilityInfo::swapSuccEdgesProbabilities(const BasicBlock *Src) {
assert(Src->getTerminator()->getNumSuccessors() == 2);
- auto It0 = Probs.find(std::make_pair(Src, 0));
- if (It0 == Probs.end())
- return; // No probability is set for edges from Src
- auto It1 = Probs.find(std::make_pair(Src, 1));
- assert(It1 != Probs.end());
- std::swap(It0->second, It1->second);
+ if (BranchProbability *P = const_cast<BranchProbability *>(getEdges(Src)))
+ std::swap(P[0], P[1]);
}
raw_ostream &
@@ -1431,24 +1443,10 @@ BranchProbabilityInfo::printEdgeProbability(raw_ostream &OS,
void BranchProbabilityInfo::eraseBlock(const BasicBlock *BB) {
LLVM_DEBUG(dbgs() << "eraseBlock " << BB->getName() << "\n");
-
- // Note that we cannot use successors of BB because the terminator of BB may
- // have changed when eraseBlock is called as a BasicBlockCallbackVH callback.
- // Instead we remove prob data for the block by iterating successors by their
- // indices from 0 till the last which exists. There could not be prob data for
- // a pair (BB, N) if there is no data for (BB, N-1) because the data is always
- // set for all successors from 0 to M at once by the method
- // setEdgeProbability().
- Handles.erase(BasicBlockCallbackVH(BB, this));
- for (unsigned I = 0;; ++I) {
- auto MapI = Probs.find(std::make_pair(BB, I));
- if (MapI == Probs.end()) {
- assert(Probs.count(std::make_pair(BB, I + 1)) == 0 &&
- "Must be no more successors");
- return;
- }
- Probs.erase(MapI);
- }
+ assert(BB->getParent() == LastF);
+ assert(BlockNumberEpoch == LastF->getBlockNumberEpoch());
+ if (EdgeStarts.size() > BB->getNumber())
+ EdgeStarts[BB->getNumber()] = 0;
}
void BranchProbabilityInfo::calculate(const Function &F, const LoopInfo &LoopI,
@@ -1458,6 +1456,9 @@ void BranchProbabilityInfo::calculate(const Function &F, const LoopInfo &LoopI,
LLVM_DEBUG(dbgs() << "---- Branch Probability Info : " << F.getName()
<< " ----\n\n");
LastF = &F; // Store the last function we ran on for printing.
+ BlockNumberEpoch = F.getBlockNumberEpoch();
+ Probs.clear();
+ EdgeStarts.clear();
BPIConstruction(*this).calculate(F, LoopI, TLI, DT, PDT);
if (PrintBranchProb && (PrintBranchProbFuncName.empty() ||
@@ -1490,8 +1491,6 @@ bool BranchProbabilityInfoWrapperPass::runOnFunction(Function &F) {
return false;
}
-void BranchProbabilityInfoWrapperPass::releaseMemory() { BPI.releaseMemory(); }
-
void BranchProbabilityInfoWrapperPass::print(raw_ostream &OS,
const Module *) const {
BPI.print(OS);
diff --git a/llvm/test/Transforms/JumpThreading/thread-prob-8.ll b/llvm/test/Transforms/JumpThreading/thread-prob-8.ll
index b63c789515966..d23121ff8a928 100644
--- a/llvm/test/Transforms/JumpThreading/thread-prob-8.ll
+++ b/llvm/test/Transforms/JumpThreading/thread-prob-8.ll
@@ -8,7 +8,6 @@
; prob[L0->2] + prob[L0->3]
; CHECK: Computing probabilities for entry
-; CHECK: eraseBlock L0
; CHECK-NOT: set edge L0 -> 0 successor probability to 0x12492492 / 0x80000000 = 14.29%
; CHECK-NOT: set edge L0 -> 1 successor probability to 0x24924925 / 0x80000000 = 28.57%
; CHECK-NOT: set edge L0 -> 2 successor probability to 0x24924925 / 0x80000000 = 28.57%
``````````
</details>
https://github.com/llvm/llvm-project/pull/186658
More information about the llvm-commits
mailing list