[llvm] [SLP][REVEC] Make ShuffleCostEstimator and ShuffleInstructionBuilder can vectorize vector instructions. (PR #99499)

[Han-Kuan Chen via llvm-commits]

https://github.com/HanKuanChen updated https://github.com/llvm/llvm-project/pull/99499

>From 078dfc2e81ceec38849c45013c058abbf3fdbe9e Mon Sep 17 00:00:00 2001
From: Han-Kuan Chen <hankuan.chen at sifive.com>
Date: Fri, 28 Jun 2024 00:24:24 -0700
Subject: [PATCH 1/3] [SLP][REVEC] Pre-commit test.

---
 llvm/test/Transforms/SLPVectorizer/revec.ll | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/llvm/test/Transforms/SLPVectorizer/revec.ll b/llvm/test/Transforms/SLPVectorizer/revec.ll
index c2dc6d0ab73b7..65018a034e6fb 100644
--- a/llvm/test/Transforms/SLPVectorizer/revec.ll
+++ b/llvm/test/Transforms/SLPVectorizer/revec.ll
@@ -58,3 +58,18 @@ entry:
   store <8 x i16> %4, ptr %5, align 2
   ret void
 }
+
+define void @test3(ptr %in, ptr %out) {
+entry:
+  %0 = load <8 x float>, ptr %in, align 4
+  %1 = fmul <8 x float> %0, zeroinitializer
+  %2 = fmul <8 x float> %0, zeroinitializer
+  %3 = fadd <8 x float> zeroinitializer, %1
+  %4 = fadd <8 x float> %0, %2
+  %5 = fcmp ogt <8 x float> %3, zeroinitializer
+  %6 = fcmp ogt <8 x float> %4, zeroinitializer
+  %7 = getelementptr i1, ptr %out, i64 8
+  store <8 x i1> %5, ptr %out, align 1
+  store <8 x i1> %6, ptr %7, align 1
+  ret void
+}

>From 3dd721740c643ed3ec8fcae97ad7e17738e3bc0d Mon Sep 17 00:00:00 2001
From: Han-Kuan Chen <hankuan.chen at sifive.com>
Date: Fri, 28 Jun 2024 00:33:11 -0700
Subject: [PATCH 2/3] [SLP][REVEC] Make ShuffleCostEstimator::gather support
 vector instructions.

When REVEC is enabled, we need to expand vector types into scalar types.
---
 .../Transforms/Vectorize/SLPVectorizer.cpp    | 20 ++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index d8c3bae06e932..9dacfb71ca580 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -9131,12 +9131,30 @@ class BoUpSLP::ShuffleCostEstimator : public BaseShuffleAnalysis {
         }
         Vals.push_back(Constant::getNullValue(V->getType()));
       }
+      if (auto *VecTy = dyn_cast<FixedVectorType>(Vals.front()->getType())) {
+        // When REVEC is enabled, we need to expand vector types into scalar
+        // types.
+        unsigned VecTyNumElements = VecTy->getNumElements();
+        SmallVector<Constant *> NewVals;
+        NewVals.reserve(VL.size() * VecTyNumElements);
+        for (Constant *V : Vals)
+          for (unsigned I = 0; I != VecTyNumElements; ++I) {
+            Type *ScalarTy = V->getType()->getScalarType();
+            if (isa<PoisonValue>(V))
+              NewVals.push_back(PoisonValue::get(ScalarTy));
+            else if (isa<UndefValue>(V))
+              NewVals.push_back(UndefValue::get(ScalarTy));
+            else
+              NewVals.push_back(Constant::getNullValue(ScalarTy));
+          }
+        Vals.swap(NewVals);
+      }
       return ConstantVector::get(Vals);
     }
     return ConstantVector::getSplat(
         ElementCount::getFixed(
             cast<FixedVectorType>(Root->getType())->getNumElements()),
-        getAllOnesValue(*R.DL, ScalarTy));
+        getAllOnesValue(*R.DL, ScalarTy->getScalarType()));
   }
   InstructionCost createFreeze(InstructionCost Cost) { return Cost; }
   /// Finalize emission of the shuffles.

>From 40306f1f9edd0dbf076d492fa41176c8a788b2e3 Mon Sep 17 00:00:00 2001
From: Han-Kuan Chen <hankuan.chen at sifive.com>
Date: Fri, 28 Jun 2024 00:56:08 -0700
Subject: [PATCH 3/3] [SLP][REVEC] Make BoUpSLP::gather support vector
 instructions.

---
 .../Transforms/Vectorize/SLPVectorizer.cpp    | 22 ++++++++++++++-----
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index 9dacfb71ca580..e50a0c242a61a 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -11636,8 +11636,8 @@ Value *BoUpSLP::gather(ArrayRef<Value *> VL, Value *Root, Type *ScalarTy) {
                                       Type *Ty) {
     Value *Scalar = V;
     if (Scalar->getType() != Ty) {
-      assert(Scalar->getType()->isIntegerTy() && Ty->isIntegerTy() &&
-             "Expected integer types only.");
+      assert(Scalar->getType()->isIntOrIntVectorTy() &&
+             Ty->isIntOrIntVectorTy() && "Expected integer types only.");
       Value *V = Scalar;
       if (auto *CI = dyn_cast<CastInst>(Scalar);
           isa_and_nonnull<SExtInst, ZExtInst>(CI)) {
@@ -11650,10 +11650,20 @@ Value *BoUpSLP::gather(ArrayRef<Value *> VL, Value *Root, Type *ScalarTy) {
           V, Ty, !isKnownNonNegative(Scalar, SimplifyQuery(*DL)));
     }
 
-    Vec = Builder.CreateInsertElement(Vec, Scalar, Builder.getInt32(Pos));
-    auto *InsElt = dyn_cast<InsertElementInst>(Vec);
-    if (!InsElt)
-      return Vec;
+    Instruction *InsElt;
+    if (auto *VecTy = dyn_cast<FixedVectorType>(Scalar->getType())) {
+      Vec = InsElt = Builder.CreateInsertVector(
+          Vec->getType(), Vec, V,
+          Builder.getInt64(Pos * VecTy->getNumElements()));
+      auto *II = dyn_cast<IntrinsicInst>(InsElt);
+      if (!II || II->getIntrinsicID() != Intrinsic::vector_insert)
+        return Vec;
+    } else {
+      Vec = Builder.CreateInsertElement(Vec, Scalar, Builder.getInt32(Pos));
+      InsElt = dyn_cast<InsertElementInst>(Vec);
+      if (!InsElt)
+        return Vec;
+    }
     GatherShuffleExtractSeq.insert(InsElt);
     CSEBlocks.insert(InsElt->getParent());
     // Add to our 'need-to-extract' list.