[llvm] [SLP][REVEC] The vectorized result for ShuffleVector may not be ShuffleVectorInst. (PR #116940)

[Han-Kuan Chen via llvm-commits]

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

>From 06678789bed8e9c5f6c58c42f311151790edd877 Mon Sep 17 00:00:00 2001
From: Han-Kuan Chen <hankuan.chen at sifive.com>
Date: Wed, 20 Nov 2024 00:32:20 -0800
Subject: [PATCH 1/3] [SLP][REVEC] Pre-commit test.

---
 .../Transforms/SLPVectorizer/revec-shufflevector.ll   | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/llvm/test/Transforms/SLPVectorizer/revec-shufflevector.ll b/llvm/test/Transforms/SLPVectorizer/revec-shufflevector.ll
index 8091c218addfab..e3133df4ecc47b 100644
--- a/llvm/test/Transforms/SLPVectorizer/revec-shufflevector.ll
+++ b/llvm/test/Transforms/SLPVectorizer/revec-shufflevector.ll
@@ -105,3 +105,14 @@ entry:
   store <4 x i32> %2, ptr %6, align 4
   ret void
 }
+
+define void @test5(ptr %out) {
+entry:
+  %0 = shufflevector <8 x i32> zeroinitializer, <8 x i32> poison, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+  %1 = shufflevector <8 x i32> zeroinitializer, <8 x i32> poison, <4 x i32> <i32 4, i32 5, i32 6, i32 7>
+  %2 = getelementptr inbounds i32, ptr %out, i64 0
+  %3 = getelementptr inbounds i32, ptr %out, i64 4
+  store <4 x i32> %0, ptr %2, align 4
+  store <4 x i32> %1, ptr %3, align 4
+  ret void
+}

>From f449be8c187cdcab62fd4621771f005a52b5b084 Mon Sep 17 00:00:00 2001
From: Han-Kuan Chen <hankuan.chen at sifive.com>
Date: Wed, 20 Nov 2024 00:49:51 -0800
Subject: [PATCH 2/3] [SLP][REVEC] The vectorized result for ShuffleVector may
 not be ShuffleVectorInst.

---
 .../Transforms/Vectorize/SLPVectorizer.cpp    | 21 +++++++++++--------
 .../SLPVectorizer/revec-shufflevector.ll      |  5 +++++
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index dc0dffd9fcbf81..22f66da190c41f 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -15713,16 +15713,19 @@ Value *BoUpSLP::vectorizeTree(TreeEntry *E, bool PostponedPHIs) {
           LLVM_DEBUG(dbgs() << "SLP: Diamond merged for " << *VL0 << ".\n");
           return E->VectorizedValue;
         }
-        assert(isa<ShuffleVectorInst>(Src) &&
-               "Not supported shufflevector usage.");
-        auto *SVSrc = cast<ShuffleVectorInst>(Src);
-        assert(isa<PoisonValue>(SVSrc->getOperand(1)) &&
-               "Not supported shufflevector usage.");
         SmallVector<int> ThisMask(calculateShufflevectorMask(E->Scalars));
-        SmallVector<int> NewMask(ThisMask.size());
-        transform(ThisMask, NewMask.begin(),
-                  [&SVSrc](int Mask) { return SVSrc->getShuffleMask()[Mask]; });
-        V = Builder.CreateShuffleVector(SVSrc->getOperand(0), NewMask);
+        if (isa<ShuffleVectorInst>(Src)) {
+          auto *SVSrc = cast<ShuffleVectorInst>(Src);
+          assert(isa<PoisonValue>(SVSrc->getOperand(1)) &&
+                 "Not supported shufflevector usage.");
+          SmallVector<int> NewMask(ThisMask.size());
+          transform(ThisMask, NewMask.begin(), [&SVSrc](int Mask) {
+            return SVSrc->getShuffleMask()[Mask];
+          });
+          V = Builder.CreateShuffleVector(SVSrc->getOperand(0), NewMask);
+        } else {
+          V = Builder.CreateShuffleVector(Src, ThisMask);
+        }
         propagateIRFlags(V, E->Scalars, VL0);
         if (auto *I = dyn_cast<Instruction>(V))
           V = propagateMetadata(I, E->Scalars);
diff --git a/llvm/test/Transforms/SLPVectorizer/revec-shufflevector.ll b/llvm/test/Transforms/SLPVectorizer/revec-shufflevector.ll
index e3133df4ecc47b..a2673d81068d8d 100644
--- a/llvm/test/Transforms/SLPVectorizer/revec-shufflevector.ll
+++ b/llvm/test/Transforms/SLPVectorizer/revec-shufflevector.ll
@@ -107,6 +107,11 @@ entry:
 }
 
 define void @test5(ptr %out) {
+; CHECK-LABEL: @test5(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    store <8 x i32> zeroinitializer, ptr [[OUT:%.*]], align 4
+; CHECK-NEXT:    ret void
+;
 entry:
   %0 = shufflevector <8 x i32> zeroinitializer, <8 x i32> poison, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   %1 = shufflevector <8 x i32> zeroinitializer, <8 x i32> poison, <4 x i32> <i32 4, i32 5, i32 6, i32 7>

>From 9dd249cd76648d909c00fb0416362bc19298f7c3 Mon Sep 17 00:00:00 2001
From: Han-Kuan Chen <hankuan.chen at sifive.com>
Date: Wed, 20 Nov 2024 06:48:40 -0800
Subject: [PATCH 3/3] apply comment

---
 llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index 22f66da190c41f..26dea50ae0b50e 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -15714,8 +15714,7 @@ Value *BoUpSLP::vectorizeTree(TreeEntry *E, bool PostponedPHIs) {
           return E->VectorizedValue;
         }
         SmallVector<int> ThisMask(calculateShufflevectorMask(E->Scalars));
-        if (isa<ShuffleVectorInst>(Src)) {
-          auto *SVSrc = cast<ShuffleVectorInst>(Src);
+        if (auto *SVSrc = dyn_cast<ShuffleVectorInst>(Src)) {
           assert(isa<PoisonValue>(SVSrc->getOperand(1)) &&
                  "Not supported shufflevector usage.");
           SmallVector<int> NewMask(ThisMask.size());