1//===- IndirectBrExpandPass.cpp - Expand indirectbr to switch -------------===//
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/// \file
9///
10/// Implements an expansion pass to turn `indirectbr` instructions in the IR
11/// into `switch` instructions. This works by enumerating the basic blocks in
12/// a dense range of integers, replacing each `blockaddr` constant with the
13/// corresponding integer constant, and then building a switch that maps from
14/// the integers to the actual blocks. All of the indirectbr instructions in the
15/// function are redirected to this common switch.
16///
17/// While this is generically useful if a target is unable to codegen
18/// `indirectbr` natively, it is primarily useful when there is some desire to
19/// get the builtin non-jump-table lowering of a switch even when the input
20/// source contained an explicit indirect branch construct.
21///
22/// Note that it doesn't make any sense to enable this pass unless a target also
23/// disables jump-table lowering of switches. Doing that is likely to pessimize
24/// the code.
25///
26//===----------------------------------------------------------------------===//
27
28#include "llvm/ADT/STLExtras.h"
29#include "llvm/ADT/Sequence.h"
30#include "llvm/ADT/SmallVector.h"
31#include "llvm/Analysis/DomTreeUpdater.h"
32#include "llvm/CodeGen/IndirectBrExpand.h"
33#include "llvm/CodeGen/TargetPassConfig.h"
34#include "llvm/CodeGen/TargetSubtargetInfo.h"
35#include "llvm/IR/BasicBlock.h"
36#include "llvm/IR/Constants.h"
37#include "llvm/IR/Dominators.h"
38#include "llvm/IR/Function.h"
39#include "llvm/IR/Instructions.h"
40#include "llvm/InitializePasses.h"
41#include "llvm/Pass.h"
42#include "llvm/Support/ErrorHandling.h"
43#include "llvm/Target/TargetMachine.h"
44#include <optional>
45
46using namespace llvm;
47
48#define DEBUG_TYPE "indirectbr-expand"
49
50namespace {
51
52class IndirectBrExpandLegacyPass : public FunctionPass {
53public:
54 static char ID; // Pass identification, replacement for typeid
55
56 IndirectBrExpandLegacyPass() : FunctionPass(ID) {
57 initializeIndirectBrExpandLegacyPassPass(*PassRegistry::getPassRegistry());
58 }
59
60 void getAnalysisUsage(AnalysisUsage &AU) const override {
61 AU.addPreserved<DominatorTreeWrapperPass>();
62 }
63
64 bool runOnFunction(Function &F) override;
65};
66
67} // end anonymous namespace
68
69static bool runImpl(Function &F, const TargetLowering *TLI,
70 DomTreeUpdater *DTU);
71
72PreservedAnalyses IndirectBrExpandPass::run(Function &F,
73 FunctionAnalysisManager &FAM) {
74 auto *STI = TM->getSubtargetImpl(F);
75 if (!STI->enableIndirectBrExpand())
76 return PreservedAnalyses::all();
77
78 auto *TLI = STI->getTargetLowering();
79 auto *DT = FAM.getCachedResult<DominatorTreeAnalysis>(IR&: F);
80 DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
81
82 bool Changed = runImpl(F, TLI, DTU: DT ? &DTU : nullptr);
83 if (!Changed)
84 return PreservedAnalyses::all();
85 PreservedAnalyses PA;
86 PA.preserve<DominatorTreeAnalysis>();
87 return PA;
88}
89
90char IndirectBrExpandLegacyPass::ID = 0;
91
92INITIALIZE_PASS_BEGIN(IndirectBrExpandLegacyPass, DEBUG_TYPE,
93 "Expand indirectbr instructions", false, false)
94INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
95INITIALIZE_PASS_END(IndirectBrExpandLegacyPass, DEBUG_TYPE,
96 "Expand indirectbr instructions", false, false)
97
98FunctionPass *llvm::createIndirectBrExpandPass() {
99 return new IndirectBrExpandLegacyPass();
100}
101
102bool runImpl(Function &F, const TargetLowering *TLI, DomTreeUpdater *DTU) {
103 auto &DL = F.getDataLayout();
104
105 SmallVector<IndirectBrInst *, 1> IndirectBrs;
106
107 // Set of all potential successors for indirectbr instructions.
108 SmallPtrSet<BasicBlock *, 4> IndirectBrSuccs;
109
110 // Build a list of indirectbrs that we want to rewrite.
111 for (BasicBlock &BB : F)
112 if (auto *IBr = dyn_cast<IndirectBrInst>(Val: BB.getTerminator())) {
113 // Handle the degenerate case of no successors by replacing the indirectbr
114 // with unreachable as there is no successor available.
115 if (IBr->getNumSuccessors() == 0) {
116 (void)new UnreachableInst(F.getContext(), IBr->getIterator());
117 IBr->eraseFromParent();
118 continue;
119 }
120
121 IndirectBrs.push_back(Elt: IBr);
122 for (BasicBlock *SuccBB : IBr->successors())
123 IndirectBrSuccs.insert(Ptr: SuccBB);
124 }
125
126 if (IndirectBrs.empty())
127 return false;
128
129 // If we need to replace any indirectbrs we need to establish integer
130 // constants that will correspond to each of the basic blocks in the function
131 // whose address escapes. We do that here and rewrite all the blockaddress
132 // constants to just be those integer constants cast to a pointer type.
133 SmallVector<BasicBlock *, 4> BBs;
134
135 for (BasicBlock &BB : F) {
136 // Skip blocks that aren't successors to an indirectbr we're going to
137 // rewrite.
138 if (!IndirectBrSuccs.count(Ptr: &BB))
139 continue;
140
141 auto IsBlockAddressUse = [&](const Use &U) {
142 return isa<BlockAddress>(Val: U.getUser());
143 };
144 auto BlockAddressUseIt = llvm::find_if(Range: BB.uses(), P: IsBlockAddressUse);
145 if (BlockAddressUseIt == BB.use_end())
146 continue;
147
148 assert(std::find_if(std::next(BlockAddressUseIt), BB.use_end(),
149 IsBlockAddressUse) == BB.use_end() &&
150 "There should only ever be a single blockaddress use because it is "
151 "a constant and should be uniqued.");
152
153 auto *BA = cast<BlockAddress>(Val: BlockAddressUseIt->getUser());
154
155 // Skip if the constant was formed but ended up not being used (due to DCE
156 // or whatever).
157 if (!BA->isConstantUsed())
158 continue;
159
160 // Compute the index we want to use for this basic block. We can't use zero
161 // because null can be compared with block addresses.
162 int BBIndex = BBs.size() + 1;
163 BBs.push_back(Elt: &BB);
164
165 auto *ITy = cast<IntegerType>(Val: DL.getIntPtrType(BA->getType()));
166 ConstantInt *BBIndexC = ConstantInt::get(Ty: ITy, V: BBIndex);
167
168 // Now rewrite the blockaddress to an integer constant based on the index.
169 // FIXME: This part doesn't properly recognize other uses of blockaddress
170 // expressions, for instance, where they are used to pass labels to
171 // asm-goto. This part of the pass needs a rework.
172 BA->replaceAllUsesWith(V: ConstantExpr::getIntToPtr(C: BBIndexC, Ty: BA->getType()));
173 }
174
175 if (BBs.empty()) {
176 // There are no blocks whose address is taken, so any indirectbr instruction
177 // cannot get a valid input and we can replace all of them with unreachable.
178 SmallVector<DominatorTree::UpdateType, 8> Updates;
179 if (DTU)
180 Updates.reserve(N: IndirectBrSuccs.size());
181 for (auto *IBr : IndirectBrs) {
182 if (DTU) {
183 for (BasicBlock *SuccBB : IBr->successors())
184 Updates.push_back(Elt: {DominatorTree::Delete, IBr->getParent(), SuccBB});
185 }
186 (void)new UnreachableInst(F.getContext(), IBr->getIterator());
187 IBr->eraseFromParent();
188 }
189 if (DTU) {
190 assert(Updates.size() == IndirectBrSuccs.size() &&
191 "Got unexpected update count.");
192 DTU->applyUpdates(Updates);
193 }
194 return true;
195 }
196
197 BasicBlock *SwitchBB;
198 Value *SwitchValue;
199
200 // Compute a common integer type across all the indirectbr instructions.
201 IntegerType *CommonITy = nullptr;
202 for (auto *IBr : IndirectBrs) {
203 auto *ITy =
204 cast<IntegerType>(Val: DL.getIntPtrType(IBr->getAddress()->getType()));
205 if (!CommonITy || ITy->getBitWidth() > CommonITy->getBitWidth())
206 CommonITy = ITy;
207 }
208
209 auto GetSwitchValue = [CommonITy](IndirectBrInst *IBr) {
210 return CastInst::CreatePointerCast(S: IBr->getAddress(), Ty: CommonITy,
211 Name: Twine(IBr->getAddress()->getName()) +
212 ".switch_cast",
213 InsertBefore: IBr->getIterator());
214 };
215
216 SmallVector<DominatorTree::UpdateType, 8> Updates;
217
218 if (IndirectBrs.size() == 1) {
219 // If we only have one indirectbr, we can just directly replace it within
220 // its block.
221 IndirectBrInst *IBr = IndirectBrs[0];
222 SwitchBB = IBr->getParent();
223 SwitchValue = GetSwitchValue(IBr);
224 if (DTU) {
225 Updates.reserve(N: IndirectBrSuccs.size());
226 for (BasicBlock *SuccBB : IBr->successors())
227 Updates.push_back(Elt: {DominatorTree::Delete, IBr->getParent(), SuccBB});
228 assert(Updates.size() == IndirectBrSuccs.size() &&
229 "Got unexpected update count.");
230 }
231 IBr->eraseFromParent();
232 } else {
233 // Otherwise we need to create a new block to hold the switch across BBs,
234 // jump to that block instead of each indirectbr, and phi together the
235 // values for the switch.
236 SwitchBB = BasicBlock::Create(Context&: F.getContext(), Name: "switch_bb", Parent: &F);
237 auto *SwitchPN = PHINode::Create(Ty: CommonITy, NumReservedValues: IndirectBrs.size(),
238 NameStr: "switch_value_phi", InsertBefore: SwitchBB);
239 SwitchValue = SwitchPN;
240
241 // Now replace the indirectbr instructions with direct branches to the
242 // switch block and fill out the PHI operands.
243 if (DTU)
244 Updates.reserve(N: IndirectBrs.size() + 2 * IndirectBrSuccs.size());
245 for (auto *IBr : IndirectBrs) {
246 SwitchPN->addIncoming(V: GetSwitchValue(IBr), BB: IBr->getParent());
247 BranchInst::Create(IfTrue: SwitchBB, InsertBefore: IBr->getIterator());
248 if (DTU) {
249 Updates.push_back(Elt: {DominatorTree::Insert, IBr->getParent(), SwitchBB});
250 for (BasicBlock *SuccBB : IBr->successors())
251 Updates.push_back(Elt: {DominatorTree::Delete, IBr->getParent(), SuccBB});
252 }
253 IBr->eraseFromParent();
254 }
255 }
256
257 // Now build the switch in the block. The block will have no terminator
258 // already.
259 auto *SI = SwitchInst::Create(Value: SwitchValue, Default: BBs[0], NumCases: BBs.size(), InsertBefore: SwitchBB);
260
261 // Add a case for each block.
262 for (int i : llvm::seq<int>(Begin: 1, End: BBs.size()))
263 SI->addCase(OnVal: ConstantInt::get(Ty: CommonITy, V: i + 1), Dest: BBs[i]);
264
265 if (DTU) {
266 // If there were multiple indirectbr's, they may have common successors,
267 // but in the dominator tree, we only track unique edges.
268 SmallPtrSet<BasicBlock *, 8> UniqueSuccessors;
269 Updates.reserve(N: Updates.size() + BBs.size());
270 for (BasicBlock *BB : BBs) {
271 if (UniqueSuccessors.insert(Ptr: BB).second)
272 Updates.push_back(Elt: {DominatorTree::Insert, SwitchBB, BB});
273 }
274 DTU->applyUpdates(Updates);
275 }
276
277 return true;
278}
279
280bool IndirectBrExpandLegacyPass::runOnFunction(Function &F) {
281 auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
282 if (!TPC)
283 return false;
284
285 auto &TM = TPC->getTM<TargetMachine>();
286 auto &STI = *TM.getSubtargetImpl(F);
287 if (!STI.enableIndirectBrExpand())
288 return false;
289 auto *TLI = STI.getTargetLowering();
290
291 std::optional<DomTreeUpdater> DTU;
292 if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
293 DTU.emplace(args&: DTWP->getDomTree(), args: DomTreeUpdater::UpdateStrategy::Lazy);
294
295 return runImpl(F, TLI, DTU: DTU ? &*DTU : nullptr);
296}
297