Skip to content

Commit 19b7d6a

Browse files
alexander-shaposhnikovalexander-shaposhnikov
committed
Add a pass to convert jump tables to switches
1 parent a6161a2 commit 19b7d6a

19 files changed

+734
-0
lines changed

‎llvm/include/llvm/Transforms/Scalar/JumpTableToSwitch.h

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
//===- JumpTableToSwitch.h - ------------------------------------*- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#ifndef LLVM_TRANSFORMS_SCALAR_JUMP_TABLE_TO_SWITCH_H
10+
#define LLVM_TRANSFORMS_SCALAR_JUMP_TABLE_TO_SWITCH_H
11+
12+
#include "llvm/IR/PassManager.h"
13+
14+
namespace llvm {
15+
16+
class Function;
17+
18+
struct JumpTableToSwitchPass : PassInfoMixin<JumpTableToSwitchPass> {
19+
/// Run the pass over the function.
20+
PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
21+
};
22+
} // end namespace llvm
23+
24+
#endif // LLVM_TRANSFORMS_SCALAR_JUMP_TABLE_TO_SWITCH_H

‎llvm/lib/Passes/PassBuilder.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,7 @@
198198
#include "llvm/Transforms/Scalar/InferAddressSpaces.h"
199199
#include "llvm/Transforms/Scalar/InferAlignment.h"
200200
#include "llvm/Transforms/Scalar/InstSimplifyPass.h"
201+
#include "llvm/Transforms/Scalar/JumpTableToSwitch.h"
201202
#include "llvm/Transforms/Scalar/JumpThreading.h"
202203
#include "llvm/Transforms/Scalar/LICM.h"
203204
#include "llvm/Transforms/Scalar/LoopAccessAnalysisPrinter.h"

