1//===- UnifyLoopExits.cpp - Redirect exiting edges to one block -*- C++ -*-===//
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// For each natural loop with multiple exit blocks, this pass creates a new
10// block N such that all exiting blocks now branch to N, and then control flow
11// is redistributed to all the original exit blocks.
12//
13// Limitation: This assumes that all terminators in the CFG are direct branches
14// (the "br" instruction). The presence of any other control flow
15// such as indirectbr, switch or callbr will cause an assert.
16//
17//===----------------------------------------------------------------------===//
18
19#include "llvm/Transforms/Utils/UnifyLoopExits.h"
20#include "llvm/ADT/MapVector.h"
21#include "llvm/Analysis/DomTreeUpdater.h"
22#include "llvm/Analysis/LoopInfo.h"
23#include "llvm/IR/Constants.h"
24#include "llvm/IR/Dominators.h"
25#include "llvm/InitializePasses.h"
26#include "llvm/Support/CommandLine.h"
27#include "llvm/Transforms/Utils.h"
28#include "llvm/Transforms/Utils/BasicBlockUtils.h"
29
30#define DEBUG_TYPE "unify-loop-exits"
31
32using namespace llvm;
33
34static cl::opt<unsigned> MaxBooleansInControlFlowHub(
35 "max-booleans-in-control-flow-hub", cl::init(Val: 32), cl::Hidden,
36 cl::desc("Set the maximum number of outgoing blocks for using a boolean "
37 "value to record the exiting block in CreateControlFlowHub."));
38
39namespace {
40struct UnifyLoopExitsLegacyPass : public FunctionPass {
41 static char ID;
42 UnifyLoopExitsLegacyPass() : FunctionPass(ID) {
43 initializeUnifyLoopExitsLegacyPassPass(*PassRegistry::getPassRegistry());
44 }
45
46 void getAnalysisUsage(AnalysisUsage &AU) const override {
47 AU.addRequired<LoopInfoWrapperPass>();
48 AU.addRequired<DominatorTreeWrapperPass>();
49 AU.addPreserved<LoopInfoWrapperPass>();
50 AU.addPreserved<DominatorTreeWrapperPass>();
51 }
52
53 bool runOnFunction(Function &F) override;
54};
55} // namespace
56
57char UnifyLoopExitsLegacyPass::ID = 0;
58
59FunctionPass *llvm::createUnifyLoopExitsPass() {
60 return new UnifyLoopExitsLegacyPass();
61}
62
63INITIALIZE_PASS_BEGIN(UnifyLoopExitsLegacyPass, "unify-loop-exits",
64 "Fixup each natural loop to have a single exit block",
65 false /* Only looks at CFG */, false /* Analysis Pass */)
66INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
67INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
68INITIALIZE_PASS_END(UnifyLoopExitsLegacyPass, "unify-loop-exits",
69 "Fixup each natural loop to have a single exit block",
70 false /* Only looks at CFG */, false /* Analysis Pass */)
71
72// The current transform introduces new control flow paths which may break the
73// SSA requirement that every def must dominate all its uses. For example,
74// consider a value D defined inside the loop that is used by some instruction
75// U outside the loop. It follows that D dominates U, since the original
76// program has valid SSA form. After merging the exits, all paths from D to U
77// now flow through the unified exit block. In addition, there may be other
78// paths that do not pass through D, but now reach the unified exit
79// block. Thus, D no longer dominates U.
80//
81// Restore the dominance by creating a phi for each such D at the new unified
82// loop exit. But when doing this, ignore any uses U that are in the new unified
83// loop exit, since those were introduced specially when the block was created.
84//
85// The use of SSAUpdater seems like overkill for this operation. The location
86// for creating the new PHI is well-known, and also the set of incoming blocks
87// to the new PHI.
88static void restoreSSA(const DominatorTree &DT, const Loop *L,
89 const SetVector<BasicBlock *> &Incoming,
90 BasicBlock *LoopExitBlock) {
91 using InstVector = SmallVector<Instruction *, 8>;
92 using IIMap = MapVector<Instruction *, InstVector>;
93 IIMap ExternalUsers;
94 for (auto *BB : L->blocks()) {
95 for (auto &I : *BB) {
96 for (auto &U : I.uses()) {
97 auto UserInst = cast<Instruction>(Val: U.getUser());
98 auto UserBlock = UserInst->getParent();
99 if (UserBlock == LoopExitBlock)
100 continue;
101 if (L->contains(BB: UserBlock))
102 continue;
103 LLVM_DEBUG(dbgs() << "added ext use for " << I.getName() << "("
104 << BB->getName() << ")"
105 << ": " << UserInst->getName() << "("
106 << UserBlock->getName() << ")"
107 << "\n");
108 ExternalUsers[&I].push_back(Elt: UserInst);
109 }
110 }
111 }
112
113 for (const auto &II : ExternalUsers) {
114 // For each Def used outside the loop, create NewPhi in
115 // LoopExitBlock. NewPhi receives Def only along exiting blocks that
116 // dominate it, while the remaining values are undefined since those paths
117 // didn't exist in the original CFG.
118 auto Def = II.first;
119 LLVM_DEBUG(dbgs() << "externally used: " << Def->getName() << "\n");
120 auto NewPhi =
121 PHINode::Create(Ty: Def->getType(), NumReservedValues: Incoming.size(),
122 NameStr: Def->getName() + ".moved", InsertBefore: LoopExitBlock->begin());
123 for (auto *In : Incoming) {
124 LLVM_DEBUG(dbgs() << "predecessor " << In->getName() << ": ");
125 if (Def->getParent() == In || DT.dominates(Def, BB: In)) {
126 LLVM_DEBUG(dbgs() << "dominated\n");
127 NewPhi->addIncoming(V: Def, BB: In);
128 } else {
129 LLVM_DEBUG(dbgs() << "not dominated\n");
130 NewPhi->addIncoming(V: PoisonValue::get(T: Def->getType()), BB: In);
131 }
132 }
133
134 LLVM_DEBUG(dbgs() << "external users:");
135 for (auto *U : II.second) {
136 LLVM_DEBUG(dbgs() << " " << U->getName());
137 U->replaceUsesOfWith(From: Def, To: NewPhi);
138 }
139 LLVM_DEBUG(dbgs() << "\n");
140 }
141}
142
143static bool unifyLoopExits(DominatorTree &DT, LoopInfo &LI, Loop *L) {
144 // To unify the loop exits, we need a list of the exiting blocks as
145 // well as exit blocks. The functions for locating these lists both
146 // traverse the entire loop body. It is more efficient to first
147 // locate the exiting blocks and then examine their successors to
148 // locate the exit blocks.
149 SetVector<BasicBlock *> ExitingBlocks;
150 SetVector<BasicBlock *> Exits;
151
152 // We need SetVectors, but the Loop API takes a vector, so we use a temporary.
153 SmallVector<BasicBlock *, 8> Temp;
154 L->getExitingBlocks(ExitingBlocks&: Temp);
155 for (auto *BB : Temp) {
156 ExitingBlocks.insert(X: BB);
157 for (auto *S : successors(BB)) {
158 auto SL = LI.getLoopFor(BB: S);
159 // A successor is not an exit if it is directly or indirectly in the
160 // current loop.
161 if (SL == L || L->contains(L: SL))
162 continue;
163 Exits.insert(X: S);
164 }
165 }
166
167 LLVM_DEBUG(
168 dbgs() << "Found exit blocks:";
169 for (auto Exit : Exits) {
170 dbgs() << " " << Exit->getName();
171 }
172 dbgs() << "\n";
173
174 dbgs() << "Found exiting blocks:";
175 for (auto EB : ExitingBlocks) {
176 dbgs() << " " << EB->getName();
177 }
178 dbgs() << "\n";);
179
180 if (Exits.size() <= 1) {
181 LLVM_DEBUG(dbgs() << "loop does not have multiple exits; nothing to do\n");
182 return false;
183 }
184
185 SmallVector<BasicBlock *, 8> GuardBlocks;
186 DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
187 auto LoopExitBlock =
188 CreateControlFlowHub(DTU: &DTU, GuardBlocks, Predecessors: ExitingBlocks, Successors: Exits, Prefix: "loop.exit",
189 MaxControlFlowBooleans: MaxBooleansInControlFlowHub.getValue());
190
191 restoreSSA(DT, L, Incoming: ExitingBlocks, LoopExitBlock);
192
193#if defined(EXPENSIVE_CHECKS)
194 assert(DT.verify(DominatorTree::VerificationLevel::Full));
195#else
196 assert(DT.verify(DominatorTree::VerificationLevel::Fast));
197#endif // EXPENSIVE_CHECKS
198 L->verifyLoop();
199
200 // The guard blocks were created outside the loop, so they need to become
201 // members of the parent loop.
202 if (auto ParentLoop = L->getParentLoop()) {
203 for (auto *G : GuardBlocks) {
204 ParentLoop->addBasicBlockToLoop(NewBB: G, LI);
205 }
206 ParentLoop->verifyLoop();
207 }
208
209#if defined(EXPENSIVE_CHECKS)
210 LI.verify(DT);
211#endif // EXPENSIVE_CHECKS
212
213 return true;
214}
215
216static bool runImpl(LoopInfo &LI, DominatorTree &DT) {
217
218 bool Changed = false;
219 auto Loops = LI.getLoopsInPreorder();
220 for (auto *L : Loops) {
221 LLVM_DEBUG(dbgs() << "Loop: " << L->getHeader()->getName() << " (depth: "
222 << LI.getLoopDepth(L->getHeader()) << ")\n");
223 Changed |= unifyLoopExits(DT, LI, L);
224 }
225 return Changed;
226}
227
228bool UnifyLoopExitsLegacyPass::runOnFunction(Function &F) {
229 LLVM_DEBUG(dbgs() << "===== Unifying loop exits in function " << F.getName()
230 << "\n");
231 auto &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
232 auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
233
234 assert(hasOnlySimpleTerminator(F) && "Unsupported block terminator.");
235
236 return runImpl(LI, DT);
237}
238
239namespace llvm {
240
241PreservedAnalyses UnifyLoopExitsPass::run(Function &F,
242 FunctionAnalysisManager &AM) {
243 auto &LI = AM.getResult<LoopAnalysis>(IR&: F);
244 auto &DT = AM.getResult<DominatorTreeAnalysis>(IR&: F);
245
246 if (!runImpl(LI, DT))
247 return PreservedAnalyses::all();
248 PreservedAnalyses PA;
249 PA.preserve<LoopAnalysis>();
250 PA.preserve<DominatorTreeAnalysis>();
251 return PA;
252}
253} // namespace llvm
254