[LoopUnroll] Allow loops with multiple exiting blocks where loop latch is not necessary one of them.

[Whitney]

Summary

Currently LoopUnrollPass already allow loops with multiple exiting blocks, but it is only allowed when the loop latch is one of the exiting blocks.
When the loop latch is not an exiting block, then only single exiting block is supported.
When possible, the single loop latch or the single exiting block terminator is optimized to an unconditional branch in the unrolled loop.

This patch allows loops with multiple exiting blocks even if the loop latch is not one of them. However, the optimization of exiting block terminator to unconditional branch is not done when there exists more than one exiting block.

[Whitney]

Obtain ContinueOnTrue and LoopExit before modifying loop latch branch.

[etiotto]

llvm/test/Transforms/LoopUnroll/nonlatchcondbr.ll

217

Would be better IMO to have conditional branch that isn't foldable here.

[efriedma]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

797

We should probably rewrite the latch here even if ExitingBI is null.

[Meinersbur]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

310

[typo] conditional

[Whitney]

Addressed review comments.

[efriedma]

LGTM

llvm/lib/Transforms/Utils/LoopUnroll.cpp

LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
   // Loops with indirectbr cannot be cloned.
   if (!L->isSafeToClone()) {
     LLVM_DEBUG(dbgs() << "  Can't unroll; Loop body cannot be cloned.\n");
     return LoopUnrollResult::Unmodified;
   }
 
   // The current loop unroll pass can unroll loops that have
   // (1) single latch; and
