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

(On Diff #237182)

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 ↗

(On Diff #237182)

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/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);
   }