‎llvm/lib/Passes/PassBuilderPipelines.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@
9191
#include "llvm/Transforms/Scalar/IndVarSimplify.h"
9292
#include "llvm/Transforms/Scalar/InferAlignment.h"
9393
#include "llvm/Transforms/Scalar/InstSimplifyPass.h"
94+
#include "llvm/Transforms/Scalar/JumpTableToSwitch.h"
9495
#include "llvm/Transforms/Scalar/JumpThreading.h"
9596
#include "llvm/Transforms/Scalar/LICM.h"
9697
#include "llvm/Transforms/Scalar/LoopDeletion.h"
@@ -558,6 +559,7 @@ PassBuilder::buildFunctionSimplificationPipeline(OptimizationLevel Level,
558559
// Optimize based on known information about branches, and cleanup afterward.
559560
FPM.addPass(JumpThreadingPass());
560561
FPM.addPass(CorrelatedValuePropagationPass());
562+
FPM.addPass(JumpTableToSwitchPass());
561563

562564
FPM.addPass(
563565
SimplifyCFGPass(SimplifyCFGOptions().convertSwitchRangeToICmp(true)));

‎llvm/lib/Passes/PassRegistry.def

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -348,6 +348,7 @@ FUNCTION_PASS("interleaved-load-combine", InterleavedLoadCombinePass(TM))
348348
FUNCTION_PASS("invalidate<all>", InvalidateAllAnalysesPass())
349349
FUNCTION_PASS("irce", IRCEPass())
350350
FUNCTION_PASS("jump-threading", JumpThreadingPass())
351+
FUNCTION_PASS("jump-table-to-switch", JumpTableToSwitchPass());
351352
FUNCTION_PASS("kcfi", KCFIPass())
352353
FUNCTION_PASS("lcssa", LCSSAPass())
353354
FUNCTION_PASS("libcalls-shrinkwrap", LibCallsShrinkWrapPass())

‎llvm/lib/Transforms/Scalar/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ add_llvm_component_library(LLVMScalarOpts
2525
InferAlignment.cpp
2626
InstSimplifyPass.cpp
2727
JumpThreading.cpp
28+
JumpTableToSwitch.cpp
2829
LICM.cpp
2930
LoopAccessAnalysisPrinter.cpp
3031
LoopBoundSplit.cpp

‎llvm/lib/Transforms/Scalar/JumpTableToSwitch.cpp

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
//===- JumpTableToSwitch.cpp ----------------------------------------------===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#include "llvm/Transforms/Scalar/JumpTableToSwitch.h"
10+
#include "llvm/ADT/DenseMap.h"
11+
#include "llvm/ADT/SmallSet.h"
12+
#include "llvm/Analysis/ConstantFolding.h"
13+
#include "llvm/Analysis/DomTreeUpdater.h"
14+
#include "llvm/Analysis/OptimizationRemarkEmitter.h"
15+
#include "llvm/Analysis/PostDominators.h"
16+
#include "llvm/Analysis/TargetLibraryInfo.h"
17+
#include "llvm/Analysis/TargetTransformInfo.h"
18+
#include "llvm/Analysis/ValueTracking.h"
19+
#include "llvm/IR/IRBuilder.h"
20+
#include "llvm/IR/IntrinsicInst.h"
21+
#include "llvm/Support/CommandLine.h"
22+
#include "llvm/Support/Debug.h"
23+
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
24+
#include "llvm/Transforms/Utils/Cloning.h"
25+
#include "llvm/Transforms/Utils/Local.h"
26+
27+
using namespace llvm;
28+
using namespace PatternMatch;
29+
30+
static cl::opt<unsigned>
31+
JumpTableSizeThreshold("jump-table-to-switch-size-threshold", cl::Hidden,
32+
cl::desc("Only split jump tables with size less or "
33+
"equal than JumpTableSizeThreshold."),
34+
cl::init(10));
35+
36+
static cl::opt<unsigned> FunctionSizeThreshold(
37+
"jump-table-to-switch-function-size-threshold", cl::Hidden,
38+
cl::desc("Only split jump tables containing functions whose sizes are less "
39+
"or equal than this threshold."),
40+
cl::init(50));
41+
42+
#define DEBUG_TYPE "jump-table-to-switch"
43+
44+
namespace {
45+
struct JumpTableTy {
46+
Value *Index;
47+
SmallVector<Function *, 10> Funcs;
48+
};
49+
} // anonymous namespace
50+
51+
static std::optional<JumpTableTy> parseJumpTable(GetElementPtrInst *GEP,
52+
PointerType *PtrTy) {
53+
Constant *Ptr = dyn_cast<Constant>(GEP->getPointerOperand());
54+
if (!Ptr)
55+
return std::nullopt;
56+
57+
GlobalVariable *GV = dyn_cast<GlobalVariable>(Ptr);
58+
if (!GV || !GV->isConstant() || !GV->hasDefinitiveInitializer())
59+
return std::nullopt;
60+
61+
Function &F = *GEP->getParent()->getParent();
62+
const DataLayout &DL = F.getParent()->getDataLayout();
63+
const unsigned BitWidth =
64+
DL.getIndexSizeInBits(GEP->getPointerAddressSpace());
65+
MapVector<Value *, APInt> VariableOffsets;
66+
APInt ConstantOffset(BitWidth, 0);
67+
if (!GEP->collectOffset(DL, BitWidth, VariableOffsets, ConstantOffset))
68+
return std::nullopt;
69+
if (VariableOffsets.size() != 1)
70+
return std::nullopt;
71+
// TODO: consider supporting more general patterns
72+
if (!ConstantOffset.isZero())
73+
return std::nullopt;
74+
const uint64_t StrideBytes = VariableOffsets.front().second.getZExtValue();
75+
const uint64_t JumpTableSizeBytes = DL.getTypeAllocSize(GV->getValueType());
76+
if (JumpTableSizeBytes % StrideBytes != 0)
77+
return std::nullopt;
78+
const uint64_t N = JumpTableSizeBytes / StrideBytes;
79+
if (N > JumpTableSizeThreshold)
80+
return std::nullopt;
81+
82+
JumpTableTy JumpTable;
83+
JumpTable.Index = VariableOffsets.front().first;
84+
JumpTable.Funcs.assign(N, nullptr);
85+
const unsigned PtrSizeBits = DL.getPointerTypeSizeInBits(PtrTy);
86+
for (uint64_t Index = 0; Index < N; ++Index) {
87+
APInt Offset(PtrSizeBits, Index * StrideBytes);
88+
Offset += ConstantOffset;
89+
Constant *C = ConstantFoldLoadFromConst(
90+
cast<Constant>(GV->getInitializer()), PtrTy, Offset, DL);
91+
auto *Func = dyn_cast_or_null<Function>(C);
92+
if (!Func || Func->isDeclaration() ||
93+
Func->getInstructionCount() > FunctionSizeThreshold)
94+
return std::nullopt;
95+
JumpTable.Funcs[Index] = Func;
96+
}
97+
return JumpTable;
98+
}
99+
100+
static BasicBlock *expandToSwitch(CallBase *CB, const JumpTableTy &JT,
101+
DomTreeUpdater &DTU,
102+
OptimizationRemarkEmitter &ORE) {
103+
const bool IsVoid = CB->getType() == Type::getVoidTy(CB->getContext());
104+
105+
SmallVector<DominatorTree::UpdateType, 8> DTUpdates;
106+
BasicBlock *BB = CB->getParent();
107+
BasicBlock *Tail = SplitBlock(BB, CB, &DTU, nullptr, nullptr,
108+
BB->getName() + Twine(".tail"));
109+
DTUpdates.push_back({DominatorTree::Delete, BB, Tail});
110+
BB->getTerminator()->eraseFromParent();
111+
112+
Function &F = *BB->getParent();
113+
BasicBlock *BBUnreachable = BasicBlock::Create(
114+
F.getContext(), "default.switch.case.unreachable", &F, Tail);
115+
IRBuilder<> BuilderUnreachable(BBUnreachable);
116+
BuilderUnreachable.CreateUnreachable();
117+
118+
IRBuilder<> Builder(BB);
119+
SwitchInst *Switch = Builder.CreateSwitch(JT.Index, BBUnreachable);
120+
DTUpdates.push_back({DominatorTree::Insert, BB, BBUnreachable});
121+
122+
IRBuilder<> BuilderTail(CB);
123+
PHINode *PHI =
124+
IsVoid ? nullptr : BuilderTail.CreatePHI(CB->getType(), JT.Funcs.size());
125+
126+
for (auto [Index, Func] : llvm::enumerate(JT.Funcs)) {
127+
BasicBlock *B = BasicBlock::Create(Func->getContext(),
128+
"call." + Twine(Index), &F, Tail);
129+
DTUpdates.push_back({DominatorTree::Insert, BB, B});
130+
DTUpdates.push_back({DominatorTree::Insert, B, Tail});
131+
132+
CallBase *Call = cast<CallBase>(CB->clone());
133+
Call->setCalledFunction(Func);
134+
Call->insertInto(B, B->end());
135+
Switch->addCase(
136+
cast<ConstantInt>(ConstantInt::get(JT.Index->getType(), Index)), B);
137+
BranchInst::Create(Tail, B);
138+
if (PHI)
139+
PHI->addIncoming(Call, B);
140+
}
141+
DTU.applyUpdates(DTUpdates);
142+
ORE.emit([&]() {
143+
return OptimizationRemark(DEBUG_TYPE, "ReplacedJumpTableWithSwitch", CB)
144+
<< "expanded indirect call into switch";
145+
});
146+
if (PHI)
147+
CB->replaceAllUsesWith(PHI);
148+
CB->eraseFromParent();
149+
return Tail;
150+
}
151+
152+
PreservedAnalyses JumpTableToSwitchPass::run(Function &F,
153+
FunctionAnalysisManager &AM) {
154+
OptimizationRemarkEmitter &ORE =
155+
AM.getResult<OptimizationRemarkEmitterAnalysis>(F);
156+
DominatorTree *DT = AM.getCachedResult<DominatorTreeAnalysis>(F);
157+
PostDominatorTree *PDT = AM.getCachedResult<PostDominatorTreeAnalysis>(F);
158+
DomTreeUpdater DTU(DT, PDT, DomTreeUpdater::UpdateStrategy::Lazy);
159+
bool Changed = false;
160+
for (BasicBlock &BB : make_early_inc_range(F)) {
161+
BasicBlock *CurrentBB = &BB;
162+
while (CurrentBB) {
163+
BasicBlock *SplittedOutTail = nullptr;
164+
for (Instruction &I : make_early_inc_range(*CurrentBB)) {
165+
auto *Call = dyn_cast<CallInst>(&I);
166+
if (!Call || Call->getCalledFunction() || Call->isMustTailCall())
167+
continue;
168+
auto *L = dyn_cast<LoadInst>(Call->getCalledOperand());
169+
// Skip atomic or volatile loads.
170+
if (!L || !L->isSimple())
171+
continue;
172+
auto *GEP = dyn_cast<GetElementPtrInst>(L->getPointerOperand());
173+
if (!GEP)
174+
continue;
175+
auto *PtrTy = dyn_cast<PointerType>(L->getType());
176+
assert(PtrTy && "call operand must be a pointer");
177+
std::optional<JumpTableTy> JumpTable = parseJumpTable(GEP, PtrTy);
178+
if (!JumpTable)
179+
continue;
180+
SplittedOutTail = expandToSwitch(Call, *JumpTable, DTU, ORE);
181+
Changed = true;
182+
break;
183+
}
184+
CurrentBB = SplittedOutTail ? SplittedOutTail : nullptr;
185+
}
186+
}
187+
188+
if (!Changed)
189+
return PreservedAnalyses::all();
190+
191+
PreservedAnalyses PA;
192+
if (DT)
193+
PA.preserve<DominatorTreeAnalysis>();
194+
if (PDT)
195+
PA.preserve<PostDominatorTreeAnalysis>();
196+
return PA;
197+
}

‎llvm/test/Other/new-pm-defaults.ll

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,7 @@
151151
; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
152152
; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
153153
; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
154+
; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
154155
; CHECK-O-NEXT: Running pass: SimplifyCFGPass
155156
; CHECK-O-NEXT: Running pass: InstCombinePass
156157
; CHECK-O23SZ-NEXT: Running pass: AggressiveInstCombinePass

‎llvm/test/Other/new-pm-thinlto-postlink-defaults.ll

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@
9090
; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
9191
; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
9292
; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
93+
; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
9394
; CHECK-O-NEXT: Running pass: SimplifyCFGPass
9495
; CHECK-O-NEXT: Running pass: InstCombinePass
9596
; CHECK-O23SZ-NEXT: Running pass: AggressiveInstCombinePass

‎llvm/test/Other/new-pm-thinlto-postlink-pgo-defaults.ll

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@
7878
; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
7979
; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
8080
; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
81+
; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
8182
; CHECK-O-NEXT: Running pass: SimplifyCFGPass
8283
; CHECK-O-NEXT: Running pass: InstCombinePass
8384
; CHECK-O23SZ-NEXT: Running pass: AggressiveInstCombinePass

‎llvm/test/Other/new-pm-thinlto-postlink-samplepgo-defaults.ll

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@
8686
; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
8787
; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
8888
; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
89+
; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
8990
; CHECK-O-NEXT: Running pass: SimplifyCFGPass
9091
; CHECK-O-NEXT: Running pass: InstCombinePass
9192
; CHECK-O23SZ-NEXT: Running pass: AggressiveInstCombinePass

‎llvm/test/Other/new-pm-thinlto-prelink-defaults.ll

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,7 @@
121121
; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
122122
; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
123123
; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
124+
; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
124125
; CHECK-O-NEXT: Running pass: SimplifyCFGPass
125126
; CHECK-O-NEXT: Running pass: InstCombinePass
126127
; CHECK-O23SZ-NEXT: Running pass: AggressiveInstCombinePass

‎llvm/test/Other/new-pm-thinlto-prelink-pgo-defaults.ll

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,7 @@
118118
; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
119119
; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
120120
; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
121+
; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
121122
; CHECK-O-NEXT: Running pass: SimplifyCFGPass
122123
; CHECK-O-NEXT: Running pass: InstCombinePass
123124
; CHECK-O-NEXT: Running analysis: BlockFrequencyAnalysis on foo

‎llvm/test/Other/new-pm-thinlto-prelink-samplepgo-defaults.ll

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@
9090
; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
9191
; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
9292
; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
93+
; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
9394
; CHECK-O-NEXT: Running pass: SimplifyCFGPass
9495
; CHECK-O-NEXT: Running pass: InstCombinePass
9596
; CHECK-O23SZ-NEXT: Running pass: AggressiveInstCombinePass

0 commit comments

Comments
 (0)