lib/Transforms/IPO/BlockExtractor.cpp

   1 //===- BlockExtractor.cpp - Extracts blocks into their own functions ------===//
   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 //
   9 // This pass extracts the specified basic blocks from the module into their
  10 // own functions.
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "llvm/ADT/STLExtras.h"
  15 #include "llvm/ADT/Statistic.h"
  16 #include "llvm/IR/Instructions.h"
  17 #include "llvm/IR/Module.h"
  18 #include "llvm/Pass.h"
  19 #include "llvm/Support/CommandLine.h"
  20 #include "llvm/Support/Debug.h"
  21 #include "llvm/Support/MemoryBuffer.h"
  22 #include "llvm/Transforms/IPO.h"
  23 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
  24 #include "llvm/Transforms/Utils/CodeExtractor.h"
  25 using namespace llvm;
  26
  27 #define DEBUG_TYPE "block-extractor"
  28
  29 STATISTIC(NumExtracted, "Number of basic blocks extracted");
  30
  31 static cl::opt<std::string> BlockExtractorFile(
  32     "extract-blocks-file", cl::value_desc("filename"),
  33     cl::desc("A file containing list of basic blocks to extract"), cl::Hidden);
  34
  35 cl::opt<bool> BlockExtractorEraseFuncs("extract-blocks-erase-funcs",
  36                                        cl::desc("Erase the existing functions"),
  37                                        cl::Hidden);
  38
  39 namespace {
  40 class BlockExtractor : public ModulePass {
  41   SmallVector<BasicBlock *, 16> Blocks;
  42   bool EraseFunctions;
  43   SmallVector<std::pair<std::string, std::string>, 32> BlocksByName;
  44
  45 public:
  46   static char ID;
  47   BlockExtractor(const SmallVectorImpl<BasicBlock *> &BlocksToExtract,
  48                  bool EraseFunctions)
  49       : ModulePass(ID), Blocks(BlocksToExtract.begin(), BlocksToExtract.end()),
  50         EraseFunctions(EraseFunctions) {
  51     if (!BlockExtractorFile.empty())
  52       loadFile();
  53   }
  54   BlockExtractor() : BlockExtractor(SmallVector<BasicBlock *, 0>(), false) {}
  55   bool runOnModule(Module &M) override;
  56
  57 private:
  58   void loadFile();
  59   void splitLandingPadPreds(Function &F);
  60 };
  61 } // end anonymous namespace
  62
  63 char BlockExtractor::ID = 0;
  64 INITIALIZE_PASS(BlockExtractor, "extract-blocks",
  65                 "Extract basic blocks from module", false, false)
  66
  67 ModulePass *llvm::createBlockExtractorPass() { return new BlockExtractor(); }
  68 ModulePass *llvm::createBlockExtractorPass(
  69     const SmallVectorImpl<BasicBlock *> &BlocksToExtract, bool EraseFunctions) {
  70   return new BlockExtractor(BlocksToExtract, EraseFunctions);
  71 }
  72
  73 /// Gets all of the blocks specified in the input file.
  74 void BlockExtractor::loadFile() {
  75   auto ErrOrBuf = MemoryBuffer::getFile(BlockExtractorFile);
  76   if (ErrOrBuf.getError())
  77     report_fatal_error("BlockExtractor couldn't load the file.");
  78   // Read the file.
  79   auto &Buf = *ErrOrBuf;
  80   SmallVector<StringRef, 16> Lines;
  81   Buf->getBuffer().split(Lines, '\n', /*MaxSplit=*/-1,
  82                          /*KeepEmpty=*/false);
  83   for (const auto &Line : Lines) {
  84     auto FBPair = Line.split(' ');
  85     BlocksByName.push_back({FBPair.first, FBPair.second});
  86   }
  87 }
  88
  89 /// Extracts the landing pads to make sure all of them have only one
  90 /// predecessor.
  91 void BlockExtractor::splitLandingPadPreds(Function &F) {
  92   for (BasicBlock &BB : F) {
  93     for (Instruction &I : BB) {
  94       if (!isa<InvokeInst>(&I))
  95         continue;
  96       InvokeInst *II = cast<InvokeInst>(&I);
  97       BasicBlock *Parent = II->getParent();
  98       BasicBlock *LPad = II->getUnwindDest();
  99
 100       // Look through the landing pad's predecessors. If one of them ends in an
 101       // 'invoke', then we want to split the landing pad.
 102       bool Split = false;
 103       for (auto PredBB : predecessors(LPad)) {
 104         if (PredBB->isLandingPad() && PredBB != Parent &&
 105             isa<InvokeInst>(Parent->getTerminator())) {
 106           Split = true;
 107           break;
 108         }
 109       }
 110
 111       if (!Split)
 112         continue;
 113
 114       SmallVector<BasicBlock *, 2> NewBBs;
 115       SplitLandingPadPredecessors(LPad, Parent, ".1", ".2", NewBBs);
 116     }
 117   }
 118 }
 119
 120 bool BlockExtractor::runOnModule(Module &M) {
 121
 122   bool Changed = false;
 123
 124   // Get all the functions.
 125   SmallVector<Function *, 4> Functions;
 126   for (Function &F : M) {
 127     splitLandingPadPreds(F);
 128     Functions.push_back(&F);
 129   }
 130
 131   // Get all the blocks specified in the input file.
 132   for (const auto &BInfo : BlocksByName) {
 133     Function *F = M.getFunction(BInfo.first);
 134     if (!F)
 135       report_fatal_error("Invalid function name specified in the input file");
 136     auto Res = llvm::find_if(*F, [&](const BasicBlock &BB) {
 137       return BB.getName().equals(BInfo.second);
 138     });
 139     if (Res == F->end())
 140       report_fatal_error("Invalid block name specified in the input file");
 141     Blocks.push_back(&*Res);
 142   }
 143
 144   // Extract basic blocks.
 145   for (BasicBlock *BB : Blocks) {
 146     // Check if the module contains BB.
 147     if (BB->getParent()->getParent() != &M)
 148       report_fatal_error("Invalid basic block");
 149     LLVM_DEBUG(dbgs() << "BlockExtractor: Extracting "
 150                       << BB->getParent()->getName() << ":" << BB->getName()
 151                       << "\n");
 152     SmallVector<BasicBlock *, 2> BlocksToExtractVec;
 153     BlocksToExtractVec.push_back(BB);
 154     if (const InvokeInst *II = dyn_cast<InvokeInst>(BB->getTerminator()))
 155       BlocksToExtractVec.push_back(II->getUnwindDest());
 156     CodeExtractor(BlocksToExtractVec).extractCodeRegion();
 157     ++NumExtracted;
 158     Changed = true;
 159   }
 160
 161   // Erase the functions.
 162   if (EraseFunctions || BlockExtractorEraseFuncs) {
 163     for (Function *F : Functions) {
 164       LLVM_DEBUG(dbgs() << "BlockExtractor: Trying to delete " << F->getName()
 165                         << "\n");
 166       F->deleteBody();
 167     }
 168     // Set linkage as ExternalLinkage to avoid erasing unreachable functions.
 169     for (Function &F : M)
 170       F.setLinkage(GlobalValue::ExternalLinkage);
 171     Changed = true;
 172   }
 173
 174   return Changed;
 175 }