[llvm] [ConstantFolding] Fold vector.partial.reduce.add constants (PR #212112)
Jeong Jihyeon via llvm-commits
llvm-commits at lists.llvm.org
Fri Sep 11 01:06:58 PDT 2026
https://github.com/JihyeonJeong129 updated https://github.com/llvm/llvm-project/pull/212112
>From 53c65bb4e411f75f58d4226f0fdc7cbd52596d80 Mon Sep 17 00:00:00 2001
From: Jihyeon Jeong <jh.jeong129 at gmail.com>
Date: Tue, 8 Sep 2026 13:21:36 +0000
Subject: [PATCH 1/5] [ConstantFolding] Add baseline tests for
vector.partial.reduce.add
---
.../InstSimplify/ConstProp/vecreduce.ll | 97 +++++++++++++++++++
1 file changed, 97 insertions(+)
diff --git a/llvm/test/Transforms/InstSimplify/ConstProp/vecreduce.ll b/llvm/test/Transforms/InstSimplify/ConstProp/vecreduce.ll
index 479b3f8ea4128..520b8906c9c2b 100644
--- a/llvm/test/Transforms/InstSimplify/ConstProp/vecreduce.ll
+++ b/llvm/test/Transforms/InstSimplify/ConstProp/vecreduce.ll
@@ -928,3 +928,100 @@ define i32 @umax_poison_elt() {
%x = call i32 @llvm.vector.reduce.umax.v8i32(<8 x i32> <i32 1, i32 1, i32 poison, i32 1, i32 1, i32 poison, i32 1, i32 1>)
ret i32 %x
}
+
+define <4 x i32> @partial_reduce_add_constants() {
+; CHECK-LABEL: @partial_reduce_add_constants(
+; CHECK-NEXT: [[X:%.*]] = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(<4 x i32> <i32 100, i32 200, i32 300, i32 400>, <16 x i32> <i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15, i32 16>)
+; CHECK-NEXT: ret <4 x i32> [[X]]
+;
+ %x = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(
+ <4 x i32> <i32 100, i32 200, i32 300, i32 400>,
+ <16 x i32> <i32 1, i32 2, i32 3, i32 4,
+ i32 5, i32 6, i32 7, i32 8,
+ i32 9, i32 10, i32 11, i32 12,
+ i32 13, i32 14, i32 15, i32 16>)
+ ret <4 x i32> %x
+}
+
+define <4 x i32> @partial_reduce_add_nonconstant_acc(<4 x i32> %acc) {
+; CHECK-LABEL: @partial_reduce_add_nonconstant_acc(
+; CHECK-NEXT: [[X:%.*]] = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(<4 x i32> [[ACC:%.*]], <16 x i32> <i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15, i32 16>)
+; CHECK-NEXT: ret <4 x i32> [[X]]
+;
+ %x = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(
+ <4 x i32> %acc,
+ <16 x i32> <i32 1, i32 2, i32 3, i32 4,
+ i32 5, i32 6, i32 7, i32 8,
+ i32 9, i32 10, i32 11, i32 12,
+ i32 13, i32 14, i32 15, i32 16>)
+ ret <4 x i32> %x
+}
+
+define <4 x i32> @partial_reduce_add_nonconstant_input(<16 x i32> %input) {
+; CHECK-LABEL: @partial_reduce_add_nonconstant_input(
+; CHECK-NEXT: [[X:%.*]] = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(<4 x i32> <i32 100, i32 200, i32 300, i32 400>, <16 x i32> [[INPUT:%.*]])
+; CHECK-NEXT: ret <4 x i32> [[X]]
+;
+ %x = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(
+ <4 x i32> <i32 100, i32 200, i32 300, i32 400>,
+ <16 x i32> %input)
+ ret <4 x i32> %x
+}
+
+define <4 x i32> @partial_reduce_add_poison_element() {
+; CHECK-LABEL: @partial_reduce_add_poison_element(
+; CHECK-NEXT: [[X:%.*]] = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(<4 x i32> <i32 100, i32 200, i32 300, i32 400>, <16 x i32> <i32 1, i32 2, i32 3, i32 4, i32 5, i32 poison, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15, i32 16>)
+; CHECK-NEXT: ret <4 x i32> [[X]]
+;
+ %x = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(
+ <4 x i32> <i32 100, i32 200, i32 300, i32 400>,
+ <16 x i32> <i32 1, i32 2, i32 3, i32 4,
+ i32 5, i32 poison, i32 7, i32 8,
+ i32 9, i32 10, i32 11, i32 12,
+ i32 13, i32 14, i32 15, i32 16>)
+ ret <4 x i32> %x
+}
+
+define <4 x i32> @partial_reduce_add_ratio_one() {
+; CHECK-LABEL: @partial_reduce_add_ratio_one(
+; CHECK-NEXT: [[X:%.*]] = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v4i32(<4 x i32> <i32 100, i32 200, i32 300, i32 400>, <4 x i32> <i32 1, i32 2, i32 3, i32 4>)
+; CHECK-NEXT: ret <4 x i32> [[X]]
+;
+ %x = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v4i32(
+ <4 x i32> <i32 100, i32 200, i32 300, i32 400>,
+ <4 x i32> <i32 1, i32 2, i32 3, i32 4>)
+ ret <4 x i32> %x
+}
+
+define <2 x i32> @partial_reduce_add_ratio_two() {
+; CHECK-LABEL: @partial_reduce_add_ratio_two(
+; CHECK-NEXT: [[X:%.*]] = call <2 x i32> @llvm.vector.partial.reduce.add.v2i32.v4i32(<2 x i32> <i32 100, i32 200>, <4 x i32> <i32 1, i32 2, i32 3, i32 4>)
+; CHECK-NEXT: ret <2 x i32> [[X]]
+;
+ %x = call <2 x i32> @llvm.vector.partial.reduce.add.v2i32.v4i32(
+ <2 x i32> <i32 100, i32 200>,
+ <4 x i32> <i32 1, i32 2, i32 3, i32 4>)
+ ret <2 x i32> %x
+}
+
+define <2 x i32> @partial_reduce_add_negative() {
+; CHECK-LABEL: @partial_reduce_add_negative(
+; CHECK-NEXT: [[X:%.*]] = call <2 x i32> @llvm.vector.partial.reduce.add.v2i32.v4i32(<2 x i32> <i32 -100, i32 -200>, <4 x i32> <i32 1, i32 2, i32 3, i32 4>)
+; CHECK-NEXT: ret <2 x i32> [[X]]
+;
+ %x = call <2 x i32> @llvm.vector.partial.reduce.add.v2i32.v4i32(
+ <2 x i32> <i32 -100, i32 -200>,
+ <4 x i32> <i32 1, i32 2, i32 3, i32 4>)
+ ret <2 x i32> %x
+}
+
+define <2 x i8> @partial_reduce_add_wrap() {
+; CHECK-LABEL: @partial_reduce_add_wrap(
+; CHECK-NEXT: [[X:%.*]] = call <2 x i8> @llvm.vector.partial.reduce.add.v2i8.v4i8(<2 x i8> <i8 127, i8 126>, <4 x i8> <i8 1, i8 1, i8 2, i8 4>)
+; CHECK-NEXT: ret <2 x i8> [[X]]
+;
+ %x = call <2 x i8> @llvm.vector.partial.reduce.add.v2i8.v4i8(
+ <2 x i8> <i8 127, i8 126>,
+ <4 x i8> <i8 1, i8 1, i8 2, i8 4>)
+ ret <2 x i8> %x
+}
>From 4ec0f4e1fb1857accda0c47874f4d9aefeb16d73 Mon Sep 17 00:00:00 2001
From: Jihyeon Jeong <jh.jeong129 at gmail.com>
Date: Tue, 8 Sep 2026 13:22:53 +0000
Subject: [PATCH 2/5] [ConstantFolding] Fold vector.partial.reduce.add
constants
---
llvm/lib/Analysis/ConstantFolding.cpp | 47 +++++++++++++++++++
.../InstSimplify/ConstProp/vecreduce.ll | 18 +++----
2 files changed, 53 insertions(+), 12 deletions(-)
diff --git a/llvm/lib/Analysis/ConstantFolding.cpp b/llvm/lib/Analysis/ConstantFolding.cpp
index aca2556cd48eb..9a3135f084d51 100644
--- a/llvm/lib/Analysis/ConstantFolding.cpp
+++ b/llvm/lib/Analysis/ConstantFolding.cpp
@@ -1777,6 +1777,7 @@ static bool canConstantFoldIntrinsic(Intrinsic::ID ID, bool IsStrictFP) {
case Intrinsic::vector_reduce_smax:
case Intrinsic::vector_reduce_umin:
case Intrinsic::vector_reduce_umax:
+ case Intrinsic::vector_partial_reduce_add:
case Intrinsic::vector_extract:
case Intrinsic::vector_insert:
case Intrinsic::vector_interleave2:
@@ -2413,6 +2414,50 @@ Constant *constantFoldVectorReduce(Intrinsic::ID IID, Constant *Op) {
return ConstantInt::get(Op->getContext(), Acc);
}
+/// Fold a vector partial reduction add using the deterministic grouping
+/// chosen by TargetLowering::expandPartialReduceMLA. Although the
+/// LangRef leaves the grouping unspecified, input element I is accumulated
+/// into result lane I % NumAccElts, with each accumulator element seeding
+/// its corresponding result lane. Returns nullptr if any element cannot be
+/// folded.
+static Constant *constantFoldVectorPartialReduceAdd(Constant *Acc,
+ Constant *Input,
+ const DataLayout &DL) {
+ auto *AccTy = cast<FixedVectorType>(Acc->getType());
+ auto *InputTy = cast<FixedVectorType>(Input->getType());
+
+ unsigned NumAccElts = AccTy->getNumElements();
+ unsigned NumInputElts = InputTy->getNumElements();
+
+ SmallVector<Constant *> ResultElts;
+ ResultElts.reserve(NumAccElts);
+
+ for (unsigned I = 0; I < NumAccElts; ++I) {
+ Constant *AccElt = Acc->getAggregateElement(I);
+
+ if (!AccElt)
+ return nullptr;
+
+ ResultElts.push_back(AccElt);
+ }
+
+ for (unsigned I = 0; I < NumInputElts; ++I) {
+ Constant *InputElt = Input->getAggregateElement(I);
+ if (!InputElt)
+ return nullptr;
+
+ unsigned ResultIdx = I % NumAccElts;
+ Constant *Folded = ConstantFoldBinaryOpOperands(
+ Instruction::Add, ResultElts[ResultIdx], InputElt, DL);
+ if (!Folded)
+ return nullptr;
+
+ ResultElts[ResultIdx] = Folded;
+ }
+
+ return ConstantVector::get(ResultElts);
+}
+
/// Attempt to fold an SSE floating point to integer conversion of a constant
/// floating point. If roundTowardZero is false, the default IEEE rounding is
/// used (toward nearest, ties to even). This matches the behavior of the
@@ -4453,6 +4498,8 @@ static Constant *ConstantFoldFixedVectorCall(
}
return ConstantVector::get(Result);
}
+ case Intrinsic::vector_partial_reduce_add:
+ return constantFoldVectorPartialReduceAdd(Operands[0], Operands[1], DL);
case Intrinsic::wasm_dot: {
unsigned NumElements =
cast<FixedVectorType>(Operands[0]->getType())->getNumElements();
diff --git a/llvm/test/Transforms/InstSimplify/ConstProp/vecreduce.ll b/llvm/test/Transforms/InstSimplify/ConstProp/vecreduce.ll
index 520b8906c9c2b..314d7165dce0f 100644
--- a/llvm/test/Transforms/InstSimplify/ConstProp/vecreduce.ll
+++ b/llvm/test/Transforms/InstSimplify/ConstProp/vecreduce.ll
@@ -931,8 +931,7 @@ define i32 @umax_poison_elt() {
define <4 x i32> @partial_reduce_add_constants() {
; CHECK-LABEL: @partial_reduce_add_constants(
-; CHECK-NEXT: [[X:%.*]] = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(<4 x i32> <i32 100, i32 200, i32 300, i32 400>, <16 x i32> <i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15, i32 16>)
-; CHECK-NEXT: ret <4 x i32> [[X]]
+; CHECK-NEXT: ret <4 x i32> <i32 128, i32 232, i32 336, i32 440>
;
%x = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(
<4 x i32> <i32 100, i32 200, i32 300, i32 400>,
@@ -970,8 +969,7 @@ define <4 x i32> @partial_reduce_add_nonconstant_input(<16 x i32> %input) {
define <4 x i32> @partial_reduce_add_poison_element() {
; CHECK-LABEL: @partial_reduce_add_poison_element(
-; CHECK-NEXT: [[X:%.*]] = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(<4 x i32> <i32 100, i32 200, i32 300, i32 400>, <16 x i32> <i32 1, i32 2, i32 3, i32 4, i32 5, i32 poison, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15, i32 16>)
-; CHECK-NEXT: ret <4 x i32> [[X]]
+; CHECK-NEXT: ret <4 x i32> <i32 128, i32 poison, i32 336, i32 440>
;
%x = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(
<4 x i32> <i32 100, i32 200, i32 300, i32 400>,
@@ -984,8 +982,7 @@ define <4 x i32> @partial_reduce_add_poison_element() {
define <4 x i32> @partial_reduce_add_ratio_one() {
; CHECK-LABEL: @partial_reduce_add_ratio_one(
-; CHECK-NEXT: [[X:%.*]] = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v4i32(<4 x i32> <i32 100, i32 200, i32 300, i32 400>, <4 x i32> <i32 1, i32 2, i32 3, i32 4>)
-; CHECK-NEXT: ret <4 x i32> [[X]]
+; CHECK-NEXT: ret <4 x i32> <i32 101, i32 202, i32 303, i32 404>
;
%x = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v4i32(
<4 x i32> <i32 100, i32 200, i32 300, i32 400>,
@@ -995,8 +992,7 @@ define <4 x i32> @partial_reduce_add_ratio_one() {
define <2 x i32> @partial_reduce_add_ratio_two() {
; CHECK-LABEL: @partial_reduce_add_ratio_two(
-; CHECK-NEXT: [[X:%.*]] = call <2 x i32> @llvm.vector.partial.reduce.add.v2i32.v4i32(<2 x i32> <i32 100, i32 200>, <4 x i32> <i32 1, i32 2, i32 3, i32 4>)
-; CHECK-NEXT: ret <2 x i32> [[X]]
+; CHECK-NEXT: ret <2 x i32> <i32 104, i32 206>
;
%x = call <2 x i32> @llvm.vector.partial.reduce.add.v2i32.v4i32(
<2 x i32> <i32 100, i32 200>,
@@ -1006,8 +1002,7 @@ define <2 x i32> @partial_reduce_add_ratio_two() {
define <2 x i32> @partial_reduce_add_negative() {
; CHECK-LABEL: @partial_reduce_add_negative(
-; CHECK-NEXT: [[X:%.*]] = call <2 x i32> @llvm.vector.partial.reduce.add.v2i32.v4i32(<2 x i32> <i32 -100, i32 -200>, <4 x i32> <i32 1, i32 2, i32 3, i32 4>)
-; CHECK-NEXT: ret <2 x i32> [[X]]
+; CHECK-NEXT: ret <2 x i32> <i32 -96, i32 -194>
;
%x = call <2 x i32> @llvm.vector.partial.reduce.add.v2i32.v4i32(
<2 x i32> <i32 -100, i32 -200>,
@@ -1017,8 +1012,7 @@ define <2 x i32> @partial_reduce_add_negative() {
define <2 x i8> @partial_reduce_add_wrap() {
; CHECK-LABEL: @partial_reduce_add_wrap(
-; CHECK-NEXT: [[X:%.*]] = call <2 x i8> @llvm.vector.partial.reduce.add.v2i8.v4i8(<2 x i8> <i8 127, i8 126>, <4 x i8> <i8 1, i8 1, i8 2, i8 4>)
-; CHECK-NEXT: ret <2 x i8> [[X]]
+; CHECK-NEXT: ret <2 x i8> <i8 -126, i8 -125>
;
%x = call <2 x i8> @llvm.vector.partial.reduce.add.v2i8.v4i8(
<2 x i8> <i8 127, i8 126>,
>From feeebd7ad41fdeb832c5b2758e09b0172ebfed1f Mon Sep 17 00:00:00 2001
From: Jihyeon Jeong <jh.jeong129 at gmail.com>
Date: Wed, 9 Sep 2026 15:46:16 +0000
Subject: [PATCH 3/5] [ConstantFolding] Add regression test for mixed
fixed/scalable partial reduction
---
.../Transforms/InstSimplify/ConstProp/vecreduce.ll | 10 ++++++++++
1 file changed, 10 insertions(+)
diff --git a/llvm/test/Transforms/InstSimplify/ConstProp/vecreduce.ll b/llvm/test/Transforms/InstSimplify/ConstProp/vecreduce.ll
index 314d7165dce0f..6c2ee928d7c0e 100644
--- a/llvm/test/Transforms/InstSimplify/ConstProp/vecreduce.ll
+++ b/llvm/test/Transforms/InstSimplify/ConstProp/vecreduce.ll
@@ -967,6 +967,16 @@ define <4 x i32> @partial_reduce_add_nonconstant_input(<16 x i32> %input) {
ret <4 x i32> %x
}
+define <4 x i32> @partial_reduce_add_mixed_fixed_scalable() {
+; CHECK-LABEL: @partial_reduce_add_mixed_fixed_scalable(
+; CHECK-NEXT: [[X:%.*]] = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.nxv8i32(<4 x i32> zeroinitializer, <vscale x 8 x i32> zeroinitializer)
+; CHECK-NEXT: ret <4 x i32> [[X]]
+;
+ %x = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.nxv8i32(
+ <4 x i32> zeroinitializer, <vscale x 8 x i32> zeroinitializer)
+ ret <4 x i32> %x
+}
+
define <4 x i32> @partial_reduce_add_poison_element() {
; CHECK-LABEL: @partial_reduce_add_poison_element(
; CHECK-NEXT: ret <4 x i32> <i32 128, i32 poison, i32 336, i32 440>
>From a75ff4e702cc6873e24b910f8bf605ccb8766593 Mon Sep 17 00:00:00 2001
From: Jihyeon Jeong <jh.jeong129 at gmail.com>
Date: Wed, 9 Sep 2026 15:47:00 +0000
Subject: [PATCH 4/5] [ConstantFolding] Avoid crash on scalable partial
reduction input
---
llvm/lib/Analysis/ConstantFolding.cpp | 5 ++++-
1 file changed, 4 insertions(+), 1 deletion(-)
diff --git a/llvm/lib/Analysis/ConstantFolding.cpp b/llvm/lib/Analysis/ConstantFolding.cpp
index 9a3135f084d51..4e470177f9ff5 100644
--- a/llvm/lib/Analysis/ConstantFolding.cpp
+++ b/llvm/lib/Analysis/ConstantFolding.cpp
@@ -2424,7 +2424,10 @@ static Constant *constantFoldVectorPartialReduceAdd(Constant *Acc,
Constant *Input,
const DataLayout &DL) {
auto *AccTy = cast<FixedVectorType>(Acc->getType());
- auto *InputTy = cast<FixedVectorType>(Input->getType());
+ // A fixed result type does not guarantee a fixed input type.
+ auto *InputTy = dyn_cast<FixedVectorType>(Input->getType());
+ if (!InputTy)
+ return nullptr;
unsigned NumAccElts = AccTy->getNumElements();
unsigned NumInputElts = InputTy->getNumElements();
>From 168c56f52342268ab678a721c11649a87bf2b43f Mon Sep 17 00:00:00 2001
From: Jihyeon Jeong <jh.jeong129 at gmail.com>
Date: Thu, 10 Sep 2026 23:11:02 +0000
Subject: [PATCH 5/5] [ConstantFolding] Simplify partial reduction result
initialization
---
llvm/lib/Analysis/ConstantFolding.cpp | 11 +++--------
1 file changed, 3 insertions(+), 8 deletions(-)
diff --git a/llvm/lib/Analysis/ConstantFolding.cpp b/llvm/lib/Analysis/ConstantFolding.cpp
index 4e470177f9ff5..473f87c8d5e06 100644
--- a/llvm/lib/Analysis/ConstantFolding.cpp
+++ b/llvm/lib/Analysis/ConstantFolding.cpp
@@ -2432,16 +2432,11 @@ static Constant *constantFoldVectorPartialReduceAdd(Constant *Acc,
unsigned NumAccElts = AccTy->getNumElements();
unsigned NumInputElts = InputTy->getNumElements();
- SmallVector<Constant *> ResultElts;
- ResultElts.reserve(NumAccElts);
-
+ SmallVector<Constant *> ResultElts(NumAccElts);
for (unsigned I = 0; I < NumAccElts; ++I) {
- Constant *AccElt = Acc->getAggregateElement(I);
-
- if (!AccElt)
+ ResultElts[I] = Acc->getAggregateElement(I);
+ if (!ResultElts[I])
return nullptr;
-
- ResultElts.push_back(AccElt);
}
for (unsigned I = 0; I < NumInputElts; ++I) {
More information about the llvm-commits
mailing list