llvm/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/Transforms/IPO/BlockExtractor.h"
  15 #include "llvm/ADT/STLExtras.h"
  16 #include "llvm/ADT/Statistic.h"
  17 #include "llvm/IR/Instructions.h"
  18 #include "llvm/IR/Module.h"
  19 #include "llvm/IR/PassManager.h"
  20 #include "llvm/InitializePasses.h"
  21 #include "llvm/Pass.h"
  22 #include "llvm/Support/CommandLine.h"
  23 #include "llvm/Support/Debug.h"
  24 #include "llvm/Support/MemoryBuffer.h"
  25 #include "llvm/Transforms/IPO.h"
  26 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
  27 #include "llvm/Transforms/Utils/CodeExtractor.h"
  28
  29 using namespace llvm;
  30
  31 #define DEBUG_TYPE "block-extractor"
  32
  33 STATISTIC(NumExtracted, "Number of basic blocks extracted");
  34
  35 static cl::opt<std::string> BlockExtractorFile(
  36     "extract-blocks-file", cl::value_desc("filename"),
  37     cl::desc("A file containing list of basic blocks to extract"), cl::Hidden);
  38
  39 static cl::opt<bool>
  40     BlockExtractorEraseFuncs("extract-blocks-erase-funcs",
  41                              cl::desc("Erase the existing functions"),
  42                              cl::Hidden);
  43 namespace {
  44 class BlockExtractor {
  45 public:
  46   BlockExtractor(bool EraseFunctions) : EraseFunctions(EraseFunctions) {}
  47   bool runOnModule(Module &M);
  48   void init(const SmallVectorImpl<SmallVector<BasicBlock *, 16>>
  49                 &GroupsOfBlocksToExtract) {
  50     for (const SmallVectorImpl<BasicBlock *> &GroupOfBlocks :
  51          GroupsOfBlocksToExtract) {
  52       SmallVector<BasicBlock *, 16> NewGroup;
  53       NewGroup.append(GroupOfBlocks.begin(), GroupOfBlocks.end());
  54       GroupsOfBlocks.emplace_back(NewGroup);
  55     }
  56     if (!BlockExtractorFile.empty())
  57       loadFile();
  58   }
  59
  60 private:
  61   SmallVector<SmallVector<BasicBlock *, 16>, 4> GroupsOfBlocks;
  62   bool EraseFunctions;
  63   /// Map a function name to groups of blocks.
  64   SmallVector<std::pair<std::string, SmallVector<std::string, 4>>, 4>
  65       BlocksByName;
  66
  67   void loadFile();
  68   void splitLandingPadPreds(Function &F);
  69 };
  70
  71 class BlockExtractorLegacyPass : public ModulePass {
  72   BlockExtractor BE;
  73   bool runOnModule(Module &M) override;
  74
  75 public:
  76   static char ID;
  77   BlockExtractorLegacyPass(const SmallVectorImpl<BasicBlock *> &BlocksToExtract,
  78                            bool EraseFunctions)
  79       : ModulePass(ID), BE(EraseFunctions) {
  80     // We want one group per element of the input list.
  81     SmallVector<SmallVector<BasicBlock *, 16>, 4> MassagedGroupsOfBlocks;
  82     for (BasicBlock *BB : BlocksToExtract) {
  83       SmallVector<BasicBlock *, 16> NewGroup;
  84       NewGroup.push_back(BB);
  85       MassagedGroupsOfBlocks.push_back(NewGroup);
  86     }
  87     BE.init(MassagedGroupsOfBlocks);
  88   }
  89
  90   BlockExtractorLegacyPass(const SmallVectorImpl<SmallVector<BasicBlock *, 16>>
  91                                &GroupsOfBlocksToExtract,
  92                            bool EraseFunctions)
  93       : ModulePass(ID), BE(EraseFunctions) {
  94     BE.init(GroupsOfBlocksToExtract);
  95   }
  96
  97   BlockExtractorLegacyPass()
  98       : BlockExtractorLegacyPass(SmallVector<BasicBlock *, 0>(), false) {}
  99 };
 100
 101 } // end anonymous namespace
 102
 103 char BlockExtractorLegacyPass::ID = 0;
 104 INITIALIZE_PASS(BlockExtractorLegacyPass, "extract-blocks",
 105                 "Extract basic blocks from module", false, false)
 106
 107 ModulePass *llvm::createBlockExtractorPass() {
 108   return new BlockExtractorLegacyPass();
 109 }
 110 ModulePass *llvm::createBlockExtractorPass(
 111     const SmallVectorImpl<BasicBlock *> &BlocksToExtract, bool EraseFunctions) {
 112   return new BlockExtractorLegacyPass(BlocksToExtract, EraseFunctions);
 113 }
 114 ModulePass *llvm::createBlockExtractorPass(
 115     const SmallVectorImpl<SmallVector<BasicBlock *, 16>>
 116         &GroupsOfBlocksToExtract,
 117     bool EraseFunctions) {
 118   return new BlockExtractorLegacyPass(GroupsOfBlocksToExtract, EraseFunctions);
 119 }
 120
 121 /// Gets all of the blocks specified in the input file.
 122 void BlockExtractor::loadFile() {
 123   auto ErrOrBuf = MemoryBuffer::getFile(BlockExtractorFile);
 124   if (ErrOrBuf.getError())
 125     report_fatal_error("BlockExtractor couldn't load the file.");
 126   // Read the file.
 127   auto &Buf = *ErrOrBuf;
 128   SmallVector<StringRef, 16> Lines;
 129   Buf->getBuffer().split(Lines, '\n', /*MaxSplit=*/-1,
 130                          /*KeepEmpty=*/false);
 131   for (const auto &Line : Lines) {
 132     SmallVector<StringRef, 4> LineSplit;
 133     Line.split(LineSplit, ' ', /*MaxSplit=*/-1,
 134                /*KeepEmpty=*/false);
 135     if (LineSplit.empty())
 136       continue;
 137     if (LineSplit.size()!=2)
 138       report_fatal_error("Invalid line format, expecting lines like: 'funcname bb1[;bb2..]'");
 139     SmallVector<StringRef, 4> BBNames;
 140     LineSplit[1].split(BBNames, ';', /*MaxSplit=*/-1,
 141                        /*KeepEmpty=*/false);
 142     if (BBNames.empty())
 143       report_fatal_error("Missing bbs name");
 144     BlocksByName.push_back(
 145         {std::string(LineSplit[0]), {BBNames.begin(), BBNames.end()}});
 146   }
 147 }
 148
 149 /// Extracts the landing pads to make sure all of them have only one
 150 /// predecessor.
 151 void BlockExtractor::splitLandingPadPreds(Function &F) {
 152   for (BasicBlock &BB : F) {
 153     for (Instruction &I : BB) {
 154       if (!isa<InvokeInst>(&I))
 155         continue;
 156       InvokeInst *II = cast<InvokeInst>(&I);
 157       BasicBlock *Parent = II->getParent();
 158       BasicBlock *LPad = II->getUnwindDest();
 159
 160       // Look through the landing pad's predecessors. If one of them ends in an
 161       // 'invoke', then we want to split the landing pad.
 162       bool Split = false;
 163       for (auto PredBB : predecessors(LPad)) {
 164         if (PredBB->isLandingPad() && PredBB != Parent &&
 165             isa<InvokeInst>(Parent->getTerminator())) {
 166           Split = true;
 167           break;
 168         }
 169       }
 170
 171       if (!Split)
 172         continue;
 173
 174       SmallVector<BasicBlock *, 2> NewBBs;
 175       SplitLandingPadPredecessors(LPad, Parent, ".1", ".2", NewBBs);
 176     }
 177   }
 178 }
 179
 180 bool BlockExtractor::runOnModule(Module &M) {
 181
 182   bool Changed = false;
 183
 184   // Get all the functions.
 185   SmallVector<Function *, 4> Functions;
 186   for (Function &F : M) {
 187     splitLandingPadPreds(F);
 188     Functions.push_back(&F);
 189   }
 190
 191   // Get all the blocks specified in the input file.
 192   unsigned NextGroupIdx = GroupsOfBlocks.size();
 193   GroupsOfBlocks.resize(NextGroupIdx + BlocksByName.size());
 194   for (const auto &BInfo : BlocksByName) {
 195     Function *F = M.getFunction(BInfo.first);
 196     if (!F)
 197       report_fatal_error("Invalid function name specified in the input file");
 198     for (const auto &BBInfo : BInfo.second) {
 199       auto Res = llvm::find_if(*F, [&](const BasicBlock &BB) {
 200         return BB.getName().equals(BBInfo);
 201       });
 202       if (Res == F->end())
 203         report_fatal_error("Invalid block name specified in the input file");
 204       GroupsOfBlocks[NextGroupIdx].push_back(&*Res);
 205     }
 206     ++NextGroupIdx;
 207   }
 208
 209   // Extract each group of basic blocks.
 210   for (auto &BBs : GroupsOfBlocks) {
 211     SmallVector<BasicBlock *, 32> BlocksToExtractVec;
 212     for (BasicBlock *BB : BBs) {
 213       // Check if the module contains BB.
 214       if (BB->getParent()->getParent() != &M)
 215         report_fatal_error("Invalid basic block");
 216       LLVM_DEBUG(dbgs() << "BlockExtractor: Extracting "
 217                         << BB->getParent()->getName() << ":" << BB->getName()
 218                         << "\n");
 219       BlocksToExtractVec.push_back(BB);
 220       if (const InvokeInst *II = dyn_cast<InvokeInst>(BB->getTerminator()))
 221         BlocksToExtractVec.push_back(II->getUnwindDest());
 222       ++NumExtracted;
 223       Changed = true;
 224     }
 225     CodeExtractorAnalysisCache CEAC(*BBs[0]->getParent());
 226     Function *F = CodeExtractor(BlocksToExtractVec).extractCodeRegion(CEAC);
 227     if (F)
 228       LLVM_DEBUG(dbgs() << "Extracted group '" << (*BBs.begin())->getName()
 229                         << "' in: " << F->getName() << '\n');
 230     else
 231       LLVM_DEBUG(dbgs() << "Failed to extract for group '"
 232                         << (*BBs.begin())->getName() << "'\n");
 233   }
 234
 235   // Erase the functions.
 236   if (EraseFunctions || BlockExtractorEraseFuncs) {
 237     for (Function *F : Functions) {
 238       LLVM_DEBUG(dbgs() << "BlockExtractor: Trying to delete " << F->getName()
 239                         << "\n");
 240       F->deleteBody();
 241     }
 242     // Set linkage as ExternalLinkage to avoid erasing unreachable functions.
 243     for (Function &F : M)
 244       F.setLinkage(GlobalValue::ExternalLinkage);
 245     Changed = true;
 246   }
 247
 248   return Changed;
 249 }
 250
 251 bool BlockExtractorLegacyPass::runOnModule(Module &M) {
 252   return BE.runOnModule(M);
 253 }
 254
 255 PreservedAnalyses BlockExtractorPass::run(Module &M,
 256                                           ModuleAnalysisManager &AM) {
 257   BlockExtractor BE(false);
 258   BE.init(SmallVector<SmallVector<BasicBlock *, 16>, 0>());
 259   return BE.runOnModule(M) ? PreservedAnalyses::none()
 260                            : PreservedAnalyses::all();
 261 }