[llvm] [InstCombine] Transform `vector.reduce.add` and `splat` into multiplication (PR #161020)

[Gábor Spaits via llvm-commits]

================
@@ -3761,6 +3762,39 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
             return replaceInstUsesWith(CI, Res);
           }
       }
+
+      // Handle the case where a splat is summarized. In that case we have a
+      // multpilication. For example: %2 = insertelement <4 x i32> poison, i32
+      // %0, i64 0 %3 = shufflevector <4 x i32> %2, poison, <4 x i32>
+      // zeroinitializer %4 = tail call i32 @llvm.vector.reduce.add.v4i32(%3)
+      // =>
+      // %2 = shl i32 %0, 2
+      if (Value *Splat = getSplatValue(Arg)) {
+        // It is only a multiplication if we add the same element over and over.
+        assert(Arg->getType()->isVectorTy() &&
+               "The vector.reduce.add intrinsic's argument must be a vector!");
+        ElementCount ReducedVectorElementCount =
+            static_cast<VectorType *>(Arg->getType())->getElementCount();
+        if (ReducedVectorElementCount.isFixed()) {
+          unsigned VectorSize = ReducedVectorElementCount.getFixedValue();
+          Type *SplatType = Splat->getType();
+          unsigned SplatTypeWidth = SplatType->getIntegerBitWidth();
+          Value *Res;
+          // Power of two is a special case. We can just use a left shif here.
+          if (isPowerOf2_32(VectorSize)) {
+            unsigned Pow2 = Log2_32(VectorSize);
+            Res = Builder.CreateShl(
+                Splat, Constant::getIntegerValue(SplatType,
+                                                 APInt(SplatTypeWidth, Pow2)));
+            return replaceInstUsesWith(CI, Res);
+          }
+          // Otherwise just multiply.
+          Res = Builder.CreateMul(
+              Splat, Constant::getIntegerValue(
+                         SplatType, APInt(SplatTypeWidth, VectorSize)));
----------------
spaits wrote:

This is much more elegant than the current one.

https://github.com/llvm/llvm-project/pull/161020