Skip to content

[DFAJumpThreading] Preserve BFI and BPI during unfolding selects #162262

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 2 commits into from
Closed

[DFAJumpThreading] Preserve BFI and BPI during unfolding selects #162262

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cc06bc0
[DFAJumpThreading] Preserve BFI and BPI during unfolding selects
XChy Oct 7, 2025
6704873
correct comment
XChy Oct 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 56 additions & 10 deletions llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/AssumptionCache.h"
#include "llvm/Analysis/BranchProbabilityInfo.h"
#include "llvm/Analysis/CodeMetrics.h"
#include "llvm/Analysis/DomTreeUpdater.h"
#include "llvm/Analysis/LoopInfo.h"
Expand Down Expand Up @@ -141,18 +142,21 @@ class SelectInstToUnfold {
explicit operator bool() const { return SI && SIUse; }
};

void unfold(DomTreeUpdater *DTU, LoopInfo *LI, SelectInstToUnfold SIToUnfold,
void unfold(DomTreeUpdater *DTU, LoopInfo *LI, BlockFrequencyInfo *BFI,
BranchProbabilityInfo *BPI, SelectInstToUnfold SIToUnfold,
std::vector<SelectInstToUnfold> *NewSIsToUnfold,
std::vector<BasicBlock *> *NewBBs);

class DFAJumpThreading {
public:
DFAJumpThreading(AssumptionCache *AC, DominatorTree *DT, LoopInfo *LI,
BlockFrequencyInfo *BFI, BranchProbabilityInfo *BPI,
TargetTransformInfo *TTI, OptimizationRemarkEmitter *ORE)
: AC(AC), DT(DT), LI(LI), TTI(TTI), ORE(ORE) {}
: AC(AC), DT(DT), LI(LI), BFI(BFI), BPI(BPI), TTI(TTI), ORE(ORE) {}

bool run(Function &F);
bool LoopInfoBroken;
bool BFIBPIBroken;

private:
void
Expand All @@ -167,7 +171,7 @@ class DFAJumpThreading {

std::vector<SelectInstToUnfold> NewSIsToUnfold;
std::vector<BasicBlock *> NewBBs;
unfold(&DTU, LI, SIToUnfold, &NewSIsToUnfold, &NewBBs);
unfold(&DTU, LI, BFI, BPI, SIToUnfold, &NewSIsToUnfold, &NewBBs);

// Put newly discovered select instructions into the work list.
llvm::append_range(Stack, NewSIsToUnfold);
Expand All @@ -177,6 +181,8 @@ class DFAJumpThreading {
AssumptionCache *AC;
DominatorTree *DT;
LoopInfo *LI;
BlockFrequencyInfo *BFI;
BranchProbabilityInfo *BPI;
TargetTransformInfo *TTI;
OptimizationRemarkEmitter *ORE;
};
Expand All @@ -192,17 +198,32 @@ namespace {
/// created basic blocks into \p NewBBs.
///
/// TODO: merge it with CodeGenPrepare::optimizeSelectInst() if possible.
void unfold(DomTreeUpdater *DTU, LoopInfo *LI, SelectInstToUnfold SIToUnfold,
void unfold(DomTreeUpdater *DTU, LoopInfo *LI, BlockFrequencyInfo *BFI,
BranchProbabilityInfo *BPI, SelectInstToUnfold SIToUnfold,
std::vector<SelectInstToUnfold> *NewSIsToUnfold,
std::vector<BasicBlock *> *NewBBs) {
SelectInst *SI = SIToUnfold.getInst();
PHINode *SIUse = SIToUnfold.getUse();
assert(SI->hasOneUse());
// The select may come indirectly, instead of from where it is defined.
BasicBlock *StartBlock = SIUse->getIncomingBlock(*SI->use_begin());
BranchInst *StartBlockTerm =
dyn_cast<BranchInst>(StartBlock->getTerminator());
assert(StartBlockTerm);
BranchInst *StartBlockTerm = cast<BranchInst>(StartBlock->getTerminator());

uint64_t TrueWeight = 1;
uint64_t FalseWeight = 1;
// Copy probabilities from 'SI' to the created conditional branch.
SmallVector<BranchProbability, 2> SIProbs;
if (extractBranchWeights(*SI, TrueWeight, FalseWeight) &&
(TrueWeight + FalseWeight) != 0) {
SIProbs.emplace_back(BranchProbability::getBranchProbability(
TrueWeight, TrueWeight + FalseWeight));
SIProbs.emplace_back(BranchProbability::getBranchProbability(
FalseWeight, TrueWeight + FalseWeight));
}
if ((TrueWeight + FalseWeight) == 0)
TrueWeight = FalseWeight = 1;
auto FalseProb = BranchProbability::getBranchProbability(
FalseWeight, TrueWeight + FalseWeight);

if (StartBlockTerm->isUnconditional()) {
BasicBlock *EndBlock = StartBlock->getUniqueSuccessor();
Expand Down Expand Up @@ -263,6 +284,13 @@ void unfold(DomTreeUpdater *DTU, LoopInfo *LI, SelectInstToUnfold SIToUnfold,
BranchInst::Create(EndBlock, NewBlock, SI->getCondition(), StartBlock);
DTU->applyUpdates({{DominatorTree::Insert, StartBlock, EndBlock},
{DominatorTree::Insert, StartBlock, NewBlock}});

// Update BPI if exists.
if (BPI && !SIProbs.empty())
BPI->setEdgeProbability(StartBlock, SIProbs);
// Update the block frequency of NewBlock.
if (BFI)
BFI->setBlockFreq(NewBlock, BFI->getBlockFreq(StartBlock) * FalseProb);
} else {
BasicBlock *EndBlock = SIUse->getParent();
BasicBlock *NewBlockT = BasicBlock::Create(
Expand Down Expand Up @@ -336,6 +364,17 @@ void unfold(DomTreeUpdater *DTU, LoopInfo *LI, SelectInstToUnfold SIToUnfold,
StartBlockTerm->setSuccessor(SuccNum, NewBlockT);
DTU->applyUpdates({{DominatorTree::Delete, StartBlock, EndBlock},
{DominatorTree::Insert, StartBlock, NewBlockT}});
// Update BPI if exists.
if (BPI && !SIProbs.empty())
BPI->setEdgeProbability(NewBlockT, SIProbs);
// Update the block frequency of both NewBlockT and NewBlockF.
if (BFI) {
assert(BPI && "BPI should be valid if BFI exists");
auto NewBlockTFreq = BFI->getBlockFreq(StartBlock) *
BPI->getEdgeProbability(StartBlock, SuccNum);
BFI->setBlockFreq(NewBlockT, NewBlockTFreq);
BFI->setBlockFreq(NewBlockF, NewBlockTFreq * FalseProb);
}
}

// Preserve loop info
Expand Down Expand Up @@ -994,6 +1033,7 @@ struct TransformDFA {
SmallPtrSet<BasicBlock *, 16> BlocksToClean;
BlocksToClean.insert_range(successors(SwitchBlock));

// TODO: Preserve BFI/BPI during creating exit paths.
{
DomTreeUpdater DTU(*DT, DomTreeUpdater::UpdateStrategy::Lazy);
for (const ThreadingPath &TPath : SwitchPaths->getThreadingPaths()) {
Expand Down Expand Up @@ -1378,7 +1418,7 @@ bool DFAJumpThreading::run(Function &F) {

SmallVector<AllSwitchPaths, 2> ThreadableLoops;
bool MadeChanges = false;
LoopInfoBroken = false;
LoopInfoBroken = BFIBPIBroken = false;

for (BasicBlock &BB : F) {
auto *SI = dyn_cast<SwitchInst>(BB.getTerminator());
Expand Down Expand Up @@ -1431,7 +1471,7 @@ bool DFAJumpThreading::run(Function &F) {
for (AllSwitchPaths SwitchPaths : ThreadableLoops) {
TransformDFA Transform(&SwitchPaths, DT, AC, TTI, ORE, EphValues);
if (Transform.run())
MadeChanges = LoopInfoBroken = true;
MadeChanges = LoopInfoBroken = BFIBPIBroken = true;
}

#ifdef EXPENSIVE_CHECKS
Expand All @@ -1450,15 +1490,21 @@ PreservedAnalyses DFAJumpThreadingPass::run(Function &F,
AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(F);
DominatorTree &DT = AM.getResult<DominatorTreeAnalysis>(F);
LoopInfo &LI = AM.getResult<LoopAnalysis>(F);
BlockFrequencyInfo *BFI = AM.getCachedResult<BlockFrequencyAnalysis>(F);
BranchProbabilityInfo *BPI = AM.getCachedResult<BranchProbabilityAnalysis>(F);
TargetTransformInfo &TTI = AM.getResult<TargetIRAnalysis>(F);
OptimizationRemarkEmitter ORE(&F);
DFAJumpThreading ThreadImpl(&AC, &DT, &LI, &TTI, &ORE);
DFAJumpThreading ThreadImpl(&AC, &DT, &LI, BFI, BPI, &TTI, &ORE);
if (!ThreadImpl.run(F))
return PreservedAnalyses::all();

PreservedAnalyses PA;
PA.preserve<DominatorTreeAnalysis>();
if (!ThreadImpl.LoopInfoBroken)
PA.preserve<LoopAnalysis>();
if (!ThreadImpl.BFIBPIBroken) {
PA.preserve<BranchProbabilityAnalysis>();
PA.preserve<BlockFrequencyAnalysis>();
}
return PA;
}
Loading