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