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