llvm/lib/CodeGen/IndirectBrExpandPass.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-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 }