[VPlan] Also use original parent loop for exit VPBBs.

When vectorizing loops with early exits that is nested within another
one, one of the loop exits may be outside both loops, so setting adding
it to the parent loop is incorrect. Also use the original parent loop
for exit blocks.

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -411,12 +411,16 @@ void VPBasicBlock::connectToPredecessors(VPTransformState &State) {
   // Register NewBB in its loop. In innermost loops its the same for all
   // BB's.
   Loop *ParentLoop = State.CurrentParentLoop;
-  // If this block has a sole successor that is an exit block then it needs
-  // adding to the same parent loop as the exit block.
-  VPBlockBase *SuccVPBB = getSingleSuccessor();
-  if (SuccVPBB && State.Plan->isExitBlock(SuccVPBB))
-    ParentLoop =
-        State.LI->getLoopFor(cast<VPIRBasicBlock>(SuccVPBB)->getIRBasicBlock());
+  // If this block has a sole successor that is an exit block or is an exit
+  // block itself then it needs adding to the same parent loop as the exit
+  // block.
+  VPBlockBase *SuccOrExitVPB = getSingleSuccessor();
+  SuccOrExitVPB = SuccOrExitVPB ? SuccOrExitVPB : this;
+  if (State.Plan->isExitBlock(SuccOrExitVPB)) {
+    ParentLoop = State.LI->getLoopFor(
+        cast<VPIRBasicBlock>(SuccOrExitVPB)->getIRBasicBlock());
+  }
+
   if (ParentLoop && !State.LI->getLoopFor(NewBB))
     ParentLoop->addBasicBlockToLoop(NewBB, *State.LI);
 

llvm/test/Transforms/LoopVectorize/single_early_exit_with_outer_loop.ll

@@ -1,4 +1,4 @@
-; RUN: opt -S < %s -p loop-vectorize,'print<loops>' -disable-output -enable-early-exit-vectorization 2>&1 | FileCheck %s
+; RUN: opt -S < %s -passes='loop-vectorize,verify<loops>,print<loops>' -disable-output -enable-early-exit-vectorization 2>&1 | FileCheck %s
 
 declare void @init_mem(ptr, i64);
 
@@ -89,3 +89,32 @@ loop.outer.latch:
   %count.outer.next = add i64 %count.outer, %t
   br label %loop.outer
 }
+
+define i32 @early_exit_branch_to_outer_header() {
+; CHECK-LABEL: Loop info for function 'early_exit_branch_to_outer_header':
+; CHECK-NEXT:  Loop at depth 1 containing: %outer.header<header>,%loop.header,%loop.latch<exiting>,%outer.header.loopexit<latch>,%scalar.ph,%vector.ph,%vector.body,%middle.split,%middle.block<exiting>,%vector.early.exit
+; CHECK-NEXT:    Loop at depth 2 containing: %loop.header<header><exiting>,%loop.latch<latch><exiting>
+; CHECK-NEXT:    Loop at depth 2 containing: %vector.body<header><latch><exiting>
+entry:
+  %src = alloca [1024 x i8]
+  call void @init_mem(ptr %src, i64 1024)
+  br label %outer.header
+
+outer.header:
+  br label %loop.header
+
+loop.header:
+  %iv = phi i64 [ 0, %outer.header ], [ %iv.next, %loop.latch ]
+  %gep.src = getelementptr i8, ptr %src, i64 %iv
+  %l = load i8, ptr %gep.src, align 1
+  %c = icmp eq i8 %l, 0
+  br i1 %c, label %outer.header, label %loop.latch
+
+loop.latch:
+  %iv.next = add i64 %iv, 1
+  %ec = icmp eq i64 %iv.next, 1024
+  br i1 %ec, label %exit, label %loop.header
+
+exit:
+  ret i32 1
+}