-  // (2a) latch is an exiting block; or
-  // (2b) latch is unconditional and there exists a single exiting block.
+  // (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.
   BasicBlock *Header = L->getHeader();
   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 || !ExitingBI) {
-    LLVM_DEBUG(dbgs() << "  Can't unroll; loop not terminated by a conditional "
-                         "branch in latch or a single exiting block.\n");
-    return LoopUnrollResult::Unmodified;
-  }
-  if (LatchBI->isConditional() && LatchBI != ExitingBI) {
+  if (!LatchBI || (LatchBI->isConditional() && !LatchIsExiting)) {
     LLVM_DEBUG(
         dbgs() << "Can't unroll; a conditional latch must exit the loop");
     return LoopUnrollResult::Unmodified;
   }
-  LLVM_DEBUG(dbgs() << "  Exiting Block = " << ExitingBI->getParent()->getName()
-                    << "\n");
+  LLVM_DEBUG({
+    if (ExitingBI)
+      dbgs() << "  Exiting Block = " << ExitingBI->getParent()->getName()
+             << "\n";
+    else
+      dbgs() << "  No single exiting block\n";
+  });
 
   if (Header->hasAddressTaken()) {
     // The loop-rotate pass can be helpful to avoid this in many cases.
     LLVM_DEBUG(
         dbgs() << "  Won't unroll loop: address of header block is taken.\n");
     return LoopUnrollResult::Unmodified;
   }
 
if (SE) {
     if (ULO.ForgetAllSCEV)
       SE->forgetAllLoops();
     else
       SE->forgetTopmostLoop(L);
   }
 
   if (!LatchIsExiting)
     ++NumUnrolledNotLatch;
-  bool ContinueOnTrue = L->contains(ExitingBI->getSuccessor(0));
-  BasicBlock *LoopExit = ExitingBI->getSuccessor(ContinueOnTrue);
+  Optional<bool> ContinueOnTrue = None;
+  BasicBlock *LoopExit = nullptr;
+  if (ExitingBI) {
+    ContinueOnTrue = L->contains(ExitingBI->getSuccessor(0));
+    LoopExit = ExitingBI->getSuccessor(*ContinueOnTrue);
+  }
 
   // For the first iteration of the loop, we should use the precloned values for
   // 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 *> ExitingSucc;
   std::vector<BasicBlock *> Latches;
   Headers.push_back(Header);
   Latches.push_back(LatchBlock);
+  if (ExitingBI) {
     ExitingBlocks.push_back(ExitingBI->getParent());
-  ExitingSucc.push_back(ExitingBI->getSuccessor(!ContinueOnTrue));
+    ExitingSucc.push_back(ExitingBI->getSuccessor(!(*ContinueOnTrue)));
+  }
 
   // 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);
         if (*BB == ExitingSucc[0])
           ExitingSucc.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 (CompletelyUnroll) {
           InVal = LastValueMap[InVal];
       }
       assert(Latches.back() == LastValueMap[LatchBlock] && "bad last latch");
       PN->addIncoming(InVal, Latches.back());
     }
   }
 
   auto setDest = [](BasicBlock *Src, BasicBlock *Dest, BasicBlock *BlockInLoop,
-                    bool NeedConditional, bool ContinueOnTrue,
+                    bool NeedConditional, Optional<bool> ContinueOnTrue,
                     bool IsDestLoopExit) {
     auto *Term = cast<BranchInst>(Src->getTerminator());
     if (NeedConditional) {
       // Update the conditional branch's successor for the following
       // iteration.
-      Term->setSuccessor(!ContinueOnTrue, Dest);
+      assert(ContinueOnTrue.hasValue() &&
+             "Expecting valid ContinueOnTrue when NeedConditional is true");
+      Term->setSuccessor(!(*ContinueOnTrue), Dest);
     } else {
       // Remove phi operands at this loop exit
       if (!IsDestLoopExit) {
         BasicBlock *BB = Src;
         for (BasicBlock *Succ : successors(BB)) {
           // Preserve the incoming value from BB if we are jumping to the block
           // in the current loop.
           if (Succ == BlockInLoop)
for (unsigned i = 0, e = ExitingBlocks.size(); i != e; ++i) {
         NeedConditional = false;
 
       // Conditional branches from non-latch exiting block have successors
       // either in the same loop iteration or outside the loop. The branches are
       // already correct.
       if (NeedConditional)
         continue;
       setDest(ExitingBlocks[i], ExitingSucc[i], ExitingSucc[i], NeedConditional,
-              ContinueOnTrue, false);
+              None, false);
     }
 
     // When completely unrolling, the last latch becomes unreachable.
     if (CompletelyUnroll) {
       BranchInst *Term = cast<BranchInst>(Latches.back()->getTerminator());
       new UnreachableInst(Term->getContext(), Term);
       Term->eraseFromParent();
     }
   }
 
   // Update dominators of blocks we might reach through exits.
   // Immediate dominator of such block might change, because we add more
   // routes which can lead to the exit: we can now reach it from the copied
   // iterations too.
   if (DT && ULO.Count > 1) {
     for (auto *BB : OriginalLoopBlocks) {
       auto *BBDomNode = DT->getNode(BB);
       SmallVector<BasicBlock *, 16> ChildrenToUpdate;
       for (auto *ChildDomNode : BBDomNode->getChildren()) {
         auto *ChildBB = ChildDomNode->getBlock();
         if (!L->contains(ChildBB))
           ChildrenToUpdate.push_back(ChildBB);
       }
       BasicBlock *NewIDom;
-      BasicBlock *&TermBlock = ExitingBlocks[0];
-      auto &TermBlocks = ExitingBlocks;
-      if (BB == TermBlock) {
+      if (ExitingBI && BB == ExitingBlocks[0]) {
         // The latch is special because we emit unconditional branches in
         // some cases where the original loop contained a conditional branch.
         // Since the latch is always at the bottom of the loop, if the latch
         // dominated an exit before unrolling, the new dominator of that exit
         // must also be a latch.  Specifically, the dominator is the first
         // latch which ends in a conditional branch, or the last latch if
         // there is no such latch.
         // For loops exiting from non latch exiting block, we limit the
-        // supported loops to have a single exiting block.
-        NewIDom = TermBlocks.back();
-        for (unsigned i = 0, e = TermBlocks.size(); i != e; ++i) {
-          Instruction *Term = TermBlocks[i]->getTerminator();
+        // branch simplification to single exiting block loops.
+        NewIDom = ExitingBlocks.back();
+        for (unsigned i = 0, e = ExitingBlocks.size(); i != e; ++i) {
+          Instruction *Term = ExitingBlocks[i]->getTerminator();
           if (isa<BranchInst>(Term) && cast<BranchInst>(Term)->isConditional()) {
-            NewIDom = DT->findNearestCommonDominator(TermBlocks[i], Latches[i]);
+            NewIDom =
+                DT->findNearestCommonDominator(ExitingBlocks[i], Latches[i]);
             break;
           }
         }
       } else {
         // The new idom of the block will be the nearest common dominator
         // of all copies of the previous idom. This is equivalent to the
         // nearest common dominator of the previous idom and the first latch,
         // which dominates all copies of the previous idom.

llvm/test/Transforms/LoopUnroll/nonlatchcondbr.ll

for.body.for.body_crit_edge:
   %arrayidx.phi.trans.insert = getelementptr inbounds i32, i32* %A, i64 %inc
   %.pre = load i32, i32* %arrayidx.phi.trans.insert, align 4
   br label %for.header
 
 for.end:
   ret void
 }
 
+; Check that loop unroll pass correctly handle loops with
+; (1) multiple exiting blocks; and
+; (2) loop latch is not an exiting block.
+
+define void @test3(i32* noalias %A, i1 %cond) {
+; CHECK-LABEL: @test3(
+; 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:    [[TMP1:%.*]] = phi i32 [ [[TMP0]], [[ENTRY:%.*]] ], [ [[DOTPRE_3:%.*]], [[FOR_BODY_FOR_BODY_CRIT_EDGE_3:%.*]] ]
+; CHECK-NEXT:    [[I:%.*]] = phi i64 [ 0, [[ENTRY]] ], [ [[INC_3:%.*]], [[FOR_BODY_FOR_BODY_CRIT_EDGE_3]] ]
+; CHECK-NEXT:    call void @bar(i32 [[TMP1]])
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[FOR_BODY:%.*]], label [[FOR_END:%.*]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[INC:%.*]] = add nuw nsw i64 [[I]], 1
+; CHECK-NEXT:    br i1 true, label [[FOR_BODY_FOR_BODY_CRIT_EDGE:%.*]], label [[FOR_END]]
+; CHECK:       for.body.for.body_crit_edge:
+; CHECK-NEXT:    [[ARRAYIDX_PHI_TRANS_INSERT:%.*]] = getelementptr inbounds i32, i32* [[A]], i64 [[INC]]
+; 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:    [[INC_1:%.*]] = add nuw nsw i64 [[INC]], 1
+; CHECK-NEXT:    br i1 true, label [[FOR_BODY_FOR_BODY_CRIT_EDGE_1:%.*]], label [[FOR_END]]
+; CHECK:       for.body.for.body_crit_edge.1:
+; CHECK-NEXT:    [[ARRAYIDX_PHI_TRANS_INSERT_1:%.*]] = getelementptr inbounds i32, i32* [[A]], i64 [[INC_1]]
+; 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:    [[INC_2:%.*]] = add nuw nsw i64 [[INC_1]], 1
+; CHECK-NEXT:    br i1 true, label [[FOR_BODY_FOR_BODY_CRIT_EDGE_2:%.*]], label [[FOR_END]]
+; CHECK:       for.body.for.body_crit_edge.2:
+; CHECK-NEXT:    [[ARRAYIDX_PHI_TRANS_INSERT_2:%.*]] = getelementptr inbounds i32, i32* [[A]], i64 [[INC_2]]
+; 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:    [[INC_3]] = add nuw nsw i64 [[INC_2]], 1
+; CHECK-NEXT:    br i1 false, label [[FOR_BODY_FOR_BODY_CRIT_EDGE_3]], label [[FOR_END]]
+; CHECK:       for.body.for.body_crit_edge.3:
+; CHECK-NEXT:    [[ARRAYIDX_PHI_TRANS_INSERT_3:%.*]] = getelementptr inbounds i32, i32* [[A]], i64 [[INC_3]]
+; CHECK-NEXT:    [[DOTPRE_3]] = load i32, i32* [[ARRAYIDX_PHI_TRANS_INSERT_3]], align 4
+; CHECK-NEXT:    br label [[FOR_HEADER]], !llvm.loop !2
+;
+entry:
+  %0 = load i32, i32* %A, align 4
+  call void @bar(i32 %0)
+  br label %for.header
+
+for.header:
+  %1 = phi i32 [ %0, %entry ], [ %.pre, %for.body.for.body_crit_edge ]
+  %i = phi i64 [ 0, %entry ], [ %inc, %for.body.for.body_crit_edge ]
+  %arrayidx = getelementptr inbounds i32, i32* %A, i64 %i
+  call void @bar(i32 %1)
+  br i1 %cond, label %for.body, label %for.end
+
+for.body:
+  %inc = add nsw i64 %i, 1
+  %cmp = icmp slt i64 %inc, 4
+  br i1 %cmp, label %for.body.for.body_crit_edge, label %for.end
+
+for.body.for.body_crit_edge:
+  %arrayidx.phi.trans.insert = getelementptr inbounds i32, i32* %A, i64 %inc
+  %.pre = load i32, i32* %arrayidx.phi.trans.insert, align 4
+  br label %for.header
+
+for.end:
+  ret void
+}
+
 declare void @bar(i32)
 declare i1 @foo(i64)