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

[Florian Hahn via llvm-commits]

Author: Florian Hahn
Date: 2025-05-16T21:12:39+01:00
New Revision: ba93685ea228523c1f717f569e08185098b7defb

URL: https://github.com/llvm/llvm-project/commit/ba93685ea228523c1f717f569e08185098b7defb
DIFF: https://github.com/llvm/llvm-project/commit/ba93685ea228523c1f717f569e08185098b7defb.diff

LOG: [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.

Added: 
    

Modified: 
    llvm/lib/Transforms/Vectorize/VPlan.cpp
    llvm/test/Transforms/LoopVectorize/single_early_exit_with_outer_loop.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp
index c36feb0e2fdef..06b738afbd221 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.cpp
+++ b/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);
 

diff  --git a/llvm/test/Transforms/LoopVectorize/single_early_exit_with_outer_loop.ll b/llvm/test/Transforms/LoopVectorize/single_early_exit_with_outer_loop.ll
index 51cfc72752014..1aa63550e0316 100644
--- a/llvm/test/Transforms/LoopVectorize/single_early_exit_with_outer_loop.ll
+++ b/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
+}