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 
46 using namespace llvm;
47 
48 #define DEBUG_TYPE "indirectbr-expand"
49 
50 namespace {
51 
52 class IndirectBrExpandLegacyPass : public FunctionPass {
53 public:
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 
69 static bool runImpl(Function &F, const TargetLowering *TLI,
70                     DomTreeUpdater *DTU);
71 
72 PreservedAnalyses 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>(F);
80   DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
81 
82   bool Changed = runImpl(F, TLI, DT ? &DTU : nullptr);
83   if (!Changed)
84     return PreservedAnalyses::all();
85   PreservedAnalyses PA;
86   PA.preserve<DominatorTreeAnalysis>();
87   return PA;
88 }
89 
90 char IndirectBrExpandLegacyPass::ID = 0;
91 
92 INITIALIZE_PASS_BEGIN(IndirectBrExpandLegacyPass, DEBUG_TYPE,
93                       "Expand indirectbr instructions", false, false)
94 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
95 INITIALIZE_PASS_END(IndirectBrExpandLegacyPass, DEBUG_TYPE,
96                     "Expand indirectbr instructions", false, false)
97 
98 FunctionPass *llvm::createIndirectBrExpandPass() {
99   return new IndirectBrExpandLegacyPass();
100 }
101 
102 bool runImpl(Function &F, const TargetLowering *TLI, DomTreeUpdater *DTU) {
103   auto &DL = F.getParent()->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>(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);
117         IBr->eraseFromParent();
118         continue;
119       }
120 
121       IndirectBrs.push_back(IBr);
122       for (BasicBlock *SuccBB : IBr->successors())
123         IndirectBrSuccs.insert(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(&BB))
139       continue;
140 
141     auto IsBlockAddressUse = [&](const Use &U) {
142       return isa<BlockAddress>(U.getUser());
143     };
144     auto BlockAddressUseIt = llvm::find_if(BB.uses(), 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>(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(&BB);
164 
165     auto *ITy = cast<IntegerType>(DL.getIntPtrType(BA->getType()));
166     ConstantInt *BBIndexC = ConstantInt::get(ITy, 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(ConstantExpr::getIntToPtr(BBIndexC, 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(IndirectBrSuccs.size());
181     for (auto *IBr : IndirectBrs) {
182       if (DTU) {
183         for (BasicBlock *SuccBB : IBr->successors())
184           Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
185       }
186       (void)new UnreachableInst(F.getContext(), IBr);
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>(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(
211         IBr->getAddress(), CommonITy,
212         Twine(IBr->getAddress()->getName()) + ".switch_cast", IBr);
213   };
214 
215   SmallVector<DominatorTree::UpdateType, 8> Updates;
216 
217   if (IndirectBrs.size() == 1) {
218     // If we only have one indirectbr, we can just directly replace it within
219     // its block.
220     IndirectBrInst *IBr = IndirectBrs[0];
221     SwitchBB = IBr->getParent();
222     SwitchValue = GetSwitchValue(IBr);
223     if (DTU) {
224       Updates.reserve(IndirectBrSuccs.size());
225       for (BasicBlock *SuccBB : IBr->successors())
226         Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
227       assert(Updates.size() == IndirectBrSuccs.size() &&
228              "Got unexpected update count.");
229     }
230     IBr->eraseFromParent();
231   } else {
232     // Otherwise we need to create a new block to hold the switch across BBs,
233     // jump to that block instead of each indirectbr, and phi together the
234     // values for the switch.
235     SwitchBB = BasicBlock::Create(F.getContext(), "switch_bb", &F);
236     auto *SwitchPN = PHINode::Create(CommonITy, IndirectBrs.size(),
237                                      "switch_value_phi", SwitchBB);
238     SwitchValue = SwitchPN;
239 
240     // Now replace the indirectbr instructions with direct branches to the
241     // switch block and fill out the PHI operands.
242     if (DTU)
243       Updates.reserve(IndirectBrs.size() + 2 * IndirectBrSuccs.size());
244     for (auto *IBr : IndirectBrs) {
245       SwitchPN->addIncoming(GetSwitchValue(IBr), IBr->getParent());
246       BranchInst::Create(SwitchBB, IBr);
247       if (DTU) {
248         Updates.push_back({DominatorTree::Insert, IBr->getParent(), SwitchBB});
249         for (BasicBlock *SuccBB : IBr->successors())
250           Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
251       }
252       IBr->eraseFromParent();
253     }
254   }
255 
256   // Now build the switch in the block. The block will have no terminator
257   // already.
258   auto *SI = SwitchInst::Create(SwitchValue, BBs[0], BBs.size(), SwitchBB);
259 
260   // Add a case for each block.
261   for (int i : llvm::seq<int>(1, BBs.size()))
262     SI->addCase(ConstantInt::get(CommonITy, i + 1), BBs[i]);
263 
264   if (DTU) {
265     // If there were multiple indirectbr's, they may have common successors,
266     // but in the dominator tree, we only track unique edges.
267     SmallPtrSet<BasicBlock *, 8> UniqueSuccessors;
268     Updates.reserve(Updates.size() + BBs.size());
269     for (BasicBlock *BB : BBs) {
270       if (UniqueSuccessors.insert(BB).second)
271         Updates.push_back({DominatorTree::Insert, SwitchBB, BB});
272     }
273     DTU->applyUpdates(Updates);
274   }
275 
276   return true;
277 }
278 
279 bool IndirectBrExpandLegacyPass::runOnFunction(Function &F) {
280   auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
281   if (!TPC)
282     return false;
283 
284   auto &TM = TPC->getTM<TargetMachine>();
285   auto &STI = *TM.getSubtargetImpl(F);
286   if (!STI.enableIndirectBrExpand())
287     return false;
288   auto *TLI = STI.getTargetLowering();
289 
290   std::optional<DomTreeUpdater> DTU;
291   if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
292     DTU.emplace(DTWP->getDomTree(), DomTreeUpdater::UpdateStrategy::Lazy);
293 
294   return runImpl(F, TLI, DTU ? &*DTU : nullptr);
295 }
296