[LoopUnroll] Fold all exits based on known trip count/multiple

[nikic]

Summary

Fold all exits based on known trip count/multiple information from SCEV. Previously only the latch exit or the single exit were folded.

This doesn't yet eliminate ULO.TripCount and ULO.TripMultiple entirely: They're still used to a) decide whether runtime unrolling should be performed and b) for ORE remarks. However, the core unrolling logic is independent of them now.

[reames]

I think there's a problem with this code, but I'm surprised the existing test didn't catch it. Have you run this with explicit dom tree and loop info validation enabled?

Here's the problem I think I see. If we have an exit block which exits only on iteration X, and we unroll by count Y where Y > X, I believe we'll end up making the exit block unreachable. (This can't happen if there's only a single exit since that exit must be reached.) The problem with an unreachable exit is that there can be arbitrarily complicated loops reachable only through that path, and updating loop info in that case is hard. Am I missing something here?

llvm/lib/Transforms/Utils/LoopUnroll.cpp

747

This bit doesn't make sense to me, can you explain?

[nikic]

In D104203#2820498, @reames wrote:

I think there's a problem with this code, but I'm surprised the existing test didn't catch it. Have you run this with explicit dom tree and loop info validation enabled?

Here's the problem I think I see. If we have an exit block which exits only on iteration X, and we unroll by count Y where Y > X, I believe we'll end up making the exit block unreachable. (This can't happen if there's only a single exit since that exit must be reached.) The problem with an unreachable exit is that there can be arbitrarily complicated loops reachable only through that path, and updating loop info in that case is hard. Am I missing something here?

There's two bits of subtlety here: 1. For non-latch blocks, we don't fold branches towards exit (only towards loop) and 2. for a complete unroll, we always report the last iteration as exiting. These two properties ensure that we don't run into any of the hard LI invalidation scenarios, because we always retain both a branch to all exits (on the last iteration) and all loop blocks, even if we know that one or the other may be unreachable.

llvm/lib/Transforms/Utils/LoopUnroll.cpp

747

The runtime loop unrolling code may insert a prologue, in which case the computed trip counts for the original loop may no longer be correct.

[nikic]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

747

I should mention that runtime unrolling is going to need some more work, in particular I think that the decision whether to use a runtime unroll or not should be made by the caller, not this utility. Once that is done, it might make sense to compute the trip count information after doing the runtime unroll transform, in which case we could still use the trip multiple logic below this for the non-latch exits. If we teach SCEV to recognize the trip multiple from the runtime unroll pattern, then we shouldn't need to specially handle runtime unrolling here at all, though I'm not confident that we can detect the trip multiple sufficiently reliably (once folds get involved).

I think that's all followup work though, and the option used here is the conservative one of only folding the latch exit as we did before.

[reames]

LGTM, concerns addressed by responses. The bit about which exits we fold is particularly subtle. Might be worth adding a comment to the code.

llvm/lib/Transforms/Utils/LoopUnroll.cpp

LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
   const unsigned MaxTripCount = SE->getSmallConstantMaxTripCount(L);
   const bool MaxOrZero = SE->isBackedgeTakenCountMaxOrZero(L);
 
   // Effectively "DCE" unrolled iterations that are beyond the max tripcount
   // and will never be executed.
   if (MaxTripCount && ULO.Count > MaxTripCount)
     ULO.Count = MaxTripCount;
 
+  struct ExitInfo {
+    unsigned TripCount;
+    unsigned TripMultiple;
+    unsigned BreakoutTrip;
+    bool ExitOnTrue;
+    SmallVector<BasicBlock *> ExitingBlocks;
+  };
+  DenseMap<BasicBlock *, ExitInfo> ExitInfos;
+  SmallVector<BasicBlock *, 4> ExitingBlocks;
+  L->getExitingBlocks(ExitingBlocks);
+  for (auto *ExitingBlock : ExitingBlocks) {
+    // The folding code is not prepared to deal with non-branch instructions
+    // right now.
+    auto *BI = dyn_cast<BranchInst>(ExitingBlock->getTerminator());
+    if (!BI)
+      continue;
+
+    ExitInfo &Info = ExitInfos.try_emplace(ExitingBlock).first->second;
+    Info.TripCount = SE->getSmallConstantTripCount(L, ExitingBlock);
+    Info.TripMultiple = SE->getSmallConstantTripMultiple(L, ExitingBlock);
+    if (Info.TripCount != 0) {
+      Info.BreakoutTrip = Info.TripCount % ULO.Count;
+      Info.TripMultiple = 0;
+    } else {
+      Info.BreakoutTrip = Info.TripMultiple =
+          (unsigned)GreatestCommonDivisor64(ULO.Count, Info.TripMultiple);
+    }
+    Info.ExitOnTrue = !L->contains(BI->getSuccessor(0));
+    Info.ExitingBlocks.push_back(ExitingBlock);
+  }
+
   // Are we eliminating the loop control altogether?  Note that we can know
   // we're eliminating the backedge without knowing exactly which iteration
   // of the unrolled body exits.
   const bool CompletelyUnroll = ULO.Count == MaxTripCount;
 
   const bool PreserveOnlyFirst = CompletelyUnroll && MaxOrZero;
 
   // We assume a run-time trip count if the compiler cannot
LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
   // (2a) latch is unconditional; or
   // (2b) latch is conditional and is an exiting block
   // FIXME: The implementation can be extended to work with more complicated
   // cases, e.g. loops with multiple latches.
   BranchInst *LatchBI = dyn_cast<BranchInst>(LatchBlock->getTerminator());
 
   // A conditional branch which exits the loop, which can be optimized to an
   // unconditional branch in the unrolled loop in some cases.
-  BranchInst *ExitingBI = nullptr;
   bool LatchIsExiting = L->isLoopExiting(LatchBlock);
-  if (LatchIsExiting)
-    ExitingBI = LatchBI;
-  else if (BasicBlock *ExitingBlock = L->getExitingBlock())
-    ExitingBI = dyn_cast<BranchInst>(ExitingBlock->getTerminator());
   if (!LatchBI || (LatchBI->isConditional() && !LatchIsExiting)) {
     LLVM_DEBUG(
         dbgs() << "Can't unroll; a conditional latch must exit the loop");
     return LoopUnrollResult::Unmodified;
   }
-  LLVM_DEBUG({
-    if (ExitingBI)
-      dbgs() << "  Exiting Block = " << ExitingBI->getParent()->getName()
-             << "\n";
-    else
-      dbgs() << "  No single exiting block\n";
-  });
-
-  // Warning: ExactTripCount is the exact trip count for the block ending in
-  // ExitingBI, not neccessarily an exact exit count *for the loop*.  The
-  // distinction comes when we have an exiting latch, but the loop exits
-  // through another exit first.
-  const unsigned ExactTripCount = ExitingBI ?
-    SE->getSmallConstantTripCount(L,ExitingBI->getParent()) : 0;
 
   // Loops containing convergent instructions must have a count that divides
   // their TripMultiple.
   LLVM_DEBUG(
       {
         bool HasConvergent = false;
         for (auto &BB : L->blocks())
           for (auto &I : *BB)
if (RuntimeTripCount && ULO.TripMultiple % ULO.Count != 0 &&
     else {
       LLVM_DEBUG(dbgs() << "Won't unroll; remainder loop could not be "
                            "generated when assuming runtime trip count\n");
       return LoopUnrollResult::Unmodified;
     }
   }
 
   // If we know the trip count, we know the multiple...
+  // TODO: This is only used for the ORE code, remove it.
   unsigned BreakoutTrip = 0;
   if (ULO.TripCount != 0) {
     BreakoutTrip = ULO.TripCount % ULO.Count;
     ULO.TripMultiple = 0;
   } else {
     // Figure out what multiple to use.
     BreakoutTrip = ULO.TripMultiple =
         (unsigned)GreatestCommonDivisor64(ULO.Count, ULO.TripMultiple);
LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
   // PHI nodes.  Insert associations now.
   ValueToValueMapTy LastValueMap;
   std::vector<PHINode*> OrigPHINode;
   for (BasicBlock::iterator I = Header->begin(); isa<PHINode>(I); ++I) {
     OrigPHINode.push_back(cast<PHINode>(I));
   }
 
   std::vector<BasicBlock *> Headers;
-  std::vector<BasicBlock *> ExitingBlocks;
   std::vector<BasicBlock *> Latches;
   Headers.push_back(Header);
   Latches.push_back(LatchBlock);
-  if (ExitingBI)
-    ExitingBlocks.push_back(ExitingBI->getParent());
 
   // The current on-the-fly SSA update requires blocks to be processed in
   // reverse postorder so that LastValueMap contains the correct value at each
   // exit.
   LoopBlocksDFS DFS(L);
   DFS.perform(LI);
 
   // Stash the DFS iterators before adding blocks to the loop.
for (LoopBlocksDFS::RPOIterator BB = BlockBegin; BB != BlockEnd; ++BB) {
       // we can insert the proper branches later.
       if (*BB == Header)
         Headers.push_back(New);
       if (*BB == LatchBlock)
         Latches.push_back(New);
 
       // Keep track of the exiting block and its successor block contained in
       // the loop for the current iteration.
-      if (ExitingBI)
-        if (*BB == ExitingBlocks[0])
-          ExitingBlocks.push_back(New);
+      auto ExitInfoIt = ExitInfos.find(*BB);
+      if (ExitInfoIt != ExitInfos.end())
+        ExitInfoIt->second.ExitingBlocks.push_back(New);
 
       NewBlocks.push_back(New);
       UnrolledLoopBlocks.push_back(New);
 
       // Update DomTree: since we just copy the loop body, and each copy has a
       // dedicated entry block (copy of the header block), this header's copy
       // dominates all copied blocks. That means, dominance relations in the
       // copied body are the same as in the original body.
if (ULO.Count > 1) {
     }
   }
 
   assert(!UnrollVerifyDomtree ||
          DT->verify(DominatorTree::VerificationLevel::Fast));
 
   DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
 
-  if (ExitingBI) {
-    auto SetDest = [&](BasicBlock *Src, bool WillExit, bool ExitOnTrue) {
-      auto *Term = cast<BranchInst>(Src->getTerminator());
-      const unsigned Idx = ExitOnTrue ^ WillExit;
-      BasicBlock *Dest = Term->getSuccessor(Idx);
-      BasicBlock *DeadSucc = Term->getSuccessor(1-Idx);
 
-      // Remove predecessors from all non-Dest successors.
-      DeadSucc->removePredecessor(Src, /* KeepOneInputPHIs */ true);
 
-      // Replace the conditional branch with an unconditional one.
-      BranchInst::Create(Dest, Term);
-      Term->eraseFromParent();
 
-      DTU.applyUpdates({{DominatorTree::Delete, Src, DeadSucc}});
-    };
 
-    auto WillExit = [&](unsigned i, unsigned j) -> Optional<bool> {
-      if (CompletelyUnroll) {
-        if (PreserveOnlyFirst) {
-          if (i == 0)
-            return None;
-          return j == 0;
-        }
-        // Complete (but possibly inexact) unrolling
-        if (j == 0)
-          return true;
-        // Warning: ExactTripCount is the trip count of the exiting
-        // block which ends in ExitingBI, not neccessarily the loop.
-        if (ExactTripCount && j != ExactTripCount)
-          return false;
-        return None;
-      }
+  auto WillExit = [&](const ExitInfo &Info, unsigned i, unsigned j,
+                      bool IsLatch) -> Optional<bool> {
+      if (Info.TripCount && j != Info.TripCount)
 
-      if (RuntimeTripCount && j != 0)
+    if (RuntimeTripCount) {
+      // If runtime unrolling inserts a prologue, information about non-latch
+      // exits may be stale.
+      if (IsLatch && j != 0)
         return false;
+      return None;
+    }
 
-      if (j != BreakoutTrip &&
-          (ULO.TripMultiple == 0 || j % ULO.TripMultiple != 0)) {
-        // If we know the trip count or a multiple of it, we can safely use an
-        // unconditional branch for some iterations.
-        return false;
-      }
-      return None;
-    };
+    if (j != Info.BreakoutTrip &&
+        (Info.TripMultiple == 0 || j % Info.TripMultiple != 0)) {
 
-    // Fold branches for iterations where we know that they will exit or not
-    // exit.
-    bool ExitOnTrue = !L->contains(ExitingBI->getSuccessor(0));
-    for (unsigned i = 0, e = ExitingBlocks.size(); i != e; ++i) {
+  for (const auto &Pair : ExitInfos) {
+    const ExitInfo &Info = Pair.second;
+    for (unsigned i = 0, e = Info.ExitingBlocks.size(); i != e; ++i) {
       // The branch destination.
       unsigned j = (i + 1) % e;
-      Optional<bool> KnownWillExit = WillExit(i, j);
+      bool IsLatch = Pair.first == LatchBlock;
+      Optional<bool> KnownWillExit = WillExit(Info, i, j, IsLatch);
       if (!KnownWillExit)
         continue;
 
-      // TODO: Also fold known-exiting branches for non-latch exits.
-      if (*KnownWillExit && !LatchIsExiting)
+      // We don't fold known-exiting branches for non-latch exits here,
+      // because this ensures that both all loop blocks and all exit blocks
+      // remain reachable in the CFG.
+      // TODO: We could fold these branches, but it would require much more
+      // sophisticated updates to LoopInfo.
+      if (*KnownWillExit && !IsLatch)
         continue;
 
-      SetDest(ExitingBlocks[i], *KnownWillExit, ExitOnTrue);
+      SetDest(Info.ExitingBlocks[i], *KnownWillExit, Info.ExitOnTrue);
     }
   }
 
-
   // When completely unrolling, the last latch becomes unreachable.
   if (!LatchIsExiting && CompletelyUnroll)
     changeToUnreachable(Latches.back()->getTerminator(), /* UseTrap */ false,
                         PreserveLCSSA, &DTU);
 
   // Merge adjacent basic blocks, if possible.
   for (BasicBlock *Latch : Latches) {
     BranchInst *Term = dyn_cast<BranchInst>(Latch->getTerminator());

llvm/test/Transforms/LoopUnroll/multiple-exits.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -loop-unroll -S < %s | FileCheck %s
 
 declare void @bar()
 
 define void @test1() {
 ; CHECK-LABEL: @test1(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    call void @bar()
-; CHECK-NEXT:    br i1 true, label [[LATCH:%.*]], label [[EXIT:%.*]]
+; CHECK-NEXT:    br label [[LATCH:%.*]]
 ; CHECK:       latch:
 ; CHECK-NEXT:    call void @bar()
 ; CHECK-NEXT:    call void @bar()
-; CHECK-NEXT:    br i1 true, label [[LATCH_1:%.*]], label [[EXIT]]
+; CHECK-NEXT:    br label [[LATCH_1:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ; CHECK:       latch.1:
 ; CHECK-NEXT:    call void @bar()
 ; CHECK-NEXT:    call void @bar()
-; CHECK-NEXT:    br i1 true, label [[LATCH_2:%.*]], label [[EXIT]]
+; CHECK-NEXT:    br label [[LATCH_2:%.*]]
 ; CHECK:       latch.2:
 ; CHECK-NEXT:    call void @bar()
 ; CHECK-NEXT:    call void @bar()
-; CHECK-NEXT:    br i1 true, label [[LATCH_3:%.*]], label [[EXIT]]
+; CHECK-NEXT:    br label [[LATCH_3:%.*]]
 ; CHECK:       latch.3:
 ; CHECK-NEXT:    call void @bar()
 ; CHECK-NEXT:    call void @bar()
-; CHECK-NEXT:    br i1 true, label [[LATCH_4:%.*]], label [[EXIT]]
+; CHECK-NEXT:    br label [[LATCH_4:%.*]]
 ; CHECK:       latch.4:
 ; CHECK-NEXT:    call void @bar()
 ; CHECK-NEXT:    call void @bar()
-; CHECK-NEXT:    br i1 true, label [[LATCH_5:%.*]], label [[EXIT]]
+; CHECK-NEXT:    br label [[LATCH_5:%.*]]
 ; CHECK:       latch.5:
 ; CHECK-NEXT:    call void @bar()
 ; CHECK-NEXT:    call void @bar()
-; CHECK-NEXT:    br i1 true, label [[LATCH_6:%.*]], label [[EXIT]]
+; CHECK-NEXT:    br label [[LATCH_6:%.*]]
 ; CHECK:       latch.6:
 ; CHECK-NEXT:    call void @bar()
 ; CHECK-NEXT:    call void @bar()
-; CHECK-NEXT:    br i1 true, label [[LATCH_7:%.*]], label [[EXIT]]
+; CHECK-NEXT:    br label [[LATCH_7:%.*]]
 ; CHECK:       latch.7:
 ; CHECK-NEXT:    call void @bar()
 ; CHECK-NEXT:    call void @bar()
-; CHECK-NEXT:    br i1 true, label [[LATCH_8:%.*]], label [[EXIT]]
+; CHECK-NEXT:    br label [[LATCH_8:%.*]]
 ; CHECK:       latch.8:
 ; CHECK-NEXT:    call void @bar()
 ; CHECK-NEXT:    call void @bar()
-; CHECK-NEXT:    br i1 true, label [[LATCH_9:%.*]], label [[EXIT]]
+; CHECK-NEXT:    br label [[LATCH_9:%.*]]
 ; CHECK:       latch.9:
 ; CHECK-NEXT:    call void @bar()
 ; CHECK-NEXT:    call void @bar()
-; CHECK-NEXT:    br i1 false, label [[LATCH_10:%.*]], label [[EXIT]]
+; CHECK-NEXT:    br i1 false, label [[LATCH_10:%.*]], label [[EXIT:%.*]]
 ; CHECK:       latch.10:
 ; CHECK-NEXT:    call void @bar()
 ; CHECK-NEXT:    br label [[EXIT]]
 ;
 entry:
   br label %loop
 loop:
   %iv = phi i64 [0, %entry], [%iv.next, %latch]

llvm/test/Transforms/LoopUnroll/nonlatchcondbr.ll

 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = load i32, i32* [[A:%.*]], align 4
 ; CHECK-NEXT:    call void @bar(i32 [[TMP0]])
 ; CHECK-NEXT:    br label [[FOR_HEADER:%.*]]
 ; CHECK:       for.header:
 ; CHECK-NEXT:    call void @bar(i32 [[TMP0]])
 ; CHECK-NEXT:    br i1 [[COND:%.*]], label [[FOR_BODY:%.*]], label [[FOR_END:%.*]]
 ; CHECK:       for.body:
-; CHECK-NEXT:    br i1 true, label [[FOR_BODY_FOR_BODY_CRIT_EDGE:%.*]], label [[FOR_END]]
+; CHECK-NEXT:    br label [[FOR_BODY_FOR_BODY_CRIT_EDGE:%.*]]
 ; CHECK:       for.body.for.body_crit_edge:
 ; CHECK-NEXT:    [[ARRAYIDX_PHI_TRANS_INSERT:%.*]] = getelementptr inbounds i32, i32* [[A]], i64 1
 ; CHECK-NEXT:    [[DOTPRE:%.*]] = load i32, i32* [[ARRAYIDX_PHI_TRANS_INSERT]], align 4
 ; CHECK-NEXT:    call void @bar(i32 [[DOTPRE]])
 ; CHECK-NEXT:    br i1 [[COND]], label [[FOR_BODY_1:%.*]], label [[FOR_END]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    ret void
 ; CHECK:       for.body.1:
-; CHECK-NEXT:    br i1 true, label [[FOR_BODY_FOR_BODY_CRIT_EDGE_1:%.*]], label [[FOR_END]]
+; CHECK-NEXT:    br label [[FOR_BODY_FOR_BODY_CRIT_EDGE_1:%.*]]
 ; CHECK:       for.body.for.body_crit_edge.1:
 ; CHECK-NEXT:    [[ARRAYIDX_PHI_TRANS_INSERT_1:%.*]] = getelementptr inbounds i32, i32* [[A]], i64 2
 ; CHECK-NEXT:    [[DOTPRE_1:%.*]] = load i32, i32* [[ARRAYIDX_PHI_TRANS_INSERT_1]], align 4
 ; CHECK-NEXT:    call void @bar(i32 [[DOTPRE_1]])
 ; CHECK-NEXT:    br i1 [[COND]], label [[FOR_BODY_2:%.*]], label [[FOR_END]]
 ; CHECK:       for.body.2:
-; CHECK-NEXT:    br i1 true, label [[FOR_BODY_FOR_BODY_CRIT_EDGE_2:%.*]], label [[FOR_END]]
+; CHECK-NEXT:    br label [[FOR_BODY_FOR_BODY_CRIT_EDGE_2:%.*]]
 ; CHECK:       for.body.for.body_crit_edge.2:
 ; CHECK-NEXT:    [[ARRAYIDX_PHI_TRANS_INSERT_2:%.*]] = getelementptr inbounds i32, i32* [[A]], i64 3
 ; CHECK-NEXT:    [[DOTPRE_2:%.*]] = load i32, i32* [[ARRAYIDX_PHI_TRANS_INSERT_2]], align 4
 ; CHECK-NEXT:    call void @bar(i32 [[DOTPRE_2]])
 ; CHECK-NEXT:    br i1 [[COND]], label [[FOR_BODY_3:%.*]], label [[FOR_END]]
 ; CHECK:       for.body.3:
 ; CHECK-NEXT:    br i1 false, label [[FOR_BODY_FOR_BODY_CRIT_EDGE_3:%.*]], label [[FOR_END]]
 ; CHECK:       for.body.for.body_crit_edge.3: