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);
init(const SmallVectorImpl<SmallVector<BasicBlock *,16>> & GroupsOfBlocksToExtract)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;
BlockExtractorLegacyPass(const SmallVectorImpl<BasicBlock * > & BlocksToExtract,bool EraseFunctions)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
BlockExtractorLegacyPass(const SmallVectorImpl<SmallVector<BasicBlock *,16>> & GroupsOfBlocksToExtract,bool EraseFunctions)90 BlockExtractorLegacyPass(const SmallVectorImpl<SmallVector<BasicBlock *, 16>>
91 &GroupsOfBlocksToExtract,
92 bool EraseFunctions)
93 : ModulePass(ID), BE(EraseFunctions) {
94 BE.init(GroupsOfBlocksToExtract);
95 }
96
BlockExtractorLegacyPass()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
createBlockExtractorPass()107 ModulePass *llvm::createBlockExtractorPass() {
108 return new BlockExtractorLegacyPass();
109 }
createBlockExtractorPass(const SmallVectorImpl<BasicBlock * > & BlocksToExtract,bool EraseFunctions)110 ModulePass *llvm::createBlockExtractorPass(
111 const SmallVectorImpl<BasicBlock *> &BlocksToExtract, bool EraseFunctions) {
112 return new BlockExtractorLegacyPass(BlocksToExtract, EraseFunctions);
113 }
createBlockExtractorPass(const SmallVectorImpl<SmallVector<BasicBlock *,16>> & GroupsOfBlocksToExtract,bool EraseFunctions)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.
loadFile()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.
splitLandingPadPreds(Function & F)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
runOnModule(Module & M)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
runOnModule(Module & M)251 bool BlockExtractorLegacyPass::runOnModule(Module &M) {
252 return BE.runOnModule(M);
253 }
254
run(Module & M,ModuleAnalysisManager & AM)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 }
262