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 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 72 FunctionPass *llvm::createIndirectBrExpandPass() { 73 return new IndirectBrExpandPass(); 74 } 75 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: 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 BA->replaceAllUsesWith(ConstantExpr::getIntToPtr(BBIndexC, BA->getType())); 158 } 159 160 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 for (auto *IBr : IndirectBrs) { 164 (void)new UnreachableInst(F.getContext(), IBr); 165 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 for (auto *IBr : IndirectBrs) { 176 auto *ITy = 177 cast<IntegerType>(DL.getIntPtrType(IBr->getAddress()->getType())); 178 if (!CommonITy || ITy->getBitWidth() > CommonITy->getBitWidth()) 179 CommonITy = ITy; 180 } 181 182 auto GetSwitchValue = [DL, CommonITy](IndirectBrInst *IBr) { 183 return CastInst::CreatePointerCast( 184 IBr->getAddress(), CommonITy, 185 Twine(IBr->getAddress()->getName()) + ".switch_cast", IBr); 186 }; 187 188 if (IndirectBrs.size() == 1) { 189 // If we only have one indirectbr, we can just directly replace it within 190 // its block. 191 SwitchBB = IndirectBrs[0]->getParent(); 192 SwitchValue = GetSwitchValue(IndirectBrs[0]); 193 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 SwitchBB = BasicBlock::Create(F.getContext(), "switch_bb", &F); 199 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 for (auto *IBr : IndirectBrs) { 206 SwitchPN->addIncoming(GetSwitchValue(IBr), IBr->getParent()); 207 BranchInst::Create(SwitchBB, IBr); 208 IBr->eraseFromParent(); 209 } 210 } 211 212 // Now build the switch in the block. The block will have no terminator 213 // already. 214 auto *SI = SwitchInst::Create(SwitchValue, BBs[0], BBs.size(), SwitchBB); 215 216 // Add a case for each block. 217 for (int i : llvm::seq<int>(1, BBs.size())) 218 SI->addCase(ConstantInt::get(CommonITy, i + 1), BBs[i]); 219 220 return true; 221 } 222