296 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8/// \file9///10/// Implements an expansion pass to turn `indirectbr` instructions in the IR11/// into `switch` instructions. This works by enumerating the basic blocks in12/// a dense range of integers, replacing each `blockaddr` constant with the13/// corresponding integer constant, and then building a switch that maps from14/// the integers to the actual blocks. All of the indirectbr instructions in the15/// function are redirected to this common switch.16///17/// While this is generically useful if a target is unable to codegen18/// `indirectbr` natively, it is primarily useful when there is some desire to19/// get the builtin non-jump-table lowering of a switch even when the input20/// source contained an explicit indirect branch construct.21///22/// Note that it doesn't make any sense to enable this pass unless a target also23/// disables jump-table lowering of switches. Doing that is likely to pessimize24/// 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 46using namespace llvm;47 48#define DEBUG_TYPE "indirectbr-expand"49 50namespace {51 52class IndirectBrExpandLegacyPass : public FunctionPass {53public:54 static char ID; // Pass identification, replacement for typeid55 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 namespace68 69static bool runImpl(Function &F, const TargetLowering *TLI,70 DomTreeUpdater *DTU);71 72PreservedAnalyses 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 90char IndirectBrExpandLegacyPass::ID = 0;91 92INITIALIZE_PASS_BEGIN(IndirectBrExpandLegacyPass, DEBUG_TYPE,93 "Expand indirectbr instructions", false, false)94INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)95INITIALIZE_PASS_END(IndirectBrExpandLegacyPass, DEBUG_TYPE,96 "Expand indirectbr instructions", false, false)97 98FunctionPass *llvm::createIndirectBrExpandPass() {99 return new IndirectBrExpandLegacyPass();100}101 102bool runImpl(Function &F, const TargetLowering *TLI, DomTreeUpdater *DTU) {103 auto &DL = F.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 indirectbr114 // with unreachable as there is no successor available.115 if (IBr->getNumSuccessors() == 0) {116 (void)new UnreachableInst(F.getContext(), IBr->getIterator());117 IBr->eraseFromParent();118 continue;119 }120 121 IndirectBrs.push_back(IBr);122 IndirectBrSuccs.insert_range(IBr->successors());123 }124 125 if (IndirectBrs.empty())126 return false;127 128 // If we need to replace any indirectbrs we need to establish integer129 // constants that will correspond to each of the basic blocks in the function130 // whose address escapes. We do that here and rewrite all the blockaddress131 // constants to just be those integer constants cast to a pointer type.132 SmallVector<BasicBlock *, 4> BBs;133 134 for (BasicBlock &BB : F) {135 // Skip blocks that aren't successors to an indirectbr we're going to136 // rewrite.137 if (!IndirectBrSuccs.count(&BB))138 continue;139 140 auto IsBlockAddressUse = [&](const Use &U) {141 return isa<BlockAddress>(U.getUser());142 };143 auto BlockAddressUseIt = llvm::find_if(BB.uses(), IsBlockAddressUse);144 if (BlockAddressUseIt == BB.use_end())145 continue;146 147 assert(std::none_of(std::next(BlockAddressUseIt), BB.use_end(),148 IsBlockAddressUse) &&149 "There should only ever be a single blockaddress use because it is "150 "a constant and should be uniqued.");151 152 auto *BA = cast<BlockAddress>(BlockAddressUseIt->getUser());153 154 // Skip if the constant was formed but ended up not being used (due to DCE155 // or whatever).156 if (!BA->isConstantUsed())157 continue;158 159 // Compute the index we want to use for this basic block. We can't use zero160 // because null can be compared with block addresses.161 int BBIndex = BBs.size() + 1;162 BBs.push_back(&BB);163 164 auto *ITy = cast<IntegerType>(DL.getIntPtrType(BA->getType()));165 ConstantInt *BBIndexC = ConstantInt::get(ITy, BBIndex);166 167 // Now rewrite the blockaddress to an integer constant based on the index.168 // FIXME: This part doesn't properly recognize other uses of blockaddress169 // expressions, for instance, where they are used to pass labels to170 // asm-goto. This part of the pass needs a rework.171 BA->replaceAllUsesWith(ConstantExpr::getIntToPtr(BBIndexC, BA->getType()));172 }173 174 if (BBs.empty()) {175 // There are no blocks whose address is taken, so any indirectbr instruction176 // cannot get a valid input and we can replace all of them with unreachable.177 SmallVector<DominatorTree::UpdateType, 8> Updates;178 if (DTU)179 Updates.reserve(IndirectBrSuccs.size());180 for (auto *IBr : IndirectBrs) {181 if (DTU) {182 for (BasicBlock *SuccBB : IBr->successors())183 Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});184 }185 (void)new UnreachableInst(F.getContext(), IBr->getIterator());186 IBr->eraseFromParent();187 }188 if (DTU) {189 assert(Updates.size() == IndirectBrSuccs.size() &&190 "Got unexpected update count.");191 DTU->applyUpdates(Updates);192 }193 return true;194 }195 196 BasicBlock *SwitchBB;197 Value *SwitchValue;198 199 // Compute a common integer type across all the indirectbr instructions.200 IntegerType *CommonITy = nullptr;201 for (auto *IBr : IndirectBrs) {202 auto *ITy =203 cast<IntegerType>(DL.getIntPtrType(IBr->getAddress()->getType()));204 if (!CommonITy || ITy->getBitWidth() > CommonITy->getBitWidth())205 CommonITy = ITy;206 }207 208 auto GetSwitchValue = [CommonITy](IndirectBrInst *IBr) {209 return CastInst::CreatePointerCast(IBr->getAddress(), CommonITy,210 Twine(IBr->getAddress()->getName()) +211 ".switch_cast",212 IBr->getIterator());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 within219 // 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 the234 // 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 the241 // 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->getIterator());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 terminator257 // 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 279bool 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