contrib/llvm/lib/CodeGen/IndirectBrExpandPass.cpp

   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 }