Line data Source code
1 : //===- IndirectBrExpandPass.cpp - Expand indirectbr to switch -------------===//
2 : //
3 : // The LLVM Compiler Infrastructure
4 : //
5 : // This file is distributed under the University of Illinois Open Source
6 : // License. See LICENSE.TXT for details.
7 : //
8 : //===----------------------------------------------------------------------===//
9 : /// \file
10 : ///
11 : /// Implements an expansion pass to turn `indirectbr` instructions in the IR
12 : /// into `switch` instructions. This works by enumerating the basic blocks in
13 : /// a dense range of integers, replacing each `blockaddr` constant with the
14 : /// corresponding integer constant, and then building a switch that maps from
15 : /// the integers to the actual blocks. All of the indirectbr instructions in the
16 : /// function are redirected to this common switch.
17 : ///
18 : /// While this is generically useful if a target is unable to codegen
19 : /// `indirectbr` natively, it is primarily useful when there is some desire to
20 : /// get the builtin non-jump-table lowering of a switch even when the input
21 : /// source contained an explicit indirect branch construct.
22 : ///
23 : /// Note that it doesn't make any sense to enable this pass unless a target also
24 : /// disables jump-table lowering of switches. Doing that is likely to pessimize
25 : /// the code.
26 : ///
27 : //===----------------------------------------------------------------------===//
28 :
29 : #include "llvm/ADT/STLExtras.h"
30 : #include "llvm/ADT/Sequence.h"
31 : #include "llvm/ADT/SmallVector.h"
32 : #include "llvm/CodeGen/TargetPassConfig.h"
33 : #include "llvm/CodeGen/TargetSubtargetInfo.h"
34 : #include "llvm/IR/BasicBlock.h"
35 : #include "llvm/IR/Function.h"
36 : #include "llvm/IR/IRBuilder.h"
37 : #include "llvm/IR/InstIterator.h"
38 : #include "llvm/IR/Instruction.h"
39 : #include "llvm/IR/Instructions.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 :
58 14776 : IndirectBrExpandPass() : FunctionPass(ID) {
59 14776 : initializeIndirectBrExpandPassPass(*PassRegistry::getPassRegistry());
60 14776 : }
61 :
62 : bool runOnFunction(Function &F) override;
63 : };
64 :
65 : } // end anonymous namespace
66 :
67 : char IndirectBrExpandPass::ID = 0;
68 :
69 117926 : INITIALIZE_PASS(IndirectBrExpandPass, DEBUG_TYPE,
70 : "Expand indirectbr instructions", false, false)
71 :
72 14775 : FunctionPass *llvm::createIndirectBrExpandPass() {
73 14775 : return new IndirectBrExpandPass();
74 : }
75 :
76 313963 : bool IndirectBrExpandPass::runOnFunction(Function &F) {
77 313963 : auto &DL = F.getParent()->getDataLayout();
78 313963 : auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
79 313963 : if (!TPC)
80 : return false;
81 :
82 313963 : auto &TM = TPC->getTM<TargetMachine>();
83 313963 : auto &STI = *TM.getSubtargetImpl(F);
84 313963 : if (!STI.enableIndirectBrExpand())
85 : return false;
86 13 : 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 97 : for (BasicBlock &BB : F)
95 84 : 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 12 : if (IBr->getNumSuccessors() == 0) {
99 0 : (void)new UnreachableInst(F.getContext(), IBr);
100 0 : IBr->eraseFromParent();
101 0 : continue;
102 : }
103 :
104 12 : IndirectBrs.push_back(IBr);
105 84 : for (BasicBlock *SuccBB : IBr->successors())
106 60 : IndirectBrSuccs.insert(SuccBB);
107 : }
108 :
109 13 : 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 75 : for (BasicBlock &BB : F) {
119 : // Skip blocks that aren't successors to an indirectbr we're going to
120 : // rewrite.
121 68 : if (!IndirectBrSuccs.count(&BB))
122 : continue;
123 :
124 : auto IsBlockAddressUse = [&](const Use &U) {
125 0 : return isa<BlockAddress>(U.getUser());
126 : };
127 51 : auto BlockAddressUseIt = llvm::find_if(BB.uses(), IsBlockAddressUse);
128 51 : 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 47 : 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 47 : 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 47 : int BBIndex = BBs.size() + 1;
146 47 : BBs.push_back(&BB);
147 :
148 47 : auto *ITy = cast<IntegerType>(DL.getIntPtrType(BA->getType()));
149 47 : ConstantInt *BBIndexC = ConstantInt::get(ITy, BBIndex);
150 :
151 : // Now rewrite the blockaddress to an integer constant based on the index.
152 : // FIXME: We could potentially preserve the uses as arguments to inline asm.
153 : // This would allow some uses such as diagnostic information in crashes to
154 : // have higher quality even when this transform is enabled, but would break
155 : // users that round-trip blockaddresses through inline assembly and then
156 : // back into an indirectbr.
157 47 : BA->replaceAllUsesWith(ConstantExpr::getIntToPtr(BBIndexC, BA->getType()));
158 : }
159 :
160 7 : if (BBs.empty()) {
161 : // There are no blocks whose address is taken, so any indirectbr instruction
162 : // cannot get a valid input and we can replace all of them with unreachable.
163 2 : for (auto *IBr : IndirectBrs) {
164 2 : (void)new UnreachableInst(F.getContext(), IBr);
165 1 : IBr->eraseFromParent();
166 : }
167 : return true;
168 : }
169 :
170 : BasicBlock *SwitchBB;
171 : Value *SwitchValue;
172 :
173 : // Compute a common integer type across all the indirectbr instructions.
174 : IntegerType *CommonITy = nullptr;
175 17 : for (auto *IBr : IndirectBrs) {
176 : auto *ITy =
177 11 : cast<IntegerType>(DL.getIntPtrType(IBr->getAddress()->getType()));
178 11 : if (!CommonITy || ITy->getBitWidth() > CommonITy->getBitWidth())
179 : CommonITy = ITy;
180 : }
181 :
182 6 : auto GetSwitchValue = [DL, CommonITy](IndirectBrInst *IBr) {
183 : return CastInst::CreatePointerCast(
184 : IBr->getAddress(), CommonITy,
185 : Twine(IBr->getAddress()->getName()) + ".switch_cast", IBr);
186 6 : };
187 :
188 6 : if (IndirectBrs.size() == 1) {
189 : // If we only have one indirectbr, we can just directly replace it within
190 : // its block.
191 1 : SwitchBB = IndirectBrs[0]->getParent();
192 1 : SwitchValue = GetSwitchValue(IndirectBrs[0]);
193 1 : IndirectBrs[0]->eraseFromParent();
194 : } else {
195 : // Otherwise we need to create a new block to hold the switch across BBs,
196 : // jump to that block instead of each indirectbr, and phi together the
197 : // values for the switch.
198 10 : SwitchBB = BasicBlock::Create(F.getContext(), "switch_bb", &F);
199 10 : auto *SwitchPN = PHINode::Create(CommonITy, IndirectBrs.size(),
200 : "switch_value_phi", SwitchBB);
201 : SwitchValue = SwitchPN;
202 :
203 : // Now replace the indirectbr instructions with direct branches to the
204 : // switch block and fill out the PHI operands.
205 15 : for (auto *IBr : IndirectBrs) {
206 10 : SwitchPN->addIncoming(GetSwitchValue(IBr), IBr->getParent());
207 10 : BranchInst::Create(SwitchBB, IBr);
208 10 : IBr->eraseFromParent();
209 : }
210 : }
211 :
212 : // Now build the switch in the block. The block will have no terminator
213 : // already.
214 12 : auto *SI = SwitchInst::Create(SwitchValue, BBs[0], BBs.size(), SwitchBB);
215 :
216 : // Add a case for each block.
217 47 : for (int i : llvm::seq<int>(1, BBs.size()))
218 82 : SI->addCase(ConstantInt::get(CommonITy, i + 1), BBs[i]);
219 :
220 : return true;
221 : }
|