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 or switch will cause an assert.
16// The callbr and switch terminators are supported by creating
17// intermediate target blocks that unconditionally branch to the
18// original target blocks. These intermediate target blocks can then
19// be redirected through the ControlFlowHub as usual.
20//
21//===----------------------------------------------------------------------===//
22
23#include "llvm/Transforms/Utils/UnifyLoopExits.h"
24#include "llvm/ADT/DenseMap.h"
25#include "llvm/ADT/MapVector.h"
26#include "llvm/Analysis/DomTreeUpdater.h"
27#include "llvm/Analysis/LoopInfo.h"
28#include "llvm/IR/Constants.h"
29#include "llvm/IR/Dominators.h"
30#include "llvm/IR/Instructions.h"
31#include "llvm/InitializePasses.h"
32#include "llvm/Support/CommandLine.h"
33#include "llvm/Transforms/Utils.h"
34#include "llvm/Transforms/Utils/BasicBlockUtils.h"
35#include "llvm/Transforms/Utils/ControlFlowUtils.h"
36
37#define DEBUG_TYPE "unify-loop-exits"
38
39using namespace llvm;
40
41static cl::opt<unsigned> MaxBooleansInControlFlowHub(
42 "max-booleans-in-control-flow-hub", cl::init(Val: 32), cl::Hidden,
43 cl::desc("Set the maximum number of outgoing blocks for using a boolean "
44 "value to record the exiting block in the ControlFlowHub."));
45
46namespace {
47struct UnifyLoopExitsLegacyPass : public FunctionPass {
48 static char ID;
49 UnifyLoopExitsLegacyPass() : FunctionPass(ID) {
50 initializeUnifyLoopExitsLegacyPassPass(*PassRegistry::getPassRegistry());
51 }
52
53 void getAnalysisUsage(AnalysisUsage &AU) const override {
54 AU.addRequired<LoopInfoWrapperPass>();
55 AU.addRequired<DominatorTreeWrapperPass>();
56 AU.addPreserved<LoopInfoWrapperPass>();
57 AU.addPreserved<DominatorTreeWrapperPass>();
58 }
59
60 bool runOnFunction(Function &F) override;
61};
62} // namespace
63
64char UnifyLoopExitsLegacyPass::ID = 0;
65
66FunctionPass *llvm::createUnifyLoopExitsPass() {
67 return new UnifyLoopExitsLegacyPass();
68}
69
70INITIALIZE_PASS_BEGIN(UnifyLoopExitsLegacyPass, "unify-loop-exits",
71 "Fixup each natural loop to have a single exit block",
72 false /* Only looks at CFG */, false /* Analysis Pass */)
73INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
74INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
75INITIALIZE_PASS_END(UnifyLoopExitsLegacyPass, "unify-loop-exits",
76 "Fixup each natural loop to have a single exit block",
77 false /* Only looks at CFG */, false /* Analysis Pass */)
78
79// The current transform introduces new control flow paths which may break the
80// SSA requirement that every def must dominate all its uses. For example,
81// consider a value D defined inside the loop that is used by some instruction
82// U outside the loop. It follows that D dominates U, since the original
83// program has valid SSA form. After merging the exits, all paths from D to U
84// now flow through the unified exit block. In addition, there may be other
85// paths that do not pass through D, but now reach the unified exit
86// block. Thus, D no longer dominates U.
87//
88// Restore the dominance by creating a phi for each such D at the new unified
89// loop exit. But when doing this, ignore any uses U that are in the new unified
90// loop exit, since those were introduced specially when the block was created.
91//
92// The use of SSAUpdater seems like overkill for this operation. The location
93// for creating the new PHI is well-known, and also the set of incoming blocks
94// to the new PHI.
95static void restoreSSA(const DominatorTree &DT, const Loop *L,
96 SmallVectorImpl<BasicBlock *> &Incoming,
97 BasicBlock *LoopExitBlock) {
98 using InstVector = SmallVector<Instruction *, 8>;
99 using IIMap = MapVector<Instruction *, InstVector>;
100 IIMap ExternalUsers;
101 for (auto *BB : L->blocks()) {
102 for (auto &I : *BB) {
103 for (auto &U : I.uses()) {
104 auto UserInst = cast<Instruction>(Val: U.getUser());
105 auto UserBlock = UserInst->getParent();
106 if (UserBlock == LoopExitBlock)
107 continue;
108 if (L->contains(BB: UserBlock))
109 continue;
110 LLVM_DEBUG(dbgs() << "added ext use for " << I.getName() << "("
111 << BB->getName() << ")"
112 << ": " << UserInst->getName() << "("
113 << UserBlock->getName() << ")"
114 << "\n");
115 ExternalUsers[&I].push_back(Elt: UserInst);
116 }
117 }
118 }
119
120 for (const auto &II : ExternalUsers) {
121 // For each Def used outside the loop, create NewPhi in
122 // LoopExitBlock. NewPhi receives Def only along exiting blocks that
123 // dominate it, while the remaining values are undefined since those paths
124 // didn't exist in the original CFG.
125 auto Def = II.first;
126 LLVM_DEBUG(dbgs() << "externally used: " << Def->getName() << "\n");
127 auto NewPhi =
128 PHINode::Create(Ty: Def->getType(), NumReservedValues: Incoming.size(),
129 NameStr: Def->getName() + ".moved", InsertBefore: LoopExitBlock->begin());
130 for (auto *In : Incoming) {
131 LLVM_DEBUG(dbgs() << "predecessor " << In->getName() << ": ");
132 if (Def->getParent() == In || DT.dominates(Def, BB: In)) {
133 LLVM_DEBUG(dbgs() << "dominated\n");
134 NewPhi->addIncoming(V: Def, BB: In);
135 } else {
136 LLVM_DEBUG(dbgs() << "not dominated\n");
137 NewPhi->addIncoming(V: PoisonValue::get(T: Def->getType()), BB: In);
138 }
139 }
140
141 LLVM_DEBUG(dbgs() << "external users:");
142 for (auto *U : II.second) {
143 LLVM_DEBUG(dbgs() << " " << U->getName());
144 U->replaceUsesOfWith(From: Def, To: NewPhi);
145 }
146 LLVM_DEBUG(dbgs() << "\n");
147 }
148}
149
150static bool unifyLoopExits(DominatorTree &DT, LoopInfo &LI, Loop *L) {
151 // To unify the loop exits, we need a list of the exiting blocks as
152 // well as exit blocks. The functions for locating these lists both
153 // traverse the entire loop body. It is more efficient to first
154 // locate the exiting blocks and then examine their successors to
155 // locate the exit blocks.
156 SmallVector<BasicBlock *, 8> ExitingBlocks;
157 L->getExitingBlocks(ExitingBlocks);
158
159 // No exit blocks, so nothing to do. Just return.
160 if (ExitingBlocks.empty())
161 return false;
162
163 DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
164 SmallVector<BasicBlock *, 8> MultiBrTargetBlocksToFix;
165
166 // Redirect exiting edges through a control flow hub.
167 ControlFlowHub CHub;
168 bool Changed = false;
169
170 unsigned NumExitingBlocks = ExitingBlocks.size();
171 for (unsigned I = 0; I < NumExitingBlocks; ++I) {
172 BasicBlock *BB = ExitingBlocks[I];
173 Instruction *Term = BB->getTerminator();
174 if (UncondBrInst *Branch = dyn_cast<UncondBrInst>(Val: Term)) {
175 BasicBlock *Succ0 = Branch->getSuccessor(i: 0);
176 Succ0 = L->contains(BB: Succ0) ? nullptr : Succ0;
177 CHub.addBranch(BB, Succ0);
178
179 LLVM_DEBUG(dbgs() << "Added exiting branch: " << printBasicBlock(BB)
180 << " -> " << printBasicBlock(Succ0) << '\n');
181 } else if (CondBrInst *Branch = dyn_cast<CondBrInst>(Val: Term)) {
182 BasicBlock *Succ0 = Branch->getSuccessor(i: 0);
183 Succ0 = L->contains(BB: Succ0) ? nullptr : Succ0;
184
185 BasicBlock *Succ1 = Branch->getSuccessor(i: 1);
186 Succ1 = L->contains(BB: Succ1) ? nullptr : Succ1;
187 CHub.addBranch(BB, Succ0, Succ1);
188
189 LLVM_DEBUG(dbgs() << "Added exiting branch: " << printBasicBlock(BB)
190 << " -> " << printBasicBlock(Succ0)
191 << (Succ0 && Succ1 ? " " : "") << printBasicBlock(Succ1)
192 << '\n');
193 } else if (isa<CallBrInst>(Val: Term) || isa<SwitchInst>(Val: Term)) {
194 SmallDenseMap<BasicBlock *, BasicBlock *> BrTargets;
195 for (unsigned J = 0; J < Term->getNumSuccessors(); ++J) {
196 BasicBlock *Succ = Term->getSuccessor(Idx: J);
197 if (L->contains(BB: Succ))
198 continue;
199 bool UpdatedLI;
200 auto It = BrTargets.find(Val: Succ);
201 BasicBlock *ExistingTarget =
202 (It != BrTargets.end()) ? It->second : nullptr;
203 BasicBlock *NewSucc = SplitMultiBrEdge(MultiBrBlock: BB, Succ, SuccIdx: J, BrTarget: ExistingTarget,
204 DTU: &DTU, CI: nullptr, LI: &LI, UpdatedLI: &UpdatedLI);
205
206 if (!ExistingTarget) {
207 // SplitMultiBrEdge modifies the CFG because it creates an
208 // intermediate block. So we need to set the changed flag no matter
209 // what the ControlFlowHub is going to do later.
210 Changed = true;
211 // Even if the terminator and Succ do not have a common parent loop,
212 // we need to add the new target block to the parent loop of the
213 // current loop.
214 if (!UpdatedLI)
215 MultiBrTargetBlocksToFix.push_back(Elt: NewSucc);
216 // ExitingBlocks is later used to restore SSA, so we need to make sure
217 // that the blocks used for phi nodes in the guard blocks match the
218 // predecessors of the guard blocks, which, in the case of callbr or
219 // switch terminator, are the new intermediate target blocks instead
220 // of themselves. If only one exiting block is generated, the
221 // branching block itself is overwritten, while further blocks are
222 // appended as additional exiting blocks.
223 if (BrTargets.empty())
224 ExitingBlocks[I] = NewSucc;
225 else
226 ExitingBlocks.push_back(Elt: NewSucc);
227 CHub.addBranch(BB: NewSucc, Succ0: Succ);
228 BrTargets[Succ] = NewSucc;
229 }
230 LLVM_DEBUG(dbgs() << "Added exiting branch: "
231 << printBasicBlock(NewSucc) << " -> "
232 << printBasicBlock(Succ) << '\n');
233 }
234 } else {
235 reportFatalUsageError(
236 reason: "unsupported block terminator: unify-loop-exits "
237 "only supports br, callbr, and switch instructions");
238 }
239 }
240
241 SmallVector<BasicBlock *, 8> GuardBlocks;
242 BasicBlock *LoopExitBlock;
243 bool ChangedCFG;
244 std::tie(args&: LoopExitBlock, args&: ChangedCFG) = CHub.finalize(
245 DTU: &DTU, GuardBlocks, Prefix: "loop.exit", MaxControlFlowBooleans: MaxBooleansInControlFlowHub.getValue());
246 ChangedCFG |= Changed;
247 if (!ChangedCFG)
248 return false;
249
250 restoreSSA(DT, L, Incoming&: ExitingBlocks, LoopExitBlock);
251
252#if defined(EXPENSIVE_CHECKS)
253 assert(DT.verify(DominatorTree::VerificationLevel::Full));
254#else
255 assert(DT.verify(DominatorTree::VerificationLevel::Fast));
256#endif // EXPENSIVE_CHECKS
257 L->verifyLoop();
258
259 // The guard blocks were created outside the loop, so they need to become
260 // members of the parent loop.
261 // Same goes for the callbr/switch target blocks. Although we try to add them
262 // to the smallest common parent loop of the branching block and the
263 // corresponding original target block, there might not have been such a loop,
264 // in which case the newly created target blocks are not part of any
265 // loop. For nested loops, this might result in them leading to a loop with
266 // multiple entry points.
267 if (auto *ParentLoop = L->getParentLoop()) {
268 for (auto *G : GuardBlocks) {
269 ParentLoop->addBasicBlockToLoop(NewBB: G, LI);
270 }
271 for (auto *C : MultiBrTargetBlocksToFix) {
272 ParentLoop->addBasicBlockToLoop(NewBB: C, LI);
273 }
274 ParentLoop->verifyLoop();
275 }
276
277#if defined(EXPENSIVE_CHECKS)
278 LI.verify(DT);
279#endif // EXPENSIVE_CHECKS
280
281 return true;
282}
283
284static bool runImpl(LoopInfo &LI, DominatorTree &DT) {
285
286 bool Changed = false;
287 auto Loops = LI.getLoopsInPreorder();
288 for (auto *L : Loops) {
289 LLVM_DEBUG(dbgs() << "Processing loop:\n"; L->print(dbgs()));
290 Changed |= unifyLoopExits(DT, LI, L);
291 }
292 return Changed;
293}
294
295bool UnifyLoopExitsLegacyPass::runOnFunction(Function &F) {
296 LLVM_DEBUG(dbgs() << "===== Unifying loop exits in function " << F.getName()
297 << "\n");
298 auto &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
299 auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
300
301 return runImpl(LI, DT);
302}
303
304namespace llvm {
305
306PreservedAnalyses UnifyLoopExitsPass::run(Function &F,
307 FunctionAnalysisManager &AM) {
308 LLVM_DEBUG(dbgs() << "===== Unifying loop exits in function " << F.getName()
309 << "\n");
310 auto &LI = AM.getResult<LoopAnalysis>(IR&: F);
311 auto &DT = AM.getResult<DominatorTreeAnalysis>(IR&: F);
312
313 if (!runImpl(LI, DT))
314 return PreservedAnalyses::all();
315 PreservedAnalyses PA;
316 PA.preserve<LoopAnalysis>();
317 PA.preserve<DominatorTreeAnalysis>();
318 return PA;
319}
320} // namespace llvm
321