[llvm] f080b55 - [ConstantFolding] Fold vector.partial.reduce.add constants (#212112)

via llvm-commits llvm-commits at lists.llvm.org
Thu Sep 17 09:55:24 PDT 2026


Author: Jeong Jihyeon
Date: 2026-09-17T16:55:17Z
New Revision: f080b55f9a31e724faa0300b01785cbfd56be349

URL: https://github.com/llvm/llvm-project/commit/f080b55f9a31e724faa0300b01785cbfd56be349
DIFF: https://github.com/llvm/llvm-project/commit/f080b55f9a31e724faa0300b01785cbfd56be349.diff

LOG: [ConstantFolding] Fold vector.partial.reduce.add constants (#212112)

This patch adds constant folding support for `llvm.vector.partial.reduce.add`.

The intrinsic leaves the grouping of input elements into result lanes
unspecified. This implementation uses the deterministic grouping
selected by
the generic lowering in `TargetLowering::expandPartialReduceMLA`: input
element `I` is accumulated into result lane `I % NumAccElts`.

Tests cover:
* Constant accumulator and input vectors
* Non-constant accumulator and input operands
* Poison and undef elements
* Different reduction ratios
* Negative values
* Integer wraparound

Fixes #211558

Added: 
    

Modified: 
    llvm/lib/Analysis/ConstantFolding.cpp
    llvm/test/Transforms/InstSimplify/ConstProp/vecreduce.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Analysis/ConstantFolding.cpp b/llvm/lib/Analysis/ConstantFolding.cpp
index 9b29d5fa6d593..a822ab9c5a748 100644
--- a/llvm/lib/Analysis/ConstantFolding.cpp
+++ b/llvm/lib/Analysis/ConstantFolding.cpp
@@ -1779,6 +1779,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:
@@ -2415,6 +2416,48 @@ 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());
+  // 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();
+
+  SmallVector<Constant *> ResultElts(NumAccElts);
+  for (unsigned I = 0; I < NumAccElts; ++I) {
+    ResultElts[I] = Acc->getAggregateElement(I);
+    if (!ResultElts[I])
+      return nullptr;
+  }
+
+  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
@@ -4464,6 +4507,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 479b3f8ea4128..6c2ee928d7c0e 100644
--- a/llvm/test/Transforms/InstSimplify/ConstProp/vecreduce.ll
+++ b/llvm/test/Transforms/InstSimplify/ConstProp/vecreduce.ll
@@ -928,3 +928,104 @@ 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:    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>,
+  <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_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>
+;
+  %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:    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>,
+  <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:    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>,
+  <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:    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>,
+  <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:    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>,
+  <4 x i8> <i8 1, i8 1, i8 2, i8 4>)
+  ret <2 x i8> %x
+}


        


More information about the llvm-commits mailing list