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:
BlockExtractor(bool EraseFunctions)46 BlockExtractor(bool EraseFunctions) : EraseFunctions(EraseFunctions) {}
47 bool runOnModule(Module &M);
48 void
init(const std::vector<std::vector<BasicBlock * >> & GroupsOfBlocksToExtract)49 init(const std::vector<std::vector<BasicBlock *>> &GroupsOfBlocksToExtract) {
50 GroupsOfBlocks = GroupsOfBlocksToExtract;
51 if (!BlockExtractorFile.empty())
52 loadFile();
53 }
54
55 private:
56 std::vector<std::vector<BasicBlock *>> GroupsOfBlocks;
57 bool EraseFunctions;
58 /// Map a function name to groups of blocks.
59 SmallVector<std::pair<std::string, SmallVector<std::string, 4>>, 4>
60 BlocksByName;
61
62 void loadFile();
63 void splitLandingPadPreds(Function &F);
64 };
65
66 } // end anonymous namespace
67
68 /// Gets all of the blocks specified in the input file.
loadFile()69 void BlockExtractor::loadFile() {
70 auto ErrOrBuf = MemoryBuffer::getFile(BlockExtractorFile);
71 if (ErrOrBuf.getError())
72 report_fatal_error("BlockExtractor couldn't load the file.");
73 // Read the file.
74 auto &Buf = *ErrOrBuf;
75 SmallVector<StringRef, 16> Lines;
76 Buf->getBuffer().split(Lines, '\n', /*MaxSplit=*/-1,
77 /*KeepEmpty=*/false);
78 for (const auto &Line : Lines) {
79 SmallVector<StringRef, 4> LineSplit;
80 Line.split(LineSplit, ' ', /*MaxSplit=*/-1,
81 /*KeepEmpty=*/false);
82 if (LineSplit.empty())
83 continue;
84 if (LineSplit.size()!=2)
85 report_fatal_error("Invalid line format, expecting lines like: 'funcname bb1[;bb2..]'",
86 /*GenCrashDiag=*/false);
87 SmallVector<StringRef, 4> BBNames;
88 LineSplit[1].split(BBNames, ';', /*MaxSplit=*/-1,
89 /*KeepEmpty=*/false);
90 if (BBNames.empty())
91 report_fatal_error("Missing bbs name");
92 BlocksByName.push_back(
93 {std::string(LineSplit[0]), {BBNames.begin(), BBNames.end()}});
94 }
95 }
96
97 /// Extracts the landing pads to make sure all of them have only one
98 /// predecessor.
splitLandingPadPreds(Function & F)99 void BlockExtractor::splitLandingPadPreds(Function &F) {
100 for (BasicBlock &BB : F) {
101 for (Instruction &I : BB) {
102 if (!isa<InvokeInst>(&I))
103 continue;
104 InvokeInst *II = cast<InvokeInst>(&I);
105 BasicBlock *Parent = II->getParent();
106 BasicBlock *LPad = II->getUnwindDest();
107
108 // Look through the landing pad's predecessors. If one of them ends in an
109 // 'invoke', then we want to split the landing pad.
110 bool Split = false;
111 for (auto *PredBB : predecessors(LPad)) {
112 if (PredBB->isLandingPad() && PredBB != Parent &&
113 isa<InvokeInst>(Parent->getTerminator())) {
114 Split = true;
115 break;
116 }
117 }
118
119 if (!Split)
120 continue;
121
122 SmallVector<BasicBlock *, 2> NewBBs;
123 SplitLandingPadPredecessors(LPad, Parent, ".1", ".2", NewBBs);
124 }
125 }
126 }
127
runOnModule(Module & M)128 bool BlockExtractor::runOnModule(Module &M) {
129 bool Changed = false;
130
131 // Get all the functions.
132 SmallVector<Function *, 4> Functions;
133 for (Function &F : M) {
134 splitLandingPadPreds(F);
135 Functions.push_back(&F);
136 }
137
138 // Get all the blocks specified in the input file.
139 unsigned NextGroupIdx = GroupsOfBlocks.size();
140 GroupsOfBlocks.resize(NextGroupIdx + BlocksByName.size());
141 for (const auto &BInfo : BlocksByName) {
142 Function *F = M.getFunction(BInfo.first);
143 if (!F)
144 report_fatal_error("Invalid function name specified in the input file",
145 /*GenCrashDiag=*/false);
146 for (const auto &BBInfo : BInfo.second) {
147 auto Res = llvm::find_if(*F, [&](const BasicBlock &BB) {
148 return BB.getName().equals(BBInfo);
149 });
150 if (Res == F->end())
151 report_fatal_error("Invalid block name specified in the input file",
152 /*GenCrashDiag=*/false);
153 GroupsOfBlocks[NextGroupIdx].push_back(&*Res);
154 }
155 ++NextGroupIdx;
156 }
157
158 // Extract each group of basic blocks.
159 for (auto &BBs : GroupsOfBlocks) {
160 SmallVector<BasicBlock *, 32> BlocksToExtractVec;
161 for (BasicBlock *BB : BBs) {
162 // Check if the module contains BB.
163 if (BB->getParent()->getParent() != &M)
164 report_fatal_error("Invalid basic block", /*GenCrashDiag=*/false);
165 LLVM_DEBUG(dbgs() << "BlockExtractor: Extracting "
166 << BB->getParent()->getName() << ":" << BB->getName()
167 << "\n");
168 BlocksToExtractVec.push_back(BB);
169 if (const InvokeInst *II = dyn_cast<InvokeInst>(BB->getTerminator()))
170 BlocksToExtractVec.push_back(II->getUnwindDest());
171 ++NumExtracted;
172 Changed = true;
173 }
174 CodeExtractorAnalysisCache CEAC(*BBs[0]->getParent());
175 Function *F = CodeExtractor(BlocksToExtractVec).extractCodeRegion(CEAC);
176 if (F)
177 LLVM_DEBUG(dbgs() << "Extracted group '" << (*BBs.begin())->getName()
178 << "' in: " << F->getName() << '\n');
179 else
180 LLVM_DEBUG(dbgs() << "Failed to extract for group '"
181 << (*BBs.begin())->getName() << "'\n");
182 }
183
184 // Erase the functions.
185 if (EraseFunctions || BlockExtractorEraseFuncs) {
186 for (Function *F : Functions) {
187 LLVM_DEBUG(dbgs() << "BlockExtractor: Trying to delete " << F->getName()
188 << "\n");
189 F->deleteBody();
190 }
191 // Set linkage as ExternalLinkage to avoid erasing unreachable functions.
192 for (Function &F : M)
193 F.setLinkage(GlobalValue::ExternalLinkage);
194 Changed = true;
195 }
196
197 return Changed;
198 }
199
BlockExtractorPass(std::vector<std::vector<BasicBlock * >> && GroupsOfBlocks,bool EraseFunctions)200 BlockExtractorPass::BlockExtractorPass(
201 std::vector<std::vector<BasicBlock *>> &&GroupsOfBlocks,
202 bool EraseFunctions)
203 : GroupsOfBlocks(GroupsOfBlocks), EraseFunctions(EraseFunctions) {}
204
run(Module & M,ModuleAnalysisManager & AM)205 PreservedAnalyses BlockExtractorPass::run(Module &M,
206 ModuleAnalysisManager &AM) {
207 BlockExtractor BE(EraseFunctions);
208 BE.init(GroupsOfBlocks);
209 return BE.runOnModule(M) ? PreservedAnalyses::none()
210 : PreservedAnalyses::all();
211 }
212