[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 ↗

(On Diff #237182)

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

llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp

326–327

I'm surprised nobody else has run into this.

llvm/lib/Transforms/Utils/LoopRotationUtils.cpp

684–685

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

326–327

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

266–267

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

266–267

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/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 *MSSAU = Options.MSSAU;
   if (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);
         }
 
         if (!LoopPreds.empty()) {
           assert(!DestBB->isEHPad() && "We don't split edges to EH pads!");
           BasicBlock *NewExitBB = SplitBlockPredecessors(
               DestBB, LoopPreds, "split", DT, LI, MSSAU, Options.PreserveLCSSA);
           if (Options.PreserveLCSSA)
             createPHIsForSplitLoopExit(LoopPreds, NewExitBB, DestBB);

llvm/lib/Transforms/Utils/LoopRotationUtils.cpp

for (auto &DbgValue : DbgValues) {
                                                 ValueAsMetadata::get(NewVal)));
     }
   }
 }
 
 // Assuming both header and latch are exiting, look for a phi which is only
 // used outside the loop (via a LCSSA phi) in the exit from the header.
 // This means that rotating the loop can remove the phi.
-static bool profitableToRotateLoopExitingLatch(Loop *L) {
-  BasicBlock *Header = L->getHeader();
-  BranchInst *BI = dyn_cast<BranchInst>(Header->getTerminator());
-  assert(BI && BI->isConditional() && "need header with conditional exit");
-  BasicBlock *HeaderExit = BI->getSuccessor(0);
-  if (L->contains(HeaderExit))
-    HeaderExit = BI->getSuccessor(1);
-
+static bool profitableToRotateLoopExitingLatch(BasicBlock *Header,
+                                               BasicBlock *HeaderExit) {
   for (auto &Phi : Header->phis()) {
     // Look for uses of this phi in the loop/via exits other than the header.
     if (llvm::any_of(Phi.users(), [HeaderExit](const User *U) {
           return cast<Instruction>(U)->getParent() != HeaderExit;
         }))
       continue;
     return true;
   }

 /// happen once across multiple invocations of the LoopRotate pass.
 ///
 /// If -loop-rotate-multi is enabled we can do multiple rotations in one go
 /// so to reach a suitable (non-deoptimizing) exit.
 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;
 
   bool Rotated = false;
   do {
     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 Rotated;
 
     // 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 Rotated;
 
     // If the loop latch already contains a branch that leaves the loop then the
     // loop is already rotated.
     if (!OrigLatch)
       return Rotated;
 
+    // Find new Loop header. NewHeader is OrigHeader's one and only successor
+    // that is inside the loop. OrigHeader's other successor is outside the
+    // loop. Otherwise the 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.
+        // TODO: Allow multi-edge to NewHeader.
+        if (NewHeader)
+          return Rotated;
+        NewHeader = SuccBB;
+      } else {
+        // TODO: Support multiple exits.
+        if (Exit && Exit != SuccBB)
+          return Rotated;
+        Exit = SuccBB;
+      }
+    }
+
     // 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 &&
-        !profitableToRotateLoopExitingLatch(L) &&
+        !profitableToRotateLoopExitingLatch(OrigHeader, Exit) &&
         !canRotateDeoptimizingLatchExit(L))
       return Rotated;
 
     // 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);
do {
     // 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;
 
do {
 
     // 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());
+          OrigPreheader, NewHeader, 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));
+      SmallSetVector<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());
+            ExitPred, Exit, 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);
     }
 
     assert(L->getLoopPreheader() && "Invalid loop preheader after loop rotation");
     assert(L->getLoopLatch() && "Invalid loop latch after loop rotation");
 
     if (MSSAU && VerifyMemorySSA)
       MSSAU->getMemorySSA()->verifyMemorySSA();
 
     // Now that the CFG and DomTree are in a consistent state again, try to merge
     // the OrigHeader block into OrigLatch.  This will succeed if they are
     // connected by an unconditional branch.  This is just a cleanup so the
     // emitted code isn't too gross in this common case.
     DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
     BasicBlock *PredBB = OrigHeader->getUniquePredecessor();
     bool DidMerge = MergeBlockIntoPredecessor(OrigHeader, &DTU, LI, MSSAU);
     if (DidMerge)
       RemoveRedundantDbgInstrs(PredBB);
 
     if (MSSAU && VerifyMemorySSA)
       MSSAU->getMemorySSA()->verifyMemorySSA();
 
     LLVM_DEBUG(dbgs() << "LoopRotation: into "; L->dump());
 
     ++NumRotated;
 
     Rotated = true;
     SimplifiedLatch = false;

llvm/test/Transforms/LoopRotate/switch.ll

 
 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 ]