Skip to content

[DFAJumpThreading] Unify equivalent states #162447

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

Merged
merged 3 commits into from
Oct 10, 2025
Merged

[DFAJumpThreading] Unify equivalent states #162447

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f6eb2f8
Precommit testcases
XChy Oct 8, 2025
eb7f705
[DFAJumpThreading] Unify equivalent states
XChy Oct 8, 2025
f62d953
update tests
XChy Oct 8, 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
138 changes: 85 additions & 53 deletions llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -384,6 +384,22 @@ inline raw_ostream &operator<<(raw_ostream &OS, const PathType &Path) {
return OS;
}

/// Helper to get the successor corresponding to a particular case value for
/// a switch statement.
static BasicBlock *getNextCaseSuccessor(SwitchInst *Switch,
const APInt &NextState) {
BasicBlock *NextCase = nullptr;
for (auto Case : Switch->cases()) {
if (Case.getCaseValue()->getValue() == NextState) {
NextCase = Case.getCaseSuccessor();
break;
}
}
if (!NextCase)
NextCase = Switch->getDefaultDest();
return NextCase;
}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I realize you only moved this code, but it would probably make sense to compute the value => successor mapping once instead of scanning the whole switch each time.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, though it may not be the bottleneck, it's worth pre-computing.


namespace {
/// ThreadingPath is a path in the control flow of a loop that can be threaded
/// by cloning necessary basic blocks and replacing conditional branches with
Expand All @@ -396,6 +412,10 @@ struct ThreadingPath {
ExitVal = V->getValue();
IsExitValSet = true;
}
void setExitValue(const APInt &V) {
ExitVal = V;
IsExitValSet = true;
}
bool isExitValueSet() const { return IsExitValSet; }

/// Determinator is the basic block that determines the next state of the DFA.
Expand Down Expand Up @@ -578,44 +598,8 @@ struct AllSwitchPaths {
BasicBlock *getSwitchBlock() { return SwitchBlock; }

void run() {
StateDefMap StateDef = getStateDefMap();
if (StateDef.empty()) {
ORE->emit([&]() {
return OptimizationRemarkMissed(DEBUG_TYPE, "SwitchNotPredictable",
Switch)
<< "Switch instruction is not predictable.";
});
return;
}

auto *SwitchPhi = cast<PHINode>(Switch->getOperand(0));
auto *SwitchPhiDefBB = SwitchPhi->getParent();
VisitedBlocks VB;
// Get paths from the determinator BBs to SwitchPhiDefBB
std::vector<ThreadingPath> PathsToPhiDef =
getPathsFromStateDefMap(StateDef, SwitchPhi, VB, MaxNumPaths);
if (SwitchPhiDefBB == SwitchBlock || PathsToPhiDef.empty()) {
TPaths = std::move(PathsToPhiDef);
return;
}

assert(MaxNumPaths >= PathsToPhiDef.size() && !PathsToPhiDef.empty());
auto PathsLimit = MaxNumPaths / PathsToPhiDef.size();
// Find and append paths from SwitchPhiDefBB to SwitchBlock.
PathsType PathsToSwitchBB =
paths(SwitchPhiDefBB, SwitchBlock, VB, /* PathDepth = */ 1, PathsLimit);
if (PathsToSwitchBB.empty())
return;

std::vector<ThreadingPath> TempList;
for (const ThreadingPath &Path : PathsToPhiDef) {
for (const PathType &PathToSw : PathsToSwitchBB) {
ThreadingPath PathCopy(Path);
PathCopy.appendExcludingFirst(PathToSw);
TempList.push_back(PathCopy);
}
}
TPaths = std::move(TempList);
findTPaths();
unifyTPaths();
}

private:
Expand Down Expand Up @@ -807,6 +791,69 @@ struct AllSwitchPaths {
return Res;
}

// Find all threadable paths.
void findTPaths() {
StateDefMap StateDef = getStateDefMap();
if (StateDef.empty()) {
ORE->emit([&]() {
return OptimizationRemarkMissed(DEBUG_TYPE, "SwitchNotPredictable",
Switch)
<< "Switch instruction is not predictable.";
});
return;
}

auto *SwitchPhi = cast<PHINode>(Switch->getOperand(0));
auto *SwitchPhiDefBB = SwitchPhi->getParent();
VisitedBlocks VB;
// Get paths from the determinator BBs to SwitchPhiDefBB
std::vector<ThreadingPath> PathsToPhiDef =
getPathsFromStateDefMap(StateDef, SwitchPhi, VB, MaxNumPaths);
if (SwitchPhiDefBB == SwitchBlock || PathsToPhiDef.empty()) {
TPaths = std::move(PathsToPhiDef);
return;
}

assert(MaxNumPaths >= PathsToPhiDef.size() && !PathsToPhiDef.empty());
auto PathsLimit = MaxNumPaths / PathsToPhiDef.size();
// Find and append paths from SwitchPhiDefBB to SwitchBlock.
PathsType PathsToSwitchBB =
paths(SwitchPhiDefBB, SwitchBlock, VB, /* PathDepth = */ 1, PathsLimit);
if (PathsToSwitchBB.empty())
return;

std::vector<ThreadingPath> TempList;
for (const ThreadingPath &Path : PathsToPhiDef) {
for (const PathType &PathToSw : PathsToSwitchBB) {
ThreadingPath PathCopy(Path);
PathCopy.appendExcludingFirst(PathToSw);
TempList.push_back(PathCopy);
}
}
TPaths = std::move(TempList);
}

// Two states are equivalent if they have the same switch destination.
// Unify the states in different threading path if the states are equivalent.
void unifyTPaths() {
llvm::SmallDenseMap<BasicBlock *, APInt> DestToState;
for (ThreadingPath &Path : TPaths) {
APInt NextState = Path.getExitValue();
BasicBlock *Dest = getNextCaseSuccessor(Switch, NextState);
auto StateIt = DestToState.find(Dest);
if (StateIt == DestToState.end()) {
DestToState.insert({Dest, NextState});
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can avoid the duplicate lookup here by using try_emplace.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the review. I will resolve these problems.

continue;
}

if (NextState != StateIt->second) {
LLVM_DEBUG(dbgs() << "Next state in " << Path << " is equivalent to "
<< StateIt->second << "\n");
Path.setExitValue(StateIt->second);
}
}
}

unsigned NumVisited = 0;
SwitchInst *Switch;
BasicBlock *SwitchBlock;
Expand Down Expand Up @@ -1339,21 +1386,6 @@ struct TransformDFA {
return It != ClonedBBs.end() ? (*It).BB : nullptr;
}

/// Helper to get the successor corresponding to a particular case value for
/// a switch statement.
BasicBlock *getNextCaseSuccessor(SwitchInst *Switch, const APInt &NextState) {
BasicBlock *NextCase = nullptr;
for (auto Case : Switch->cases()) {
if (Case.getCaseValue()->getValue() == NextState) {
NextCase = Case.getCaseSuccessor();
break;
}
}
if (!NextCase)
NextCase = Switch->getDefaultDest();
return NextCase;
}

/// Returns true if IncomingBB is a predecessor of BB.
bool isPredecessor(BasicBlock *BB, BasicBlock *IncomingBB) {
return llvm::is_contained(predecessors(BB), IncomingBB);
Expand Down
42 changes: 11 additions & 31 deletions llvm/test/Transforms/DFAJumpThreading/dfa-unfold-select.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -227,10 +227,6 @@ define i32 @test3(i32 %num) {
; CHECK-NEXT: i32 1, label [[CASE1:%.*]]
; CHECK-NEXT: i32 2, label [[CASE2:%.*]]
; CHECK-NEXT: ]
; CHECK: for.body.jt4:
; CHECK-NEXT: [[COUNT_JT4:%.*]] = phi i32 [ [[INC_JT4:%.*]], [[FOR_INC_JT4:%.*]] ]
; CHECK-NEXT: [[STATE_JT4:%.*]] = phi i32 [ [[STATE_NEXT_JT4:%.*]], [[FOR_INC_JT4]] ]
; CHECK-NEXT: br label [[FOR_INC_JT1]]
; CHECK: for.body.jt3:
; CHECK-NEXT: [[COUNT_JT3:%.*]] = phi i32 [ [[INC_JT3:%.*]], [[FOR_INC_JT3:%.*]] ]
; CHECK-NEXT: [[STATE_JT3:%.*]] = phi i32 [ [[STATE_NEXT_JT3:%.*]], [[FOR_INC_JT3]] ]
Expand Down Expand Up @@ -261,17 +257,14 @@ define i32 @test3(i32 %num) {
; CHECK: sel.2.si.unfold.false:
; CHECK-NEXT: [[DOTSI_UNFOLD_PHI1:%.*]] = phi i32 [ 4, [[SEL_2_SI_UNFOLD_TRUE_JT3]] ]
; CHECK-NEXT: br label [[SEL_3_SI_UNFOLD_FALSE]]
; CHECK: sel.2.si.unfold.false.jt4:
; CHECK: sel.2.si.unfold.false.jt3:
; CHECK-NEXT: [[DOTSI_UNFOLD_PHI1_JT4:%.*]] = phi i32 [ 4, [[SEL_2_SI_UNFOLD_TRUE:%.*]] ]
; CHECK-NEXT: br label [[SEL_3_SI_UNFOLD_FALSE_JT4:%.*]]
; CHECK-NEXT: br label [[SEL_3_SI_UNFOLD_FALSE_JT3]]
; CHECK: sel.3.si.unfold.false:
; CHECK-NEXT: [[SEL_2_SI_UNFOLD_PHI:%.*]] = phi i32 [ poison, [[SEL_2_SI_UNFOLD_TRUE]] ], [ [[DOTSI_UNFOLD_PHI1]], [[SEL_2_SI_UNFOLD_FALSE]] ]
; CHECK-NEXT: br label [[FOR_INC]]
; CHECK: sel.3.si.unfold.false.jt4:
; CHECK-NEXT: [[SEL_2_SI_UNFOLD_PHI_JT4:%.*]] = phi i32 [ [[DOTSI_UNFOLD_PHI1_JT4]], [[SEL_2_SI_UNFOLD_FALSE_JT4]] ]
; CHECK-NEXT: br label [[FOR_INC_JT4]]
; CHECK: sel.3.si.unfold.false.jt3:
; CHECK-NEXT: [[SEL_2_SI_UNFOLD_PHI_JT3:%.*]] = phi i32 [ [[DOTSI_UNFOLD_PHI_JT3]], [[SEL_2_SI_UNFOLD_TRUE_JT3]] ]
; CHECK-NEXT: [[SEL_2_SI_UNFOLD_PHI_JT3:%.*]] = phi i32 [ [[DOTSI_UNFOLD_PHI_JT3]], [[SEL_2_SI_UNFOLD_TRUE_JT3]] ], [ [[DOTSI_UNFOLD_PHI1_JT4]], [[SEL_2_SI_UNFOLD_FALSE_JT4]] ]
; CHECK-NEXT: br label [[FOR_INC_JT3]]
; CHECK: sel.1.si.unfold.true:
; CHECK-NEXT: br i1 [[CMP_1]], label [[FOR_INC]], label [[SEL_1_SI_UNFOLD_FALSE_JT2:%.*]]
Expand All @@ -289,11 +282,6 @@ define i32 @test3(i32 %num) {
; CHECK-NEXT: [[INC]] = add nsw i32 [[COUNT5]], 1
; CHECK-NEXT: [[CMP_EXIT:%.*]] = icmp slt i32 [[INC]], [[NUM:%.*]]
; CHECK-NEXT: br i1 [[CMP_EXIT]], label [[FOR_BODY]], label [[FOR_END:%.*]]
; CHECK: for.inc.jt4:
; CHECK-NEXT: [[STATE_NEXT_JT4]] = phi i32 [ [[SEL_2_SI_UNFOLD_PHI_JT4]], [[SEL_3_SI_UNFOLD_FALSE_JT4]] ]
; CHECK-NEXT: [[INC_JT4]] = add nsw i32 undef, 1
; CHECK-NEXT: [[CMP_EXIT_JT4:%.*]] = icmp slt i32 [[INC_JT4]], [[NUM]]
; CHECK-NEXT: br i1 [[CMP_EXIT_JT4]], label [[FOR_BODY_JT4:%.*]], label [[FOR_END]]
; CHECK: for.inc.jt3:
; CHECK-NEXT: [[STATE_NEXT_JT3]] = phi i32 [ [[SEL_2_SI_UNFOLD_PHI_JT3]], [[SEL_3_SI_UNFOLD_FALSE_JT3]] ]
; CHECK-NEXT: [[INC_JT3]] = add nsw i32 [[COUNT5]], 1
Expand All @@ -305,8 +293,8 @@ define i32 @test3(i32 %num) {
; CHECK-NEXT: [[CMP_EXIT_JT2:%.*]] = icmp slt i32 [[INC_JT2]], [[NUM]]
; CHECK-NEXT: br i1 [[CMP_EXIT_JT2]], label [[FOR_BODY_JT2]], label [[FOR_END]]
; CHECK: for.inc.jt1:
; CHECK-NEXT: [[COUNT4:%.*]] = phi i32 [ [[COUNT_JT4]], [[FOR_BODY_JT4]] ], [ [[COUNT_JT3]], [[FOR_BODY_JT3]] ], [ [[COUNT5]], [[SEL_1_SI_UNFOLD_TRUE_JT1]] ], [ [[COUNT]], [[FOR_BODY]] ]
; CHECK-NEXT: [[STATE_NEXT_JT1]] = phi i32 [ 1, [[FOR_BODY]] ], [ 1, [[FOR_BODY_JT3]] ], [ 1, [[FOR_BODY_JT4]] ], [ [[DOTSI_UNFOLD_PHI2_JT1]], [[SEL_1_SI_UNFOLD_TRUE_JT1]] ]
; CHECK-NEXT: [[COUNT4:%.*]] = phi i32 [ [[COUNT_JT3]], [[FOR_BODY_JT3]] ], [ [[COUNT5]], [[SEL_1_SI_UNFOLD_TRUE_JT1]] ], [ [[COUNT]], [[FOR_BODY]] ]
; CHECK-NEXT: [[STATE_NEXT_JT1]] = phi i32 [ 1, [[FOR_BODY]] ], [ 1, [[FOR_BODY_JT3]] ], [ [[DOTSI_UNFOLD_PHI2_JT1]], [[SEL_1_SI_UNFOLD_TRUE_JT1]] ]
; CHECK-NEXT: [[INC_JT1]] = add nsw i32 [[COUNT4]], 1
; CHECK-NEXT: [[CMP_EXIT_JT1:%.*]] = icmp slt i32 [[INC_JT1]], [[NUM]]
; CHECK-NEXT: br i1 [[CMP_EXIT_JT1]], label [[FOR_BODY_JT1]], label [[FOR_END]]
Expand Down Expand Up @@ -402,36 +390,28 @@ define void @pr65222(i32 %flags, i1 %cmp, i1 %tobool.not) {
; CHECK-NEXT: br label [[IF_END_JT2:%.*]]
; CHECK: cond1.si.unfold.true:
; CHECK-NEXT: br i1 [[CMP]], label [[IF_END]], label [[COND1_SI_UNFOLD_FALSE_JT1:%.*]]
; CHECK: cond1.si.unfold.true.jt3:
; CHECK: cond1.si.unfold.true.jt2:
; CHECK-NEXT: [[DOTSI_UNFOLD_PHI2:%.*]] = phi i32 [ 3, [[THEN]] ]
; CHECK-NEXT: br i1 [[CMP]], label [[IF_END_JT3:%.*]], label [[COND1_SI_UNFOLD_FALSE:%.*]]
; CHECK-NEXT: br i1 [[CMP]], label [[IF_END_JT2]], label [[COND1_SI_UNFOLD_FALSE:%.*]]
; CHECK: cond1.si.unfold.false:
; CHECK-NEXT: [[DOTSI_UNFOLD_PHI3:%.*]] = phi i32 [ 1, [[COND1_SI_UNFOLD_TRUE]] ]
; CHECK-NEXT: br label [[IF_END]]
; CHECK: cond1.si.unfold.false.jt1:
; CHECK: cond1.si.unfold.false.jt2:
; CHECK-NEXT: [[DOTSI_UNFOLD_PHI3_JT1:%.*]] = phi i32 [ 1, [[COND1_SI_UNFOLD_TRUE1:%.*]] ]
; CHECK-NEXT: br label [[IF_END_JT1:%.*]]
; CHECK-NEXT: br label [[IF_END_JT2]]
; CHECK: if.end:
; CHECK-NEXT: [[UNFOLDED:%.*]] = phi i32 [ [[FLAGS:%.*]], [[WHILE_COND]] ], [ [[COND_SI_UNFOLD_PHI]], [[TOUNFOLD_SI_UNFOLD_FALSE1]] ], [ poison, [[COND1_SI_UNFOLD_TRUE1]] ], [ [[DOTSI_UNFOLD_PHI3]], [[COND1_SI_UNFOLD_FALSE]] ]
; CHECK-NEXT: [[OTHER:%.*]] = phi i32 [ [[FLAGS]], [[WHILE_COND]] ], [ 0, [[TOUNFOLD_SI_UNFOLD_FALSE1]] ], [ 0, [[COND1_SI_UNFOLD_TRUE1]] ], [ 0, [[COND1_SI_UNFOLD_FALSE]] ]
; CHECK-NEXT: switch i32 [[UNFOLDED]], label [[UNREACHABLE:%.*]] [
; CHECK-NEXT: i32 0, label [[SW_BB:%.*]]
; CHECK-NEXT: ]
; CHECK: if.end.jt1:
; CHECK-NEXT: [[UNFOLDED_JT1:%.*]] = phi i32 [ [[DOTSI_UNFOLD_PHI3_JT1]], [[COND1_SI_UNFOLD_FALSE_JT1]] ]
; CHECK-NEXT: [[OTHER_JT1:%.*]] = phi i32 [ 0, [[COND1_SI_UNFOLD_FALSE_JT1]] ]
; CHECK-NEXT: br label [[UNREACHABLE]]
; CHECK: if.end.jt3:
; CHECK-NEXT: [[UNFOLDED_JT3:%.*]] = phi i32 [ [[DOTSI_UNFOLD_PHI2]], [[COND1_SI_UNFOLD_TRUE]] ]
; CHECK-NEXT: [[OTHER_JT3:%.*]] = phi i32 [ 0, [[COND1_SI_UNFOLD_TRUE]] ]
; CHECK-NEXT: br label [[UNREACHABLE]]
; CHECK: if.end.jt0:
; CHECK-NEXT: [[UNFOLDED_JT0:%.*]] = phi i32 [ [[COND_SI_UNFOLD_PHI_JT0]], [[TOUNFOLD_SI_UNFOLD_FALSE_JT0]] ]
; CHECK-NEXT: [[OTHER_JT0:%.*]] = phi i32 [ 0, [[TOUNFOLD_SI_UNFOLD_FALSE_JT0]] ]
; CHECK-NEXT: br label [[SW_BB]]
; CHECK: if.end.jt2:
; CHECK-NEXT: [[UNFOLDED_JT2:%.*]] = phi i32 [ [[COND_SI_UNFOLD_PHI_JT2]], [[TOUNFOLD_SI_UNFOLD_FALSE]] ]
; CHECK-NEXT: [[OTHER_JT2:%.*]] = phi i32 [ 0, [[TOUNFOLD_SI_UNFOLD_FALSE]] ]
; CHECK-NEXT: [[UNFOLDED_JT2:%.*]] = phi i32 [ [[COND_SI_UNFOLD_PHI_JT2]], [[TOUNFOLD_SI_UNFOLD_FALSE]] ], [ [[DOTSI_UNFOLD_PHI2]], [[COND1_SI_UNFOLD_TRUE]] ], [ [[DOTSI_UNFOLD_PHI3_JT1]], [[COND1_SI_UNFOLD_FALSE_JT1]] ]
; CHECK-NEXT: [[OTHER_JT2:%.*]] = phi i32 [ 0, [[TOUNFOLD_SI_UNFOLD_FALSE]] ], [ 0, [[COND1_SI_UNFOLD_TRUE]] ], [ 0, [[COND1_SI_UNFOLD_FALSE_JT1]] ]
; CHECK-NEXT: br label [[UNREACHABLE]]
; CHECK: unreachable:
; CHECK-NEXT: unreachable
Expand Down
Loading