xref: /llvm-project/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp (revision 8aaeee5fb6d7c513d065e55b4bc2ad39bd6e760a)
1 ///===- SimpleLoopUnswitch.cpp - Hoist loop-invariant control flow ---------===//
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 #include "llvm/Transforms/Scalar/SimpleLoopUnswitch.h"
10 #include "llvm/ADT/DenseMap.h"
11 #include "llvm/ADT/STLExtras.h"
12 #include "llvm/ADT/Sequence.h"
13 #include "llvm/ADT/SetVector.h"
14 #include "llvm/ADT/SmallPtrSet.h"
15 #include "llvm/ADT/SmallVector.h"
16 #include "llvm/ADT/Statistic.h"
17 #include "llvm/ADT/Twine.h"
18 #include "llvm/Analysis/AssumptionCache.h"
19 #include "llvm/Analysis/CFG.h"
20 #include "llvm/Analysis/CodeMetrics.h"
21 #include "llvm/Analysis/GuardUtils.h"
22 #include "llvm/Analysis/InstructionSimplify.h"
23 #include "llvm/Analysis/LoopAnalysisManager.h"
24 #include "llvm/Analysis/LoopInfo.h"
25 #include "llvm/Analysis/LoopIterator.h"
26 #include "llvm/Analysis/LoopPass.h"
27 #include "llvm/Analysis/MemorySSA.h"
28 #include "llvm/Analysis/MemorySSAUpdater.h"
29 #include "llvm/Analysis/MustExecute.h"
30 #include "llvm/IR/BasicBlock.h"
31 #include "llvm/IR/Constant.h"
32 #include "llvm/IR/Constants.h"
33 #include "llvm/IR/Dominators.h"
34 #include "llvm/IR/Function.h"
35 #include "llvm/IR/InstrTypes.h"
36 #include "llvm/IR/Instruction.h"
37 #include "llvm/IR/Instructions.h"
38 #include "llvm/IR/IntrinsicInst.h"
39 #include "llvm/IR/IRBuilder.h"
40 #include "llvm/IR/Use.h"
41 #include "llvm/IR/Value.h"
42 #include "llvm/InitializePasses.h"
43 #include "llvm/Pass.h"
44 #include "llvm/Support/Casting.h"
45 #include "llvm/Support/CommandLine.h"
46 #include "llvm/Support/Debug.h"
47 #include "llvm/Support/ErrorHandling.h"
48 #include "llvm/Support/GenericDomTree.h"
49 #include "llvm/Support/raw_ostream.h"
50 #include "llvm/Transforms/Scalar/SimpleLoopUnswitch.h"
51 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
52 #include "llvm/Transforms/Utils/Cloning.h"
53 #include "llvm/Transforms/Utils/LoopUtils.h"
54 #include "llvm/Transforms/Utils/ValueMapper.h"
55 #include <algorithm>
56 #include <cassert>
57 #include <iterator>
58 #include <numeric>
59 #include <utility>
60 
61 #define DEBUG_TYPE "simple-loop-unswitch"
62 
63 using namespace llvm;
64 
65 STATISTIC(NumBranches, "Number of branches unswitched");
66 STATISTIC(NumSwitches, "Number of switches unswitched");
67 STATISTIC(NumGuards, "Number of guards turned into branches for unswitching");
68 STATISTIC(NumTrivial, "Number of unswitches that are trivial");
69 STATISTIC(
70     NumCostMultiplierSkipped,
71     "Number of unswitch candidates that had their cost multiplier skipped");
72 
73 static cl::opt<bool> EnableNonTrivialUnswitch(
74     "enable-nontrivial-unswitch", cl::init(false), cl::Hidden,
75     cl::desc("Forcibly enables non-trivial loop unswitching rather than "
76              "following the configuration passed into the pass."));
77 
78 static cl::opt<int>
79     UnswitchThreshold("unswitch-threshold", cl::init(50), cl::Hidden,
80                       cl::desc("The cost threshold for unswitching a loop."));
81 
82 static cl::opt<bool> EnableUnswitchCostMultiplier(
83     "enable-unswitch-cost-multiplier", cl::init(true), cl::Hidden,
84     cl::desc("Enable unswitch cost multiplier that prohibits exponential "
85              "explosion in nontrivial unswitch."));
86 static cl::opt<int> UnswitchSiblingsToplevelDiv(
87     "unswitch-siblings-toplevel-div", cl::init(2), cl::Hidden,
88     cl::desc("Toplevel siblings divisor for cost multiplier."));
89 static cl::opt<int> UnswitchNumInitialUnscaledCandidates(
90     "unswitch-num-initial-unscaled-candidates", cl::init(8), cl::Hidden,
91     cl::desc("Number of unswitch candidates that are ignored when calculating "
92              "cost multiplier."));
93 static cl::opt<bool> UnswitchGuards(
94     "simple-loop-unswitch-guards", cl::init(true), cl::Hidden,
95     cl::desc("If enabled, simple loop unswitching will also consider "
96              "llvm.experimental.guard intrinsics as unswitch candidates."));
97 
98 /// Collect all of the loop invariant input values transitively used by the
99 /// homogeneous instruction graph from a given root.
100 ///
101 /// This essentially walks from a root recursively through loop variant operands
102 /// which have the exact same opcode and finds all inputs which are loop
103 /// invariant. For some operations these can be re-associated and unswitched out
104 /// of the loop entirely.
105 static TinyPtrVector<Value *>
106 collectHomogenousInstGraphLoopInvariants(Loop &L, Instruction &Root,
107                                          LoopInfo &LI) {
108   assert(!L.isLoopInvariant(&Root) &&
109          "Only need to walk the graph if root itself is not invariant.");
110   TinyPtrVector<Value *> Invariants;
111 
112   // Build a worklist and recurse through operators collecting invariants.
113   SmallVector<Instruction *, 4> Worklist;
114   SmallPtrSet<Instruction *, 8> Visited;
115   Worklist.push_back(&Root);
116   Visited.insert(&Root);
117   do {
118     Instruction &I = *Worklist.pop_back_val();
119     for (Value *OpV : I.operand_values()) {
120       // Skip constants as unswitching isn't interesting for them.
121       if (isa<Constant>(OpV))
122         continue;
123 
124       // Add it to our result if loop invariant.
125       if (L.isLoopInvariant(OpV)) {
126         Invariants.push_back(OpV);
127         continue;
128       }
129 
130       // If not an instruction with the same opcode, nothing we can do.
131       Instruction *OpI = dyn_cast<Instruction>(OpV);
132       if (!OpI || OpI->getOpcode() != Root.getOpcode())
133         continue;
134 
135       // Visit this operand.
136       if (Visited.insert(OpI).second)
137         Worklist.push_back(OpI);
138     }
139   } while (!Worklist.empty());
140 
141   return Invariants;
142 }
143 
144 static void replaceLoopInvariantUses(Loop &L, Value *Invariant,
145                                      Constant &Replacement) {
146   assert(!isa<Constant>(Invariant) && "Why are we unswitching on a constant?");
147 
148   // Replace uses of LIC in the loop with the given constant.
149   for (auto UI = Invariant->use_begin(), UE = Invariant->use_end(); UI != UE;) {
150     // Grab the use and walk past it so we can clobber it in the use list.
151     Use *U = &*UI++;
152     Instruction *UserI = dyn_cast<Instruction>(U->getUser());
153 
154     // Replace this use within the loop body.
155     if (UserI && L.contains(UserI))
156       U->set(&Replacement);
157   }
158 }
159 
160 /// Check that all the LCSSA PHI nodes in the loop exit block have trivial
161 /// incoming values along this edge.
162 static bool areLoopExitPHIsLoopInvariant(Loop &L, BasicBlock &ExitingBB,
163                                          BasicBlock &ExitBB) {
164   for (Instruction &I : ExitBB) {
165     auto *PN = dyn_cast<PHINode>(&I);
166     if (!PN)
167       // No more PHIs to check.
168       return true;
169 
170     // If the incoming value for this edge isn't loop invariant the unswitch
171     // won't be trivial.
172     if (!L.isLoopInvariant(PN->getIncomingValueForBlock(&ExitingBB)))
173       return false;
174   }
175   llvm_unreachable("Basic blocks should never be empty!");
176 }
177 
178 /// Insert code to test a set of loop invariant values, and conditionally branch
179 /// on them.
180 static void buildPartialUnswitchConditionalBranch(BasicBlock &BB,
181                                                   ArrayRef<Value *> Invariants,
182                                                   bool Direction,
183                                                   BasicBlock &UnswitchedSucc,
184                                                   BasicBlock &NormalSucc) {
185   IRBuilder<> IRB(&BB);
186 
187   Value *Cond = Direction ? IRB.CreateOr(Invariants) :
188     IRB.CreateAnd(Invariants);
189   IRB.CreateCondBr(Cond, Direction ? &UnswitchedSucc : &NormalSucc,
190                    Direction ? &NormalSucc : &UnswitchedSucc);
191 }
192 
193 /// Rewrite the PHI nodes in an unswitched loop exit basic block.
194 ///
195 /// Requires that the loop exit and unswitched basic block are the same, and
196 /// that the exiting block was a unique predecessor of that block. Rewrites the
197 /// PHI nodes in that block such that what were LCSSA PHI nodes become trivial
198 /// PHI nodes from the old preheader that now contains the unswitched
199 /// terminator.
200 static void rewritePHINodesForUnswitchedExitBlock(BasicBlock &UnswitchedBB,
201                                                   BasicBlock &OldExitingBB,
202                                                   BasicBlock &OldPH) {
203   for (PHINode &PN : UnswitchedBB.phis()) {
204     // When the loop exit is directly unswitched we just need to update the
205     // incoming basic block. We loop to handle weird cases with repeated
206     // incoming blocks, but expect to typically only have one operand here.
207     for (auto i : seq<int>(0, PN.getNumOperands())) {
208       assert(PN.getIncomingBlock(i) == &OldExitingBB &&
209              "Found incoming block different from unique predecessor!");
210       PN.setIncomingBlock(i, &OldPH);
211     }
212   }
213 }
214 
215 /// Rewrite the PHI nodes in the loop exit basic block and the split off
216 /// unswitched block.
217 ///
218 /// Because the exit block remains an exit from the loop, this rewrites the
219 /// LCSSA PHI nodes in it to remove the unswitched edge and introduces PHI
220 /// nodes into the unswitched basic block to select between the value in the
221 /// old preheader and the loop exit.
222 static void rewritePHINodesForExitAndUnswitchedBlocks(BasicBlock &ExitBB,
223                                                       BasicBlock &UnswitchedBB,
224                                                       BasicBlock &OldExitingBB,
225                                                       BasicBlock &OldPH,
226                                                       bool FullUnswitch) {
227   assert(&ExitBB != &UnswitchedBB &&
228          "Must have different loop exit and unswitched blocks!");
229   Instruction *InsertPt = &*UnswitchedBB.begin();
230   for (PHINode &PN : ExitBB.phis()) {
231     auto *NewPN = PHINode::Create(PN.getType(), /*NumReservedValues*/ 2,
232                                   PN.getName() + ".split", InsertPt);
233 
234     // Walk backwards over the old PHI node's inputs to minimize the cost of
235     // removing each one. We have to do this weird loop manually so that we
236     // create the same number of new incoming edges in the new PHI as we expect
237     // each case-based edge to be included in the unswitched switch in some
238     // cases.
239     // FIXME: This is really, really gross. It would be much cleaner if LLVM
240     // allowed us to create a single entry for a predecessor block without
241     // having separate entries for each "edge" even though these edges are
242     // required to produce identical results.
243     for (int i = PN.getNumIncomingValues() - 1; i >= 0; --i) {
244       if (PN.getIncomingBlock(i) != &OldExitingBB)
245         continue;
246 
247       Value *Incoming = PN.getIncomingValue(i);
248       if (FullUnswitch)
249         // No more edge from the old exiting block to the exit block.
250         PN.removeIncomingValue(i);
251 
252       NewPN->addIncoming(Incoming, &OldPH);
253     }
254 
255     // Now replace the old PHI with the new one and wire the old one in as an
256     // input to the new one.
257     PN.replaceAllUsesWith(NewPN);
258     NewPN->addIncoming(&PN, &ExitBB);
259   }
260 }
261 
262 /// Hoist the current loop up to the innermost loop containing a remaining exit.
263 ///
264 /// Because we've removed an exit from the loop, we may have changed the set of
265 /// loops reachable and need to move the current loop up the loop nest or even
266 /// to an entirely separate nest.
267 static void hoistLoopToNewParent(Loop &L, BasicBlock &Preheader,
268                                  DominatorTree &DT, LoopInfo &LI,
269                                  MemorySSAUpdater *MSSAU, ScalarEvolution *SE) {
270   // If the loop is already at the top level, we can't hoist it anywhere.
271   Loop *OldParentL = L.getParentLoop();
272   if (!OldParentL)
273     return;
274 
275   SmallVector<BasicBlock *, 4> Exits;
276   L.getExitBlocks(Exits);
277   Loop *NewParentL = nullptr;
278   for (auto *ExitBB : Exits)
279     if (Loop *ExitL = LI.getLoopFor(ExitBB))
280       if (!NewParentL || NewParentL->contains(ExitL))
281         NewParentL = ExitL;
282 
283   if (NewParentL == OldParentL)
284     return;
285 
286   // The new parent loop (if different) should always contain the old one.
287   if (NewParentL)
288     assert(NewParentL->contains(OldParentL) &&
289            "Can only hoist this loop up the nest!");
290 
291   // The preheader will need to move with the body of this loop. However,
292   // because it isn't in this loop we also need to update the primary loop map.
293   assert(OldParentL == LI.getLoopFor(&Preheader) &&
294          "Parent loop of this loop should contain this loop's preheader!");
295   LI.changeLoopFor(&Preheader, NewParentL);
296 
297   // Remove this loop from its old parent.
298   OldParentL->removeChildLoop(&L);
299 
300   // Add the loop either to the new parent or as a top-level loop.
301   if (NewParentL)
302     NewParentL->addChildLoop(&L);
303   else
304     LI.addTopLevelLoop(&L);
305 
306   // Remove this loops blocks from the old parent and every other loop up the
307   // nest until reaching the new parent. Also update all of these
308   // no-longer-containing loops to reflect the nesting change.
309   for (Loop *OldContainingL = OldParentL; OldContainingL != NewParentL;
310        OldContainingL = OldContainingL->getParentLoop()) {
311     llvm::erase_if(OldContainingL->getBlocksVector(),
312                    [&](const BasicBlock *BB) {
313                      return BB == &Preheader || L.contains(BB);
314                    });
315 
316     OldContainingL->getBlocksSet().erase(&Preheader);
317     for (BasicBlock *BB : L.blocks())
318       OldContainingL->getBlocksSet().erase(BB);
319 
320     // Because we just hoisted a loop out of this one, we have essentially
321     // created new exit paths from it. That means we need to form LCSSA PHI
322     // nodes for values used in the no-longer-nested loop.
323     formLCSSA(*OldContainingL, DT, &LI, SE);
324 
325     // We shouldn't need to form dedicated exits because the exit introduced
326     // here is the (just split by unswitching) preheader. However, after trivial
327     // unswitching it is possible to get new non-dedicated exits out of parent
328     // loop so let's conservatively form dedicated exit blocks and figure out
329     // if we can optimize later.
330     formDedicatedExitBlocks(OldContainingL, &DT, &LI, MSSAU,
331                             /*PreserveLCSSA*/ true);
332   }
333 }
334 
335 // Return the top-most loop containing ExitBB and having ExitBB as exiting block
336 // or the loop containing ExitBB, if there is no parent loop containing ExitBB
337 // as exiting block.
338 static Loop *getTopMostExitingLoop(BasicBlock *ExitBB, LoopInfo &LI) {
339   Loop *TopMost = LI.getLoopFor(ExitBB);
340   Loop *Current = TopMost;
341   while (Current) {
342     if (Current->isLoopExiting(ExitBB))
343       TopMost = Current;
344     Current = Current->getParentLoop();
345   }
346   return TopMost;
347 }
348 
349 /// Unswitch a trivial branch if the condition is loop invariant.
350 ///
351 /// This routine should only be called when loop code leading to the branch has
352 /// been validated as trivial (no side effects). This routine checks if the
353 /// condition is invariant and one of the successors is a loop exit. This
354 /// allows us to unswitch without duplicating the loop, making it trivial.
355 ///
356 /// If this routine fails to unswitch the branch it returns false.
357 ///
358 /// If the branch can be unswitched, this routine splits the preheader and
359 /// hoists the branch above that split. Preserves loop simplified form
360 /// (splitting the exit block as necessary). It simplifies the branch within
361 /// the loop to an unconditional branch but doesn't remove it entirely. Further
362 /// cleanup can be done with some simplify-cfg like pass.
363 ///
364 /// If `SE` is not null, it will be updated based on the potential loop SCEVs
365 /// invalidated by this.
366 static bool unswitchTrivialBranch(Loop &L, BranchInst &BI, DominatorTree &DT,
367                                   LoopInfo &LI, ScalarEvolution *SE,
368                                   MemorySSAUpdater *MSSAU) {
369   assert(BI.isConditional() && "Can only unswitch a conditional branch!");
370   LLVM_DEBUG(dbgs() << "  Trying to unswitch branch: " << BI << "\n");
371 
372   // The loop invariant values that we want to unswitch.
373   TinyPtrVector<Value *> Invariants;
374 
375   // When true, we're fully unswitching the branch rather than just unswitching
376   // some input conditions to the branch.
377   bool FullUnswitch = false;
378 
379   if (L.isLoopInvariant(BI.getCondition())) {
380     Invariants.push_back(BI.getCondition());
381     FullUnswitch = true;
382   } else {
383     if (auto *CondInst = dyn_cast<Instruction>(BI.getCondition()))
384       Invariants = collectHomogenousInstGraphLoopInvariants(L, *CondInst, LI);
385     if (Invariants.empty())
386       // Couldn't find invariant inputs!
387       return false;
388   }
389 
390   // Check that one of the branch's successors exits, and which one.
391   bool ExitDirection = true;
392   int LoopExitSuccIdx = 0;
393   auto *LoopExitBB = BI.getSuccessor(0);
394   if (L.contains(LoopExitBB)) {
395     ExitDirection = false;
396     LoopExitSuccIdx = 1;
397     LoopExitBB = BI.getSuccessor(1);
398     if (L.contains(LoopExitBB))
399       return false;
400   }
401   auto *ContinueBB = BI.getSuccessor(1 - LoopExitSuccIdx);
402   auto *ParentBB = BI.getParent();
403   if (!areLoopExitPHIsLoopInvariant(L, *ParentBB, *LoopExitBB))
404     return false;
405 
406   // When unswitching only part of the branch's condition, we need the exit
407   // block to be reached directly from the partially unswitched input. This can
408   // be done when the exit block is along the true edge and the branch condition
409   // is a graph of `or` operations, or the exit block is along the false edge
410   // and the condition is a graph of `and` operations.
411   if (!FullUnswitch) {
412     if (ExitDirection) {
413       if (cast<Instruction>(BI.getCondition())->getOpcode() != Instruction::Or)
414         return false;
415     } else {
416       if (cast<Instruction>(BI.getCondition())->getOpcode() != Instruction::And)
417         return false;
418     }
419   }
420 
421   LLVM_DEBUG({
422     dbgs() << "    unswitching trivial invariant conditions for: " << BI
423            << "\n";
424     for (Value *Invariant : Invariants) {
425       dbgs() << "      " << *Invariant << " == true";
426       if (Invariant != Invariants.back())
427         dbgs() << " ||";
428       dbgs() << "\n";
429     }
430   });
431 
432   // If we have scalar evolutions, we need to invalidate them including this
433   // loop, the loop containing the exit block and the topmost parent loop
434   // exiting via LoopExitBB.
435   if (SE) {
436     if (Loop *ExitL = getTopMostExitingLoop(LoopExitBB, LI))
437       SE->forgetLoop(ExitL);
438     else
439       // Forget the entire nest as this exits the entire nest.
440       SE->forgetTopmostLoop(&L);
441   }
442 
443   if (MSSAU && VerifyMemorySSA)
444     MSSAU->getMemorySSA()->verifyMemorySSA();
445 
446   // Split the preheader, so that we know that there is a safe place to insert
447   // the conditional branch. We will change the preheader to have a conditional
448   // branch on LoopCond.
449   BasicBlock *OldPH = L.getLoopPreheader();
450   BasicBlock *NewPH = SplitEdge(OldPH, L.getHeader(), &DT, &LI, MSSAU);
451 
452   // Now that we have a place to insert the conditional branch, create a place
453   // to branch to: this is the exit block out of the loop that we are
454   // unswitching. We need to split this if there are other loop predecessors.
455   // Because the loop is in simplified form, *any* other predecessor is enough.
456   BasicBlock *UnswitchedBB;
457   if (FullUnswitch && LoopExitBB->getUniquePredecessor()) {
458     assert(LoopExitBB->getUniquePredecessor() == BI.getParent() &&
459            "A branch's parent isn't a predecessor!");
460     UnswitchedBB = LoopExitBB;
461   } else {
462     UnswitchedBB =
463         SplitBlock(LoopExitBB, &LoopExitBB->front(), &DT, &LI, MSSAU);
464   }
465 
466   if (MSSAU && VerifyMemorySSA)
467     MSSAU->getMemorySSA()->verifyMemorySSA();
468 
469   // Actually move the invariant uses into the unswitched position. If possible,
470   // we do this by moving the instructions, but when doing partial unswitching
471   // we do it by building a new merge of the values in the unswitched position.
472   OldPH->getTerminator()->eraseFromParent();
473   if (FullUnswitch) {
474     // If fully unswitching, we can use the existing branch instruction.
475     // Splice it into the old PH to gate reaching the new preheader and re-point
476     // its successors.
477     OldPH->getInstList().splice(OldPH->end(), BI.getParent()->getInstList(),
478                                 BI);
479     if (MSSAU) {
480       // Temporarily clone the terminator, to make MSSA update cheaper by
481       // separating "insert edge" updates from "remove edge" ones.
482       ParentBB->getInstList().push_back(BI.clone());
483     } else {
484       // Create a new unconditional branch that will continue the loop as a new
485       // terminator.
486       BranchInst::Create(ContinueBB, ParentBB);
487     }
488     BI.setSuccessor(LoopExitSuccIdx, UnswitchedBB);
489     BI.setSuccessor(1 - LoopExitSuccIdx, NewPH);
490   } else {
491     // Only unswitching a subset of inputs to the condition, so we will need to
492     // build a new branch that merges the invariant inputs.
493     if (ExitDirection)
494       assert(cast<Instruction>(BI.getCondition())->getOpcode() ==
495                  Instruction::Or &&
496              "Must have an `or` of `i1`s for the condition!");
497     else
498       assert(cast<Instruction>(BI.getCondition())->getOpcode() ==
499                  Instruction::And &&
500              "Must have an `and` of `i1`s for the condition!");
501     buildPartialUnswitchConditionalBranch(*OldPH, Invariants, ExitDirection,
502                                           *UnswitchedBB, *NewPH);
503   }
504 
505   // Update the dominator tree with the added edge.
506   DT.insertEdge(OldPH, UnswitchedBB);
507 
508   // After the dominator tree was updated with the added edge, update MemorySSA
509   // if available.
510   if (MSSAU) {
511     SmallVector<CFGUpdate, 1> Updates;
512     Updates.push_back({cfg::UpdateKind::Insert, OldPH, UnswitchedBB});
513     MSSAU->applyInsertUpdates(Updates, DT);
514   }
515 
516   // Finish updating dominator tree and memory ssa for full unswitch.
517   if (FullUnswitch) {
518     if (MSSAU) {
519       // Remove the cloned branch instruction.
520       ParentBB->getTerminator()->eraseFromParent();
521       // Create unconditional branch now.
522       BranchInst::Create(ContinueBB, ParentBB);
523       MSSAU->removeEdge(ParentBB, LoopExitBB);
524     }
525     DT.deleteEdge(ParentBB, LoopExitBB);
526   }
527 
528   if (MSSAU && VerifyMemorySSA)
529     MSSAU->getMemorySSA()->verifyMemorySSA();
530 
531   // Rewrite the relevant PHI nodes.
532   if (UnswitchedBB == LoopExitBB)
533     rewritePHINodesForUnswitchedExitBlock(*UnswitchedBB, *ParentBB, *OldPH);
534   else
535     rewritePHINodesForExitAndUnswitchedBlocks(*LoopExitBB, *UnswitchedBB,
536                                               *ParentBB, *OldPH, FullUnswitch);
537 
538   // The constant we can replace all of our invariants with inside the loop
539   // body. If any of the invariants have a value other than this the loop won't
540   // be entered.
541   ConstantInt *Replacement = ExitDirection
542                                  ? ConstantInt::getFalse(BI.getContext())
543                                  : ConstantInt::getTrue(BI.getContext());
544 
545   // Since this is an i1 condition we can also trivially replace uses of it
546   // within the loop with a constant.
547   for (Value *Invariant : Invariants)
548     replaceLoopInvariantUses(L, Invariant, *Replacement);
549 
550   // If this was full unswitching, we may have changed the nesting relationship
551   // for this loop so hoist it to its correct parent if needed.
552   if (FullUnswitch)
553     hoistLoopToNewParent(L, *NewPH, DT, LI, MSSAU, SE);
554 
555   if (MSSAU && VerifyMemorySSA)
556     MSSAU->getMemorySSA()->verifyMemorySSA();
557 
558   LLVM_DEBUG(dbgs() << "    done: unswitching trivial branch...\n");
559   ++NumTrivial;
560   ++NumBranches;
561   return true;
562 }
563 
564 /// Unswitch a trivial switch if the condition is loop invariant.
565 ///
566 /// This routine should only be called when loop code leading to the switch has
567 /// been validated as trivial (no side effects). This routine checks if the
568 /// condition is invariant and that at least one of the successors is a loop
569 /// exit. This allows us to unswitch without duplicating the loop, making it
570 /// trivial.
571 ///
572 /// If this routine fails to unswitch the switch it returns false.
573 ///
574 /// If the switch can be unswitched, this routine splits the preheader and
575 /// copies the switch above that split. If the default case is one of the
576 /// exiting cases, it copies the non-exiting cases and points them at the new
577 /// preheader. If the default case is not exiting, it copies the exiting cases
578 /// and points the default at the preheader. It preserves loop simplified form
579 /// (splitting the exit blocks as necessary). It simplifies the switch within
580 /// the loop by removing now-dead cases. If the default case is one of those
581 /// unswitched, it replaces its destination with a new basic block containing
582 /// only unreachable. Such basic blocks, while technically loop exits, are not
583 /// considered for unswitching so this is a stable transform and the same
584 /// switch will not be revisited. If after unswitching there is only a single
585 /// in-loop successor, the switch is further simplified to an unconditional
586 /// branch. Still more cleanup can be done with some simplify-cfg like pass.
587 ///
588 /// If `SE` is not null, it will be updated based on the potential loop SCEVs
589 /// invalidated by this.
590 static bool unswitchTrivialSwitch(Loop &L, SwitchInst &SI, DominatorTree &DT,
591                                   LoopInfo &LI, ScalarEvolution *SE,
592                                   MemorySSAUpdater *MSSAU) {
593   LLVM_DEBUG(dbgs() << "  Trying to unswitch switch: " << SI << "\n");
594   Value *LoopCond = SI.getCondition();
595 
596   // If this isn't switching on an invariant condition, we can't unswitch it.
597   if (!L.isLoopInvariant(LoopCond))
598     return false;
599 
600   auto *ParentBB = SI.getParent();
601 
602   // The same check must be used both for the default and the exit cases. We
603   // should never leave edges from the switch instruction to a basic block that
604   // we are unswitching, hence the condition used to determine the default case
605   // needs to also be used to populate ExitCaseIndices, which is then used to
606   // remove cases from the switch.
607   auto IsTriviallyUnswitchableExitBlock = [&](BasicBlock &BBToCheck) {
608     // BBToCheck is not an exit block if it is inside loop L.
609     if (L.contains(&BBToCheck))
610       return false;
611     // BBToCheck is not trivial to unswitch if its phis aren't loop invariant.
612     if (!areLoopExitPHIsLoopInvariant(L, *ParentBB, BBToCheck))
613       return false;
614     // We do not unswitch a block that only has an unreachable statement, as
615     // it's possible this is a previously unswitched block. Only unswitch if
616     // either the terminator is not unreachable, or, if it is, it's not the only
617     // instruction in the block.
618     auto *TI = BBToCheck.getTerminator();
619     bool isUnreachable = isa<UnreachableInst>(TI);
620     return !isUnreachable ||
621            (isUnreachable && (BBToCheck.getFirstNonPHIOrDbg() != TI));
622   };
623 
624   SmallVector<int, 4> ExitCaseIndices;
625   for (auto Case : SI.cases())
626     if (IsTriviallyUnswitchableExitBlock(*Case.getCaseSuccessor()))
627       ExitCaseIndices.push_back(Case.getCaseIndex());
628   BasicBlock *DefaultExitBB = nullptr;
629   SwitchInstProfUpdateWrapper::CaseWeightOpt DefaultCaseWeight =
630       SwitchInstProfUpdateWrapper::getSuccessorWeight(SI, 0);
631   if (IsTriviallyUnswitchableExitBlock(*SI.getDefaultDest())) {
632     DefaultExitBB = SI.getDefaultDest();
633   } else if (ExitCaseIndices.empty())
634     return false;
635 
636   LLVM_DEBUG(dbgs() << "    unswitching trivial switch...\n");
637 
638   if (MSSAU && VerifyMemorySSA)
639     MSSAU->getMemorySSA()->verifyMemorySSA();
640 
641   // We may need to invalidate SCEVs for the outermost loop reached by any of
642   // the exits.
643   Loop *OuterL = &L;
644 
645   if (DefaultExitBB) {
646     // Clear out the default destination temporarily to allow accurate
647     // predecessor lists to be examined below.
648     SI.setDefaultDest(nullptr);
649     // Check the loop containing this exit.
650     Loop *ExitL = LI.getLoopFor(DefaultExitBB);
651     if (!ExitL || ExitL->contains(OuterL))
652       OuterL = ExitL;
653   }
654 
655   // Store the exit cases into a separate data structure and remove them from
656   // the switch.
657   SmallVector<std::tuple<ConstantInt *, BasicBlock *,
658                          SwitchInstProfUpdateWrapper::CaseWeightOpt>,
659               4> ExitCases;
660   ExitCases.reserve(ExitCaseIndices.size());
661   SwitchInstProfUpdateWrapper SIW(SI);
662   // We walk the case indices backwards so that we remove the last case first
663   // and don't disrupt the earlier indices.
664   for (unsigned Index : reverse(ExitCaseIndices)) {
665     auto CaseI = SI.case_begin() + Index;
666     // Compute the outer loop from this exit.
667     Loop *ExitL = LI.getLoopFor(CaseI->getCaseSuccessor());
668     if (!ExitL || ExitL->contains(OuterL))
669       OuterL = ExitL;
670     // Save the value of this case.
671     auto W = SIW.getSuccessorWeight(CaseI->getSuccessorIndex());
672     ExitCases.emplace_back(CaseI->getCaseValue(), CaseI->getCaseSuccessor(), W);
673     // Delete the unswitched cases.
674     SIW.removeCase(CaseI);
675   }
676 
677   if (SE) {
678     if (OuterL)
679       SE->forgetLoop(OuterL);
680     else
681       SE->forgetTopmostLoop(&L);
682   }
683 
684   // Check if after this all of the remaining cases point at the same
685   // successor.
686   BasicBlock *CommonSuccBB = nullptr;
687   if (SI.getNumCases() > 0 &&
688       std::all_of(std::next(SI.case_begin()), SI.case_end(),
689                   [&SI](const SwitchInst::CaseHandle &Case) {
690                     return Case.getCaseSuccessor() ==
691                            SI.case_begin()->getCaseSuccessor();
692                   }))
693     CommonSuccBB = SI.case_begin()->getCaseSuccessor();
694   if (!DefaultExitBB) {
695     // If we're not unswitching the default, we need it to match any cases to
696     // have a common successor or if we have no cases it is the common
697     // successor.
698     if (SI.getNumCases() == 0)
699       CommonSuccBB = SI.getDefaultDest();
700     else if (SI.getDefaultDest() != CommonSuccBB)
701       CommonSuccBB = nullptr;
702   }
703 
704   // Split the preheader, so that we know that there is a safe place to insert
705   // the switch.
706   BasicBlock *OldPH = L.getLoopPreheader();
707   BasicBlock *NewPH = SplitEdge(OldPH, L.getHeader(), &DT, &LI, MSSAU);
708   OldPH->getTerminator()->eraseFromParent();
709 
710   // Now add the unswitched switch.
711   auto *NewSI = SwitchInst::Create(LoopCond, NewPH, ExitCases.size(), OldPH);
712   SwitchInstProfUpdateWrapper NewSIW(*NewSI);
713 
714   // Rewrite the IR for the unswitched basic blocks. This requires two steps.
715   // First, we split any exit blocks with remaining in-loop predecessors. Then
716   // we update the PHIs in one of two ways depending on if there was a split.
717   // We walk in reverse so that we split in the same order as the cases
718   // appeared. This is purely for convenience of reading the resulting IR, but
719   // it doesn't cost anything really.
720   SmallPtrSet<BasicBlock *, 2> UnswitchedExitBBs;
721   SmallDenseMap<BasicBlock *, BasicBlock *, 2> SplitExitBBMap;
722   // Handle the default exit if necessary.
723   // FIXME: It'd be great if we could merge this with the loop below but LLVM's
724   // ranges aren't quite powerful enough yet.
725   if (DefaultExitBB) {
726     if (pred_empty(DefaultExitBB)) {
727       UnswitchedExitBBs.insert(DefaultExitBB);
728       rewritePHINodesForUnswitchedExitBlock(*DefaultExitBB, *ParentBB, *OldPH);
729     } else {
730       auto *SplitBB =
731           SplitBlock(DefaultExitBB, &DefaultExitBB->front(), &DT, &LI, MSSAU);
732       rewritePHINodesForExitAndUnswitchedBlocks(*DefaultExitBB, *SplitBB,
733                                                 *ParentBB, *OldPH,
734                                                 /*FullUnswitch*/ true);
735       DefaultExitBB = SplitExitBBMap[DefaultExitBB] = SplitBB;
736     }
737   }
738   // Note that we must use a reference in the for loop so that we update the
739   // container.
740   for (auto &ExitCase : reverse(ExitCases)) {
741     // Grab a reference to the exit block in the pair so that we can update it.
742     BasicBlock *ExitBB = std::get<1>(ExitCase);
743 
744     // If this case is the last edge into the exit block, we can simply reuse it
745     // as it will no longer be a loop exit. No mapping necessary.
746     if (pred_empty(ExitBB)) {
747       // Only rewrite once.
748       if (UnswitchedExitBBs.insert(ExitBB).second)
749         rewritePHINodesForUnswitchedExitBlock(*ExitBB, *ParentBB, *OldPH);
750       continue;
751     }
752 
753     // Otherwise we need to split the exit block so that we retain an exit
754     // block from the loop and a target for the unswitched condition.
755     BasicBlock *&SplitExitBB = SplitExitBBMap[ExitBB];
756     if (!SplitExitBB) {
757       // If this is the first time we see this, do the split and remember it.
758       SplitExitBB = SplitBlock(ExitBB, &ExitBB->front(), &DT, &LI, MSSAU);
759       rewritePHINodesForExitAndUnswitchedBlocks(*ExitBB, *SplitExitBB,
760                                                 *ParentBB, *OldPH,
761                                                 /*FullUnswitch*/ true);
762     }
763     // Update the case pair to point to the split block.
764     std::get<1>(ExitCase) = SplitExitBB;
765   }
766 
767   // Now add the unswitched cases. We do this in reverse order as we built them
768   // in reverse order.
769   for (auto &ExitCase : reverse(ExitCases)) {
770     ConstantInt *CaseVal = std::get<0>(ExitCase);
771     BasicBlock *UnswitchedBB = std::get<1>(ExitCase);
772 
773     NewSIW.addCase(CaseVal, UnswitchedBB, std::get<2>(ExitCase));
774   }
775 
776   // If the default was unswitched, re-point it and add explicit cases for
777   // entering the loop.
778   if (DefaultExitBB) {
779     NewSIW->setDefaultDest(DefaultExitBB);
780     NewSIW.setSuccessorWeight(0, DefaultCaseWeight);
781 
782     // We removed all the exit cases, so we just copy the cases to the
783     // unswitched switch.
784     for (const auto &Case : SI.cases())
785       NewSIW.addCase(Case.getCaseValue(), NewPH,
786                      SIW.getSuccessorWeight(Case.getSuccessorIndex()));
787   } else if (DefaultCaseWeight) {
788     // We have to set branch weight of the default case.
789     uint64_t SW = *DefaultCaseWeight;
790     for (const auto &Case : SI.cases()) {
791       auto W = SIW.getSuccessorWeight(Case.getSuccessorIndex());
792       assert(W &&
793              "case weight must be defined as default case weight is defined");
794       SW += *W;
795     }
796     NewSIW.setSuccessorWeight(0, SW);
797   }
798 
799   // If we ended up with a common successor for every path through the switch
800   // after unswitching, rewrite it to an unconditional branch to make it easy
801   // to recognize. Otherwise we potentially have to recognize the default case
802   // pointing at unreachable and other complexity.
803   if (CommonSuccBB) {
804     BasicBlock *BB = SI.getParent();
805     // We may have had multiple edges to this common successor block, so remove
806     // them as predecessors. We skip the first one, either the default or the
807     // actual first case.
808     bool SkippedFirst = DefaultExitBB == nullptr;
809     for (auto Case : SI.cases()) {
810       assert(Case.getCaseSuccessor() == CommonSuccBB &&
811              "Non-common successor!");
812       (void)Case;
813       if (!SkippedFirst) {
814         SkippedFirst = true;
815         continue;
816       }
817       CommonSuccBB->removePredecessor(BB,
818                                       /*KeepOneInputPHIs*/ true);
819     }
820     // Now nuke the switch and replace it with a direct branch.
821     SIW.eraseFromParent();
822     BranchInst::Create(CommonSuccBB, BB);
823   } else if (DefaultExitBB) {
824     assert(SI.getNumCases() > 0 &&
825            "If we had no cases we'd have a common successor!");
826     // Move the last case to the default successor. This is valid as if the
827     // default got unswitched it cannot be reached. This has the advantage of
828     // being simple and keeping the number of edges from this switch to
829     // successors the same, and avoiding any PHI update complexity.
830     auto LastCaseI = std::prev(SI.case_end());
831 
832     SI.setDefaultDest(LastCaseI->getCaseSuccessor());
833     SIW.setSuccessorWeight(
834         0, SIW.getSuccessorWeight(LastCaseI->getSuccessorIndex()));
835     SIW.removeCase(LastCaseI);
836   }
837 
838   // Walk the unswitched exit blocks and the unswitched split blocks and update
839   // the dominator tree based on the CFG edits. While we are walking unordered
840   // containers here, the API for applyUpdates takes an unordered list of
841   // updates and requires them to not contain duplicates.
842   SmallVector<DominatorTree::UpdateType, 4> DTUpdates;
843   for (auto *UnswitchedExitBB : UnswitchedExitBBs) {
844     DTUpdates.push_back({DT.Delete, ParentBB, UnswitchedExitBB});
845     DTUpdates.push_back({DT.Insert, OldPH, UnswitchedExitBB});
846   }
847   for (auto SplitUnswitchedPair : SplitExitBBMap) {
848     DTUpdates.push_back({DT.Delete, ParentBB, SplitUnswitchedPair.first});
849     DTUpdates.push_back({DT.Insert, OldPH, SplitUnswitchedPair.second});
850   }
851   DT.applyUpdates(DTUpdates);
852 
853   if (MSSAU) {
854     MSSAU->applyUpdates(DTUpdates, DT);
855     if (VerifyMemorySSA)
856       MSSAU->getMemorySSA()->verifyMemorySSA();
857   }
858 
859   assert(DT.verify(DominatorTree::VerificationLevel::Fast));
860 
861   // We may have changed the nesting relationship for this loop so hoist it to
862   // its correct parent if needed.
863   hoistLoopToNewParent(L, *NewPH, DT, LI, MSSAU, SE);
864 
865   if (MSSAU && VerifyMemorySSA)
866     MSSAU->getMemorySSA()->verifyMemorySSA();
867 
868   ++NumTrivial;
869   ++NumSwitches;
870   LLVM_DEBUG(dbgs() << "    done: unswitching trivial switch...\n");
871   return true;
872 }
873 
874 /// This routine scans the loop to find a branch or switch which occurs before
875 /// any side effects occur. These can potentially be unswitched without
876 /// duplicating the loop. If a branch or switch is successfully unswitched the
877 /// scanning continues to see if subsequent branches or switches have become
878 /// trivial. Once all trivial candidates have been unswitched, this routine
879 /// returns.
880 ///
881 /// The return value indicates whether anything was unswitched (and therefore
882 /// changed).
883 ///
884 /// If `SE` is not null, it will be updated based on the potential loop SCEVs
885 /// invalidated by this.
886 static bool unswitchAllTrivialConditions(Loop &L, DominatorTree &DT,
887                                          LoopInfo &LI, ScalarEvolution *SE,
888                                          MemorySSAUpdater *MSSAU) {
889   bool Changed = false;
890 
891   // If loop header has only one reachable successor we should keep looking for
892   // trivial condition candidates in the successor as well. An alternative is
893   // to constant fold conditions and merge successors into loop header (then we
894   // only need to check header's terminator). The reason for not doing this in
895   // LoopUnswitch pass is that it could potentially break LoopPassManager's
896   // invariants. Folding dead branches could either eliminate the current loop
897   // or make other loops unreachable. LCSSA form might also not be preserved
898   // after deleting branches. The following code keeps traversing loop header's
899   // successors until it finds the trivial condition candidate (condition that
900   // is not a constant). Since unswitching generates branches with constant
901   // conditions, this scenario could be very common in practice.
902   BasicBlock *CurrentBB = L.getHeader();
903   SmallPtrSet<BasicBlock *, 8> Visited;
904   Visited.insert(CurrentBB);
905   do {
906     // Check if there are any side-effecting instructions (e.g. stores, calls,
907     // volatile loads) in the part of the loop that the code *would* execute
908     // without unswitching.
909     if (MSSAU) // Possible early exit with MSSA
910       if (auto *Defs = MSSAU->getMemorySSA()->getBlockDefs(CurrentBB))
911         if (!isa<MemoryPhi>(*Defs->begin()) || (++Defs->begin() != Defs->end()))
912           return Changed;
913     if (llvm::any_of(*CurrentBB,
914                      [](Instruction &I) { return I.mayHaveSideEffects(); }))
915       return Changed;
916 
917     Instruction *CurrentTerm = CurrentBB->getTerminator();
918 
919     if (auto *SI = dyn_cast<SwitchInst>(CurrentTerm)) {
920       // Don't bother trying to unswitch past a switch with a constant
921       // condition. This should be removed prior to running this pass by
922       // simplify-cfg.
923       if (isa<Constant>(SI->getCondition()))
924         return Changed;
925 
926       if (!unswitchTrivialSwitch(L, *SI, DT, LI, SE, MSSAU))
927         // Couldn't unswitch this one so we're done.
928         return Changed;
929 
930       // Mark that we managed to unswitch something.
931       Changed = true;
932 
933       // If unswitching turned the terminator into an unconditional branch then
934       // we can continue. The unswitching logic specifically works to fold any
935       // cases it can into an unconditional branch to make it easier to
936       // recognize here.
937       auto *BI = dyn_cast<BranchInst>(CurrentBB->getTerminator());
938       if (!BI || BI->isConditional())
939         return Changed;
940 
941       CurrentBB = BI->getSuccessor(0);
942       continue;
943     }
944 
945     auto *BI = dyn_cast<BranchInst>(CurrentTerm);
946     if (!BI)
947       // We do not understand other terminator instructions.
948       return Changed;
949 
950     // Don't bother trying to unswitch past an unconditional branch or a branch
951     // with a constant value. These should be removed by simplify-cfg prior to
952     // running this pass.
953     if (!BI->isConditional() || isa<Constant>(BI->getCondition()))
954       return Changed;
955 
956     // Found a trivial condition candidate: non-foldable conditional branch. If
957     // we fail to unswitch this, we can't do anything else that is trivial.
958     if (!unswitchTrivialBranch(L, *BI, DT, LI, SE, MSSAU))
959       return Changed;
960 
961     // Mark that we managed to unswitch something.
962     Changed = true;
963 
964     // If we only unswitched some of the conditions feeding the branch, we won't
965     // have collapsed it to a single successor.
966     BI = cast<BranchInst>(CurrentBB->getTerminator());
967     if (BI->isConditional())
968       return Changed;
969 
970     // Follow the newly unconditional branch into its successor.
971     CurrentBB = BI->getSuccessor(0);
972 
973     // When continuing, if we exit the loop or reach a previous visited block,
974     // then we can not reach any trivial condition candidates (unfoldable
975     // branch instructions or switch instructions) and no unswitch can happen.
976   } while (L.contains(CurrentBB) && Visited.insert(CurrentBB).second);
977 
978   return Changed;
979 }
980 
981 /// Build the cloned blocks for an unswitched copy of the given loop.
982 ///
983 /// The cloned blocks are inserted before the loop preheader (`LoopPH`) and
984 /// after the split block (`SplitBB`) that will be used to select between the
985 /// cloned and original loop.
986 ///
987 /// This routine handles cloning all of the necessary loop blocks and exit
988 /// blocks including rewriting their instructions and the relevant PHI nodes.
989 /// Any loop blocks or exit blocks which are dominated by a different successor
990 /// than the one for this clone of the loop blocks can be trivially skipped. We
991 /// use the `DominatingSucc` map to determine whether a block satisfies that
992 /// property with a simple map lookup.
993 ///
994 /// It also correctly creates the unconditional branch in the cloned
995 /// unswitched parent block to only point at the unswitched successor.
996 ///
997 /// This does not handle most of the necessary updates to `LoopInfo`. Only exit
998 /// block splitting is correctly reflected in `LoopInfo`, essentially all of
999 /// the cloned blocks (and their loops) are left without full `LoopInfo`
1000 /// updates. This also doesn't fully update `DominatorTree`. It adds the cloned
1001 /// blocks to them but doesn't create the cloned `DominatorTree` structure and
1002 /// instead the caller must recompute an accurate DT. It *does* correctly
1003 /// update the `AssumptionCache` provided in `AC`.
1004 static BasicBlock *buildClonedLoopBlocks(
1005     Loop &L, BasicBlock *LoopPH, BasicBlock *SplitBB,
1006     ArrayRef<BasicBlock *> ExitBlocks, BasicBlock *ParentBB,
1007     BasicBlock *UnswitchedSuccBB, BasicBlock *ContinueSuccBB,
1008     const SmallDenseMap<BasicBlock *, BasicBlock *, 16> &DominatingSucc,
1009     ValueToValueMapTy &VMap,
1010     SmallVectorImpl<DominatorTree::UpdateType> &DTUpdates, AssumptionCache &AC,
1011     DominatorTree &DT, LoopInfo &LI, MemorySSAUpdater *MSSAU) {
1012   SmallVector<BasicBlock *, 4> NewBlocks;
1013   NewBlocks.reserve(L.getNumBlocks() + ExitBlocks.size());
1014 
1015   // We will need to clone a bunch of blocks, wrap up the clone operation in
1016   // a helper.
1017   auto CloneBlock = [&](BasicBlock *OldBB) {
1018     // Clone the basic block and insert it before the new preheader.
1019     BasicBlock *NewBB = CloneBasicBlock(OldBB, VMap, ".us", OldBB->getParent());
1020     NewBB->moveBefore(LoopPH);
1021 
1022     // Record this block and the mapping.
1023     NewBlocks.push_back(NewBB);
1024     VMap[OldBB] = NewBB;
1025 
1026     return NewBB;
1027   };
1028 
1029   // We skip cloning blocks when they have a dominating succ that is not the
1030   // succ we are cloning for.
1031   auto SkipBlock = [&](BasicBlock *BB) {
1032     auto It = DominatingSucc.find(BB);
1033     return It != DominatingSucc.end() && It->second != UnswitchedSuccBB;
1034   };
1035 
1036   // First, clone the preheader.
1037   auto *ClonedPH = CloneBlock(LoopPH);
1038 
1039   // Then clone all the loop blocks, skipping the ones that aren't necessary.
1040   for (auto *LoopBB : L.blocks())
1041     if (!SkipBlock(LoopBB))
1042       CloneBlock(LoopBB);
1043 
1044   // Split all the loop exit edges so that when we clone the exit blocks, if
1045   // any of the exit blocks are *also* a preheader for some other loop, we
1046   // don't create multiple predecessors entering the loop header.
1047   for (auto *ExitBB : ExitBlocks) {
1048     if (SkipBlock(ExitBB))
1049       continue;
1050 
1051     // When we are going to clone an exit, we don't need to clone all the
1052     // instructions in the exit block and we want to ensure we have an easy
1053     // place to merge the CFG, so split the exit first. This is always safe to
1054     // do because there cannot be any non-loop predecessors of a loop exit in
1055     // loop simplified form.
1056     auto *MergeBB = SplitBlock(ExitBB, &ExitBB->front(), &DT, &LI, MSSAU);
1057 
1058     // Rearrange the names to make it easier to write test cases by having the
1059     // exit block carry the suffix rather than the merge block carrying the
1060     // suffix.
1061     MergeBB->takeName(ExitBB);
1062     ExitBB->setName(Twine(MergeBB->getName()) + ".split");
1063 
1064     // Now clone the original exit block.
1065     auto *ClonedExitBB = CloneBlock(ExitBB);
1066     assert(ClonedExitBB->getTerminator()->getNumSuccessors() == 1 &&
1067            "Exit block should have been split to have one successor!");
1068     assert(ClonedExitBB->getTerminator()->getSuccessor(0) == MergeBB &&
1069            "Cloned exit block has the wrong successor!");
1070 
1071     // Remap any cloned instructions and create a merge phi node for them.
1072     for (auto ZippedInsts : llvm::zip_first(
1073              llvm::make_range(ExitBB->begin(), std::prev(ExitBB->end())),
1074              llvm::make_range(ClonedExitBB->begin(),
1075                               std::prev(ClonedExitBB->end())))) {
1076       Instruction &I = std::get<0>(ZippedInsts);
1077       Instruction &ClonedI = std::get<1>(ZippedInsts);
1078 
1079       // The only instructions in the exit block should be PHI nodes and
1080       // potentially a landing pad.
1081       assert(
1082           (isa<PHINode>(I) || isa<LandingPadInst>(I) || isa<CatchPadInst>(I)) &&
1083           "Bad instruction in exit block!");
1084       // We should have a value map between the instruction and its clone.
1085       assert(VMap.lookup(&I) == &ClonedI && "Mismatch in the value map!");
1086 
1087       auto *MergePN =
1088           PHINode::Create(I.getType(), /*NumReservedValues*/ 2, ".us-phi",
1089                           &*MergeBB->getFirstInsertionPt());
1090       I.replaceAllUsesWith(MergePN);
1091       MergePN->addIncoming(&I, ExitBB);
1092       MergePN->addIncoming(&ClonedI, ClonedExitBB);
1093     }
1094   }
1095 
1096   // Rewrite the instructions in the cloned blocks to refer to the instructions
1097   // in the cloned blocks. We have to do this as a second pass so that we have
1098   // everything available. Also, we have inserted new instructions which may
1099   // include assume intrinsics, so we update the assumption cache while
1100   // processing this.
1101   for (auto *ClonedBB : NewBlocks)
1102     for (Instruction &I : *ClonedBB) {
1103       RemapInstruction(&I, VMap,
1104                        RF_NoModuleLevelChanges | RF_IgnoreMissingLocals);
1105       if (auto *II = dyn_cast<IntrinsicInst>(&I))
1106         if (II->getIntrinsicID() == Intrinsic::assume)
1107           AC.registerAssumption(II);
1108     }
1109 
1110   // Update any PHI nodes in the cloned successors of the skipped blocks to not
1111   // have spurious incoming values.
1112   for (auto *LoopBB : L.blocks())
1113     if (SkipBlock(LoopBB))
1114       for (auto *SuccBB : successors(LoopBB))
1115         if (auto *ClonedSuccBB = cast_or_null<BasicBlock>(VMap.lookup(SuccBB)))
1116           for (PHINode &PN : ClonedSuccBB->phis())
1117             PN.removeIncomingValue(LoopBB, /*DeletePHIIfEmpty*/ false);
1118 
1119   // Remove the cloned parent as a predecessor of any successor we ended up
1120   // cloning other than the unswitched one.
1121   auto *ClonedParentBB = cast<BasicBlock>(VMap.lookup(ParentBB));
1122   for (auto *SuccBB : successors(ParentBB)) {
1123     if (SuccBB == UnswitchedSuccBB)
1124       continue;
1125 
1126     auto *ClonedSuccBB = cast_or_null<BasicBlock>(VMap.lookup(SuccBB));
1127     if (!ClonedSuccBB)
1128       continue;
1129 
1130     ClonedSuccBB->removePredecessor(ClonedParentBB,
1131                                     /*KeepOneInputPHIs*/ true);
1132   }
1133 
1134   // Replace the cloned branch with an unconditional branch to the cloned
1135   // unswitched successor.
1136   auto *ClonedSuccBB = cast<BasicBlock>(VMap.lookup(UnswitchedSuccBB));
1137   ClonedParentBB->getTerminator()->eraseFromParent();
1138   BranchInst::Create(ClonedSuccBB, ClonedParentBB);
1139 
1140   // If there are duplicate entries in the PHI nodes because of multiple edges
1141   // to the unswitched successor, we need to nuke all but one as we replaced it
1142   // with a direct branch.
1143   for (PHINode &PN : ClonedSuccBB->phis()) {
1144     bool Found = false;
1145     // Loop over the incoming operands backwards so we can easily delete as we
1146     // go without invalidating the index.
1147     for (int i = PN.getNumOperands() - 1; i >= 0; --i) {
1148       if (PN.getIncomingBlock(i) != ClonedParentBB)
1149         continue;
1150       if (!Found) {
1151         Found = true;
1152         continue;
1153       }
1154       PN.removeIncomingValue(i, /*DeletePHIIfEmpty*/ false);
1155     }
1156   }
1157 
1158   // Record the domtree updates for the new blocks.
1159   SmallPtrSet<BasicBlock *, 4> SuccSet;
1160   for (auto *ClonedBB : NewBlocks) {
1161     for (auto *SuccBB : successors(ClonedBB))
1162       if (SuccSet.insert(SuccBB).second)
1163         DTUpdates.push_back({DominatorTree::Insert, ClonedBB, SuccBB});
1164     SuccSet.clear();
1165   }
1166 
1167   return ClonedPH;
1168 }
1169 
1170 /// Recursively clone the specified loop and all of its children.
1171 ///
1172 /// The target parent loop for the clone should be provided, or can be null if
1173 /// the clone is a top-level loop. While cloning, all the blocks are mapped
1174 /// with the provided value map. The entire original loop must be present in
1175 /// the value map. The cloned loop is returned.
1176 static Loop *cloneLoopNest(Loop &OrigRootL, Loop *RootParentL,
1177                            const ValueToValueMapTy &VMap, LoopInfo &LI) {
1178   auto AddClonedBlocksToLoop = [&](Loop &OrigL, Loop &ClonedL) {
1179     assert(ClonedL.getBlocks().empty() && "Must start with an empty loop!");
1180     ClonedL.reserveBlocks(OrigL.getNumBlocks());
1181     for (auto *BB : OrigL.blocks()) {
1182       auto *ClonedBB = cast<BasicBlock>(VMap.lookup(BB));
1183       ClonedL.addBlockEntry(ClonedBB);
1184       if (LI.getLoopFor(BB) == &OrigL)
1185         LI.changeLoopFor(ClonedBB, &ClonedL);
1186     }
1187   };
1188 
1189   // We specially handle the first loop because it may get cloned into
1190   // a different parent and because we most commonly are cloning leaf loops.
1191   Loop *ClonedRootL = LI.AllocateLoop();
1192   if (RootParentL)
1193     RootParentL->addChildLoop(ClonedRootL);
1194   else
1195     LI.addTopLevelLoop(ClonedRootL);
1196   AddClonedBlocksToLoop(OrigRootL, *ClonedRootL);
1197 
1198   if (OrigRootL.empty())
1199     return ClonedRootL;
1200 
1201   // If we have a nest, we can quickly clone the entire loop nest using an
1202   // iterative approach because it is a tree. We keep the cloned parent in the
1203   // data structure to avoid repeatedly querying through a map to find it.
1204   SmallVector<std::pair<Loop *, Loop *>, 16> LoopsToClone;
1205   // Build up the loops to clone in reverse order as we'll clone them from the
1206   // back.
1207   for (Loop *ChildL : llvm::reverse(OrigRootL))
1208     LoopsToClone.push_back({ClonedRootL, ChildL});
1209   do {
1210     Loop *ClonedParentL, *L;
1211     std::tie(ClonedParentL, L) = LoopsToClone.pop_back_val();
1212     Loop *ClonedL = LI.AllocateLoop();
1213     ClonedParentL->addChildLoop(ClonedL);
1214     AddClonedBlocksToLoop(*L, *ClonedL);
1215     for (Loop *ChildL : llvm::reverse(*L))
1216       LoopsToClone.push_back({ClonedL, ChildL});
1217   } while (!LoopsToClone.empty());
1218 
1219   return ClonedRootL;
1220 }
1221 
1222 /// Build the cloned loops of an original loop from unswitching.
1223 ///
1224 /// Because unswitching simplifies the CFG of the loop, this isn't a trivial
1225 /// operation. We need to re-verify that there even is a loop (as the backedge
1226 /// may not have been cloned), and even if there are remaining backedges the
1227 /// backedge set may be different. However, we know that each child loop is
1228 /// undisturbed, we only need to find where to place each child loop within
1229 /// either any parent loop or within a cloned version of the original loop.
1230 ///
1231 /// Because child loops may end up cloned outside of any cloned version of the
1232 /// original loop, multiple cloned sibling loops may be created. All of them
1233 /// are returned so that the newly introduced loop nest roots can be
1234 /// identified.
1235 static void buildClonedLoops(Loop &OrigL, ArrayRef<BasicBlock *> ExitBlocks,
1236                              const ValueToValueMapTy &VMap, LoopInfo &LI,
1237                              SmallVectorImpl<Loop *> &NonChildClonedLoops) {
1238   Loop *ClonedL = nullptr;
1239 
1240   auto *OrigPH = OrigL.getLoopPreheader();
1241   auto *OrigHeader = OrigL.getHeader();
1242 
1243   auto *ClonedPH = cast<BasicBlock>(VMap.lookup(OrigPH));
1244   auto *ClonedHeader = cast<BasicBlock>(VMap.lookup(OrigHeader));
1245 
1246   // We need to know the loops of the cloned exit blocks to even compute the
1247   // accurate parent loop. If we only clone exits to some parent of the
1248   // original parent, we want to clone into that outer loop. We also keep track
1249   // of the loops that our cloned exit blocks participate in.
1250   Loop *ParentL = nullptr;
1251   SmallVector<BasicBlock *, 4> ClonedExitsInLoops;
1252   SmallDenseMap<BasicBlock *, Loop *, 16> ExitLoopMap;
1253   ClonedExitsInLoops.reserve(ExitBlocks.size());
1254   for (auto *ExitBB : ExitBlocks)
1255     if (auto *ClonedExitBB = cast_or_null<BasicBlock>(VMap.lookup(ExitBB)))
1256       if (Loop *ExitL = LI.getLoopFor(ExitBB)) {
1257         ExitLoopMap[ClonedExitBB] = ExitL;
1258         ClonedExitsInLoops.push_back(ClonedExitBB);
1259         if (!ParentL || (ParentL != ExitL && ParentL->contains(ExitL)))
1260           ParentL = ExitL;
1261       }
1262   assert((!ParentL || ParentL == OrigL.getParentLoop() ||
1263           ParentL->contains(OrigL.getParentLoop())) &&
1264          "The computed parent loop should always contain (or be) the parent of "
1265          "the original loop.");
1266 
1267   // We build the set of blocks dominated by the cloned header from the set of
1268   // cloned blocks out of the original loop. While not all of these will
1269   // necessarily be in the cloned loop, it is enough to establish that they
1270   // aren't in unreachable cycles, etc.
1271   SmallSetVector<BasicBlock *, 16> ClonedLoopBlocks;
1272   for (auto *BB : OrigL.blocks())
1273     if (auto *ClonedBB = cast_or_null<BasicBlock>(VMap.lookup(BB)))
1274       ClonedLoopBlocks.insert(ClonedBB);
1275 
1276   // Rebuild the set of blocks that will end up in the cloned loop. We may have
1277   // skipped cloning some region of this loop which can in turn skip some of
1278   // the backedges so we have to rebuild the blocks in the loop based on the
1279   // backedges that remain after cloning.
1280   SmallVector<BasicBlock *, 16> Worklist;
1281   SmallPtrSet<BasicBlock *, 16> BlocksInClonedLoop;
1282   for (auto *Pred : predecessors(ClonedHeader)) {
1283     // The only possible non-loop header predecessor is the preheader because
1284     // we know we cloned the loop in simplified form.
1285     if (Pred == ClonedPH)
1286       continue;
1287 
1288     // Because the loop was in simplified form, the only non-loop predecessor
1289     // should be the preheader.
1290     assert(ClonedLoopBlocks.count(Pred) && "Found a predecessor of the loop "
1291                                            "header other than the preheader "
1292                                            "that is not part of the loop!");
1293 
1294     // Insert this block into the loop set and on the first visit (and if it
1295     // isn't the header we're currently walking) put it into the worklist to
1296     // recurse through.
1297     if (BlocksInClonedLoop.insert(Pred).second && Pred != ClonedHeader)
1298       Worklist.push_back(Pred);
1299   }
1300 
1301   // If we had any backedges then there *is* a cloned loop. Put the header into
1302   // the loop set and then walk the worklist backwards to find all the blocks
1303   // that remain within the loop after cloning.
1304   if (!BlocksInClonedLoop.empty()) {
1305     BlocksInClonedLoop.insert(ClonedHeader);
1306 
1307     while (!Worklist.empty()) {
1308       BasicBlock *BB = Worklist.pop_back_val();
1309       assert(BlocksInClonedLoop.count(BB) &&
1310              "Didn't put block into the loop set!");
1311 
1312       // Insert any predecessors that are in the possible set into the cloned
1313       // set, and if the insert is successful, add them to the worklist. Note
1314       // that we filter on the blocks that are definitely reachable via the
1315       // backedge to the loop header so we may prune out dead code within the
1316       // cloned loop.
1317       for (auto *Pred : predecessors(BB))
1318         if (ClonedLoopBlocks.count(Pred) &&
1319             BlocksInClonedLoop.insert(Pred).second)
1320           Worklist.push_back(Pred);
1321     }
1322 
1323     ClonedL = LI.AllocateLoop();
1324     if (ParentL) {
1325       ParentL->addBasicBlockToLoop(ClonedPH, LI);
1326       ParentL->addChildLoop(ClonedL);
1327     } else {
1328       LI.addTopLevelLoop(ClonedL);
1329     }
1330     NonChildClonedLoops.push_back(ClonedL);
1331 
1332     ClonedL->reserveBlocks(BlocksInClonedLoop.size());
1333     // We don't want to just add the cloned loop blocks based on how we
1334     // discovered them. The original order of blocks was carefully built in
1335     // a way that doesn't rely on predecessor ordering. Rather than re-invent
1336     // that logic, we just re-walk the original blocks (and those of the child
1337     // loops) and filter them as we add them into the cloned loop.
1338     for (auto *BB : OrigL.blocks()) {
1339       auto *ClonedBB = cast_or_null<BasicBlock>(VMap.lookup(BB));
1340       if (!ClonedBB || !BlocksInClonedLoop.count(ClonedBB))
1341         continue;
1342 
1343       // Directly add the blocks that are only in this loop.
1344       if (LI.getLoopFor(BB) == &OrigL) {
1345         ClonedL->addBasicBlockToLoop(ClonedBB, LI);
1346         continue;
1347       }
1348 
1349       // We want to manually add it to this loop and parents.
1350       // Registering it with LoopInfo will happen when we clone the top
1351       // loop for this block.
1352       for (Loop *PL = ClonedL; PL; PL = PL->getParentLoop())
1353         PL->addBlockEntry(ClonedBB);
1354     }
1355 
1356     // Now add each child loop whose header remains within the cloned loop. All
1357     // of the blocks within the loop must satisfy the same constraints as the
1358     // header so once we pass the header checks we can just clone the entire
1359     // child loop nest.
1360     for (Loop *ChildL : OrigL) {
1361       auto *ClonedChildHeader =
1362           cast_or_null<BasicBlock>(VMap.lookup(ChildL->getHeader()));
1363       if (!ClonedChildHeader || !BlocksInClonedLoop.count(ClonedChildHeader))
1364         continue;
1365 
1366 #ifndef NDEBUG
1367       // We should never have a cloned child loop header but fail to have
1368       // all of the blocks for that child loop.
1369       for (auto *ChildLoopBB : ChildL->blocks())
1370         assert(BlocksInClonedLoop.count(
1371                    cast<BasicBlock>(VMap.lookup(ChildLoopBB))) &&
1372                "Child cloned loop has a header within the cloned outer "
1373                "loop but not all of its blocks!");
1374 #endif
1375 
1376       cloneLoopNest(*ChildL, ClonedL, VMap, LI);
1377     }
1378   }
1379 
1380   // Now that we've handled all the components of the original loop that were
1381   // cloned into a new loop, we still need to handle anything from the original
1382   // loop that wasn't in a cloned loop.
1383 
1384   // Figure out what blocks are left to place within any loop nest containing
1385   // the unswitched loop. If we never formed a loop, the cloned PH is one of
1386   // them.
1387   SmallPtrSet<BasicBlock *, 16> UnloopedBlockSet;
1388   if (BlocksInClonedLoop.empty())
1389     UnloopedBlockSet.insert(ClonedPH);
1390   for (auto *ClonedBB : ClonedLoopBlocks)
1391     if (!BlocksInClonedLoop.count(ClonedBB))
1392       UnloopedBlockSet.insert(ClonedBB);
1393 
1394   // Copy the cloned exits and sort them in ascending loop depth, we'll work
1395   // backwards across these to process them inside out. The order shouldn't
1396   // matter as we're just trying to build up the map from inside-out; we use
1397   // the map in a more stably ordered way below.
1398   auto OrderedClonedExitsInLoops = ClonedExitsInLoops;
1399   llvm::sort(OrderedClonedExitsInLoops, [&](BasicBlock *LHS, BasicBlock *RHS) {
1400     return ExitLoopMap.lookup(LHS)->getLoopDepth() <
1401            ExitLoopMap.lookup(RHS)->getLoopDepth();
1402   });
1403 
1404   // Populate the existing ExitLoopMap with everything reachable from each
1405   // exit, starting from the inner most exit.
1406   while (!UnloopedBlockSet.empty() && !OrderedClonedExitsInLoops.empty()) {
1407     assert(Worklist.empty() && "Didn't clear worklist!");
1408 
1409     BasicBlock *ExitBB = OrderedClonedExitsInLoops.pop_back_val();
1410     Loop *ExitL = ExitLoopMap.lookup(ExitBB);
1411 
1412     // Walk the CFG back until we hit the cloned PH adding everything reachable
1413     // and in the unlooped set to this exit block's loop.
1414     Worklist.push_back(ExitBB);
1415     do {
1416       BasicBlock *BB = Worklist.pop_back_val();
1417       // We can stop recursing at the cloned preheader (if we get there).
1418       if (BB == ClonedPH)
1419         continue;
1420 
1421       for (BasicBlock *PredBB : predecessors(BB)) {
1422         // If this pred has already been moved to our set or is part of some
1423         // (inner) loop, no update needed.
1424         if (!UnloopedBlockSet.erase(PredBB)) {
1425           assert(
1426               (BlocksInClonedLoop.count(PredBB) || ExitLoopMap.count(PredBB)) &&
1427               "Predecessor not mapped to a loop!");
1428           continue;
1429         }
1430 
1431         // We just insert into the loop set here. We'll add these blocks to the
1432         // exit loop after we build up the set in an order that doesn't rely on
1433         // predecessor order (which in turn relies on use list order).
1434         bool Inserted = ExitLoopMap.insert({PredBB, ExitL}).second;
1435         (void)Inserted;
1436         assert(Inserted && "Should only visit an unlooped block once!");
1437 
1438         // And recurse through to its predecessors.
1439         Worklist.push_back(PredBB);
1440       }
1441     } while (!Worklist.empty());
1442   }
1443 
1444   // Now that the ExitLoopMap gives as  mapping for all the non-looping cloned
1445   // blocks to their outer loops, walk the cloned blocks and the cloned exits
1446   // in their original order adding them to the correct loop.
1447 
1448   // We need a stable insertion order. We use the order of the original loop
1449   // order and map into the correct parent loop.
1450   for (auto *BB : llvm::concat<BasicBlock *const>(
1451            makeArrayRef(ClonedPH), ClonedLoopBlocks, ClonedExitsInLoops))
1452     if (Loop *OuterL = ExitLoopMap.lookup(BB))
1453       OuterL->addBasicBlockToLoop(BB, LI);
1454 
1455 #ifndef NDEBUG
1456   for (auto &BBAndL : ExitLoopMap) {
1457     auto *BB = BBAndL.first;
1458     auto *OuterL = BBAndL.second;
1459     assert(LI.getLoopFor(BB) == OuterL &&
1460            "Failed to put all blocks into outer loops!");
1461   }
1462 #endif
1463 
1464   // Now that all the blocks are placed into the correct containing loop in the
1465   // absence of child loops, find all the potentially cloned child loops and
1466   // clone them into whatever outer loop we placed their header into.
1467   for (Loop *ChildL : OrigL) {
1468     auto *ClonedChildHeader =
1469         cast_or_null<BasicBlock>(VMap.lookup(ChildL->getHeader()));
1470     if (!ClonedChildHeader || BlocksInClonedLoop.count(ClonedChildHeader))
1471       continue;
1472 
1473 #ifndef NDEBUG
1474     for (auto *ChildLoopBB : ChildL->blocks())
1475       assert(VMap.count(ChildLoopBB) &&
1476              "Cloned a child loop header but not all of that loops blocks!");
1477 #endif
1478 
1479     NonChildClonedLoops.push_back(cloneLoopNest(
1480         *ChildL, ExitLoopMap.lookup(ClonedChildHeader), VMap, LI));
1481   }
1482 }
1483 
1484 static void
1485 deleteDeadClonedBlocks(Loop &L, ArrayRef<BasicBlock *> ExitBlocks,
1486                        ArrayRef<std::unique_ptr<ValueToValueMapTy>> VMaps,
1487                        DominatorTree &DT, MemorySSAUpdater *MSSAU) {
1488   // Find all the dead clones, and remove them from their successors.
1489   SmallVector<BasicBlock *, 16> DeadBlocks;
1490   for (BasicBlock *BB : llvm::concat<BasicBlock *const>(L.blocks(), ExitBlocks))
1491     for (auto &VMap : VMaps)
1492       if (BasicBlock *ClonedBB = cast_or_null<BasicBlock>(VMap->lookup(BB)))
1493         if (!DT.isReachableFromEntry(ClonedBB)) {
1494           for (BasicBlock *SuccBB : successors(ClonedBB))
1495             SuccBB->removePredecessor(ClonedBB);
1496           DeadBlocks.push_back(ClonedBB);
1497         }
1498 
1499   // Remove all MemorySSA in the dead blocks
1500   if (MSSAU) {
1501     SmallSetVector<BasicBlock *, 8> DeadBlockSet(DeadBlocks.begin(),
1502                                                  DeadBlocks.end());
1503     MSSAU->removeBlocks(DeadBlockSet);
1504   }
1505 
1506   // Drop any remaining references to break cycles.
1507   for (BasicBlock *BB : DeadBlocks)
1508     BB->dropAllReferences();
1509   // Erase them from the IR.
1510   for (BasicBlock *BB : DeadBlocks)
1511     BB->eraseFromParent();
1512 }
1513 
1514 static void deleteDeadBlocksFromLoop(Loop &L,
1515                                      SmallVectorImpl<BasicBlock *> &ExitBlocks,
1516                                      DominatorTree &DT, LoopInfo &LI,
1517                                      MemorySSAUpdater *MSSAU) {
1518   // Find all the dead blocks tied to this loop, and remove them from their
1519   // successors.
1520   SmallSetVector<BasicBlock *, 8> DeadBlockSet;
1521 
1522   // Start with loop/exit blocks and get a transitive closure of reachable dead
1523   // blocks.
1524   SmallVector<BasicBlock *, 16> DeathCandidates(ExitBlocks.begin(),
1525                                                 ExitBlocks.end());
1526   DeathCandidates.append(L.blocks().begin(), L.blocks().end());
1527   while (!DeathCandidates.empty()) {
1528     auto *BB = DeathCandidates.pop_back_val();
1529     if (!DeadBlockSet.count(BB) && !DT.isReachableFromEntry(BB)) {
1530       for (BasicBlock *SuccBB : successors(BB)) {
1531         SuccBB->removePredecessor(BB);
1532         DeathCandidates.push_back(SuccBB);
1533       }
1534       DeadBlockSet.insert(BB);
1535     }
1536   }
1537 
1538   // Remove all MemorySSA in the dead blocks
1539   if (MSSAU)
1540     MSSAU->removeBlocks(DeadBlockSet);
1541 
1542   // Filter out the dead blocks from the exit blocks list so that it can be
1543   // used in the caller.
1544   llvm::erase_if(ExitBlocks,
1545                  [&](BasicBlock *BB) { return DeadBlockSet.count(BB); });
1546 
1547   // Walk from this loop up through its parents removing all of the dead blocks.
1548   for (Loop *ParentL = &L; ParentL; ParentL = ParentL->getParentLoop()) {
1549     for (auto *BB : DeadBlockSet)
1550       ParentL->getBlocksSet().erase(BB);
1551     llvm::erase_if(ParentL->getBlocksVector(),
1552                    [&](BasicBlock *BB) { return DeadBlockSet.count(BB); });
1553   }
1554 
1555   // Now delete the dead child loops. This raw delete will clear them
1556   // recursively.
1557   llvm::erase_if(L.getSubLoopsVector(), [&](Loop *ChildL) {
1558     if (!DeadBlockSet.count(ChildL->getHeader()))
1559       return false;
1560 
1561     assert(llvm::all_of(ChildL->blocks(),
1562                         [&](BasicBlock *ChildBB) {
1563                           return DeadBlockSet.count(ChildBB);
1564                         }) &&
1565            "If the child loop header is dead all blocks in the child loop must "
1566            "be dead as well!");
1567     LI.destroy(ChildL);
1568     return true;
1569   });
1570 
1571   // Remove the loop mappings for the dead blocks and drop all the references
1572   // from these blocks to others to handle cyclic references as we start
1573   // deleting the blocks themselves.
1574   for (auto *BB : DeadBlockSet) {
1575     // Check that the dominator tree has already been updated.
1576     assert(!DT.getNode(BB) && "Should already have cleared domtree!");
1577     LI.changeLoopFor(BB, nullptr);
1578     // Drop all uses of the instructions to make sure we won't have dangling
1579     // uses in other blocks.
1580     for (auto &I : *BB)
1581       if (!I.use_empty())
1582         I.replaceAllUsesWith(UndefValue::get(I.getType()));
1583     BB->dropAllReferences();
1584   }
1585 
1586   // Actually delete the blocks now that they've been fully unhooked from the
1587   // IR.
1588   for (auto *BB : DeadBlockSet)
1589     BB->eraseFromParent();
1590 }
1591 
1592 /// Recompute the set of blocks in a loop after unswitching.
1593 ///
1594 /// This walks from the original headers predecessors to rebuild the loop. We
1595 /// take advantage of the fact that new blocks can't have been added, and so we
1596 /// filter by the original loop's blocks. This also handles potentially
1597 /// unreachable code that we don't want to explore but might be found examining
1598 /// the predecessors of the header.
1599 ///
1600 /// If the original loop is no longer a loop, this will return an empty set. If
1601 /// it remains a loop, all the blocks within it will be added to the set
1602 /// (including those blocks in inner loops).
1603 static SmallPtrSet<const BasicBlock *, 16> recomputeLoopBlockSet(Loop &L,
1604                                                                  LoopInfo &LI) {
1605   SmallPtrSet<const BasicBlock *, 16> LoopBlockSet;
1606 
1607   auto *PH = L.getLoopPreheader();
1608   auto *Header = L.getHeader();
1609 
1610   // A worklist to use while walking backwards from the header.
1611   SmallVector<BasicBlock *, 16> Worklist;
1612 
1613   // First walk the predecessors of the header to find the backedges. This will
1614   // form the basis of our walk.
1615   for (auto *Pred : predecessors(Header)) {
1616     // Skip the preheader.
1617     if (Pred == PH)
1618       continue;
1619 
1620     // Because the loop was in simplified form, the only non-loop predecessor
1621     // is the preheader.
1622     assert(L.contains(Pred) && "Found a predecessor of the loop header other "
1623                                "than the preheader that is not part of the "
1624                                "loop!");
1625 
1626     // Insert this block into the loop set and on the first visit and, if it
1627     // isn't the header we're currently walking, put it into the worklist to
1628     // recurse through.
1629     if (LoopBlockSet.insert(Pred).second && Pred != Header)
1630       Worklist.push_back(Pred);
1631   }
1632 
1633   // If no backedges were found, we're done.
1634   if (LoopBlockSet.empty())
1635     return LoopBlockSet;
1636 
1637   // We found backedges, recurse through them to identify the loop blocks.
1638   while (!Worklist.empty()) {
1639     BasicBlock *BB = Worklist.pop_back_val();
1640     assert(LoopBlockSet.count(BB) && "Didn't put block into the loop set!");
1641 
1642     // No need to walk past the header.
1643     if (BB == Header)
1644       continue;
1645 
1646     // Because we know the inner loop structure remains valid we can use the
1647     // loop structure to jump immediately across the entire nested loop.
1648     // Further, because it is in loop simplified form, we can directly jump
1649     // to its preheader afterward.
1650     if (Loop *InnerL = LI.getLoopFor(BB))
1651       if (InnerL != &L) {
1652         assert(L.contains(InnerL) &&
1653                "Should not reach a loop *outside* this loop!");
1654         // The preheader is the only possible predecessor of the loop so
1655         // insert it into the set and check whether it was already handled.
1656         auto *InnerPH = InnerL->getLoopPreheader();
1657         assert(L.contains(InnerPH) && "Cannot contain an inner loop block "
1658                                       "but not contain the inner loop "
1659                                       "preheader!");
1660         if (!LoopBlockSet.insert(InnerPH).second)
1661           // The only way to reach the preheader is through the loop body
1662           // itself so if it has been visited the loop is already handled.
1663           continue;
1664 
1665         // Insert all of the blocks (other than those already present) into
1666         // the loop set. We expect at least the block that led us to find the
1667         // inner loop to be in the block set, but we may also have other loop
1668         // blocks if they were already enqueued as predecessors of some other
1669         // outer loop block.
1670         for (auto *InnerBB : InnerL->blocks()) {
1671           if (InnerBB == BB) {
1672             assert(LoopBlockSet.count(InnerBB) &&
1673                    "Block should already be in the set!");
1674             continue;
1675           }
1676 
1677           LoopBlockSet.insert(InnerBB);
1678         }
1679 
1680         // Add the preheader to the worklist so we will continue past the
1681         // loop body.
1682         Worklist.push_back(InnerPH);
1683         continue;
1684       }
1685 
1686     // Insert any predecessors that were in the original loop into the new
1687     // set, and if the insert is successful, add them to the worklist.
1688     for (auto *Pred : predecessors(BB))
1689       if (L.contains(Pred) && LoopBlockSet.insert(Pred).second)
1690         Worklist.push_back(Pred);
1691   }
1692 
1693   assert(LoopBlockSet.count(Header) && "Cannot fail to add the header!");
1694 
1695   // We've found all the blocks participating in the loop, return our completed
1696   // set.
1697   return LoopBlockSet;
1698 }
1699 
1700 /// Rebuild a loop after unswitching removes some subset of blocks and edges.
1701 ///
1702 /// The removal may have removed some child loops entirely but cannot have
1703 /// disturbed any remaining child loops. However, they may need to be hoisted
1704 /// to the parent loop (or to be top-level loops). The original loop may be
1705 /// completely removed.
1706 ///
1707 /// The sibling loops resulting from this update are returned. If the original
1708 /// loop remains a valid loop, it will be the first entry in this list with all
1709 /// of the newly sibling loops following it.
1710 ///
1711 /// Returns true if the loop remains a loop after unswitching, and false if it
1712 /// is no longer a loop after unswitching (and should not continue to be
1713 /// referenced).
1714 static bool rebuildLoopAfterUnswitch(Loop &L, ArrayRef<BasicBlock *> ExitBlocks,
1715                                      LoopInfo &LI,
1716                                      SmallVectorImpl<Loop *> &HoistedLoops) {
1717   auto *PH = L.getLoopPreheader();
1718 
1719   // Compute the actual parent loop from the exit blocks. Because we may have
1720   // pruned some exits the loop may be different from the original parent.
1721   Loop *ParentL = nullptr;
1722   SmallVector<Loop *, 4> ExitLoops;
1723   SmallVector<BasicBlock *, 4> ExitsInLoops;
1724   ExitsInLoops.reserve(ExitBlocks.size());
1725   for (auto *ExitBB : ExitBlocks)
1726     if (Loop *ExitL = LI.getLoopFor(ExitBB)) {
1727       ExitLoops.push_back(ExitL);
1728       ExitsInLoops.push_back(ExitBB);
1729       if (!ParentL || (ParentL != ExitL && ParentL->contains(ExitL)))
1730         ParentL = ExitL;
1731     }
1732 
1733   // Recompute the blocks participating in this loop. This may be empty if it
1734   // is no longer a loop.
1735   auto LoopBlockSet = recomputeLoopBlockSet(L, LI);
1736 
1737   // If we still have a loop, we need to re-set the loop's parent as the exit
1738   // block set changing may have moved it within the loop nest. Note that this
1739   // can only happen when this loop has a parent as it can only hoist the loop
1740   // *up* the nest.
1741   if (!LoopBlockSet.empty() && L.getParentLoop() != ParentL) {
1742     // Remove this loop's (original) blocks from all of the intervening loops.
1743     for (Loop *IL = L.getParentLoop(); IL != ParentL;
1744          IL = IL->getParentLoop()) {
1745       IL->getBlocksSet().erase(PH);
1746       for (auto *BB : L.blocks())
1747         IL->getBlocksSet().erase(BB);
1748       llvm::erase_if(IL->getBlocksVector(), [&](BasicBlock *BB) {
1749         return BB == PH || L.contains(BB);
1750       });
1751     }
1752 
1753     LI.changeLoopFor(PH, ParentL);
1754     L.getParentLoop()->removeChildLoop(&L);
1755     if (ParentL)
1756       ParentL->addChildLoop(&L);
1757     else
1758       LI.addTopLevelLoop(&L);
1759   }
1760 
1761   // Now we update all the blocks which are no longer within the loop.
1762   auto &Blocks = L.getBlocksVector();
1763   auto BlocksSplitI =
1764       LoopBlockSet.empty()
1765           ? Blocks.begin()
1766           : std::stable_partition(
1767                 Blocks.begin(), Blocks.end(),
1768                 [&](BasicBlock *BB) { return LoopBlockSet.count(BB); });
1769 
1770   // Before we erase the list of unlooped blocks, build a set of them.
1771   SmallPtrSet<BasicBlock *, 16> UnloopedBlocks(BlocksSplitI, Blocks.end());
1772   if (LoopBlockSet.empty())
1773     UnloopedBlocks.insert(PH);
1774 
1775   // Now erase these blocks from the loop.
1776   for (auto *BB : make_range(BlocksSplitI, Blocks.end()))
1777     L.getBlocksSet().erase(BB);
1778   Blocks.erase(BlocksSplitI, Blocks.end());
1779 
1780   // Sort the exits in ascending loop depth, we'll work backwards across these
1781   // to process them inside out.
1782   llvm::stable_sort(ExitsInLoops, [&](BasicBlock *LHS, BasicBlock *RHS) {
1783     return LI.getLoopDepth(LHS) < LI.getLoopDepth(RHS);
1784   });
1785 
1786   // We'll build up a set for each exit loop.
1787   SmallPtrSet<BasicBlock *, 16> NewExitLoopBlocks;
1788   Loop *PrevExitL = L.getParentLoop(); // The deepest possible exit loop.
1789 
1790   auto RemoveUnloopedBlocksFromLoop =
1791       [](Loop &L, SmallPtrSetImpl<BasicBlock *> &UnloopedBlocks) {
1792         for (auto *BB : UnloopedBlocks)
1793           L.getBlocksSet().erase(BB);
1794         llvm::erase_if(L.getBlocksVector(), [&](BasicBlock *BB) {
1795           return UnloopedBlocks.count(BB);
1796         });
1797       };
1798 
1799   SmallVector<BasicBlock *, 16> Worklist;
1800   while (!UnloopedBlocks.empty() && !ExitsInLoops.empty()) {
1801     assert(Worklist.empty() && "Didn't clear worklist!");
1802     assert(NewExitLoopBlocks.empty() && "Didn't clear loop set!");
1803 
1804     // Grab the next exit block, in decreasing loop depth order.
1805     BasicBlock *ExitBB = ExitsInLoops.pop_back_val();
1806     Loop &ExitL = *LI.getLoopFor(ExitBB);
1807     assert(ExitL.contains(&L) && "Exit loop must contain the inner loop!");
1808 
1809     // Erase all of the unlooped blocks from the loops between the previous
1810     // exit loop and this exit loop. This works because the ExitInLoops list is
1811     // sorted in increasing order of loop depth and thus we visit loops in
1812     // decreasing order of loop depth.
1813     for (; PrevExitL != &ExitL; PrevExitL = PrevExitL->getParentLoop())
1814       RemoveUnloopedBlocksFromLoop(*PrevExitL, UnloopedBlocks);
1815 
1816     // Walk the CFG back until we hit the cloned PH adding everything reachable
1817     // and in the unlooped set to this exit block's loop.
1818     Worklist.push_back(ExitBB);
1819     do {
1820       BasicBlock *BB = Worklist.pop_back_val();
1821       // We can stop recursing at the cloned preheader (if we get there).
1822       if (BB == PH)
1823         continue;
1824 
1825       for (BasicBlock *PredBB : predecessors(BB)) {
1826         // If this pred has already been moved to our set or is part of some
1827         // (inner) loop, no update needed.
1828         if (!UnloopedBlocks.erase(PredBB)) {
1829           assert((NewExitLoopBlocks.count(PredBB) ||
1830                   ExitL.contains(LI.getLoopFor(PredBB))) &&
1831                  "Predecessor not in a nested loop (or already visited)!");
1832           continue;
1833         }
1834 
1835         // We just insert into the loop set here. We'll add these blocks to the
1836         // exit loop after we build up the set in a deterministic order rather
1837         // than the predecessor-influenced visit order.
1838         bool Inserted = NewExitLoopBlocks.insert(PredBB).second;
1839         (void)Inserted;
1840         assert(Inserted && "Should only visit an unlooped block once!");
1841 
1842         // And recurse through to its predecessors.
1843         Worklist.push_back(PredBB);
1844       }
1845     } while (!Worklist.empty());
1846 
1847     // If blocks in this exit loop were directly part of the original loop (as
1848     // opposed to a child loop) update the map to point to this exit loop. This
1849     // just updates a map and so the fact that the order is unstable is fine.
1850     for (auto *BB : NewExitLoopBlocks)
1851       if (Loop *BBL = LI.getLoopFor(BB))
1852         if (BBL == &L || !L.contains(BBL))
1853           LI.changeLoopFor(BB, &ExitL);
1854 
1855     // We will remove the remaining unlooped blocks from this loop in the next
1856     // iteration or below.
1857     NewExitLoopBlocks.clear();
1858   }
1859 
1860   // Any remaining unlooped blocks are no longer part of any loop unless they
1861   // are part of some child loop.
1862   for (; PrevExitL; PrevExitL = PrevExitL->getParentLoop())
1863     RemoveUnloopedBlocksFromLoop(*PrevExitL, UnloopedBlocks);
1864   for (auto *BB : UnloopedBlocks)
1865     if (Loop *BBL = LI.getLoopFor(BB))
1866       if (BBL == &L || !L.contains(BBL))
1867         LI.changeLoopFor(BB, nullptr);
1868 
1869   // Sink all the child loops whose headers are no longer in the loop set to
1870   // the parent (or to be top level loops). We reach into the loop and directly
1871   // update its subloop vector to make this batch update efficient.
1872   auto &SubLoops = L.getSubLoopsVector();
1873   auto SubLoopsSplitI =
1874       LoopBlockSet.empty()
1875           ? SubLoops.begin()
1876           : std::stable_partition(
1877                 SubLoops.begin(), SubLoops.end(), [&](Loop *SubL) {
1878                   return LoopBlockSet.count(SubL->getHeader());
1879                 });
1880   for (auto *HoistedL : make_range(SubLoopsSplitI, SubLoops.end())) {
1881     HoistedLoops.push_back(HoistedL);
1882     HoistedL->setParentLoop(nullptr);
1883 
1884     // To compute the new parent of this hoisted loop we look at where we
1885     // placed the preheader above. We can't lookup the header itself because we
1886     // retained the mapping from the header to the hoisted loop. But the
1887     // preheader and header should have the exact same new parent computed
1888     // based on the set of exit blocks from the original loop as the preheader
1889     // is a predecessor of the header and so reached in the reverse walk. And
1890     // because the loops were all in simplified form the preheader of the
1891     // hoisted loop can't be part of some *other* loop.
1892     if (auto *NewParentL = LI.getLoopFor(HoistedL->getLoopPreheader()))
1893       NewParentL->addChildLoop(HoistedL);
1894     else
1895       LI.addTopLevelLoop(HoistedL);
1896   }
1897   SubLoops.erase(SubLoopsSplitI, SubLoops.end());
1898 
1899   // Actually delete the loop if nothing remained within it.
1900   if (Blocks.empty()) {
1901     assert(SubLoops.empty() &&
1902            "Failed to remove all subloops from the original loop!");
1903     if (Loop *ParentL = L.getParentLoop())
1904       ParentL->removeChildLoop(llvm::find(*ParentL, &L));
1905     else
1906       LI.removeLoop(llvm::find(LI, &L));
1907     LI.destroy(&L);
1908     return false;
1909   }
1910 
1911   return true;
1912 }
1913 
1914 /// Helper to visit a dominator subtree, invoking a callable on each node.
1915 ///
1916 /// Returning false at any point will stop walking past that node of the tree.
1917 template <typename CallableT>
1918 void visitDomSubTree(DominatorTree &DT, BasicBlock *BB, CallableT Callable) {
1919   SmallVector<DomTreeNode *, 4> DomWorklist;
1920   DomWorklist.push_back(DT[BB]);
1921 #ifndef NDEBUG
1922   SmallPtrSet<DomTreeNode *, 4> Visited;
1923   Visited.insert(DT[BB]);
1924 #endif
1925   do {
1926     DomTreeNode *N = DomWorklist.pop_back_val();
1927 
1928     // Visit this node.
1929     if (!Callable(N->getBlock()))
1930       continue;
1931 
1932     // Accumulate the child nodes.
1933     for (DomTreeNode *ChildN : *N) {
1934       assert(Visited.insert(ChildN).second &&
1935              "Cannot visit a node twice when walking a tree!");
1936       DomWorklist.push_back(ChildN);
1937     }
1938   } while (!DomWorklist.empty());
1939 }
1940 
1941 static void unswitchNontrivialInvariants(
1942     Loop &L, Instruction &TI, ArrayRef<Value *> Invariants,
1943     SmallVectorImpl<BasicBlock *> &ExitBlocks, DominatorTree &DT, LoopInfo &LI,
1944     AssumptionCache &AC, function_ref<void(bool, ArrayRef<Loop *>)> UnswitchCB,
1945     ScalarEvolution *SE, MemorySSAUpdater *MSSAU) {
1946   auto *ParentBB = TI.getParent();
1947   BranchInst *BI = dyn_cast<BranchInst>(&TI);
1948   SwitchInst *SI = BI ? nullptr : cast<SwitchInst>(&TI);
1949 
1950   // We can only unswitch switches, conditional branches with an invariant
1951   // condition, or combining invariant conditions with an instruction.
1952   assert((SI || (BI && BI->isConditional())) &&
1953          "Can only unswitch switches and conditional branch!");
1954   bool FullUnswitch = SI || BI->getCondition() == Invariants[0];
1955   if (FullUnswitch)
1956     assert(Invariants.size() == 1 &&
1957            "Cannot have other invariants with full unswitching!");
1958   else
1959     assert(isa<Instruction>(BI->getCondition()) &&
1960            "Partial unswitching requires an instruction as the condition!");
1961 
1962   if (MSSAU && VerifyMemorySSA)
1963     MSSAU->getMemorySSA()->verifyMemorySSA();
1964 
1965   // Constant and BBs tracking the cloned and continuing successor. When we are
1966   // unswitching the entire condition, this can just be trivially chosen to
1967   // unswitch towards `true`. However, when we are unswitching a set of
1968   // invariants combined with `and` or `or`, the combining operation determines
1969   // the best direction to unswitch: we want to unswitch the direction that will
1970   // collapse the branch.
1971   bool Direction = true;
1972   int ClonedSucc = 0;
1973   if (!FullUnswitch) {
1974     if (cast<Instruction>(BI->getCondition())->getOpcode() != Instruction::Or) {
1975       assert(cast<Instruction>(BI->getCondition())->getOpcode() ==
1976                  Instruction::And &&
1977              "Only `or` and `and` instructions can combine invariants being "
1978              "unswitched.");
1979       Direction = false;
1980       ClonedSucc = 1;
1981     }
1982   }
1983 
1984   BasicBlock *RetainedSuccBB =
1985       BI ? BI->getSuccessor(1 - ClonedSucc) : SI->getDefaultDest();
1986   SmallSetVector<BasicBlock *, 4> UnswitchedSuccBBs;
1987   if (BI)
1988     UnswitchedSuccBBs.insert(BI->getSuccessor(ClonedSucc));
1989   else
1990     for (auto Case : SI->cases())
1991       if (Case.getCaseSuccessor() != RetainedSuccBB)
1992         UnswitchedSuccBBs.insert(Case.getCaseSuccessor());
1993 
1994   assert(!UnswitchedSuccBBs.count(RetainedSuccBB) &&
1995          "Should not unswitch the same successor we are retaining!");
1996 
1997   // The branch should be in this exact loop. Any inner loop's invariant branch
1998   // should be handled by unswitching that inner loop. The caller of this
1999   // routine should filter out any candidates that remain (but were skipped for
2000   // whatever reason).
2001   assert(LI.getLoopFor(ParentBB) == &L && "Branch in an inner loop!");
2002 
2003   // Compute the parent loop now before we start hacking on things.
2004   Loop *ParentL = L.getParentLoop();
2005   // Get blocks in RPO order for MSSA update, before changing the CFG.
2006   LoopBlocksRPO LBRPO(&L);
2007   if (MSSAU)
2008     LBRPO.perform(&LI);
2009 
2010   // Compute the outer-most loop containing one of our exit blocks. This is the
2011   // furthest up our loopnest which can be mutated, which we will use below to
2012   // update things.
2013   Loop *OuterExitL = &L;
2014   for (auto *ExitBB : ExitBlocks) {
2015     Loop *NewOuterExitL = LI.getLoopFor(ExitBB);
2016     if (!NewOuterExitL) {
2017       // We exited the entire nest with this block, so we're done.
2018       OuterExitL = nullptr;
2019       break;
2020     }
2021     if (NewOuterExitL != OuterExitL && NewOuterExitL->contains(OuterExitL))
2022       OuterExitL = NewOuterExitL;
2023   }
2024 
2025   // At this point, we're definitely going to unswitch something so invalidate
2026   // any cached information in ScalarEvolution for the outer most loop
2027   // containing an exit block and all nested loops.
2028   if (SE) {
2029     if (OuterExitL)
2030       SE->forgetLoop(OuterExitL);
2031     else
2032       SE->forgetTopmostLoop(&L);
2033   }
2034 
2035   // If the edge from this terminator to a successor dominates that successor,
2036   // store a map from each block in its dominator subtree to it. This lets us
2037   // tell when cloning for a particular successor if a block is dominated by
2038   // some *other* successor with a single data structure. We use this to
2039   // significantly reduce cloning.
2040   SmallDenseMap<BasicBlock *, BasicBlock *, 16> DominatingSucc;
2041   for (auto *SuccBB : llvm::concat<BasicBlock *const>(
2042            makeArrayRef(RetainedSuccBB), UnswitchedSuccBBs))
2043     if (SuccBB->getUniquePredecessor() ||
2044         llvm::all_of(predecessors(SuccBB), [&](BasicBlock *PredBB) {
2045           return PredBB == ParentBB || DT.dominates(SuccBB, PredBB);
2046         }))
2047       visitDomSubTree(DT, SuccBB, [&](BasicBlock *BB) {
2048         DominatingSucc[BB] = SuccBB;
2049         return true;
2050       });
2051 
2052   // Split the preheader, so that we know that there is a safe place to insert
2053   // the conditional branch. We will change the preheader to have a conditional
2054   // branch on LoopCond. The original preheader will become the split point
2055   // between the unswitched versions, and we will have a new preheader for the
2056   // original loop.
2057   BasicBlock *SplitBB = L.getLoopPreheader();
2058   BasicBlock *LoopPH = SplitEdge(SplitBB, L.getHeader(), &DT, &LI, MSSAU);
2059 
2060   // Keep track of the dominator tree updates needed.
2061   SmallVector<DominatorTree::UpdateType, 4> DTUpdates;
2062 
2063   // Clone the loop for each unswitched successor.
2064   SmallVector<std::unique_ptr<ValueToValueMapTy>, 4> VMaps;
2065   VMaps.reserve(UnswitchedSuccBBs.size());
2066   SmallDenseMap<BasicBlock *, BasicBlock *, 4> ClonedPHs;
2067   for (auto *SuccBB : UnswitchedSuccBBs) {
2068     VMaps.emplace_back(new ValueToValueMapTy());
2069     ClonedPHs[SuccBB] = buildClonedLoopBlocks(
2070         L, LoopPH, SplitBB, ExitBlocks, ParentBB, SuccBB, RetainedSuccBB,
2071         DominatingSucc, *VMaps.back(), DTUpdates, AC, DT, LI, MSSAU);
2072   }
2073 
2074   // Drop metadata if we may break its semantics by moving this instr into the
2075   // split block.
2076   if (TI.getMetadata(LLVMContext::MD_make_implicit)) {
2077     // It is only legal to preserve make.implicit metadata if we are guaranteed
2078     // to reach implicit null check block after following this branch.
2079     ICFLoopSafetyInfo SafetyInfo;
2080     SafetyInfo.computeLoopSafetyInfo(&L);
2081     if (!SafetyInfo.isGuaranteedToExecute(TI, &DT, &L))
2082       TI.setMetadata(LLVMContext::MD_make_implicit, nullptr);
2083   }
2084 
2085   // The stitching of the branched code back together depends on whether we're
2086   // doing full unswitching or not with the exception that we always want to
2087   // nuke the initial terminator placed in the split block.
2088   SplitBB->getTerminator()->eraseFromParent();
2089   if (FullUnswitch) {
2090     // Splice the terminator from the original loop and rewrite its
2091     // successors.
2092     SplitBB->getInstList().splice(SplitBB->end(), ParentBB->getInstList(), TI);
2093 
2094     // Keep a clone of the terminator for MSSA updates.
2095     Instruction *NewTI = TI.clone();
2096     ParentBB->getInstList().push_back(NewTI);
2097 
2098     // First wire up the moved terminator to the preheaders.
2099     if (BI) {
2100       BasicBlock *ClonedPH = ClonedPHs.begin()->second;
2101       BI->setSuccessor(ClonedSucc, ClonedPH);
2102       BI->setSuccessor(1 - ClonedSucc, LoopPH);
2103       DTUpdates.push_back({DominatorTree::Insert, SplitBB, ClonedPH});
2104     } else {
2105       assert(SI && "Must either be a branch or switch!");
2106 
2107       // Walk the cases and directly update their successors.
2108       assert(SI->getDefaultDest() == RetainedSuccBB &&
2109              "Not retaining default successor!");
2110       SI->setDefaultDest(LoopPH);
2111       for (auto &Case : SI->cases())
2112         if (Case.getCaseSuccessor() == RetainedSuccBB)
2113           Case.setSuccessor(LoopPH);
2114         else
2115           Case.setSuccessor(ClonedPHs.find(Case.getCaseSuccessor())->second);
2116 
2117       // We need to use the set to populate domtree updates as even when there
2118       // are multiple cases pointing at the same successor we only want to
2119       // remove and insert one edge in the domtree.
2120       for (BasicBlock *SuccBB : UnswitchedSuccBBs)
2121         DTUpdates.push_back(
2122             {DominatorTree::Insert, SplitBB, ClonedPHs.find(SuccBB)->second});
2123     }
2124 
2125     if (MSSAU) {
2126       DT.applyUpdates(DTUpdates);
2127       DTUpdates.clear();
2128 
2129       // Remove all but one edge to the retained block and all unswitched
2130       // blocks. This is to avoid having duplicate entries in the cloned Phis,
2131       // when we know we only keep a single edge for each case.
2132       MSSAU->removeDuplicatePhiEdgesBetween(ParentBB, RetainedSuccBB);
2133       for (BasicBlock *SuccBB : UnswitchedSuccBBs)
2134         MSSAU->removeDuplicatePhiEdgesBetween(ParentBB, SuccBB);
2135 
2136       for (auto &VMap : VMaps)
2137         MSSAU->updateForClonedLoop(LBRPO, ExitBlocks, *VMap,
2138                                    /*IgnoreIncomingWithNoClones=*/true);
2139       MSSAU->updateExitBlocksForClonedLoop(ExitBlocks, VMaps, DT);
2140 
2141       // Remove all edges to unswitched blocks.
2142       for (BasicBlock *SuccBB : UnswitchedSuccBBs)
2143         MSSAU->removeEdge(ParentBB, SuccBB);
2144     }
2145 
2146     // Now unhook the successor relationship as we'll be replacing
2147     // the terminator with a direct branch. This is much simpler for branches
2148     // than switches so we handle those first.
2149     if (BI) {
2150       // Remove the parent as a predecessor of the unswitched successor.
2151       assert(UnswitchedSuccBBs.size() == 1 &&
2152              "Only one possible unswitched block for a branch!");
2153       BasicBlock *UnswitchedSuccBB = *UnswitchedSuccBBs.begin();
2154       UnswitchedSuccBB->removePredecessor(ParentBB,
2155                                           /*KeepOneInputPHIs*/ true);
2156       DTUpdates.push_back({DominatorTree::Delete, ParentBB, UnswitchedSuccBB});
2157     } else {
2158       // Note that we actually want to remove the parent block as a predecessor
2159       // of *every* case successor. The case successor is either unswitched,
2160       // completely eliminating an edge from the parent to that successor, or it
2161       // is a duplicate edge to the retained successor as the retained successor
2162       // is always the default successor and as we'll replace this with a direct
2163       // branch we no longer need the duplicate entries in the PHI nodes.
2164       SwitchInst *NewSI = cast<SwitchInst>(NewTI);
2165       assert(NewSI->getDefaultDest() == RetainedSuccBB &&
2166              "Not retaining default successor!");
2167       for (auto &Case : NewSI->cases())
2168         Case.getCaseSuccessor()->removePredecessor(
2169             ParentBB,
2170             /*KeepOneInputPHIs*/ true);
2171 
2172       // We need to use the set to populate domtree updates as even when there
2173       // are multiple cases pointing at the same successor we only want to
2174       // remove and insert one edge in the domtree.
2175       for (BasicBlock *SuccBB : UnswitchedSuccBBs)
2176         DTUpdates.push_back({DominatorTree::Delete, ParentBB, SuccBB});
2177     }
2178 
2179     // After MSSAU update, remove the cloned terminator instruction NewTI.
2180     ParentBB->getTerminator()->eraseFromParent();
2181 
2182     // Create a new unconditional branch to the continuing block (as opposed to
2183     // the one cloned).
2184     BranchInst::Create(RetainedSuccBB, ParentBB);
2185   } else {
2186     assert(BI && "Only branches have partial unswitching.");
2187     assert(UnswitchedSuccBBs.size() == 1 &&
2188            "Only one possible unswitched block for a branch!");
2189     BasicBlock *ClonedPH = ClonedPHs.begin()->second;
2190     // When doing a partial unswitch, we have to do a bit more work to build up
2191     // the branch in the split block.
2192     buildPartialUnswitchConditionalBranch(*SplitBB, Invariants, Direction,
2193                                           *ClonedPH, *LoopPH);
2194     DTUpdates.push_back({DominatorTree::Insert, SplitBB, ClonedPH});
2195 
2196     if (MSSAU) {
2197       DT.applyUpdates(DTUpdates);
2198       DTUpdates.clear();
2199 
2200       // Perform MSSA cloning updates.
2201       for (auto &VMap : VMaps)
2202         MSSAU->updateForClonedLoop(LBRPO, ExitBlocks, *VMap,
2203                                    /*IgnoreIncomingWithNoClones=*/true);
2204       MSSAU->updateExitBlocksForClonedLoop(ExitBlocks, VMaps, DT);
2205     }
2206   }
2207 
2208   // Apply the updates accumulated above to get an up-to-date dominator tree.
2209   DT.applyUpdates(DTUpdates);
2210 
2211   // Now that we have an accurate dominator tree, first delete the dead cloned
2212   // blocks so that we can accurately build any cloned loops. It is important to
2213   // not delete the blocks from the original loop yet because we still want to
2214   // reference the original loop to understand the cloned loop's structure.
2215   deleteDeadClonedBlocks(L, ExitBlocks, VMaps, DT, MSSAU);
2216 
2217   // Build the cloned loop structure itself. This may be substantially
2218   // different from the original structure due to the simplified CFG. This also
2219   // handles inserting all the cloned blocks into the correct loops.
2220   SmallVector<Loop *, 4> NonChildClonedLoops;
2221   for (std::unique_ptr<ValueToValueMapTy> &VMap : VMaps)
2222     buildClonedLoops(L, ExitBlocks, *VMap, LI, NonChildClonedLoops);
2223 
2224   // Now that our cloned loops have been built, we can update the original loop.
2225   // First we delete the dead blocks from it and then we rebuild the loop
2226   // structure taking these deletions into account.
2227   deleteDeadBlocksFromLoop(L, ExitBlocks, DT, LI, MSSAU);
2228 
2229   if (MSSAU && VerifyMemorySSA)
2230     MSSAU->getMemorySSA()->verifyMemorySSA();
2231 
2232   SmallVector<Loop *, 4> HoistedLoops;
2233   bool IsStillLoop = rebuildLoopAfterUnswitch(L, ExitBlocks, LI, HoistedLoops);
2234 
2235   if (MSSAU && VerifyMemorySSA)
2236     MSSAU->getMemorySSA()->verifyMemorySSA();
2237 
2238   // This transformation has a high risk of corrupting the dominator tree, and
2239   // the below steps to rebuild loop structures will result in hard to debug
2240   // errors in that case so verify that the dominator tree is sane first.
2241   // FIXME: Remove this when the bugs stop showing up and rely on existing
2242   // verification steps.
2243   assert(DT.verify(DominatorTree::VerificationLevel::Fast));
2244 
2245   if (BI) {
2246     // If we unswitched a branch which collapses the condition to a known
2247     // constant we want to replace all the uses of the invariants within both
2248     // the original and cloned blocks. We do this here so that we can use the
2249     // now updated dominator tree to identify which side the users are on.
2250     assert(UnswitchedSuccBBs.size() == 1 &&
2251            "Only one possible unswitched block for a branch!");
2252     BasicBlock *ClonedPH = ClonedPHs.begin()->second;
2253 
2254     // When considering multiple partially-unswitched invariants
2255     // we cant just go replace them with constants in both branches.
2256     //
2257     // For 'AND' we infer that true branch ("continue") means true
2258     // for each invariant operand.
2259     // For 'OR' we can infer that false branch ("continue") means false
2260     // for each invariant operand.
2261     // So it happens that for multiple-partial case we dont replace
2262     // in the unswitched branch.
2263     bool ReplaceUnswitched = FullUnswitch || (Invariants.size() == 1);
2264 
2265     ConstantInt *UnswitchedReplacement =
2266         Direction ? ConstantInt::getTrue(BI->getContext())
2267                   : ConstantInt::getFalse(BI->getContext());
2268     ConstantInt *ContinueReplacement =
2269         Direction ? ConstantInt::getFalse(BI->getContext())
2270                   : ConstantInt::getTrue(BI->getContext());
2271     for (Value *Invariant : Invariants)
2272       for (auto UI = Invariant->use_begin(), UE = Invariant->use_end();
2273            UI != UE;) {
2274         // Grab the use and walk past it so we can clobber it in the use list.
2275         Use *U = &*UI++;
2276         Instruction *UserI = dyn_cast<Instruction>(U->getUser());
2277         if (!UserI)
2278           continue;
2279 
2280         // Replace it with the 'continue' side if in the main loop body, and the
2281         // unswitched if in the cloned blocks.
2282         if (DT.dominates(LoopPH, UserI->getParent()))
2283           U->set(ContinueReplacement);
2284         else if (ReplaceUnswitched &&
2285                  DT.dominates(ClonedPH, UserI->getParent()))
2286           U->set(UnswitchedReplacement);
2287       }
2288   }
2289 
2290   // We can change which blocks are exit blocks of all the cloned sibling
2291   // loops, the current loop, and any parent loops which shared exit blocks
2292   // with the current loop. As a consequence, we need to re-form LCSSA for
2293   // them. But we shouldn't need to re-form LCSSA for any child loops.
2294   // FIXME: This could be made more efficient by tracking which exit blocks are
2295   // new, and focusing on them, but that isn't likely to be necessary.
2296   //
2297   // In order to reasonably rebuild LCSSA we need to walk inside-out across the
2298   // loop nest and update every loop that could have had its exits changed. We
2299   // also need to cover any intervening loops. We add all of these loops to
2300   // a list and sort them by loop depth to achieve this without updating
2301   // unnecessary loops.
2302   auto UpdateLoop = [&](Loop &UpdateL) {
2303 #ifndef NDEBUG
2304     UpdateL.verifyLoop();
2305     for (Loop *ChildL : UpdateL) {
2306       ChildL->verifyLoop();
2307       assert(ChildL->isRecursivelyLCSSAForm(DT, LI) &&
2308              "Perturbed a child loop's LCSSA form!");
2309     }
2310 #endif
2311     // First build LCSSA for this loop so that we can preserve it when
2312     // forming dedicated exits. We don't want to perturb some other loop's
2313     // LCSSA while doing that CFG edit.
2314     formLCSSA(UpdateL, DT, &LI, SE);
2315 
2316     // For loops reached by this loop's original exit blocks we may
2317     // introduced new, non-dedicated exits. At least try to re-form dedicated
2318     // exits for these loops. This may fail if they couldn't have dedicated
2319     // exits to start with.
2320     formDedicatedExitBlocks(&UpdateL, &DT, &LI, MSSAU, /*PreserveLCSSA*/ true);
2321   };
2322 
2323   // For non-child cloned loops and hoisted loops, we just need to update LCSSA
2324   // and we can do it in any order as they don't nest relative to each other.
2325   //
2326   // Also check if any of the loops we have updated have become top-level loops
2327   // as that will necessitate widening the outer loop scope.
2328   for (Loop *UpdatedL :
2329        llvm::concat<Loop *>(NonChildClonedLoops, HoistedLoops)) {
2330     UpdateLoop(*UpdatedL);
2331     if (!UpdatedL->getParentLoop())
2332       OuterExitL = nullptr;
2333   }
2334   if (IsStillLoop) {
2335     UpdateLoop(L);
2336     if (!L.getParentLoop())
2337       OuterExitL = nullptr;
2338   }
2339 
2340   // If the original loop had exit blocks, walk up through the outer most loop
2341   // of those exit blocks to update LCSSA and form updated dedicated exits.
2342   if (OuterExitL != &L)
2343     for (Loop *OuterL = ParentL; OuterL != OuterExitL;
2344          OuterL = OuterL->getParentLoop())
2345       UpdateLoop(*OuterL);
2346 
2347 #ifndef NDEBUG
2348   // Verify the entire loop structure to catch any incorrect updates before we
2349   // progress in the pass pipeline.
2350   LI.verify(DT);
2351 #endif
2352 
2353   // Now that we've unswitched something, make callbacks to report the changes.
2354   // For that we need to merge together the updated loops and the cloned loops
2355   // and check whether the original loop survived.
2356   SmallVector<Loop *, 4> SibLoops;
2357   for (Loop *UpdatedL : llvm::concat<Loop *>(NonChildClonedLoops, HoistedLoops))
2358     if (UpdatedL->getParentLoop() == ParentL)
2359       SibLoops.push_back(UpdatedL);
2360   UnswitchCB(IsStillLoop, SibLoops);
2361 
2362   if (MSSAU && VerifyMemorySSA)
2363     MSSAU->getMemorySSA()->verifyMemorySSA();
2364 
2365   if (BI)
2366     ++NumBranches;
2367   else
2368     ++NumSwitches;
2369 }
2370 
2371 /// Recursively compute the cost of a dominator subtree based on the per-block
2372 /// cost map provided.
2373 ///
2374 /// The recursive computation is memozied into the provided DT-indexed cost map
2375 /// to allow querying it for most nodes in the domtree without it becoming
2376 /// quadratic.
2377 static int
2378 computeDomSubtreeCost(DomTreeNode &N,
2379                       const SmallDenseMap<BasicBlock *, int, 4> &BBCostMap,
2380                       SmallDenseMap<DomTreeNode *, int, 4> &DTCostMap) {
2381   // Don't accumulate cost (or recurse through) blocks not in our block cost
2382   // map and thus not part of the duplication cost being considered.
2383   auto BBCostIt = BBCostMap.find(N.getBlock());
2384   if (BBCostIt == BBCostMap.end())
2385     return 0;
2386 
2387   // Lookup this node to see if we already computed its cost.
2388   auto DTCostIt = DTCostMap.find(&N);
2389   if (DTCostIt != DTCostMap.end())
2390     return DTCostIt->second;
2391 
2392   // If not, we have to compute it. We can't use insert above and update
2393   // because computing the cost may insert more things into the map.
2394   int Cost = std::accumulate(
2395       N.begin(), N.end(), BBCostIt->second, [&](int Sum, DomTreeNode *ChildN) {
2396         return Sum + computeDomSubtreeCost(*ChildN, BBCostMap, DTCostMap);
2397       });
2398   bool Inserted = DTCostMap.insert({&N, Cost}).second;
2399   (void)Inserted;
2400   assert(Inserted && "Should not insert a node while visiting children!");
2401   return Cost;
2402 }
2403 
2404 /// Turns a llvm.experimental.guard intrinsic into implicit control flow branch,
2405 /// making the following replacement:
2406 ///
2407 ///   --code before guard--
2408 ///   call void (i1, ...) @llvm.experimental.guard(i1 %cond) [ "deopt"() ]
2409 ///   --code after guard--
2410 ///
2411 /// into
2412 ///
2413 ///   --code before guard--
2414 ///   br i1 %cond, label %guarded, label %deopt
2415 ///
2416 /// guarded:
2417 ///   --code after guard--
2418 ///
2419 /// deopt:
2420 ///   call void (i1, ...) @llvm.experimental.guard(i1 false) [ "deopt"() ]
2421 ///   unreachable
2422 ///
2423 /// It also makes all relevant DT and LI updates, so that all structures are in
2424 /// valid state after this transform.
2425 static BranchInst *
2426 turnGuardIntoBranch(IntrinsicInst *GI, Loop &L,
2427                     SmallVectorImpl<BasicBlock *> &ExitBlocks,
2428                     DominatorTree &DT, LoopInfo &LI, MemorySSAUpdater *MSSAU) {
2429   SmallVector<DominatorTree::UpdateType, 4> DTUpdates;
2430   LLVM_DEBUG(dbgs() << "Turning " << *GI << " into a branch.\n");
2431   BasicBlock *CheckBB = GI->getParent();
2432 
2433   if (MSSAU && VerifyMemorySSA)
2434      MSSAU->getMemorySSA()->verifyMemorySSA();
2435 
2436   // Remove all CheckBB's successors from DomTree. A block can be seen among
2437   // successors more than once, but for DomTree it should be added only once.
2438   SmallPtrSet<BasicBlock *, 4> Successors;
2439   for (auto *Succ : successors(CheckBB))
2440     if (Successors.insert(Succ).second)
2441       DTUpdates.push_back({DominatorTree::Delete, CheckBB, Succ});
2442 
2443   Instruction *DeoptBlockTerm =
2444       SplitBlockAndInsertIfThen(GI->getArgOperand(0), GI, true);
2445   BranchInst *CheckBI = cast<BranchInst>(CheckBB->getTerminator());
2446   // SplitBlockAndInsertIfThen inserts control flow that branches to
2447   // DeoptBlockTerm if the condition is true.  We want the opposite.
2448   CheckBI->swapSuccessors();
2449 
2450   BasicBlock *GuardedBlock = CheckBI->getSuccessor(0);
2451   GuardedBlock->setName("guarded");
2452   CheckBI->getSuccessor(1)->setName("deopt");
2453   BasicBlock *DeoptBlock = CheckBI->getSuccessor(1);
2454 
2455   // We now have a new exit block.
2456   ExitBlocks.push_back(CheckBI->getSuccessor(1));
2457 
2458   if (MSSAU)
2459     MSSAU->moveAllAfterSpliceBlocks(CheckBB, GuardedBlock, GI);
2460 
2461   GI->moveBefore(DeoptBlockTerm);
2462   GI->setArgOperand(0, ConstantInt::getFalse(GI->getContext()));
2463 
2464   // Add new successors of CheckBB into DomTree.
2465   for (auto *Succ : successors(CheckBB))
2466     DTUpdates.push_back({DominatorTree::Insert, CheckBB, Succ});
2467 
2468   // Now the blocks that used to be CheckBB's successors are GuardedBlock's
2469   // successors.
2470   for (auto *Succ : Successors)
2471     DTUpdates.push_back({DominatorTree::Insert, GuardedBlock, Succ});
2472 
2473   // Make proper changes to DT.
2474   DT.applyUpdates(DTUpdates);
2475   // Inform LI of a new loop block.
2476   L.addBasicBlockToLoop(GuardedBlock, LI);
2477 
2478   if (MSSAU) {
2479     MemoryDef *MD = cast<MemoryDef>(MSSAU->getMemorySSA()->getMemoryAccess(GI));
2480     MSSAU->moveToPlace(MD, DeoptBlock, MemorySSA::BeforeTerminator);
2481     if (VerifyMemorySSA)
2482       MSSAU->getMemorySSA()->verifyMemorySSA();
2483   }
2484 
2485   ++NumGuards;
2486   return CheckBI;
2487 }
2488 
2489 /// Cost multiplier is a way to limit potentially exponential behavior
2490 /// of loop-unswitch. Cost is multipied in proportion of 2^number of unswitch
2491 /// candidates available. Also accounting for the number of "sibling" loops with
2492 /// the idea to account for previous unswitches that already happened on this
2493 /// cluster of loops. There was an attempt to keep this formula simple,
2494 /// just enough to limit the worst case behavior. Even if it is not that simple
2495 /// now it is still not an attempt to provide a detailed heuristic size
2496 /// prediction.
2497 ///
2498 /// TODO: Make a proper accounting of "explosion" effect for all kinds of
2499 /// unswitch candidates, making adequate predictions instead of wild guesses.
2500 /// That requires knowing not just the number of "remaining" candidates but
2501 /// also costs of unswitching for each of these candidates.
2502 static int CalculateUnswitchCostMultiplier(
2503     Instruction &TI, Loop &L, LoopInfo &LI, DominatorTree &DT,
2504     ArrayRef<std::pair<Instruction *, TinyPtrVector<Value *>>>
2505         UnswitchCandidates) {
2506 
2507   // Guards and other exiting conditions do not contribute to exponential
2508   // explosion as soon as they dominate the latch (otherwise there might be
2509   // another path to the latch remaining that does not allow to eliminate the
2510   // loop copy on unswitch).
2511   BasicBlock *Latch = L.getLoopLatch();
2512   BasicBlock *CondBlock = TI.getParent();
2513   if (DT.dominates(CondBlock, Latch) &&
2514       (isGuard(&TI) ||
2515        llvm::count_if(successors(&TI), [&L](BasicBlock *SuccBB) {
2516          return L.contains(SuccBB);
2517        }) <= 1)) {
2518     NumCostMultiplierSkipped++;
2519     return 1;
2520   }
2521 
2522   auto *ParentL = L.getParentLoop();
2523   int SiblingsCount = (ParentL ? ParentL->getSubLoopsVector().size()
2524                                : std::distance(LI.begin(), LI.end()));
2525   // Count amount of clones that all the candidates might cause during
2526   // unswitching. Branch/guard counts as 1, switch counts as log2 of its cases.
2527   int UnswitchedClones = 0;
2528   for (auto Candidate : UnswitchCandidates) {
2529     Instruction *CI = Candidate.first;
2530     BasicBlock *CondBlock = CI->getParent();
2531     bool SkipExitingSuccessors = DT.dominates(CondBlock, Latch);
2532     if (isGuard(CI)) {
2533       if (!SkipExitingSuccessors)
2534         UnswitchedClones++;
2535       continue;
2536     }
2537     int NonExitingSuccessors = llvm::count_if(
2538         successors(CondBlock), [SkipExitingSuccessors, &L](BasicBlock *SuccBB) {
2539           return !SkipExitingSuccessors || L.contains(SuccBB);
2540         });
2541     UnswitchedClones += Log2_32(NonExitingSuccessors);
2542   }
2543 
2544   // Ignore up to the "unscaled candidates" number of unswitch candidates
2545   // when calculating the power-of-two scaling of the cost. The main idea
2546   // with this control is to allow a small number of unswitches to happen
2547   // and rely more on siblings multiplier (see below) when the number
2548   // of candidates is small.
2549   unsigned ClonesPower =
2550       std::max(UnswitchedClones - (int)UnswitchNumInitialUnscaledCandidates, 0);
2551 
2552   // Allowing top-level loops to spread a bit more than nested ones.
2553   int SiblingsMultiplier =
2554       std::max((ParentL ? SiblingsCount
2555                         : SiblingsCount / (int)UnswitchSiblingsToplevelDiv),
2556                1);
2557   // Compute the cost multiplier in a way that won't overflow by saturating
2558   // at an upper bound.
2559   int CostMultiplier;
2560   if (ClonesPower > Log2_32(UnswitchThreshold) ||
2561       SiblingsMultiplier > UnswitchThreshold)
2562     CostMultiplier = UnswitchThreshold;
2563   else
2564     CostMultiplier = std::min(SiblingsMultiplier * (1 << ClonesPower),
2565                               (int)UnswitchThreshold);
2566 
2567   LLVM_DEBUG(dbgs() << "  Computed multiplier  " << CostMultiplier
2568                     << " (siblings " << SiblingsMultiplier << " * clones "
2569                     << (1 << ClonesPower) << ")"
2570                     << " for unswitch candidate: " << TI << "\n");
2571   return CostMultiplier;
2572 }
2573 
2574 static bool
2575 unswitchBestCondition(Loop &L, DominatorTree &DT, LoopInfo &LI,
2576                       AssumptionCache &AC, TargetTransformInfo &TTI,
2577                       function_ref<void(bool, ArrayRef<Loop *>)> UnswitchCB,
2578                       ScalarEvolution *SE, MemorySSAUpdater *MSSAU) {
2579   // Collect all invariant conditions within this loop (as opposed to an inner
2580   // loop which would be handled when visiting that inner loop).
2581   SmallVector<std::pair<Instruction *, TinyPtrVector<Value *>>, 4>
2582       UnswitchCandidates;
2583 
2584   // Whether or not we should also collect guards in the loop.
2585   bool CollectGuards = false;
2586   if (UnswitchGuards) {
2587     auto *GuardDecl = L.getHeader()->getParent()->getParent()->getFunction(
2588         Intrinsic::getName(Intrinsic::experimental_guard));
2589     if (GuardDecl && !GuardDecl->use_empty())
2590       CollectGuards = true;
2591   }
2592 
2593   for (auto *BB : L.blocks()) {
2594     if (LI.getLoopFor(BB) != &L)
2595       continue;
2596 
2597     if (CollectGuards)
2598       for (auto &I : *BB)
2599         if (isGuard(&I)) {
2600           auto *Cond = cast<IntrinsicInst>(&I)->getArgOperand(0);
2601           // TODO: Support AND, OR conditions and partial unswitching.
2602           if (!isa<Constant>(Cond) && L.isLoopInvariant(Cond))
2603             UnswitchCandidates.push_back({&I, {Cond}});
2604         }
2605 
2606     if (auto *SI = dyn_cast<SwitchInst>(BB->getTerminator())) {
2607       // We can only consider fully loop-invariant switch conditions as we need
2608       // to completely eliminate the switch after unswitching.
2609       if (!isa<Constant>(SI->getCondition()) &&
2610           L.isLoopInvariant(SI->getCondition()) && !BB->getUniqueSuccessor())
2611         UnswitchCandidates.push_back({SI, {SI->getCondition()}});
2612       continue;
2613     }
2614 
2615     auto *BI = dyn_cast<BranchInst>(BB->getTerminator());
2616     if (!BI || !BI->isConditional() || isa<Constant>(BI->getCondition()) ||
2617         BI->getSuccessor(0) == BI->getSuccessor(1))
2618       continue;
2619 
2620     if (L.isLoopInvariant(BI->getCondition())) {
2621       UnswitchCandidates.push_back({BI, {BI->getCondition()}});
2622       continue;
2623     }
2624 
2625     Instruction &CondI = *cast<Instruction>(BI->getCondition());
2626     if (CondI.getOpcode() != Instruction::And &&
2627       CondI.getOpcode() != Instruction::Or)
2628       continue;
2629 
2630     TinyPtrVector<Value *> Invariants =
2631         collectHomogenousInstGraphLoopInvariants(L, CondI, LI);
2632     if (Invariants.empty())
2633       continue;
2634 
2635     UnswitchCandidates.push_back({BI, std::move(Invariants)});
2636   }
2637 
2638   // If we didn't find any candidates, we're done.
2639   if (UnswitchCandidates.empty())
2640     return false;
2641 
2642   // Check if there are irreducible CFG cycles in this loop. If so, we cannot
2643   // easily unswitch non-trivial edges out of the loop. Doing so might turn the
2644   // irreducible control flow into reducible control flow and introduce new
2645   // loops "out of thin air". If we ever discover important use cases for doing
2646   // this, we can add support to loop unswitch, but it is a lot of complexity
2647   // for what seems little or no real world benefit.
2648   LoopBlocksRPO RPOT(&L);
2649   RPOT.perform(&LI);
2650   if (containsIrreducibleCFG<const BasicBlock *>(RPOT, LI))
2651     return false;
2652 
2653   SmallVector<BasicBlock *, 4> ExitBlocks;
2654   L.getUniqueExitBlocks(ExitBlocks);
2655 
2656   // We cannot unswitch if exit blocks contain a cleanuppad instruction as we
2657   // don't know how to split those exit blocks.
2658   // FIXME: We should teach SplitBlock to handle this and remove this
2659   // restriction.
2660   for (auto *ExitBB : ExitBlocks)
2661     if (isa<CleanupPadInst>(ExitBB->getFirstNonPHI())) {
2662       dbgs() << "Cannot unswitch because of cleanuppad in exit block\n";
2663       return false;
2664     }
2665 
2666   LLVM_DEBUG(
2667       dbgs() << "Considering " << UnswitchCandidates.size()
2668              << " non-trivial loop invariant conditions for unswitching.\n");
2669 
2670   // Given that unswitching these terminators will require duplicating parts of
2671   // the loop, so we need to be able to model that cost. Compute the ephemeral
2672   // values and set up a data structure to hold per-BB costs. We cache each
2673   // block's cost so that we don't recompute this when considering different
2674   // subsets of the loop for duplication during unswitching.
2675   SmallPtrSet<const Value *, 4> EphValues;
2676   CodeMetrics::collectEphemeralValues(&L, &AC, EphValues);
2677   SmallDenseMap<BasicBlock *, int, 4> BBCostMap;
2678 
2679   // Compute the cost of each block, as well as the total loop cost. Also, bail
2680   // out if we see instructions which are incompatible with loop unswitching
2681   // (convergent, noduplicate, or cross-basic-block tokens).
2682   // FIXME: We might be able to safely handle some of these in non-duplicated
2683   // regions.
2684   int LoopCost = 0;
2685   for (auto *BB : L.blocks()) {
2686     int Cost = 0;
2687     for (auto &I : *BB) {
2688       if (EphValues.count(&I))
2689         continue;
2690 
2691       if (I.getType()->isTokenTy() && I.isUsedOutsideOfBlock(BB))
2692         return false;
2693       if (auto *CB = dyn_cast<CallBase>(&I))
2694         if (CB->isConvergent() || CB->cannotDuplicate())
2695           return false;
2696 
2697       Cost += TTI.getUserCost(&I, TargetTransformInfo::TCK_CodeSize);
2698     }
2699     assert(Cost >= 0 && "Must not have negative costs!");
2700     LoopCost += Cost;
2701     assert(LoopCost >= 0 && "Must not have negative loop costs!");
2702     BBCostMap[BB] = Cost;
2703   }
2704   LLVM_DEBUG(dbgs() << "  Total loop cost: " << LoopCost << "\n");
2705 
2706   // Now we find the best candidate by searching for the one with the following
2707   // properties in order:
2708   //
2709   // 1) An unswitching cost below the threshold
2710   // 2) The smallest number of duplicated unswitch candidates (to avoid
2711   //    creating redundant subsequent unswitching)
2712   // 3) The smallest cost after unswitching.
2713   //
2714   // We prioritize reducing fanout of unswitch candidates provided the cost
2715   // remains below the threshold because this has a multiplicative effect.
2716   //
2717   // This requires memoizing each dominator subtree to avoid redundant work.
2718   //
2719   // FIXME: Need to actually do the number of candidates part above.
2720   SmallDenseMap<DomTreeNode *, int, 4> DTCostMap;
2721   // Given a terminator which might be unswitched, computes the non-duplicated
2722   // cost for that terminator.
2723   auto ComputeUnswitchedCost = [&](Instruction &TI, bool FullUnswitch) {
2724     BasicBlock &BB = *TI.getParent();
2725     SmallPtrSet<BasicBlock *, 4> Visited;
2726 
2727     int Cost = LoopCost;
2728     for (BasicBlock *SuccBB : successors(&BB)) {
2729       // Don't count successors more than once.
2730       if (!Visited.insert(SuccBB).second)
2731         continue;
2732 
2733       // If this is a partial unswitch candidate, then it must be a conditional
2734       // branch with a condition of either `or` or `and`. In that case, one of
2735       // the successors is necessarily duplicated, so don't even try to remove
2736       // its cost.
2737       if (!FullUnswitch) {
2738         auto &BI = cast<BranchInst>(TI);
2739         if (cast<Instruction>(BI.getCondition())->getOpcode() ==
2740             Instruction::And) {
2741           if (SuccBB == BI.getSuccessor(1))
2742             continue;
2743         } else {
2744           assert(cast<Instruction>(BI.getCondition())->getOpcode() ==
2745                      Instruction::Or &&
2746                  "Only `and` and `or` conditions can result in a partial "
2747                  "unswitch!");
2748           if (SuccBB == BI.getSuccessor(0))
2749             continue;
2750         }
2751       }
2752 
2753       // This successor's domtree will not need to be duplicated after
2754       // unswitching if the edge to the successor dominates it (and thus the
2755       // entire tree). This essentially means there is no other path into this
2756       // subtree and so it will end up live in only one clone of the loop.
2757       if (SuccBB->getUniquePredecessor() ||
2758           llvm::all_of(predecessors(SuccBB), [&](BasicBlock *PredBB) {
2759             return PredBB == &BB || DT.dominates(SuccBB, PredBB);
2760           })) {
2761         Cost -= computeDomSubtreeCost(*DT[SuccBB], BBCostMap, DTCostMap);
2762         assert(Cost >= 0 &&
2763                "Non-duplicated cost should never exceed total loop cost!");
2764       }
2765     }
2766 
2767     // Now scale the cost by the number of unique successors minus one. We
2768     // subtract one because there is already at least one copy of the entire
2769     // loop. This is computing the new cost of unswitching a condition.
2770     // Note that guards always have 2 unique successors that are implicit and
2771     // will be materialized if we decide to unswitch it.
2772     int SuccessorsCount = isGuard(&TI) ? 2 : Visited.size();
2773     assert(SuccessorsCount > 1 &&
2774            "Cannot unswitch a condition without multiple distinct successors!");
2775     return Cost * (SuccessorsCount - 1);
2776   };
2777   Instruction *BestUnswitchTI = nullptr;
2778   int BestUnswitchCost = 0;
2779   ArrayRef<Value *> BestUnswitchInvariants;
2780   for (auto &TerminatorAndInvariants : UnswitchCandidates) {
2781     Instruction &TI = *TerminatorAndInvariants.first;
2782     ArrayRef<Value *> Invariants = TerminatorAndInvariants.second;
2783     BranchInst *BI = dyn_cast<BranchInst>(&TI);
2784     int CandidateCost = ComputeUnswitchedCost(
2785         TI, /*FullUnswitch*/ !BI || (Invariants.size() == 1 &&
2786                                      Invariants[0] == BI->getCondition()));
2787     // Calculate cost multiplier which is a tool to limit potentially
2788     // exponential behavior of loop-unswitch.
2789     if (EnableUnswitchCostMultiplier) {
2790       int CostMultiplier =
2791           CalculateUnswitchCostMultiplier(TI, L, LI, DT, UnswitchCandidates);
2792       assert(
2793           (CostMultiplier > 0 && CostMultiplier <= UnswitchThreshold) &&
2794           "cost multiplier needs to be in the range of 1..UnswitchThreshold");
2795       CandidateCost *= CostMultiplier;
2796       LLVM_DEBUG(dbgs() << "  Computed cost of " << CandidateCost
2797                         << " (multiplier: " << CostMultiplier << ")"
2798                         << " for unswitch candidate: " << TI << "\n");
2799     } else {
2800       LLVM_DEBUG(dbgs() << "  Computed cost of " << CandidateCost
2801                         << " for unswitch candidate: " << TI << "\n");
2802     }
2803 
2804     if (!BestUnswitchTI || CandidateCost < BestUnswitchCost) {
2805       BestUnswitchTI = &TI;
2806       BestUnswitchCost = CandidateCost;
2807       BestUnswitchInvariants = Invariants;
2808     }
2809   }
2810   assert(BestUnswitchTI && "Failed to find loop unswitch candidate");
2811 
2812   if (BestUnswitchCost >= UnswitchThreshold) {
2813     LLVM_DEBUG(dbgs() << "Cannot unswitch, lowest cost found: "
2814                       << BestUnswitchCost << "\n");
2815     return false;
2816   }
2817 
2818   // If the best candidate is a guard, turn it into a branch.
2819   if (isGuard(BestUnswitchTI))
2820     BestUnswitchTI = turnGuardIntoBranch(cast<IntrinsicInst>(BestUnswitchTI), L,
2821                                          ExitBlocks, DT, LI, MSSAU);
2822 
2823   LLVM_DEBUG(dbgs() << "  Unswitching non-trivial (cost = "
2824                     << BestUnswitchCost << ") terminator: " << *BestUnswitchTI
2825                     << "\n");
2826   unswitchNontrivialInvariants(L, *BestUnswitchTI, BestUnswitchInvariants,
2827                                ExitBlocks, DT, LI, AC, UnswitchCB, SE, MSSAU);
2828   return true;
2829 }
2830 
2831 /// Unswitch control flow predicated on loop invariant conditions.
2832 ///
2833 /// This first hoists all branches or switches which are trivial (IE, do not
2834 /// require duplicating any part of the loop) out of the loop body. It then
2835 /// looks at other loop invariant control flows and tries to unswitch those as
2836 /// well by cloning the loop if the result is small enough.
2837 ///
2838 /// The `DT`, `LI`, `AC`, `TTI` parameters are required analyses that are also
2839 /// updated based on the unswitch.
2840 /// The `MSSA` analysis is also updated if valid (i.e. its use is enabled).
2841 ///
2842 /// If either `NonTrivial` is true or the flag `EnableNonTrivialUnswitch` is
2843 /// true, we will attempt to do non-trivial unswitching as well as trivial
2844 /// unswitching.
2845 ///
2846 /// The `UnswitchCB` callback provided will be run after unswitching is
2847 /// complete, with the first parameter set to `true` if the provided loop
2848 /// remains a loop, and a list of new sibling loops created.
2849 ///
2850 /// If `SE` is non-null, we will update that analysis based on the unswitching
2851 /// done.
2852 static bool unswitchLoop(Loop &L, DominatorTree &DT, LoopInfo &LI,
2853                          AssumptionCache &AC, TargetTransformInfo &TTI,
2854                          bool NonTrivial,
2855                          function_ref<void(bool, ArrayRef<Loop *>)> UnswitchCB,
2856                          ScalarEvolution *SE, MemorySSAUpdater *MSSAU) {
2857   assert(L.isRecursivelyLCSSAForm(DT, LI) &&
2858          "Loops must be in LCSSA form before unswitching.");
2859 
2860   // Must be in loop simplified form: we need a preheader and dedicated exits.
2861   if (!L.isLoopSimplifyForm())
2862     return false;
2863 
2864   // Try trivial unswitch first before loop over other basic blocks in the loop.
2865   if (unswitchAllTrivialConditions(L, DT, LI, SE, MSSAU)) {
2866     // If we unswitched successfully we will want to clean up the loop before
2867     // processing it further so just mark it as unswitched and return.
2868     UnswitchCB(/*CurrentLoopValid*/ true, {});
2869     return true;
2870   }
2871 
2872   // If we're not doing non-trivial unswitching, we're done. We both accept
2873   // a parameter but also check a local flag that can be used for testing
2874   // a debugging.
2875   if (!NonTrivial && !EnableNonTrivialUnswitch)
2876     return false;
2877 
2878   // For non-trivial unswitching, because it often creates new loops, we rely on
2879   // the pass manager to iterate on the loops rather than trying to immediately
2880   // reach a fixed point. There is no substantial advantage to iterating
2881   // internally, and if any of the new loops are simplified enough to contain
2882   // trivial unswitching we want to prefer those.
2883 
2884   // Try to unswitch the best invariant condition. We prefer this full unswitch to
2885   // a partial unswitch when possible below the threshold.
2886   if (unswitchBestCondition(L, DT, LI, AC, TTI, UnswitchCB, SE, MSSAU))
2887     return true;
2888 
2889   // No other opportunities to unswitch.
2890   return false;
2891 }
2892 
2893 PreservedAnalyses SimpleLoopUnswitchPass::run(Loop &L, LoopAnalysisManager &AM,
2894                                               LoopStandardAnalysisResults &AR,
2895                                               LPMUpdater &U) {
2896   Function &F = *L.getHeader()->getParent();
2897   (void)F;
2898 
2899   LLVM_DEBUG(dbgs() << "Unswitching loop in " << F.getName() << ": " << L
2900                     << "\n");
2901 
2902   // Save the current loop name in a variable so that we can report it even
2903   // after it has been deleted.
2904   std::string LoopName = std::string(L.getName());
2905 
2906   auto UnswitchCB = [&L, &U, &LoopName](bool CurrentLoopValid,
2907                                         ArrayRef<Loop *> NewLoops) {
2908     // If we did a non-trivial unswitch, we have added new (cloned) loops.
2909     if (!NewLoops.empty())
2910       U.addSiblingLoops(NewLoops);
2911 
2912     // If the current loop remains valid, we should revisit it to catch any
2913     // other unswitch opportunities. Otherwise, we need to mark it as deleted.
2914     if (CurrentLoopValid)
2915       U.revisitCurrentLoop();
2916     else
2917       U.markLoopAsDeleted(L, LoopName);
2918   };
2919 
2920   Optional<MemorySSAUpdater> MSSAU;
2921   if (AR.MSSA) {
2922     MSSAU = MemorySSAUpdater(AR.MSSA);
2923     if (VerifyMemorySSA)
2924       AR.MSSA->verifyMemorySSA();
2925   }
2926   if (!unswitchLoop(L, AR.DT, AR.LI, AR.AC, AR.TTI, NonTrivial, UnswitchCB,
2927                     &AR.SE, MSSAU.hasValue() ? MSSAU.getPointer() : nullptr))
2928     return PreservedAnalyses::all();
2929 
2930   if (AR.MSSA && VerifyMemorySSA)
2931     AR.MSSA->verifyMemorySSA();
2932 
2933   // Historically this pass has had issues with the dominator tree so verify it
2934   // in asserts builds.
2935   assert(AR.DT.verify(DominatorTree::VerificationLevel::Fast));
2936 
2937   auto PA = getLoopPassPreservedAnalyses();
2938   if (AR.MSSA)
2939     PA.preserve<MemorySSAAnalysis>();
2940   return PA;
2941 }
2942 
2943 namespace {
2944 
2945 class SimpleLoopUnswitchLegacyPass : public LoopPass {
2946   bool NonTrivial;
2947 
2948 public:
2949   static char ID; // Pass ID, replacement for typeid
2950 
2951   explicit SimpleLoopUnswitchLegacyPass(bool NonTrivial = false)
2952       : LoopPass(ID), NonTrivial(NonTrivial) {
2953     initializeSimpleLoopUnswitchLegacyPassPass(
2954         *PassRegistry::getPassRegistry());
2955   }
2956 
2957   bool runOnLoop(Loop *L, LPPassManager &LPM) override;
2958 
2959   void getAnalysisUsage(AnalysisUsage &AU) const override {
2960     AU.addRequired<AssumptionCacheTracker>();
2961     AU.addRequired<TargetTransformInfoWrapperPass>();
2962     if (EnableMSSALoopDependency) {
2963       AU.addRequired<MemorySSAWrapperPass>();
2964       AU.addPreserved<MemorySSAWrapperPass>();
2965     }
2966     getLoopAnalysisUsage(AU);
2967   }
2968 };
2969 
2970 } // end anonymous namespace
2971 
2972 bool SimpleLoopUnswitchLegacyPass::runOnLoop(Loop *L, LPPassManager &LPM) {
2973   if (skipLoop(L))
2974     return false;
2975 
2976   Function &F = *L->getHeader()->getParent();
2977 
2978   LLVM_DEBUG(dbgs() << "Unswitching loop in " << F.getName() << ": " << *L
2979                     << "\n");
2980 
2981   auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
2982   auto &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
2983   auto &AC = getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
2984   auto &TTI = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
2985   MemorySSA *MSSA = nullptr;
2986   Optional<MemorySSAUpdater> MSSAU;
2987   if (EnableMSSALoopDependency) {
2988     MSSA = &getAnalysis<MemorySSAWrapperPass>().getMSSA();
2989     MSSAU = MemorySSAUpdater(MSSA);
2990   }
2991 
2992   auto *SEWP = getAnalysisIfAvailable<ScalarEvolutionWrapperPass>();
2993   auto *SE = SEWP ? &SEWP->getSE() : nullptr;
2994 
2995   auto UnswitchCB = [&L, &LPM](bool CurrentLoopValid,
2996                                ArrayRef<Loop *> NewLoops) {
2997     // If we did a non-trivial unswitch, we have added new (cloned) loops.
2998     for (auto *NewL : NewLoops)
2999       LPM.addLoop(*NewL);
3000 
3001     // If the current loop remains valid, re-add it to the queue. This is
3002     // a little wasteful as we'll finish processing the current loop as well,
3003     // but it is the best we can do in the old PM.
3004     if (CurrentLoopValid)
3005       LPM.addLoop(*L);
3006     else
3007       LPM.markLoopAsDeleted(*L);
3008   };
3009 
3010   if (MSSA && VerifyMemorySSA)
3011     MSSA->verifyMemorySSA();
3012 
3013   bool Changed = unswitchLoop(*L, DT, LI, AC, TTI, NonTrivial, UnswitchCB, SE,
3014                               MSSAU.hasValue() ? MSSAU.getPointer() : nullptr);
3015 
3016   if (MSSA && VerifyMemorySSA)
3017     MSSA->verifyMemorySSA();
3018 
3019   // Historically this pass has had issues with the dominator tree so verify it
3020   // in asserts builds.
3021   assert(DT.verify(DominatorTree::VerificationLevel::Fast));
3022 
3023   return Changed;
3024 }
3025 
3026 char SimpleLoopUnswitchLegacyPass::ID = 0;
3027 INITIALIZE_PASS_BEGIN(SimpleLoopUnswitchLegacyPass, "simple-loop-unswitch",
3028                       "Simple unswitch loops", false, false)
3029 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
3030 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
3031 INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
3032 INITIALIZE_PASS_DEPENDENCY(LoopPass)
3033 INITIALIZE_PASS_DEPENDENCY(MemorySSAWrapperPass)
3034 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
3035 INITIALIZE_PASS_END(SimpleLoopUnswitchLegacyPass, "simple-loop-unswitch",
3036                     "Simple unswitch loops", false, false)
3037 
3038 Pass *llvm::createSimpleLoopUnswitchLegacyPass(bool NonTrivial) {
3039   return new SimpleLoopUnswitchLegacyPass(NonTrivial);
3040 }
3041