1//===- LoopExtractor.cpp - Extract each loop into a new function ----------===//
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// A pass wrapper around the ExtractLoop() scalar transformation to extract each
10// top-level loop into its own new function. If the loop is the ONLY loop in a
11// given function, it is not touched. This is a pass most useful for debugging
12// via bugpoint.
13//
14//===----------------------------------------------------------------------===//
15
16#include "llvm/Transforms/IPO/LoopExtractor.h"
17#include "llvm/ADT/Statistic.h"
18#include "llvm/Analysis/AssumptionCache.h"
19#include "llvm/Analysis/LoopInfo.h"
20#include "llvm/IR/Dominators.h"
21#include "llvm/IR/Instructions.h"
22#include "llvm/IR/Module.h"
23#include "llvm/IR/PassManager.h"
24#include "llvm/InitializePasses.h"
25#include "llvm/Pass.h"
26#include "llvm/Transforms/IPO.h"
27#include "llvm/Transforms/Utils.h"
28#include "llvm/Transforms/Utils/CodeExtractor.h"
29using namespace llvm;
30
31#define DEBUG_TYPE "loop-extract"
32
33STATISTIC(NumExtracted, "Number of loops extracted");
34
35namespace {
36struct LoopExtractorLegacyPass : public ModulePass {
37 static char ID; // Pass identification, replacement for typeid
38
39 unsigned NumLoops;
40
41 explicit LoopExtractorLegacyPass(unsigned NumLoops = ~0)
42 : ModulePass(ID), NumLoops(NumLoops) {}
43
44 bool runOnModule(Module &M) override;
45
46 void getAnalysisUsage(AnalysisUsage &AU) const override {
47 AU.addRequiredID(ID&: BreakCriticalEdgesID);
48 AU.addRequired<DominatorTreeWrapperPass>();
49 AU.addRequired<LoopInfoWrapperPass>();
50 AU.addPreserved<LoopInfoWrapperPass>();
51 AU.addRequiredID(ID&: LoopSimplifyID);
52 AU.addUsedIfAvailable<AssumptionCacheTracker>();
53 }
54};
55
56struct LoopExtractor {
57 explicit LoopExtractor(
58 unsigned NumLoops,
59 function_ref<DominatorTree &(Function &)> LookupDomTree,
60 function_ref<LoopInfo &(Function &)> LookupLoopInfo,
61 function_ref<AssumptionCache *(Function &)> LookupAssumptionCache)
62 : NumLoops(NumLoops), LookupDomTree(LookupDomTree),
63 LookupLoopInfo(LookupLoopInfo),
64 LookupAssumptionCache(LookupAssumptionCache) {}
65 bool runOnModule(Module &M);
66
67private:
68 // The number of natural loops to extract from the program into functions.
69 unsigned NumLoops;
70
71 function_ref<DominatorTree &(Function &)> LookupDomTree;
72 function_ref<LoopInfo &(Function &)> LookupLoopInfo;
73 function_ref<AssumptionCache *(Function &)> LookupAssumptionCache;
74
75 bool runOnFunction(Function &F);
76
77 bool extractLoops(Loop::iterator From, Loop::iterator To, LoopInfo &LI,
78 DominatorTree &DT);
79 bool extractLoop(Loop *L, LoopInfo &LI, DominatorTree &DT);
80};
81} // namespace
82
83char LoopExtractorLegacyPass::ID = 0;
84INITIALIZE_PASS_BEGIN(LoopExtractorLegacyPass, "loop-extract",
85 "Extract loops into new functions", false, false)
86INITIALIZE_PASS_DEPENDENCY(BreakCriticalEdges)
87INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
88INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
89INITIALIZE_PASS_DEPENDENCY(LoopSimplify)
90INITIALIZE_PASS_END(LoopExtractorLegacyPass, "loop-extract",
91 "Extract loops into new functions", false, false)
92
93namespace {
94 /// SingleLoopExtractor - For bugpoint.
95struct SingleLoopExtractor : public LoopExtractorLegacyPass {
96 static char ID; // Pass identification, replacement for typeid
97 SingleLoopExtractor() : LoopExtractorLegacyPass(1) {}
98};
99} // End anonymous namespace
100
101char SingleLoopExtractor::ID = 0;
102INITIALIZE_PASS(SingleLoopExtractor, "loop-extract-single",
103 "Extract at most one loop into a new function", false, false)
104
105// createLoopExtractorPass - This pass extracts all natural loops from the
106// program into a function if it can.
107//
108Pass *llvm::createLoopExtractorPass() { return new LoopExtractorLegacyPass(); }
109
110bool LoopExtractorLegacyPass::runOnModule(Module &M) {
111 if (skipModule(M))
112 return false;
113
114 bool Changed = false;
115 auto LookupDomTree = [this](Function &F) -> DominatorTree & {
116 return this->getAnalysis<DominatorTreeWrapperPass>(F).getDomTree();
117 };
118 auto LookupLoopInfo = [this, &Changed](Function &F) -> LoopInfo & {
119 return this->getAnalysis<LoopInfoWrapperPass>(F, Changed: &Changed).getLoopInfo();
120 };
121 auto LookupACT = [this](Function &F) -> AssumptionCache * {
122 if (auto *ACT = this->getAnalysisIfAvailable<AssumptionCacheTracker>())
123 return ACT->lookupAssumptionCache(F);
124 return nullptr;
125 };
126 return LoopExtractor(NumLoops, LookupDomTree, LookupLoopInfo, LookupACT)
127 .runOnModule(M) ||
128 Changed;
129}
130
131bool LoopExtractor::runOnModule(Module &M) {
132 if (M.empty())
133 return false;
134
135 if (!NumLoops)
136 return false;
137
138 bool Changed = false;
139
140 // The end of the function list may change (new functions will be added at the
141 // end), so we run from the first to the current last.
142 auto I = M.begin(), E = --M.end();
143 while (true) {
144 Function &F = *I;
145
146 Changed |= runOnFunction(F);
147 if (!NumLoops)
148 break;
149
150 // If this is the last function.
151 if (I == E)
152 break;
153
154 ++I;
155 }
156 return Changed;
157}
158
159bool LoopExtractor::runOnFunction(Function &F) {
160 // Do not modify `optnone` functions.
161 if (F.hasOptNone())
162 return false;
163
164 if (F.empty())
165 return false;
166
167 bool Changed = false;
168 LoopInfo &LI = LookupLoopInfo(F);
169
170 // If there are no loops in the function.
171 if (LI.empty())
172 return Changed;
173
174 DominatorTree &DT = LookupDomTree(F);
175
176 // If there is more than one top-level loop in this function, extract all of
177 // the loops.
178 if (std::next(x: LI.begin()) != LI.end())
179 return Changed | extractLoops(From: LI.begin(), To: LI.end(), LI, DT);
180
181 // Otherwise there is exactly one top-level loop.
182 Loop *TLL = *LI.begin();
183
184 // If the loop is in LoopSimplify form, then extract it only if this function
185 // is more than a minimal wrapper around the loop.
186 if (TLL->isLoopSimplifyForm()) {
187 bool ShouldExtractLoop = false;
188
189 // Extract the loop if the entry block doesn't branch to the loop header.
190 auto *EntryTI = dyn_cast<UncondBrInst>(Val: F.getEntryBlock().getTerminator());
191 if (EntryTI && EntryTI->getSuccessor() != TLL->getHeader()) {
192 ShouldExtractLoop = true;
193 } else {
194 // Check to see if any exits from the loop are more than just return
195 // blocks.
196 SmallVector<BasicBlock *, 8> ExitBlocks;
197 TLL->getExitBlocks(ExitBlocks);
198 for (auto *ExitBlock : ExitBlocks)
199 if (!isa<ReturnInst>(Val: ExitBlock->getTerminator())) {
200 ShouldExtractLoop = true;
201 break;
202 }
203 }
204
205 if (ShouldExtractLoop)
206 return Changed | extractLoop(L: TLL, LI, DT);
207 }
208
209 // Okay, this function is a minimal container around the specified loop.
210 // If we extract the loop, we will continue to just keep extracting it
211 // infinitely... so don't extract it. However, if the loop contains any
212 // sub-loops, extract them.
213 return Changed | extractLoops(From: TLL->begin(), To: TLL->end(), LI, DT);
214}
215
216bool LoopExtractor::extractLoops(Loop::iterator From, Loop::iterator To,
217 LoopInfo &LI, DominatorTree &DT) {
218 bool Changed = false;
219 SmallVector<Loop *, 8> Loops;
220
221 // Save the list of loops, as it may change.
222 Loops.assign(in_start: From, in_end: To);
223 for (Loop *L : Loops) {
224 // If LoopSimplify form is not available, stay out of trouble.
225 if (!L->isLoopSimplifyForm())
226 continue;
227
228 Changed |= extractLoop(L, LI, DT);
229 if (!NumLoops)
230 break;
231 }
232 return Changed;
233}
234
235bool LoopExtractor::extractLoop(Loop *L, LoopInfo &LI, DominatorTree &DT) {
236 assert(NumLoops != 0);
237 Function &Func = *L->getHeader()->getParent();
238 AssumptionCache *AC = LookupAssumptionCache(Func);
239 CodeExtractorAnalysisCache CEAC(Func);
240 CodeExtractor Extractor(L->getBlocks(), &DT, false, nullptr, nullptr, AC);
241 if (Extractor.extractCodeRegion(CEAC)) {
242 LI.erase(L);
243 --NumLoops;
244 ++NumExtracted;
245 return true;
246 }
247 return false;
248}
249
250// createSingleLoopExtractorPass - This pass extracts one natural loop from the
251// program into a function if it can. This is used by bugpoint.
252//
253Pass *llvm::createSingleLoopExtractorPass() {
254 return new SingleLoopExtractor();
255}
256
257PreservedAnalyses LoopExtractorPass::run(Module &M, ModuleAnalysisManager &AM) {
258 auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(IR&: M).getManager();
259 auto LookupDomTree = [&FAM](Function &F) -> DominatorTree & {
260 return FAM.getResult<DominatorTreeAnalysis>(IR&: F);
261 };
262 auto LookupLoopInfo = [&FAM](Function &F) -> LoopInfo & {
263 return FAM.getResult<LoopAnalysis>(IR&: F);
264 };
265 auto LookupAssumptionCache = [&FAM](Function &F) -> AssumptionCache * {
266 return FAM.getCachedResult<AssumptionAnalysis>(IR&: F);
267 };
268 if (!LoopExtractor(NumLoops, LookupDomTree, LookupLoopInfo,
269 LookupAssumptionCache)
270 .runOnModule(M))
271 return PreservedAnalyses::all();
272
273 PreservedAnalyses PA;
274 PA.preserve<LoopAnalysis>();
275 return PA;
276}
277
278void LoopExtractorPass::printPipeline(
279 raw_ostream &OS, function_ref<StringRef(StringRef)> MapClassName2PassName) {
280 static_cast<PassInfoMixin<LoopExtractorPass> *>(this)->printPipeline(
281 OS, MapClassName2PassName);
282 OS << '<';
283 if (NumLoops == 1)
284 OS << "single";
285 OS << '>';
286}
287