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