[LoopRotate] Add support for rotating loops with switch exit

[nikic]

Summary

Adds basic support for rotation loops with a switch terminator. This mainly involves making sure that multi edges to the exit block are handled correctly.

This does not yet fix https://bugs.llvm.org/show_bug.cgi?id=44461, because we require that there is only a single edge from the original header into the loop, while that case would require handling the multi edge case. That part is blocked on D72519.

[nikic]

Add tests

[efriedma]

llvm/include/llvm/Analysis/LoopInfoImpl.h

215

Move the new check out to avoid the extra contains() call?

llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp

283

I'm surprised nobody else has run into this.

llvm/lib/Transforms/Utils/LoopRotationUtils.cpp

489

This is going to cause non-deterministic code generation.

[nikic]

Extract condition. Use SetVector to avoid non-determinism.

[nikic]

llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp

283

Probably this is the first time PreserveLCSSA + MergeIdenticalEdges are used together.

[asbirlea]

Thank you, this looks like a great missed optimization.

llvm/lib/Transforms/Utils/LoopRotationUtils.cpp

222

I'm missing why not keep the check&exit if Term is a branch and it's unconditional.

[Gankra]

@nikic Is this patch still something that can/should be rebased and landed? (Investigating whether it can fix some issues we're seeing in Firefox related to Rust upgrading to llvm 10)

[nikic]

I'm not sure. It might make more sense to make SimplifyCFG be less aggressive about forming switches in this case. Pretty sure that converting if (a == C || B == C2) into a switch is pretty much a strict regression in optimization power, because a lot more passes support branches than switches.

[glandium]

FWIW, when applying this patch on a LLVM10-based rustc, the rustc build fails with:

PHI node has multiple entries for the same basic block with different incoming values!
  %35 = phi { i64*, i32 }* [ %83, %77 ], [ %84, %77 ], [ %33, %2 ]
label %77
  %83 = getelementptr { i64*, i32 }, { i64*, i32 }* %35, i64 -1
  %84 = getelementptr { i64*, i32 }, { i64*, i32 }* %35, i64 -1
in function _ZN80_$LT$alloc..vec..Vec$LT$T$GT$$u20$as$u20$alloc..vec..SpecExtend$LT$T$C$I$GT$$GT$9from_iter17h62b0d0798a862a64E
LLVM ERROR: Broken function found, compilation aborted!
error: could not compile `rustc_typeck`.

[xbolva00]

Reping

[nikic]

This is currently blocked on D72519.

[nikic]

Make this patch independent of D72519, by not handling multi-edges from the header into the loop yet. This means this patch doesn't address my motivating case anymore, but it would be a one line change to support it after D72519.

[nikic]

llvm/lib/Transforms/Utils/LoopRotationUtils.cpp

222

Unconditional branches will be rejected by the next check (!L->isLoopExiting(OrigHeader)) already, so it's not necessary to explicitly check for them. As this code no longer stores the BranchInst, I found it nicer to omit the check here.

llvm/include/llvm/Analysis/LoopInfoImpl.h

 /// A latch block is a block that contains a branch back to the header.
 template <class BlockT, class LoopT>
 BlockT *LoopBase<BlockT, LoopT>::getLoopLatch() const {
   assert(!isInvalid() && "Loop not in a valid state!");
   BlockT *Header = getHeader();
   BlockT *Latch = nullptr;
   for (const auto Pred : children<Inverse<BlockT *>>(Header)) {
     if (contains(Pred)) {
-      if (Latch)
+      if (Latch && Latch != Pred)
         return nullptr;
       Latch = Pred;
     }
   }
 
   return Latch;
 }
 

llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp

for (BasicBlock::iterator I = DestBB->begin(); isa<PHINode>(I); ++I) {
         BBIdx = PN->getBasicBlockIndex(TIBB);
       PN->setIncomingBlock(BBIdx, NewBB);
     }
   }
 
   // If there are any other edges from TIBB to DestBB, update those to go
   // through the split block, making those edges non-critical as well (and
   // reducing the number of phi entries in the DestBB if relevant).
