[llvm] [LV] Check full partial reduction chains in order. (PR #168036)
Florian Hahn via llvm-commits
llvm-commits at lists.llvm.org
Thu Nov 20 07:17:19 PST 2025
https://github.com/fhahn updated https://github.com/llvm/llvm-project/pull/168036
>From c47b261749dd097e177a3220f9a198965495912b Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Fri, 14 Nov 2025 09:54:07 +0000
Subject: [PATCH 1/3] [LV] Check full partial reduction chains in order.
https://github.com/llvm/llvm-project/pull/162822 added another
validation step to check if entries in a partial reduction chain have
the same scale factor. But the validation was still dependent on the
order of entries in PartialReductionChains, and would fail to reject
some cases (e.g. if the first first link matched the scale of the second
link, but the second link is invalidated later).
To fix that, group chains by their starting phi nodes, then perform the
validation for each chain, and if it fails, invalidate the whole chain
for the phi.
Fixes https://github.com/llvm/llvm-project/issues/167243.
Fixes https://github.com/llvm/llvm-project/issues/167867.
---
.../Transforms/Vectorize/LoopVectorize.cpp | 63 ++++++----
.../partial-reduce-incomplete-chains.ll | 114 ++++++++++++++++++
2 files changed, 151 insertions(+), 26 deletions(-)
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 267be2e243fc3..b808c37ac7b96 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -7992,13 +7992,13 @@ VPReplicateRecipe *VPRecipeBuilder::handleReplication(VPInstruction *VPI,
/// Find all possible partial reductions in the loop and track all of those that
/// are valid so recipes can be formed later.
void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
- // Find all possible partial reductions.
- SmallVector<std::pair<PartialReductionChain, unsigned>>
- PartialReductionChains;
- for (const auto &[Phi, RdxDesc] : Legal->getReductionVars()) {
+ // Find all possible partial reductions, grouping chains by their PHI.
+ MapVector<Instruction *,
+ SmallVector<std::pair<PartialReductionChain, unsigned>>>
+ ChainsByPhi;
+ for (const auto &[Phi, RdxDesc] : Legal->getReductionVars())
getScaledReductions(Phi, RdxDesc.getLoopExitInstr(), Range,
- PartialReductionChains);
- }
+ ChainsByPhi[Phi]);
// A partial reduction is invalid if any of its extends are used by
// something that isn't another partial reduction. This is because the
@@ -8006,8 +8006,9 @@ void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
// Build up a set of partial reduction ops for efficient use checking.
SmallPtrSet<User *, 4> PartialReductionOps;
- for (const auto &[PartialRdx, _] : PartialReductionChains)
- PartialReductionOps.insert(PartialRdx.ExtendUser);
+ for (const auto &[_, Chains] : ChainsByPhi)
+ for (const auto &[PartialRdx, _] : Chains)
+ PartialReductionOps.insert(PartialRdx.ExtendUser);
auto ExtendIsOnlyUsedByPartialReductions =
[&PartialReductionOps](Instruction *Extend) {
@@ -8018,31 +8019,41 @@ void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
// Check if each use of a chain's two extends is a partial reduction
// and only add those that don't have non-partial reduction users.
- for (auto Pair : PartialReductionChains) {
- PartialReductionChain Chain = Pair.first;
- if (ExtendIsOnlyUsedByPartialReductions(Chain.ExtendA) &&
- (!Chain.ExtendB || ExtendIsOnlyUsedByPartialReductions(Chain.ExtendB)))
- ScaledReductionMap.try_emplace(Chain.Reduction, Pair.second);
+ for (const auto &[_, Chains] : ChainsByPhi) {
+ for (const auto &[Chain, Scale] : Chains) {
+ if (ExtendIsOnlyUsedByPartialReductions(Chain.ExtendA) &&
+ (!Chain.ExtendB ||
+ ExtendIsOnlyUsedByPartialReductions(Chain.ExtendB)))
+ ScaledReductionMap.try_emplace(Chain.Reduction, Scale);
+ }
}
// Check that all partial reductions in a chain are only used by other
// partial reductions with the same scale factor. Otherwise we end up creating
// users of scaled reductions where the types of the other operands don't
// match.
- for (const auto &[Chain, Scale] : PartialReductionChains) {
- auto AllUsersPartialRdx = [ScaleVal = Scale, this](const User *U) {
- auto *UI = cast<Instruction>(U);
- if (isa<PHINode>(UI) && UI->getParent() == OrigLoop->getHeader()) {
- return all_of(UI->users(), [ScaleVal, this](const User *U) {
- auto *UI = cast<Instruction>(U);
- return ScaledReductionMap.lookup_or(UI, 0) == ScaleVal;
- });
+ for (const auto &[Phi, Chains] : ChainsByPhi) {
+ for (const auto &[Chain, Scale] : Chains) {
+ auto AllUsersPartialRdx = [ScaleVal = Scale, this](const User *U) {
+ auto *UI = cast<Instruction>(U);
+ if (isa<PHINode>(UI) && UI->getParent() == OrigLoop->getHeader()) {
+ return all_of(UI->users(), [ScaleVal, this](const User *U) {
+ auto *UI = cast<Instruction>(U);
+ return ScaledReductionMap.lookup_or(UI, 0) == ScaleVal;
+ });
+ }
+ return ScaledReductionMap.lookup_or(UI, 0) == ScaleVal ||
+ !OrigLoop->contains(UI->getParent());
+ };
+
+ // If any partial reduction entry for the phi is invalid, invalidate the
+ // whole chain.
+ if (!all_of(Chain.Reduction->users(), AllUsersPartialRdx)) {
+ for (const auto &[Chain, _] : Chains)
+ ScaledReductionMap.erase(Chain.Reduction);
+ break;
}
- return ScaledReductionMap.lookup_or(UI, 0) == ScaleVal ||
- !OrigLoop->contains(UI->getParent());
- };
- if (!all_of(Chain.Reduction->users(), AllUsersPartialRdx))
- ScaledReductionMap.erase(Chain.Reduction);
+ }
}
}
diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-incomplete-chains.ll b/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-incomplete-chains.ll
index fffab238798e3..47a5a6dc9b384 100644
--- a/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-incomplete-chains.ll
+++ b/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-incomplete-chains.ll
@@ -125,4 +125,118 @@ exit:
ret i16 %red.next
}
+
+define void @chained_sext_adds(ptr noalias %src, ptr noalias %dst) #0 {
+; CHECK-NEON-LABEL: define void @chained_sext_adds(
+; CHECK-NEON-SAME: ptr noalias [[SRC:%.*]], ptr noalias [[DST:%.*]]) #[[ATTR1]] {
+; CHECK-NEON-NEXT: [[ENTRY:.*:]]
+; CHECK-NEON-NEXT: br label %[[VECTOR_PH:.*]]
+; CHECK-NEON: [[VECTOR_PH]]:
+; CHECK-NEON-NEXT: br label %[[VECTOR_BODY:.*]]
+; CHECK-NEON: [[VECTOR_BODY]]:
+; CHECK-NEON-NEXT: [[INDEX:%.*]] = phi i64 [ 0, %[[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], %[[VECTOR_BODY]] ]
+; CHECK-NEON-NEXT: [[VEC_PHI:%.*]] = phi <4 x i32> [ zeroinitializer, %[[VECTOR_PH]] ], [ [[PARTIAL_REDUCE1:%.*]], %[[VECTOR_BODY]] ]
+; CHECK-NEON-NEXT: [[TMP0:%.*]] = getelementptr i8, ptr [[SRC]], i64 [[INDEX]]
+; CHECK-NEON-NEXT: [[WIDE_LOAD:%.*]] = load <16 x i8>, ptr [[TMP0]], align 1
+; CHECK-NEON-NEXT: [[TMP1:%.*]] = sext <16 x i8> [[WIDE_LOAD]] to <16 x i32>
+; CHECK-NEON-NEXT: [[PARTIAL_REDUCE:%.*]] = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(<4 x i32> [[VEC_PHI]], <16 x i32> [[TMP1]])
+; CHECK-NEON-NEXT: [[PARTIAL_REDUCE1]] = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(<4 x i32> [[PARTIAL_REDUCE]], <16 x i32> [[TMP1]])
+; CHECK-NEON-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 16
+; CHECK-NEON-NEXT: [[TMP2:%.*]] = icmp eq i64 [[INDEX_NEXT]], 992
+; CHECK-NEON-NEXT: br i1 [[TMP2]], label %[[MIDDLE_BLOCK:.*]], label %[[VECTOR_BODY]], !llvm.loop [[LOOP6:![0-9]+]]
+; CHECK-NEON: [[MIDDLE_BLOCK]]:
+; CHECK-NEON-NEXT: [[TMP3:%.*]] = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> [[PARTIAL_REDUCE1]])
+; CHECK-NEON-NEXT: store i32 [[TMP3]], ptr [[DST]], align 4
+; CHECK-NEON-NEXT: br label %[[SCALAR_PH:.*]]
+; CHECK-NEON: [[SCALAR_PH]]:
+; CHECK-NEON-NEXT: br label %[[LOOP:.*]]
+; CHECK-NEON: [[EXIT:.*]]:
+; CHECK-NEON-NEXT: ret void
+; CHECK-NEON: [[LOOP]]:
+; CHECK-NEON-NEXT: [[IV:%.*]] = phi i64 [ 992, %[[SCALAR_PH]] ], [ [[IV_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEON-NEXT: [[RED:%.*]] = phi i32 [ [[TMP3]], %[[SCALAR_PH]] ], [ [[ADD_1:%.*]], %[[LOOP]] ]
+; CHECK-NEON-NEXT: [[GEP_SRC:%.*]] = getelementptr i8, ptr [[SRC]], i64 [[IV]]
+; CHECK-NEON-NEXT: [[L:%.*]] = load i8, ptr [[GEP_SRC]], align 1
+; CHECK-NEON-NEXT: [[CONV8:%.*]] = sext i8 [[L]] to i32
+; CHECK-NEON-NEXT: [[ADD:%.*]] = add i32 [[RED]], [[CONV8]]
+; CHECK-NEON-NEXT: [[CONV8_1:%.*]] = sext i8 [[L]] to i32
+; CHECK-NEON-NEXT: [[ADD_1]] = add i32 [[ADD]], [[CONV8_1]]
+; CHECK-NEON-NEXT: [[GEP_DST:%.*]] = getelementptr i8, ptr [[DST]], i64 [[IV]]
+; CHECK-NEON-NEXT: store i32 [[ADD_1]], ptr [[DST]], align 4
+; CHECK-NEON-NEXT: [[IV_NEXT]] = add i64 [[IV]], 1
+; CHECK-NEON-NEXT: [[EXITCOND:%.*]] = icmp eq i64 [[IV_NEXT]], 1000
+; CHECK-NEON-NEXT: br i1 [[EXITCOND]], label %[[EXIT]], label %[[LOOP]], !llvm.loop [[LOOP7:![0-9]+]]
+;
+; CHECK-LABEL: define void @chained_sext_adds(
+; CHECK-SAME: ptr noalias [[SRC:%.*]], ptr noalias [[DST:%.*]]) #[[ATTR1]] {
+; CHECK-NEXT: [[ENTRY:.*]]:
+; CHECK-NEXT: [[TMP0:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT: [[TMP1:%.*]] = shl nuw i64 [[TMP0]], 2
+; CHECK-NEXT: [[MIN_ITERS_CHECK:%.*]] = icmp ult i64 1000, [[TMP1]]
+; CHECK-NEXT: br i1 [[MIN_ITERS_CHECK]], label %[[SCALAR_PH:.*]], label %[[VECTOR_PH:.*]]
+; CHECK: [[VECTOR_PH]]:
+; CHECK-NEXT: [[TMP2:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT: [[TMP3:%.*]] = mul nuw i64 [[TMP2]], 4
+; CHECK-NEXT: [[N_MOD_VF:%.*]] = urem i64 1000, [[TMP3]]
+; CHECK-NEXT: [[N_VEC:%.*]] = sub i64 1000, [[N_MOD_VF]]
+; CHECK-NEXT: br label %[[VECTOR_BODY:.*]]
+; CHECK: [[VECTOR_BODY]]:
+; CHECK-NEXT: [[INDEX:%.*]] = phi i64 [ 0, %[[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], %[[VECTOR_BODY]] ]
+; CHECK-NEXT: [[VEC_PHI:%.*]] = phi <vscale x 4 x i32> [ zeroinitializer, %[[VECTOR_PH]] ], [ [[TMP7:%.*]], %[[VECTOR_BODY]] ]
+; CHECK-NEXT: [[TMP4:%.*]] = getelementptr i8, ptr [[SRC]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD:%.*]] = load <vscale x 4 x i8>, ptr [[TMP4]], align 1
+; CHECK-NEXT: [[TMP5:%.*]] = sext <vscale x 4 x i8> [[WIDE_LOAD]] to <vscale x 4 x i32>
+; CHECK-NEXT: [[TMP6:%.*]] = add <vscale x 4 x i32> [[VEC_PHI]], [[TMP5]]
+; CHECK-NEXT: [[TMP7]] = add <vscale x 4 x i32> [[TMP6]], [[TMP5]]
+; CHECK-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], [[TMP3]]
+; CHECK-NEXT: [[TMP8:%.*]] = icmp eq i64 [[INDEX_NEXT]], [[N_VEC]]
+; CHECK-NEXT: br i1 [[TMP8]], label %[[MIDDLE_BLOCK:.*]], label %[[VECTOR_BODY]], !llvm.loop [[LOOP5:![0-9]+]]
+; CHECK: [[MIDDLE_BLOCK]]:
+; CHECK-NEXT: [[TMP9:%.*]] = call i32 @llvm.vector.reduce.add.nxv4i32(<vscale x 4 x i32> [[TMP7]])
+; CHECK-NEXT: store i32 [[TMP9]], ptr [[DST]], align 4
+; CHECK-NEXT: [[CMP_N:%.*]] = icmp eq i64 1000, [[N_VEC]]
+; CHECK-NEXT: br i1 [[CMP_N]], label %[[EXIT:.*]], label %[[SCALAR_PH]]
+; CHECK: [[SCALAR_PH]]:
+; CHECK-NEXT: [[BC_RESUME_VAL:%.*]] = phi i64 [ [[N_VEC]], %[[MIDDLE_BLOCK]] ], [ 0, %[[ENTRY]] ]
+; CHECK-NEXT: [[BC_MERGE_RDX:%.*]] = phi i32 [ [[TMP9]], %[[MIDDLE_BLOCK]] ], [ 0, %[[ENTRY]] ]
+; CHECK-NEXT: br label %[[LOOP:.*]]
+; CHECK: [[EXIT]]:
+; CHECK-NEXT: ret void
+; CHECK: [[LOOP]]:
+; CHECK-NEXT: [[IV:%.*]] = phi i64 [ [[BC_RESUME_VAL]], %[[SCALAR_PH]] ], [ [[IV_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEXT: [[RED:%.*]] = phi i32 [ [[BC_MERGE_RDX]], %[[SCALAR_PH]] ], [ [[ADD_1:%.*]], %[[LOOP]] ]
+; CHECK-NEXT: [[GEP_SRC:%.*]] = getelementptr i8, ptr [[SRC]], i64 [[IV]]
+; CHECK-NEXT: [[L:%.*]] = load i8, ptr [[GEP_SRC]], align 1
+; CHECK-NEXT: [[CONV8:%.*]] = sext i8 [[L]] to i32
+; CHECK-NEXT: [[ADD:%.*]] = add i32 [[RED]], [[CONV8]]
+; CHECK-NEXT: [[CONV8_1:%.*]] = sext i8 [[L]] to i32
+; CHECK-NEXT: [[ADD_1]] = add i32 [[ADD]], [[CONV8_1]]
+; CHECK-NEXT: [[GEP_DST:%.*]] = getelementptr i8, ptr [[DST]], i64 [[IV]]
+; CHECK-NEXT: store i32 [[ADD_1]], ptr [[DST]], align 4
+; CHECK-NEXT: [[IV_NEXT]] = add i64 [[IV]], 1
+; CHECK-NEXT: [[EXITCOND:%.*]] = icmp eq i64 [[IV_NEXT]], 1000
+; CHECK-NEXT: br i1 [[EXITCOND]], label %[[EXIT]], label %[[LOOP]], !llvm.loop [[LOOP6:![0-9]+]]
+;
+entry:
+ br label %loop
+
+exit: ; preds = %loop
+ ret void
+
+loop: ; preds = %loop, %entry
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
+ %red = phi i32 [ 0, %entry ], [ %add.1, %loop ]
+ %gep.src = getelementptr i8, ptr %src, i64 %iv
+ %l = load i8, ptr %gep.src, align 1
+ %conv8 = sext i8 %l to i32
+ %add = add i32 %red, %conv8
+ %conv8.1 = sext i8 %l to i32
+ %add.1 = add i32 %add, %conv8.1
+ %gep.dst = getelementptr i8, ptr %dst, i64 %iv
+ store i32 %add.1, ptr %dst, align 4
+ %iv.next = add i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1000
+ br i1 %exitcond, label %exit, label %loop
+}
+
attributes #0 = { "target-cpu"="grace" }
>From d12717c1bf2d80c933d5efbe95e73fd64c8b5b7e Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Tue, 18 Nov 2025 16:39:15 +0000
Subject: [PATCH 2/3] !fixup add comment, thanks
---
llvm/lib/Transforms/Vectorize/LoopVectorize.cpp | 4 +++-
1 file changed, 3 insertions(+), 1 deletion(-)
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index b808c37ac7b96..62a6d7874ccf8 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -7992,7 +7992,9 @@ VPReplicateRecipe *VPRecipeBuilder::handleReplication(VPInstruction *VPI,
/// Find all possible partial reductions in the loop and track all of those that
/// are valid so recipes can be formed later.
void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
- // Find all possible partial reductions, grouping chains by their PHI.
+ // Find all possible partial reductions, grouping chains by their PHI. This
+ // grouping allows invalidating the whole chain, if any link is not a valid
+ // partial reduction.
MapVector<Instruction *,
SmallVector<std::pair<PartialReductionChain, unsigned>>>
ChainsByPhi;
>From 9d7e0c20ab5497fe8cfb0107c6fe13ee9e0cac52 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Wed, 19 Nov 2025 14:23:28 +0000
Subject: [PATCH 3/3] !fixup don't check all users of phi.
---
llvm/lib/Transforms/Vectorize/LoopVectorize.cpp | 11 ++++-------
.../AArch64/partial-reduce-incomplete-chains.ll | 1 -
2 files changed, 4 insertions(+), 8 deletions(-)
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 62a6d7874ccf8..06d55b197ca33 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -8036,14 +8036,11 @@ void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
// match.
for (const auto &[Phi, Chains] : ChainsByPhi) {
for (const auto &[Chain, Scale] : Chains) {
- auto AllUsersPartialRdx = [ScaleVal = Scale, this](const User *U) {
+ auto AllUsersPartialRdx = [ScaleVal = Scale, RdxPhi = Phi,
+ this](const User *U) {
auto *UI = cast<Instruction>(U);
- if (isa<PHINode>(UI) && UI->getParent() == OrigLoop->getHeader()) {
- return all_of(UI->users(), [ScaleVal, this](const User *U) {
- auto *UI = cast<Instruction>(U);
- return ScaledReductionMap.lookup_or(UI, 0) == ScaleVal;
- });
- }
+ if (isa<PHINode>(UI) && UI->getParent() == OrigLoop->getHeader())
+ return UI == RdxPhi;
return ScaledReductionMap.lookup_or(UI, 0) == ScaleVal ||
!OrigLoop->contains(UI->getParent());
};
diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-incomplete-chains.ll b/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-incomplete-chains.ll
index 47a5a6dc9b384..2060168c531fb 100644
--- a/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-incomplete-chains.ll
+++ b/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-incomplete-chains.ll
@@ -125,7 +125,6 @@ exit:
ret i16 %red.next
}
-
define void @chained_sext_adds(ptr noalias %src, ptr noalias %dst) #0 {
; CHECK-NEON-LABEL: define void @chained_sext_adds(
; CHECK-NEON-SAME: ptr noalias [[SRC:%.*]], ptr noalias [[DST:%.*]]) #[[ATTR1]] {
More information about the llvm-commits
mailing list