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