+  unsigned NumEdges = 1;
   if (Options.MergeIdenticalEdges) {
     for (unsigned i = SuccNum+1, e = TI->getNumSuccessors(); i != e; ++i) {
       if (TI->getSuccessor(i) != DestBB) continue;
 
       // Remove an entry for TIBB from DestBB phi nodes.
       DestBB->removePredecessor(TIBB, Options.KeepOneInputPHIs);
 
       // We found another edge to DestBB, go to NewBB instead.
       TI->setSuccessor(i, NewBB);
+      ++NumEdges;
     }
   }
 
   // If we have nothing to update, just return.
   auto *DT = Options.DT;
   auto *PDT = Options.PDT;
   auto *LI = Options.LI;
   auto *MSSAU = Options.MSSAU;
if (Loop *TIL = LI->getLoopFor(TIBB)) {
       // If TIBB is in a loop and DestBB is outside of that loop, we may need
       // to update LoopSimplify form and LCSSA form.
       if (!TIL->contains(DestBB)) {
         assert(!TIL->contains(NewBB) &&
                "Split point for loop exit is contained in loop!");
 
         // Update LCSSA form in the newly created exit block.
         if (Options.PreserveLCSSA) {
-          createPHIsForSplitLoopExit(TIBB, NewBB, DestBB);
+          SmallVector<BasicBlock *, 1> Preds(NumEdges, TIBB);
+          createPHIsForSplitLoopExit(Preds, NewBB, DestBB);
         }
 
         // The only that we can break LoopSimplify form by splitting a critical
         // edge is if after the split there exists some edge from TIL to DestBB
         // *and* the only edge into DestBB from outside of TIL is that of
         // NewBB. If the first isn't true, then LoopSimplify still holds, NewBB
         // is the new exit block and it has no non-loop predecessors. If the
         // second isn't true, then DestBB was not in LoopSimplify form prior to

llvm/lib/Transforms/Utils/LoopRotationUtils.cpp

 bool LoopRotate::rotateLoop(Loop *L, bool SimplifiedLatch) {
   // If the loop has only one block then there is not much to rotate.
   if (L->getBlocks().size() == 1)
     return false;
 
   BasicBlock *OrigHeader = L->getHeader();
   BasicBlock *OrigLatch = L->getLoopLatch();
 
-  BranchInst *BI = dyn_cast<BranchInst>(OrigHeader->getTerminator());
-  if (!BI || BI->isUnconditional())
+  Instruction *Term = OrigHeader->getTerminator();
+  if (!isa<BranchInst>(Term) && !isa<SwitchInst>(Term))
     return false;
 
   // If the loop header is not one of the loop exiting blocks then
   // either this loop is already rotated or it is not
   // suitable for loop rotation transformations.
   if (!L->isLoopExiting(OrigHeader))
     return false;
 
   // If the loop latch already contains a branch that leaves the loop then the
   // loop is already rotated.
   if (!OrigLatch)
     return false;
 
   // Rotate if either the loop latch does *not* exit the loop, or if the loop
   // latch was just simplified. Or if we think it will be profitable.
   if (L->isLoopExiting(OrigLatch) && !SimplifiedLatch && IsUtilMode == false &&
       !shouldRotateLoopExitingLatch(L))
     return false;
 
+  // Find new Loop header. NewHeader is a Header's one and only successor
+  // that is inside loop.  Header's other successor is outside the
+  // loop.  Otherwise loop is not suitable for rotation.
+  BasicBlock *Exit = nullptr;
+  BasicBlock *NewHeader = nullptr;
+  for (BasicBlock *SuccBB : successors(OrigHeader)) {
+    if (L->contains(SuccBB)) {
+      // We don't have a new unique header.
+      if (NewHeader && NewHeader != SuccBB)
+        return false;
+      NewHeader = SuccBB;
+    } else {
+      // TODO: In principle, we could also support multiple exits here.
+      if (Exit && Exit != SuccBB)
+        return false;
+      Exit = SuccBB;
+    }
+  }
+
   // Check size of original header and reject loop if it is very big or we can't
   // duplicate blocks inside it.
   {
     SmallPtrSet<const Value *, 32> EphValues;
     CodeMetrics::collectEphemeralValues(L, AC, EphValues);
 
     CodeMetrics Metrics;
     Metrics.analyzeBasicBlock(OrigHeader, *TTI, EphValues);
bool LoopRotate::rotateLoop(Loop *L, bool SimplifiedLatch) {
   // infos in them.
   if (SE)
     SE->forgetTopmostLoop(L);
 
   LLVM_DEBUG(dbgs() << "LoopRotation: rotating "; L->dump());
   if (MSSAU && VerifyMemorySSA)
     MSSAU->getMemorySSA()->verifyMemorySSA();
 
-  // Find new Loop header. NewHeader is a Header's one and only successor
-  // that is inside loop.  Header's other successor is outside the
-  // loop.  Otherwise loop is not suitable for rotation.
-  BasicBlock *Exit = BI->getSuccessor(0);
-  BasicBlock *NewHeader = BI->getSuccessor(1);
-  if (L->contains(Exit))
-    std::swap(Exit, NewHeader);
-  assert(NewHeader && "Unable to determine new loop header");
-  assert(L->contains(NewHeader) && !L->contains(Exit) &&
-         "Unable to determine loop header and exit blocks");
-
   // This code assumes that the new header has exactly one predecessor.
   // Remove any single-entry PHI nodes in it.
-  assert(NewHeader->getSinglePredecessor() &&
+  assert(NewHeader->getUniquePredecessor() &&
          "New header doesn't have one pred!");
   FoldSingleEntryPHINodes(NewHeader);
 
   // Begin by walking OrigHeader and populating ValueMap with an entry for
   // each Instruction.
   BasicBlock::iterator I = OrigHeader->begin(), E = OrigHeader->end();
   ValueToValueMapTy ValueMap, ValueMapMSSA;
 
bool LoopRotate::rotateLoop(Loop *L, bool SimplifiedLatch) {
 
   // At this point, we've finished our major CFG changes.  As part of cloning
   // the loop into the preheader we've simplified instructions and the
   // duplicated conditional branch may now be branching on a constant.  If it is
   // branching on a constant and if that constant means that we enter the loop,
   // then we fold away the cond branch to an uncond branch.  This simplifies the
   // loop in cases important for nested loops, and it also means we don't have
   // to split as many edges.
-  BranchInst *PHBI = cast<BranchInst>(OrigPreheader->getTerminator());
-  assert(PHBI->isConditional() && "Should be clone of BI condbr!");
-  if (!isa<ConstantInt>(PHBI->getCondition()) ||
-      PHBI->getSuccessor(cast<ConstantInt>(PHBI->getCondition())->isZero()) !=
-          NewHeader) {
+  bool AlwaysBranchesIntoLoop = false;
+  Instruction *PHTerm = OrigPreheader->getTerminator();
+  if (auto *PHBI = dyn_cast<BranchInst>(PHTerm)) {
+    if (auto *Cond = dyn_cast<ConstantInt>(PHBI->getCondition()))
+      AlwaysBranchesIntoLoop = PHBI->getSuccessor(Cond->isZero()) == NewHeader;
+  } else if (auto *PHSI = dyn_cast<SwitchInst>(PHTerm)) {
+    if (auto *Cond = dyn_cast<ConstantInt>(PHSI->getCondition()))
+      AlwaysBranchesIntoLoop =
+          PHSI->findCaseValue(Cond)->getCaseSuccessor() == NewHeader;
+  }
+
+  if (!AlwaysBranchesIntoLoop) {
     // The conditional branch can't be folded, handle the general case.
     // Split edges as necessary to preserve LoopSimplify form.
 
     // Right now OrigPreHeader has two successors, NewHeader and ExitBlock, and
     // thus is not a preheader anymore.
     // Split the edge to form a real preheader.
     BasicBlock *NewPH = SplitCriticalEdge(
         OrigPreheader, NewHeader,
-        CriticalEdgeSplittingOptions(DT, LI, MSSAU).setPreserveLCSSA());
+        CriticalEdgeSplittingOptions(DT, LI, MSSAU)
+            .setPreserveLCSSA().setMergeIdenticalEdges());
     NewPH->setName(NewHeader->getName() + ".lr.ph");
 
     // Preserve canonical loop form, which means that 'Exit' should have only
     // one predecessor. Note that Exit could be an exit block for multiple
     // nested loops, causing both of the edges to now be critical and need to
     // be split.
-    SmallVector<BasicBlock *, 4> ExitPreds(pred_begin(Exit), pred_end(Exit));
+    SmallPtrSet<BasicBlock *, 4> ExitPreds(pred_begin(Exit), pred_end(Exit));
     bool SplitLatchEdge = false;
     for (BasicBlock *ExitPred : ExitPreds) {
       // We only need to split loop exit edges.
       Loop *PredLoop = LI->getLoopFor(ExitPred);
       if (!PredLoop || PredLoop->contains(Exit) ||
           ExitPred->getTerminator()->isIndirectTerminator())
         continue;
       SplitLatchEdge |= L->getLoopLatch() == ExitPred;
       BasicBlock *ExitSplit = SplitCriticalEdge(
           ExitPred, Exit,
-          CriticalEdgeSplittingOptions(DT, LI, MSSAU).setPreserveLCSSA());
+          CriticalEdgeSplittingOptions(DT, LI, MSSAU)
+              .setPreserveLCSSA().setMergeIdenticalEdges());
       ExitSplit->moveBefore(Exit);
     }
     assert(SplitLatchEdge &&
            "Despite splitting all preds, failed to split latch exit?");
   } else {
     // We can fold the conditional branch in the preheader, this makes things
-    // simpler. The first step is to remove the extra edge to the Exit block.
-    Exit->removePredecessor(OrigPreheader, true /*preserve LCSSA*/);
-    BranchInst *NewBI = BranchInst::Create(NewHeader, PHBI);
-    NewBI->setDebugLoc(PHBI->getDebugLoc());
-    PHBI->eraseFromParent();
+    // simpler. The first step is to remove extra edges to the Exit and
+    // NewHeader blocks. Only keep one NewHeader edge.
+    bool SeenNewHeader = false;
+    for (BasicBlock *SuccBB : successors(OrigPreheader)) {
+      if (SuccBB == NewHeader && !SeenNewHeader) {
+        SeenNewHeader = true;
+        continue;
+      }
+      SuccBB->removePredecessor(OrigPreheader, true /*preserve LCSSA*/);
+    }
+
+    BranchInst *NewBI = BranchInst::Create(NewHeader, PHTerm);
+    NewBI->setDebugLoc(PHTerm->getDebugLoc());
+    PHTerm->eraseFromParent();
 
     // With our CFG finalized, update DomTree if it is available.
     if (DT) DT->deleteEdge(OrigPreheader, Exit);
 
     // Update MSSA too, if available.
     if (MSSAU)
       MSSAU->removeEdge(OrigPreheader, Exit);
   }

llvm/test/Transforms/LoopRotate/switch.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -loop-rotate < %s | FileCheck %s
 
 define i64 @switch_multi_entry_known_entry() {
 ; CHECK-LABEL: @switch_multi_entry_known_entry(
 ; CHECK-NEXT:  start:
-; CHECK-NEXT:    br label [[HEADER:%.*]]
-; CHECK:       header:
-; CHECK-NEXT:    [[STATE:%.*]] = phi i8 [ 2, [[START:%.*]] ], [ [[NEXT_STATE:%.*]], [[LATCH:%.*]] ]
-; CHECK-NEXT:    [[COUNT:%.*]] = phi i64 [ 0, [[START]] ], [ [[INC:%.*]], [[LATCH]] ]
-; CHECK-NEXT:    switch i8 [[STATE]], label [[EXIT:%.*]] [
+; CHECK-NEXT:    br label [[LATCH:%.*]]
+; CHECK:       latch:
+; CHECK-NEXT:    [[COUNT2:%.*]] = phi i64 [ 0, [[START:%.*]] ], [ [[INC:%.*]], [[LATCH]] ], [ [[INC]], [[LATCH]] ]
+; CHECK-NEXT:    [[COUNT1:%.*]] = phi i64 [ 0, [[START]] ], [ [[INC]], [[LATCH]] ], [ [[INC]], [[LATCH]] ]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[COUNT2]], 999
+; CHECK-NEXT:    [[NEXT_STATE:%.*]] = zext i1 [[CMP]] to i8
+; CHECK-NEXT:    [[INC]] = add i64 [[COUNT1]], 1
+; CHECK-NEXT:    switch i8 [[NEXT_STATE]], label [[EXIT:%.*]] [
 ; CHECK-NEXT:    i8 0, label [[LATCH]]
 ; CHECK-NEXT:    i8 2, label [[LATCH]]
 ; CHECK-NEXT:    ]
-; CHECK:       latch:
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[COUNT]], 999
-; CHECK-NEXT:    [[NEXT_STATE]] = zext i1 [[CMP]] to i8
-; CHECK-NEXT:    [[INC]] = add i64 [[COUNT]], 1
-; CHECK-NEXT:    br label [[HEADER]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[COUNT_LCSSA:%.*]] = phi i64 [ [[COUNT]], [[HEADER]] ]
+; CHECK-NEXT:    [[COUNT_LCSSA:%.*]] = phi i64 [ [[INC]], [[LATCH]] ]
 ; CHECK-NEXT:    ret i64 [[COUNT_LCSSA]]
 ;
 start:
   br label %header
 
 header:                                           ; preds = %latch, %start
   %state = phi i8 [ 2, %start ], [ %next_state, %latch ]
   %count = phi i64 [ 0, %start ], [ %inc, %latch ]

 
 exit:                                             ; preds = %header
   ret i64 %count
 }
 
 define i64 @switch_multi_entry_unknown_entry(i8 %start_state) {
 ; CHECK-LABEL: @switch_multi_entry_unknown_entry(
 ; CHECK-NEXT:  start:
-; CHECK-NEXT:    br label [[HEADER:%.*]]
-; CHECK:       header:
-; CHECK-NEXT:    [[STATE:%.*]] = phi i8 [ [[START_STATE:%.*]], [[START:%.*]] ], [ [[NEXT_STATE:%.*]], [[LATCH:%.*]] ]
-; CHECK-NEXT:    [[COUNT:%.*]] = phi i64 [ 0, [[START]] ], [ [[INC:%.*]], [[LATCH]] ]
-; CHECK-NEXT:    switch i8 [[STATE]], label [[EXIT:%.*]] [
+; CHECK-NEXT:    switch i8 [[START_STATE:%.*]], label [[EXIT:%.*]] [
+; CHECK-NEXT:    i8 0, label [[LATCH_LR_PH:%.*]]
+; CHECK-NEXT:    i8 2, label [[LATCH_LR_PH]]
+; CHECK-NEXT:    ]
+; CHECK:       latch.lr.ph:
+; CHECK-NEXT:    br label [[LATCH:%.*]]
+; CHECK:       latch:
+; CHECK-NEXT:    [[COUNT2:%.*]] = phi i64 [ 0, [[LATCH_LR_PH]] ], [ [[INC:%.*]], [[LATCH]] ], [ [[INC]], [[LATCH]] ]
+; CHECK-NEXT:    [[COUNT1:%.*]] = phi i64 [ 0, [[LATCH_LR_PH]] ], [ [[INC]], [[LATCH]] ], [ [[INC]], [[LATCH]] ]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[COUNT2]], 999
+; CHECK-NEXT:    [[NEXT_STATE:%.*]] = zext i1 [[CMP]] to i8
+; CHECK-NEXT:    [[INC]] = add i64 [[COUNT1]], 1
+; CHECK-NEXT:    switch i8 [[NEXT_STATE]], label [[HEADER_EXIT_CRIT_EDGE:%.*]] [
 ; CHECK-NEXT:    i8 0, label [[LATCH]]
 ; CHECK-NEXT:    i8 2, label [[LATCH]]
 ; CHECK-NEXT:    ]
-; CHECK:       latch:
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[COUNT]], 999
-; CHECK-NEXT:    [[NEXT_STATE]] = zext i1 [[CMP]] to i8
-; CHECK-NEXT:    [[INC]] = add i64 [[COUNT]], 1
-; CHECK-NEXT:    br label [[HEADER]]
+; CHECK:       header.exit_crit_edge:
+; CHECK-NEXT:    [[SPLIT:%.*]] = phi i64 [ [[INC]], [[LATCH]] ]
+; CHECK-NEXT:    br label [[EXIT]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[COUNT_LCSSA:%.*]] = phi i64 [ [[COUNT]], [[HEADER]] ]
+; CHECK-NEXT:    [[COUNT_LCSSA:%.*]] = phi i64 [ [[SPLIT]], [[HEADER_EXIT_CRIT_EDGE]] ], [ 0, [[START:%.*]] ]
 ; CHECK-NEXT:    ret i64 [[COUNT_LCSSA]]
 ;
 start:
   br label %header
 
 header:                                           ; preds = %latch, %start
   %state = phi i8 [ %start_state, %start ], [ %next_state, %latch ]
   %count = phi i64 [ 0, %start ], [ %inc, %latch ]

 
 exit:                                             ; preds = %header
   ret i64 %count
 }
 
 define i64 @switch_multi_exit_known_entry() {
 ; CHECK-LABEL: @switch_multi_exit_known_entry(
 ; CHECK-NEXT:  start:
-; CHECK-NEXT:    br label [[HEADER:%.*]]
-; CHECK:       header:
-; CHECK-NEXT:    [[STATE:%.*]] = phi i8 [ 0, [[START:%.*]] ], [ [[NEXT_STATE:%.*]], [[LATCH:%.*]] ]
-; CHECK-NEXT:    [[COUNT:%.*]] = phi i64 [ 0, [[START]] ], [ [[INC:%.*]], [[LATCH]] ]
-; CHECK-NEXT:    switch i8 [[STATE]], label [[LATCH]] [
+; CHECK-NEXT:    br label [[LATCH:%.*]]
+; CHECK:       latch:
+; CHECK-NEXT:    [[COUNT1:%.*]] = phi i64 [ 0, [[START:%.*]] ], [ [[INC:%.*]], [[LATCH]] ]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[COUNT1]], 999
+; CHECK-NEXT:    [[NEXT_STATE:%.*]] = zext i1 [[CMP]] to i8
+; CHECK-NEXT:    [[INC]] = add i64 [[COUNT1]], 1
+; CHECK-NEXT:    switch i8 [[NEXT_STATE]], label [[LATCH]] [
 ; CHECK-NEXT:    i8 1, label [[EXIT:%.*]]
 ; CHECK-NEXT:    i8 2, label [[EXIT]]
 ; CHECK-NEXT:    ]
-; CHECK:       latch:
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[COUNT]], 999
-; CHECK-NEXT:    [[NEXT_STATE]] = zext i1 [[CMP]] to i8
-; CHECK-NEXT:    [[INC]] = add i64 [[COUNT]], 1
-; CHECK-NEXT:    br label [[HEADER]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[COUNT_LCSSA:%.*]] = phi i64 [ [[COUNT]], [[HEADER]] ], [ [[COUNT]], [[HEADER]] ]
+; CHECK-NEXT:    [[COUNT_LCSSA:%.*]] = phi i64 [ [[INC]], [[LATCH]] ], [ [[INC]], [[LATCH]] ]
 ; CHECK-NEXT:    ret i64 [[COUNT_LCSSA]]
 ;
 start:
   br label %header
 
 header:                                           ; preds = %latch, %start
   %state = phi i8 [ 0, %start ], [ %next_state, %latch ]
   %count = phi i64 [ 0, %start ], [ %inc, %latch ]

 
 exit:                                             ; preds = %header
   ret i64 %count
 }
 
 define i64 @switch_multi_exit_unknown_entry(i8 %start_state) {
 ; CHECK-LABEL: @switch_multi_exit_unknown_entry(
 ; CHECK-NEXT:  start:
-; CHECK-NEXT:    br label [[HEADER:%.*]]
-; CHECK:       header:
-; CHECK-NEXT:    [[STATE:%.*]] = phi i8 [ [[START_STATE:%.*]], [[START:%.*]] ], [ [[NEXT_STATE:%.*]], [[LATCH:%.*]] ]
-; CHECK-NEXT:    [[COUNT:%.*]] = phi i64 [ 0, [[START]] ], [ [[INC:%.*]], [[LATCH]] ]
-; CHECK-NEXT:    switch i8 [[STATE]], label [[LATCH]] [
+; CHECK-NEXT:    switch i8 [[START_STATE:%.*]], label [[LATCH_LR_PH:%.*]] [
 ; CHECK-NEXT:    i8 1, label [[EXIT:%.*]]
 ; CHECK-NEXT:    i8 2, label [[EXIT]]
 ; CHECK-NEXT:    ]
+; CHECK:       latch.lr.ph:
+; CHECK-NEXT:    br label [[LATCH:%.*]]
 ; CHECK:       latch:
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[COUNT]], 999
-; CHECK-NEXT:    [[NEXT_STATE]] = zext i1 [[CMP]] to i8
-; CHECK-NEXT:    [[INC]] = add i64 [[COUNT]], 1
-; CHECK-NEXT:    br label [[HEADER]]
+; CHECK-NEXT:    [[COUNT1:%.*]] = phi i64 [ 0, [[LATCH_LR_PH]] ], [ [[INC:%.*]], [[LATCH]] ]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[COUNT1]], 999
+; CHECK-NEXT:    [[NEXT_STATE:%.*]] = zext i1 [[CMP]] to i8
+; CHECK-NEXT:    [[INC]] = add i64 [[COUNT1]], 1
+; CHECK-NEXT:    switch i8 [[NEXT_STATE]], label [[LATCH]] [
+; CHECK-NEXT:    i8 1, label [[HEADER_EXIT_CRIT_EDGE:%.*]]
+; CHECK-NEXT:    i8 2, label [[HEADER_EXIT_CRIT_EDGE]]
+; CHECK-NEXT:    ]
+; CHECK:       header.exit_crit_edge:
+; CHECK-NEXT:    [[SPLIT:%.*]] = phi i64 [ [[INC]], [[LATCH]] ], [ [[INC]], [[LATCH]] ]
+; CHECK-NEXT:    br label [[EXIT]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[COUNT_LCSSA:%.*]] = phi i64 [ [[COUNT]], [[HEADER]] ], [ [[COUNT]], [[HEADER]] ]
+; CHECK-NEXT:    [[COUNT_LCSSA:%.*]] = phi i64 [ [[SPLIT]], [[HEADER_EXIT_CRIT_EDGE]] ], [ 0, [[START:%.*]] ], [ 0, [[START]] ]
 ; CHECK-NEXT:    ret i64 [[COUNT_LCSSA]]
 ;
 start:
   br label %header
 
 header:                                           ; preds = %latch, %start
   %state = phi i8 [ %start_state, %start ], [ %next_state, %latch ]
   %count = phi i64 [ 0, %start ], [ %inc, %latch ]