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/TargetPassConfig.h"
33 #include "llvm/CodeGen/TargetSubtargetInfo.h"
34 #include "llvm/IR/BasicBlock.h"
35 #include "llvm/IR/Constants.h"
36 #include "llvm/IR/Dominators.h"
37 #include "llvm/IR/Function.h"
38 #include "llvm/IR/Instructions.h"
39 #include "llvm/InitializePasses.h"
40 #include "llvm/Pass.h"
41 #include "llvm/Support/ErrorHandling.h"
42 #include "llvm/Target/TargetMachine.h"
43 #include <optional>
44
45 using namespace llvm;
46
47 #define DEBUG_TYPE "indirectbr-expand"
48
49 namespace {
50
51 class IndirectBrExpandPass : public FunctionPass {
52 const TargetLowering *TLI = nullptr;
53
54 public:
55 static char ID; // Pass identification, replacement for typeid
56
IndirectBrExpandPass()57 IndirectBrExpandPass() : FunctionPass(ID) {
58 initializeIndirectBrExpandPassPass(*PassRegistry::getPassRegistry());
59 }
60
getAnalysisUsage(AnalysisUsage & AU) const61 void getAnalysisUsage(AnalysisUsage &AU) const override {
62 AU.addPreserved<DominatorTreeWrapperPass>();
63 }
64
65 bool runOnFunction(Function &F) override;
66 };
67
68 } // end anonymous namespace
69
70 char IndirectBrExpandPass::ID = 0;
71
72 INITIALIZE_PASS_BEGIN(IndirectBrExpandPass, DEBUG_TYPE,
73 "Expand indirectbr instructions", false, false)
INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)74 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
75 INITIALIZE_PASS_END(IndirectBrExpandPass, DEBUG_TYPE,
76 "Expand indirectbr instructions", false, false)
77
78 FunctionPass *llvm::createIndirectBrExpandPass() {
79 return new IndirectBrExpandPass();
80 }
81
runOnFunction(Function & F)82 bool IndirectBrExpandPass::runOnFunction(Function &F) {
83 auto &DL = F.getParent()->getDataLayout();
84 auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
85 if (!TPC)
86 return false;
87
88 auto &TM = TPC->getTM<TargetMachine>();
89 auto &STI = *TM.getSubtargetImpl(F);
90 if (!STI.enableIndirectBrExpand())
91 return false;
92 TLI = STI.getTargetLowering();
93
94 std::optional<DomTreeUpdater> DTU;
95 if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
96 DTU.emplace(DTWP->getDomTree(), DomTreeUpdater::UpdateStrategy::Lazy);
97
98 SmallVector<IndirectBrInst *, 1> IndirectBrs;
99
100 // Set of all potential successors for indirectbr instructions.
101 SmallPtrSet<BasicBlock *, 4> IndirectBrSuccs;
102
103 // Build a list of indirectbrs that we want to rewrite.
104 for (BasicBlock &BB : F)
105 if (auto *IBr = dyn_cast<IndirectBrInst>(BB.getTerminator())) {
106 // Handle the degenerate case of no successors by replacing the indirectbr
107 // with unreachable as there is no successor available.
108 if (IBr->getNumSuccessors() == 0) {
109 (void)new UnreachableInst(F.getContext(), IBr);
110 IBr->eraseFromParent();
111 continue;
112 }
113
114 IndirectBrs.push_back(IBr);
115 for (BasicBlock *SuccBB : IBr->successors())
116 IndirectBrSuccs.insert(SuccBB);
117 }
118
119 if (IndirectBrs.empty())
120 return false;
121
122 // If we need to replace any indirectbrs we need to establish integer
123 // constants that will correspond to each of the basic blocks in the function
124 // whose address escapes. We do that here and rewrite all the blockaddress
125 // constants to just be those integer constants cast to a pointer type.
126 SmallVector<BasicBlock *, 4> BBs;
127
128 for (BasicBlock &BB : F) {
129 // Skip blocks that aren't successors to an indirectbr we're going to
130 // rewrite.
131 if (!IndirectBrSuccs.count(&BB))
132 continue;
133
134 auto IsBlockAddressUse = [&](const Use &U) {
135 return isa<BlockAddress>(U.getUser());
136 };
137 auto BlockAddressUseIt = llvm::find_if(BB.uses(), IsBlockAddressUse);
138 if (BlockAddressUseIt == BB.use_end())
139 continue;
140
141 assert(std::find_if(std::next(BlockAddressUseIt), BB.use_end(),
142 IsBlockAddressUse) == BB.use_end() &&
143 "There should only ever be a single blockaddress use because it is "
144 "a constant and should be uniqued.");
145
146 auto *BA = cast<BlockAddress>(BlockAddressUseIt->getUser());
147
148 // Skip if the constant was formed but ended up not being used (due to DCE
149 // or whatever).
150 if (!BA->isConstantUsed())
151 continue;
152
153 // Compute the index we want to use for this basic block. We can't use zero
154 // because null can be compared with block addresses.
155 int BBIndex = BBs.size() + 1;
156 BBs.push_back(&BB);
157
158 auto *ITy = cast<IntegerType>(DL.getIntPtrType(BA->getType()));
159 ConstantInt *BBIndexC = ConstantInt::get(ITy, BBIndex);
160
161 // Now rewrite the blockaddress to an integer constant based on the index.
162 // FIXME: This part doesn't properly recognize other uses of blockaddress
163 // expressions, for instance, where they are used to pass labels to
164 // asm-goto. This part of the pass needs a rework.
165 BA->replaceAllUsesWith(ConstantExpr::getIntToPtr(BBIndexC, BA->getType()));
166 }
167
168 if (BBs.empty()) {
169 // There are no blocks whose address is taken, so any indirectbr instruction
170 // cannot get a valid input and we can replace all of them with unreachable.
171 SmallVector<DominatorTree::UpdateType, 8> Updates;
172 if (DTU)
173 Updates.reserve(IndirectBrSuccs.size());
174 for (auto *IBr : IndirectBrs) {
175 if (DTU) {
176 for (BasicBlock *SuccBB : IBr->successors())
177 Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
178 }
179 (void)new UnreachableInst(F.getContext(), IBr);
180 IBr->eraseFromParent();
181 }
182 if (DTU) {
183 assert(Updates.size() == IndirectBrSuccs.size() &&
184 "Got unexpected update count.");
185 DTU->applyUpdates(Updates);
186 }
187 return true;
188 }
189
190 BasicBlock *SwitchBB;
191 Value *SwitchValue;
192
193 // Compute a common integer type across all the indirectbr instructions.
194 IntegerType *CommonITy = nullptr;
195 for (auto *IBr : IndirectBrs) {
196 auto *ITy =
197 cast<IntegerType>(DL.getIntPtrType(IBr->getAddress()->getType()));
198 if (!CommonITy || ITy->getBitWidth() > CommonITy->getBitWidth())
199 CommonITy = ITy;
200 }
201
202 auto GetSwitchValue = [CommonITy](IndirectBrInst *IBr) {
203 return CastInst::CreatePointerCast(
204 IBr->getAddress(), CommonITy,
205 Twine(IBr->getAddress()->getName()) + ".switch_cast", IBr);
206 };
207
208 SmallVector<DominatorTree::UpdateType, 8> Updates;
209
210 if (IndirectBrs.size() == 1) {
211 // If we only have one indirectbr, we can just directly replace it within
212 // its block.
213 IndirectBrInst *IBr = IndirectBrs[0];
214 SwitchBB = IBr->getParent();
215 SwitchValue = GetSwitchValue(IBr);
216 if (DTU) {
217 Updates.reserve(IndirectBrSuccs.size());
218 for (BasicBlock *SuccBB : IBr->successors())
219 Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
220 assert(Updates.size() == IndirectBrSuccs.size() &&
221 "Got unexpected update count.");
222 }
223 IBr->eraseFromParent();
224 } else {
225 // Otherwise we need to create a new block to hold the switch across BBs,
226 // jump to that block instead of each indirectbr, and phi together the
227 // values for the switch.
228 SwitchBB = BasicBlock::Create(F.getContext(), "switch_bb", &F);
229 auto *SwitchPN = PHINode::Create(CommonITy, IndirectBrs.size(),
230 "switch_value_phi", SwitchBB);
231 SwitchValue = SwitchPN;
232
233 // Now replace the indirectbr instructions with direct branches to the
234 // switch block and fill out the PHI operands.
235 if (DTU)
236 Updates.reserve(IndirectBrs.size() + 2 * IndirectBrSuccs.size());
237 for (auto *IBr : IndirectBrs) {
238 SwitchPN->addIncoming(GetSwitchValue(IBr), IBr->getParent());
239 BranchInst::Create(SwitchBB, IBr);
240 if (DTU) {
241 Updates.push_back({DominatorTree::Insert, IBr->getParent(), SwitchBB});
242 for (BasicBlock *SuccBB : IBr->successors())
243 Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
244 }
245 IBr->eraseFromParent();
246 }
247 }
248
249 // Now build the switch in the block. The block will have no terminator
250 // already.
251 auto *SI = SwitchInst::Create(SwitchValue, BBs[0], BBs.size(), SwitchBB);
252
253 // Add a case for each block.
254 for (int i : llvm::seq<int>(1, BBs.size()))
255 SI->addCase(ConstantInt::get(CommonITy, i + 1), BBs[i]);
256
257 if (DTU) {
258 // If there were multiple indirectbr's, they may have common successors,
259 // but in the dominator tree, we only track unique edges.
260 SmallPtrSet<BasicBlock *, 8> UniqueSuccessors;
261 Updates.reserve(Updates.size() + BBs.size());
262 for (BasicBlock *BB : BBs) {
263 if (UniqueSuccessors.insert(BB).second)
264 Updates.push_back({DominatorTree::Insert, SwitchBB, BB});
265 }
266 DTU->applyUpdates(Updates);
267 }
268
269 return true;
270 }
271