[llvm] Extend vector.reduce.add and splat transform to scalable vectors (PR #161101)

[Gábor Spaits via llvm-commits]

https://github.com/spaits created https://github.com/llvm/llvm-project/pull/161101

Inspired by the conversation of @nikic and @dtcxzyw in PR #161020 .

>From 7d90e64e5b5ec7ba07020471057a6a4932f6dc82 Mon Sep 17 00:00:00 2001
From: Gabor Spaits <gaborspaits1 at gmail.com>
Date: Sun, 28 Sep 2025 23:58:13 +0200
Subject: [PATCH] Extend vector.reduce.add and splat transform to scalable
 vectors

---
 .../Transforms/InstCombine/InstCombineCalls.cpp    | 14 ++++++++++----
 .../Transforms/InstCombine/vector-reductions.ll    |  7 ++++---
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index cf6d0ecab4f69..02b46b0161ad8 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3785,13 +3785,19 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
 
       // vector.reduce.add.vNiM(splat(%x)) -> mul(%x, N)
       if (Value *Splat = getSplatValue(Arg)) {
-        ElementCount VecToReduceCount =
-            cast<VectorType>(Arg->getType())->getElementCount();
+        VectorType *VecToReduceTy = cast<VectorType>(Arg->getType());
+        ElementCount VecToReduceCount = VecToReduceTy->getElementCount();
+        Value *RHS;
         if (VecToReduceCount.isFixed()) {
           unsigned VectorSize = VecToReduceCount.getFixedValue();
-          return BinaryOperator::CreateMul(
-              Splat, ConstantInt::get(Splat->getType(), VectorSize));
+          RHS = ConstantInt::get(Splat->getType(), VectorSize);
         }
+
+        RHS = Builder.CreateElementCount(Type::getInt64Ty(II->getContext()),
+                                         VecToReduceCount);
+        if (Splat->getType() != RHS->getType())
+          RHS = Builder.CreateZExtOrTrunc(RHS, Splat->getType());
+        return BinaryOperator::CreateMul(Splat, RHS);
       }
     }
     [[fallthrough]];
diff --git a/llvm/test/Transforms/InstCombine/vector-reductions.ll b/llvm/test/Transforms/InstCombine/vector-reductions.ll
index f1e0dd9bd06d7..34f0570c2698d 100644
--- a/llvm/test/Transforms/InstCombine/vector-reductions.ll
+++ b/llvm/test/Transforms/InstCombine/vector-reductions.ll
@@ -469,9 +469,10 @@ define i2 @constant_multiplied_7xi2(i2 %0) {
 
 define i32 @negative_scalable_vector(i32 %0) {
 ; CHECK-LABEL: @negative_scalable_vector(
-; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <vscale x 4 x i32> poison, i32 [[TMP0:%.*]], i64 0
-; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <vscale x 4 x i32> [[TMP2]], <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
-; CHECK-NEXT:    [[TMP4:%.*]] = tail call i32 @llvm.vector.reduce.add.nxv4i32(<vscale x 4 x i32> [[TMP3]])
+; CHECK-NEXT:    [[TMP2:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    [[DOTTR:%.*]] = trunc i64 [[TMP2]] to i32
+; CHECK-NEXT:    [[TMP3:%.*]] = shl i32 [[DOTTR]], 2
+; CHECK-NEXT:    [[TMP4:%.*]] = mul i32 [[TMP0:%.*]], [[TMP3]]
 ; CHECK-NEXT:    ret i32 [[TMP4]]
 ;
   %2 = insertelement <vscale x 4 x i32> poison, i32 %0, i64 0