https://github.com/lukel97 created https://github.com/llvm/llvm-project/pull/206244
Early exits with stores need to touch up the vector trip count used in the middle block terminator. Currently we eagerly create the branch in handlEarlyExits, but that means we need to avoid addMiddleCheck afterwards since we'll already have a branch.
We can move addMiddleCheck before handlEarlyExits to create the branch up front and then modify it for early exits with stores. This avoids the need to conditionally run addMiddleCheck based on the UncountableExitStyle, which makes it easier to compute the style purely in VPlan later.
>From c52cb4427c477ed5b3fb1562249799aa753169c0 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Sat, 27 Jun 2026 20:44:12 +0800
Subject: [PATCH] [VPlan] Run addMiddleCheck before early exit handling. NFC
Early exits with stores need to touch up the vector trip count used in the middle block terminator. Currently we eagerly create the branch in handlEarlyExits, but that means we need to avoid addMiddleCheck afterwards since we'll already have a branch.
We can move addMiddleCheck before handlEarlyExits to create the branch up front and then modify it for early exits with stores. This avoids the need to conditionally run addMiddleCheck based on the UncountableExitStyle, which makes it easier to compute the style purely in VPlan later.
---
llvm/lib/Transforms/Vectorize/LoopVectorize.cpp | 7 ++-----
llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp | 6 ++++--
.../LoopVectorize/VPlan/vplan-print-before-after-all.ll | 2 +-
3 files changed, 7 insertions(+), 8 deletions(-)
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 18e6890d13e13..a821dd24478ac 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -6562,6 +6562,8 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlan1() {
OptForSize, SCEVCheckThreshold, ORE, OrigLoop))
return nullptr;
+ RUN_VPLAN_PASS(VPlanTransforms::addMiddleCheck, *VPlan0);
+
// If we're vectorizing a loop with an uncountable exit, make sure that the
// recipes are safe to handle.
// TODO: Remove this once we can properly check the VPlan itself for both
@@ -6578,11 +6580,6 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlan1() {
return nullptr;
}
- // If we're handling uncountable exits in the scalar tail after a vector
- // loop with an in-loop mask, then the middle check has already been
- // created to compare against the actual number of lanes executed.
- if (EEStyle != UncountableExitStyle::MaskedHandleExitInScalarLoop)
- RUN_VPLAN_PASS(VPlanTransforms::addMiddleCheck, *VPlan0);
RUN_VPLAN_PASS(VPlanTransforms::createLoopRegions, *VPlan0,
getDebugLocFromInstOrOperands(Legal->getPrimaryInduction()));
if (CM.foldTailByMasking())
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index cfe21de1c7595..72151be13b241 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -4568,13 +4568,15 @@ static bool handleUncountableExitsWithSideEffects(
// against the full trip count, since we may be exiting the vector loop early.
// If we didn't take an early exit, we should get the equivalent of VF from
// the FirstActiveLane.
- VPBuilder MiddleBuilder(MiddleVPBB, MiddleVPBB->end());
+ VPBuilder MiddleBuilder(MiddleVPBB->getTerminator());
VPValue *ScalarIV = MiddleBuilder.createNaryOp(VPInstruction::ExtractLane,
{Zero, IV}, DebugLoc());
VPValue *ExitIV = MiddleBuilder.createAdd(ScalarIV, FirstActive);
VPValue *FullTC =
MiddleBuilder.createICmp(CmpInst::ICMP_EQ, ExitIV, Plan.getTripCount());
- MiddleBuilder.createNaryOp(VPInstruction::BranchOnCond, {FullTC});
+ assert(match(MiddleVPBB->getTerminator(), m_BranchOnCond()) &&
+ "Expected BranchOnCond terminator for MiddleVPBB");
+ MiddleVPBB->getTerminator()->setOperand(0, FullTC);
// Update resume phi in scalar.ph.
VPBasicBlock *ScalarPH = Plan.getScalarPreheader();
diff --git a/llvm/test/Transforms/LoopVectorize/VPlan/vplan-print-before-after-all.ll b/llvm/test/Transforms/LoopVectorize/VPlan/vplan-print-before-after-all.ll
index 86f709375286e..14590ff15c1d0 100644
--- a/llvm/test/Transforms/LoopVectorize/VPlan/vplan-print-before-after-all.ll
+++ b/llvm/test/Transforms/LoopVectorize/VPlan/vplan-print-before-after-all.ll
@@ -12,8 +12,8 @@
; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::createHeaderPhiRecipes
; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::replaceSymbolicStrides
; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::finalizeSCEVPredicates
-; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::handleEarlyExits
; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::addMiddleCheck
+; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::handleEarlyExits
; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::createLoopRegions
; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::introduceMasksAndLinearize
; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::createInLoopReductionRecipes