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