[llvm] [VPlan] Introduce reduction selects for tail folding in foldTailByMasking. NFCI (PR #192987)
Luke Lau via llvm-commits
llvm-commits at lists.llvm.org
Tue Apr 28 08:44:53 PDT 2026
https://github.com/lukel97 updated https://github.com/llvm/llvm-project/pull/192987
>From 8b956b0a53af9d9895c970d9d8bc171cdec91b0a Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Mon, 20 Apr 2026 15:31:59 +0100
Subject: [PATCH 1/4] [VPlan] Introduce reduction selects for tail folding in
foldTailByMasking. NFCI
Currently addComputeReductionResult has to check the cost model to see if the loop is tail folded, and if so then manually fix up the backedge value so any tail elements are ignored.
This PR moves this handling into foldTailByMasking itself so the plan doesn't requiring fixing up. We do this by setting the incoming value for the latch phi to the reduction phi instead of poison. A blend will be created for this automatically.
The main benefits of this are that the reduction is correct when tail folding is applied, and we don't need to worry about tail folding in as many places.
In order to preserve some of the optimizations that we get on VPInstruction::Select we need to convert the VPBlendRecipe to a select. This might not be necessessary if we decide to go down the route of https://github.com/llvm/llvm-project/pull/171851
Stacked on #192985
---
.../Transforms/Vectorize/LoopVectorize.cpp | 43 ++++++++++++-------
.../Vectorize/VPlanConstruction.cpp | 16 ++++---
.../LoopVectorize/VPlan/tail-folding.ll | 2 +-
3 files changed, 39 insertions(+), 22 deletions(-)
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 2cf693f055402..ea9677f96e46e 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -7531,9 +7531,14 @@ bool VPRecipeBuilder::replaceWithFinalIfReductionStore(
Legal->isInvariantAddressOfReduction(SI->getPointerOperand())) {
// Only create recipe for the final invariant store of the reduction.
if (Legal->isInvariantStoreOfReduction(SI)) {
+ VPValue *Val = VPI->getOperand(0), *Addr = VPI->getOperand(1);
+ // If tail folded, we need to use the blended reduction value out.
+ auto *BlendIt = find_if(Val->users(), IsaPred<VPBlendRecipe>);
+ if (BlendIt != Val->user_end())
+ Val = cast<VPBlendRecipe>(*BlendIt);
auto *Recipe = new VPReplicateRecipe(
- SI, VPI->operandsWithoutMask(), true /* IsUniform */,
- nullptr /*Mask*/, *VPI, *VPI, VPI->getDebugLoc());
+ SI, {Val, Addr}, true /* IsUniform */, nullptr /*Mask*/, *VPI, *VPI,
+ VPI->getDebugLoc());
FinalRedStoresBuilder.insert(Recipe);
}
VPI->eraseFromParent();
@@ -8022,6 +8027,7 @@ void LoopVectorizationPlanner::addReductionResultComputation(
VPBasicBlock *LatchVPBB = VectorLoopRegion->getExitingBasicBlock();
Builder.setInsertPoint(&*std::prev(std::prev(LatchVPBB->end())));
VPBasicBlock::iterator IP = MiddleVPBB->getFirstNonPhi();
+ VPValue *HeaderMask = vputils::findHeaderMask(*Plan);
for (VPRecipeBase &R :
Plan->getVectorLoopRegion()->getEntryBasicBlock()->phis()) {
VPReductionPHIRecipe *PhiR = dyn_cast<VPReductionPHIRecipe>(&R);
@@ -8034,23 +8040,28 @@ void LoopVectorizationPlanner::addReductionResultComputation(
const RecurrenceDescriptor &RdxDesc = Legal->getRecurrenceDescriptor(
cast<PHINode>(PhiR->getUnderlyingInstr()));
Type *PhiTy = TypeInfo.inferScalarType(PhiR);
- // If tail is folded by masking, introduce selects between the phi
- // and the users outside the vector region of each reduction, at the
- // beginning of the dedicated latch block.
+
+ // Convert a VPBlendRecipe backedge to a select.
+ if (auto *Blend = dyn_cast<VPBlendRecipe>(PhiR->getBackedgeValue())) {
+ if (Blend->getNumIncomingValues() == 2 &&
+ Blend->getMask(0) == HeaderMask) {
+ auto *Sel = VPBuilder(Blend).createSelect(
+ Blend->getMask(0), Blend->getIncomingValue(0),
+ Blend->getIncomingValue(1), {}, "", *Blend);
+ Blend->replaceAllUsesWith(Sel);
+ Blend->eraseFromParent();
+ }
+ }
+
auto *OrigExitingVPV = PhiR->getBackedgeValue();
auto *NewExitingVPV = PhiR->getBackedgeValue();
- if (!PhiR->isInLoop() && CM.foldTailByMasking()) {
- VPValue *Cond = vputils::findHeaderMask(*Plan);
- NewExitingVPV =
- Builder.createSelect(Cond, OrigExitingVPV, PhiR, {}, "", *PhiR);
- OrigExitingVPV->replaceUsesWithIf(NewExitingVPV, [](VPUser &U, unsigned) {
- return match(&U,
- m_VPInstruction<VPInstruction::ComputeReductionResult>());
- });
- if (CM.usePredicatedReductionSelect(RecurrenceKind))
- PhiR->setOperand(1, NewExitingVPV);
- }
+ // Remove the predicated select if the target doesn't want it.
+ VPValue *V;
+ if (!CM.usePredicatedReductionSelect(RecurrenceKind) &&
+ match(PhiR->getBackedgeValue(),
+ m_Select(m_Specific(HeaderMask), m_VPValue(V), m_Specific(PhiR))))
+ PhiR->setBackedgeValue(V);
// We want code in the middle block to appear to execute on the location of
// the scalar loop's latch terminator because: (a) it is all compiler
diff --git a/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp b/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
index f57b5298192b1..b172e79495ca0 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
@@ -1171,12 +1171,18 @@ void VPlanTransforms::foldTailByMasking(VPlan &Plan) {
for (const auto &[V, Users] : NeedsPhi) {
if (isa<VPIRValue>(V))
continue;
- // TODO: For reduction phis, use phi value instead of poison so we can
- // remove the special casing for tail folding in
- // LoopVectorizationPlanner::addReductionResultComputation
- VPValue *Poison =
+ VPValue *Backedge =
Plan.getOrAddLiveIn(PoisonValue::get(TypeInfo.inferScalarType(V)));
- VPInstruction *Phi = Builder.createScalarPhi({V, Poison});
+ VPIRFlags Flags;
+ auto *RedIt = find_if(V->users(), IsaPred<VPReductionPHIRecipe>);
+ auto *RdxPhi =
+ RedIt != V->user_end() ? cast<VPReductionPHIRecipe>(*RedIt) : nullptr;
+ if (RdxPhi && !RdxPhi->isInLoop()) {
+ Backedge = RdxPhi;
+ Flags = *RdxPhi;
+ }
+
+ VPInstruction *Phi = Builder.createScalarPhi({V, Backedge}, {}, "", Flags);
for (VPUser *U : Users)
U->replaceUsesOfWith(V, Phi);
}
diff --git a/llvm/test/Transforms/LoopVectorize/VPlan/tail-folding.ll b/llvm/test/Transforms/LoopVectorize/VPlan/tail-folding.ll
index dd8bfc8b366f8..b2c6675cd36a3 100644
--- a/llvm/test/Transforms/LoopVectorize/VPlan/tail-folding.ll
+++ b/llvm/test/Transforms/LoopVectorize/VPlan/tail-folding.ll
@@ -292,7 +292,7 @@ define i32 @reduction(ptr noalias %p, i32 %n) {
; CHECK-NEXT: Successor(s): vector.latch
; CHECK-EMPTY:
; CHECK-NEXT: vector.latch:
-; CHECK-NEXT: EMIT-SCALAR vp<[[VP8]]> = phi [ ir<%rdx.next>, vector.body.split ], [ ir<poison>, vector.body ]
+; CHECK-NEXT: EMIT-SCALAR vp<[[VP8]]> = phi [ ir<%rdx.next>, vector.body.split ], [ ir<%rdx>, vector.body ]
; CHECK-NEXT: EMIT vp<%index.next> = add nuw vp<[[VP4]]>, vp<[[VP1]]>
; CHECK-NEXT: EMIT branch-on-count vp<%index.next>, vp<[[VP2]]>
; CHECK-NEXT: No successors
>From 9265520722c12dae8a42185a3d7569376d359bac Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Mon, 20 Apr 2026 15:42:15 +0100
Subject: [PATCH 2/4] Update comment
---
llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp b/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
index b172e79495ca0..6a399b1732a7d 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
@@ -1165,7 +1165,7 @@ void VPlanTransforms::foldTailByMasking(VPlan &Plan) {
if (match(&R, m_ExtractLastPart(m_VPValue(V))))
NeedsPhi[V].push_back(&R);
- // Insert phis with a poison incoming value for past the end of the tail.
+ // Insert phis for values coming past the end of the tail.
Builder.setInsertPoint(Latch, Latch->begin());
VPTypeAnalysis TypeInfo(Plan);
for (const auto &[V, Users] : NeedsPhi) {
>From 64d477ebe8d200d22341ecddcb852a1fe9e278ad Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Tue, 28 Apr 2026 08:16:57 -0700
Subject: [PATCH 3/4] Address comments
* Rename Backedge -> TailVal
* Reuse existing Users variable
* Assert no multiple reduction phi users
---
llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp | 12 +++++++-----
1 file changed, 7 insertions(+), 5 deletions(-)
diff --git a/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp b/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
index 6a399b1732a7d..71203b379d8fd 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
@@ -1171,18 +1171,20 @@ void VPlanTransforms::foldTailByMasking(VPlan &Plan) {
for (const auto &[V, Users] : NeedsPhi) {
if (isa<VPIRValue>(V))
continue;
- VPValue *Backedge =
+ VPValue *TailVal =
Plan.getOrAddLiveIn(PoisonValue::get(TypeInfo.inferScalarType(V)));
VPIRFlags Flags;
- auto *RedIt = find_if(V->users(), IsaPred<VPReductionPHIRecipe>);
+ assert(llvm::count_if(Users, IsaPred<VPReductionPHIRecipe>) < 2 &&
+ "Value used by more than two reduction phis?");
+ auto *RedIt = find_if(Users, IsaPred<VPReductionPHIRecipe>);
auto *RdxPhi =
- RedIt != V->user_end() ? cast<VPReductionPHIRecipe>(*RedIt) : nullptr;
+ RedIt != Users.end() ? cast<VPReductionPHIRecipe>(*RedIt) : nullptr;
if (RdxPhi && !RdxPhi->isInLoop()) {
- Backedge = RdxPhi;
+ TailVal = RdxPhi;
Flags = *RdxPhi;
}
- VPInstruction *Phi = Builder.createScalarPhi({V, Backedge}, {}, "", Flags);
+ VPInstruction *Phi = Builder.createScalarPhi({V, TailVal}, {}, "", Flags);
for (VPUser *U : Users)
U->replaceUsesOfWith(V, Phi);
}
>From 5cb65b853762b232cd71ef0bc813d8d70708cadb Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Tue, 28 Apr 2026 08:44:18 -0700
Subject: [PATCH 4/4] Assert replaceWithFinalIfReductionStore is using backedge
value
---
llvm/lib/Transforms/Vectorize/LoopVectorize.cpp | 12 ++++++++----
1 file changed, 8 insertions(+), 4 deletions(-)
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 613c3af519ddc..d5cfd200fada3 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -6793,10 +6793,14 @@ bool VPRecipeBuilder::replaceWithFinalIfReductionStore(
// Only create recipe for the final invariant store of the reduction.
if (Legal->isInvariantStoreOfReduction(SI)) {
VPValue *Val = VPI->getOperand(0), *Addr = VPI->getOperand(1);
- // If tail folded, we need to use the blended reduction value out.
- auto *BlendIt = find_if(Val->users(), IsaPred<VPBlendRecipe>);
- if (BlendIt != Val->user_end())
- Val = cast<VPBlendRecipe>(*BlendIt);
+ // We need to store the exiting value of the reduction, so use the blend
+ // if tail folded.
+ if (auto *Blend = vputils::findUserOf<VPBlendRecipe>(Val))
+ Val = Blend;
+ assert(
+ vputils::findUserOf<VPReductionPHIRecipe>(Val)->getBackedgeValue() ==
+ Val &&
+ "Store isn't backedge value?");
auto *Recipe = new VPReplicateRecipe(
SI, {Val, Addr}, true /* IsUniform */, nullptr /*Mask*/, *VPI, *VPI,
VPI->getDebugLoc());
More information about the llvm-commits
mailing list