Use getSplatValue and correctly construct APInt and add i64 test

Author: spaits

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

@@ -64,6 +64,7 @@
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Support/KnownFPClass.h"
 #include "llvm/Support/MathExtras.h"
+#include "llvm/Support/TypeSize.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/InstCombine/InstCombiner.h"
 #include "llvm/Transforms/Utils/AssumeBundleBuilder.h"
@@ -3769,29 +3770,24 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
       // %4 = tail call i32 @llvm.vector.reduce.add.v4i32(%3)
       // =>
       // %2 = shl i32 %0, 2
-      Value *InputValue;
-      ArrayRef<int> Mask;
-      ConstantInt *InsertionIdx;
       assert(Arg->getType()->isVectorTy() &&
              "The vector.reduce.add intrinsic's argument must be a vector!");
 
-      if (match(Arg, m_Shuffle(m_InsertElt(m_Poison(), m_Value(InputValue),
-                                           m_ConstantInt(InsertionIdx)),
-                               m_Poison(), m_Mask(Mask)))) {
+      if (Value *Splat = getSplatValue(Arg)) {
         // It is only a multiplication if we add the same element over and over.
-        bool AllElementsAreTheSameInMask =
-            std::all_of(Mask.begin(), Mask.end(),
-                        [&Mask](int MaskElt) { return MaskElt == Mask[0]; });
-        unsigned ReducedVectorLength = Mask.size();
-
-        if (AllElementsAreTheSameInMask &&
-            InsertionIdx->getSExtValue() == Mask[0] &&
-            isPowerOf2_32(ReducedVectorLength)) {
-          unsigned Pow2 = Log2_32(ReducedVectorLength);
-          Value *Res = Builder.CreateShl(
-              InputValue, Constant::getIntegerValue(InputValue->getType(),
-                                                    APInt(32, Pow2)));
-          return replaceInstUsesWith(CI, Res);
+        ElementCount ReducedVectorElementCount =
+            static_cast<VectorType *>(Arg->getType())->getElementCount();
+        if (ReducedVectorElementCount.isFixed()) {
+          unsigned VectorSize = ReducedVectorElementCount.getFixedValue();
+          if (isPowerOf2_32(VectorSize)) {
+            unsigned Pow2 = Log2_32(VectorSize);
+            Value *Res = Builder.CreateShl(
+                Splat,
+                Constant::getIntegerValue(
+                    Splat->getType(),
+                    APInt(Splat->getType()->getIntegerBitWidth(), Pow2)));
+            return replaceInstUsesWith(CI, Res);
+          }
         }
       }
     }

llvm/test/Transforms/InstCombine/vector-reductions.ll

@@ -320,6 +320,17 @@ define i32 @constant_multiplied_at_0(i32 %0) {
   ret i32 %4
 }
 
+define i64 @constant_multiplied_at_0_64bits(i64 %0) {
+; CHECK-LABEL: @constant_multiplied_at_0_64bits(
+; CHECK-NEXT:    [[TMP2:%.*]] = shl i64 [[TMP0:%.*]], 2
+; CHECK-NEXT:    ret i64 [[TMP2]]
+;
+  %2 = insertelement <4 x i64> poison, i64 %0, i64 0
+  %3 = shufflevector <4 x i64> %2, <4 x i64> poison, <4 x i32> zeroinitializer
+  %4 = tail call i64 @llvm.vector.reduce.add.v4i64(<4 x i64> %3)
+  ret i64 %4
+}
+
 define i32 @constant_multiplied_at_0_two_pow8(i32 %0) {
 ; CHECK-LABEL: @constant_multiplied_at_0_two_pow8(
 ; CHECK-NEXT:    [[TMP2:%.*]] = shl i32 [[TMP0:%.*]], 3