[llvm] [VPlan] Account for early-exit dispatch blocks when updating LI. (PR #185618)
Florian Hahn via llvm-commits
llvm-commits at lists.llvm.org
Tue Mar 10 04:14:53 PDT 2026
https://github.com/fhahn created https://github.com/llvm/llvm-project/pull/185618
Now that we can vectorize loops with multiple early exits, we emit dispatch blocks after the middle block to go to a specific exit or continue in the dispatch chain.
With that, we need to be a bit more careful when it comes to picking the loop the dispatch block belongs to. The dispatch block will belong to the innermost loop of all exit blocks reachable from the current block.
Fixes https://github.com/llvm/llvm-project/issues/185362
>From 9fa9806b51ce80985ad574c847d72c91cb9734e8 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Tue, 10 Mar 2026 11:09:19 +0000
Subject: [PATCH] [VPlan] Account for early-exit dispatch blocks when updating
LI.
Now that we can vectorize loops with multiple early exits, we emit
dispatch blocks after the middle block to go to a specific exit or
continue in the dispatch chain.
With that, we need to be a bit more careful when it comes to picking the
loop the dispatch block belongs to. The dispatch block will belong to
the innermost loop of all exit blocks reachable from the current block.
Fixes https://github.com/llvm/llvm-project/issues/185362
---
llvm/lib/Transforms/Vectorize/VPlan.cpp | 38 +++-
.../single_early_exit_with_outer_loop.ll | 194 ++++++++++++++++++
2 files changed, 222 insertions(+), 10 deletions(-)
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp
index 0ceeb570e8b1f..eda82cca92217 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlan.cpp
@@ -390,22 +390,40 @@ BasicBlock *VPBasicBlock::createEmptyBasicBlock(VPTransformState &State) {
return NewBB;
}
+/// Determine the parent loop for \p VPBB by checking if it is an exit block
+/// or if all of its successors lead to exit blocks (following single-successor
+/// chains). Returns the innermost loop containing any of the exits if all
+/// paths lead to exits, or \p DefaultLoop otherwise.
+static Loop *getParentLoopForBlock(VPBlockBase *VPBB, VPlan &Plan, LoopInfo &LI,
+ Loop *DefaultLoop) {
+ Loop *ParentLoop = nullptr;
+ for (VPBlockBase *Succ : VPBB->getSuccessors()) {
+ while (Succ && !Plan.isExitBlock(Succ))
+ Succ = Succ->getSingleSuccessor();
+ if (!Succ)
+ return DefaultLoop;
+ Loop *L = LI.getLoopFor(cast<VPIRBasicBlock>(Succ)->getIRBasicBlock());
+ if (L && (!ParentLoop || ParentLoop->contains(L)))
+ ParentLoop = L;
+ }
+ return ParentLoop;
+}
+
void VPBasicBlock::connectToPredecessors(VPTransformState &State) {
auto &CFG = State.CFG;
BasicBlock *NewBB = CFG.VPBB2IRBB[this];
// 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 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());
- }
+ Loop *ParentLoop = nullptr;
+ if (State.Plan->isExitBlock(this))
+ ParentLoop =
+ State.LI->getLoopFor(cast<VPIRBasicBlock>(this)->getIRBasicBlock());
+ else if (getNumSuccessors() == 0)
+ ParentLoop = State.CurrentParentLoop;
+ else
+ ParentLoop = getParentLoopForBlock(this, *State.Plan, *State.LI,
+ State.CurrentParentLoop);
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 6ef55dcc334bf..b9ea300413d4b 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
@@ -119,3 +119,197 @@ loop.latch:
exit:
ret i32 1
}
+
+; Tests that when an early-exit inner loop has multiple exits and all exits
+; leave the outer loop.
+define i32 @multi_early_exit_all_leave_outer_loop(i1 %c) {
+; CHECK-LABEL: Loop info for function 'multi_early_exit_all_leave_outer_loop':
+; CHECK: Loop at depth 1 containing: %outer.header<header>,%outer.latch<latch><exiting>,%vector.ph,%vector.body<exiting>,%vector.body.interim,%middle.block
+; CHECK: {{.*}}Loop at depth 2 containing: %vector.body<header><exiting>,%vector.body.interim<latch><exiting>
+entry:
+ br label %outer.header
+
+outer.header:
+ br label %inner.header
+
+inner.header:
+ %iv = phi i64 [ 0, %outer.header ], [ %iv.next, %inner.latch ]
+ br i1 %c, label %early.exit.1, label %inner.body
+
+inner.body:
+ %trunc = trunc i64 %iv to i16
+ %cmp = icmp ult i16 %trunc, 0
+ br i1 %cmp, label %early.exit.2, label %inner.latch
+
+inner.latch:
+ %iv.next = add i64 %iv, 1
+ %exit.cond = icmp ult i64 %iv, 1
+ br i1 %exit.cond, label %inner.header, label %outer.latch
+
+outer.latch:
+ br i1 %c, label %early.exit.1, label %outer.header
+
+early.exit.1:
+ ret i32 1
+
+early.exit.2:
+ ret i32 0
+}
+
+; Tests that when an inner loop has two early exits at different loop levels
+; (one staying in the outer loop, one leaving all loops).
+define i32 @multi_early_exit_different_loop_levels(i1 %c) {
+; CHECK-LABEL: Loop info for function 'multi_early_exit_different_loop_levels':
+; CHECK: Loop at depth 1 containing: %outer.header<header>,%early.exit.outer,%outer.latch<latch>,%outer.latch.loopexit,%vector.ph,%vector.body,%vector.body.interim,%middle.block,%vector.early.exit.check<exiting>,%vector.early.exit.0
+; CHECK: {{.*}}Loop at depth 2 containing: %vector.body<header><exiting>,%vector.body.interim<latch><exiting>
+entry:
+ br label %outer.header
+
+outer.header:
+ br label %inner.header
+
+inner.header:
+ %iv = phi i64 [ 0, %outer.header ], [ %iv.next, %inner.latch ]
+ br i1 %c, label %early.exit.outer, label %inner.body
+
+inner.body:
+ %trunc = trunc i64 %iv to i16
+ %cmp = icmp ult i16 %trunc, 0
+ br i1 %cmp, label %early.exit.leave, label %inner.latch
+
+inner.latch:
+ %iv.next = add i64 %iv, 1
+ %exit.cond = icmp ult i64 %iv, 1
+ br i1 %exit.cond, label %inner.header, label %outer.latch
+
+outer.latch:
+ br label %outer.header
+
+early.exit.outer:
+ br label %outer.latch
+
+early.exit.leave:
+ ret i32 0
+}
+
+; Same as above, but the early exit order is reversed: the first early exit
+; leaves all loops and the second stays in the outer loop.
+define i32 @multi_early_exit_different_loop_levels_reversed(i1 %c) {
+; CHECK-LABEL: Loop info for function 'multi_early_exit_different_loop_levels_reversed':
+; CHECK: Loop at depth 1 containing: %outer.header<header>,%early.exit.outer,%outer.latch<latch>,%outer.latch.loopexit,%vector.ph,%vector.body,%vector.body.interim,%middle.block,%vector.early.exit.check<exiting>,%vector.early.exit.1
+; CHECK: {{.*}}Loop at depth 2 containing: %vector.body<header><exiting>,%vector.body.interim<latch><exiting>
+entry:
+ br label %outer.header
+
+outer.header:
+ br label %inner.header
+
+inner.header:
+ %iv = phi i64 [ 0, %outer.header ], [ %iv.next, %inner.latch ]
+ br i1 %c, label %early.exit.leave, label %inner.body
+
+inner.body:
+ %trunc = trunc i64 %iv to i16
+ %cmp = icmp ult i16 %trunc, 0
+ br i1 %cmp, label %early.exit.outer, label %inner.latch
+
+inner.latch:
+ %iv.next = add i64 %iv, 1
+ %exit.cond = icmp ult i64 %iv, 1
+ br i1 %exit.cond, label %inner.header, label %outer.latch
+
+outer.latch:
+ br label %outer.header
+
+early.exit.outer:
+ br label %outer.latch
+
+early.exit.leave:
+ ret i32 0
+}
+
+; Tests that when an inner loop has two early exits at different loop levels
+; (one going to a middle loop, one going to the outer loop).
+define i32 @multi_early_exit_two_different_loops(i1 %c) {
+; CHECK-LABEL: Loop info for function 'multi_early_exit_two_different_loops':
+; CHECK: Loop at depth 1 containing: %outer.header<header>,%middle.header,%early.exit.outer,%early.exit.middle,%middle.latch,%outer.latch<latch>,%middle.latch.loopexit,%outer.latch.loopexit,%vector.ph,%vector.body,%vector.body.interim,%middle.block,%vector.early.exit.check,%vector.early.exit.1,%vector.early.exit.0
+; CHECK: {{.*}}Loop at depth 2 containing: %middle.header<header>,%early.exit.middle,%middle.latch<latch><exiting>,%middle.latch.loopexit,%vector.ph,%vector.body,%vector.body.interim,%middle.block,%vector.early.exit.check<exiting>,%vector.early.exit.0
+; CHECK: {{.*}}Loop at depth 3 containing: %vector.body<header><exiting>,%vector.body.interim<latch><exiting>
+entry:
+ br label %outer.header
+
+outer.header:
+ br label %middle.header
+
+middle.header:
+ br label %inner.header
+
+inner.header:
+ %iv = phi i64 [ 0, %middle.header ], [ %iv.next, %inner.latch ]
+ br i1 %c, label %early.exit.middle, label %inner.body
+
+inner.body:
+ %trunc = trunc i64 %iv to i16
+ %cmp = icmp ult i16 %trunc, 0
+ br i1 %cmp, label %early.exit.outer, label %inner.latch
+
+inner.latch:
+ %iv.next = add i64 %iv, 1
+ %exit.cond = icmp ult i64 %iv, 1
+ br i1 %exit.cond, label %inner.header, label %middle.latch
+
+middle.latch:
+ br i1 %c, label %middle.header, label %outer.latch
+
+outer.latch:
+ br label %outer.header
+
+early.exit.middle:
+ br label %middle.latch
+
+early.exit.outer:
+ br label %outer.latch
+}
+
+; Same as above, but the early exit order is reversed: the first early exit
+; goes to the outer loop and the second goes to the middle loop.
+define i32 @multi_early_exit_two_different_loops_reversed(i1 %c) {
+; CHECK-LABEL: Loop info for function 'multi_early_exit_two_different_loops_reversed':
+; CHECK: Loop at depth 1 containing: %outer.header<header>,%middle.header,%early.exit.middle,%middle.latch,%early.exit.outer,%outer.latch<latch>,%middle.latch.loopexit,%outer.latch.loopexit,%vector.ph,%vector.body,%vector.body.interim,%middle.block,%vector.early.exit.check,%vector.early.exit.1,%vector.early.exit.0
+; CHECK: {{.*}}Loop at depth 2 containing: %middle.header<header>,%early.exit.middle,%middle.latch<latch><exiting>,%middle.latch.loopexit,%vector.ph,%vector.body,%vector.body.interim,%middle.block,%vector.early.exit.check<exiting>,%vector.early.exit.1
+; CHECK: {{.*}}Loop at depth 3 containing: %vector.body<header><exiting>,%vector.body.interim<latch><exiting>
+entry:
+ br label %outer.header
+
+outer.header:
+ br label %middle.header
+
+middle.header:
+ br label %inner.header
+
+inner.header:
+ %iv = phi i64 [ 0, %middle.header ], [ %iv.next, %inner.latch ]
+ br i1 %c, label %early.exit.outer, label %inner.body
+
+inner.body:
+ %trunc = trunc i64 %iv to i16
+ %cmp = icmp ult i16 %trunc, 0
+ br i1 %cmp, label %early.exit.middle, label %inner.latch
+
+inner.latch:
+ %iv.next = add i64 %iv, 1
+ %exit.cond = icmp ult i64 %iv, 1
+ br i1 %exit.cond, label %inner.header, label %middle.latch
+
+middle.latch:
+ br i1 %c, label %middle.header, label %outer.latch
+
+outer.latch:
+ br label %outer.header
+
+early.exit.middle:
+ br label %middle.latch
+
+early.exit.outer:
+ br label %outer.latch
+}
More information about the llvm-commits
mailing list