[llvm] [SLP][REVEC] SubVectorsMask should be transformed into vector form. (PR #120398)

[Han-Kuan Chen via llvm-commits]

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

>From 1f9f56007f7eeb97401c5adc45e73b668f8b84d6 Mon Sep 17 00:00:00 2001
From: Han-Kuan Chen <hankuan.chen at sifive.com>
Date: Wed, 18 Dec 2024 02:07:30 -0800
Subject: [PATCH 1/2] [SLP][REVEC] Pre-commit test.

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

diff --git a/llvm/test/Transforms/SLPVectorizer/RISCV/revec.ll b/llvm/test/Transforms/SLPVectorizer/RISCV/revec.ll
index 61ff4f5766d309..4267828d42e3d0 100644
--- a/llvm/test/Transforms/SLPVectorizer/RISCV/revec.ll
+++ b/llvm/test/Transforms/SLPVectorizer/RISCV/revec.ll
@@ -231,3 +231,21 @@ define ptr @test4() {
   %28 = tail call float @llvm.sqrt.f32(float %26)
   ret ptr null
 }
+
+define i32 @test5() {
+entry:
+  %div0 = fdiv <2 x double> zeroinitializer, zeroinitializer
+  %div1 = fdiv <2 x double> zeroinitializer, zeroinitializer
+  %add0 = fadd <2 x double> zeroinitializer, %div0
+  %add1 = fadd <2 x double> zeroinitializer, zeroinitializer
+  %add2 = fadd <2 x double> %div1, zeroinitializer
+  %add3 = fadd <2 x double> zeroinitializer, zeroinitializer
+  br label %for.end47
+
+for.end47:                                        ; preds = %entry
+  %add0.lcssa = phi <2 x double> [ %add0, %entry ]
+  %add1.lcssa = phi <2 x double> [ %add1, %entry ]
+  %add2.lcssa = phi <2 x double> [ %add2, %entry ]
+  %add3.lcssa = phi <2 x double> [ %add3, %entry ]
+  ret i32 0
+}

>From ec57f14b928093b929e8dfac0d595ca9cb960035 Mon Sep 17 00:00:00 2001
From: Han-Kuan Chen <hankuan.chen at sifive.com>
Date: Wed, 18 Dec 2024 02:07:39 -0800
Subject: [PATCH 2/2] [SLP][REVEC] SubVectorsMask should be transformed into
 vector form.

---
 .../Transforms/Vectorize/SLPVectorizer.cpp    |  4 ++++
 .../Transforms/SLPVectorizer/RISCV/revec.ll   | 19 +++++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index d967813075bb9f..b02bc0030d5144 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -14261,11 +14261,15 @@ class BoUpSLP::ShuffleInstructionBuilder final : public BaseShuffleAnalysis {
     IsFinalized = true;
     unsigned ScalarTyNumElements = getNumElements(ScalarTy);
     SmallVector<int> NewExtMask(ExtMask);
+    SmallVector<int> NewSubVectorsMask(SubVectorsMask);
     if (ScalarTyNumElements != 1) {
       assert(SLPReVec && "FixedVectorType is not expected.");
       transformScalarShuffleIndiciesToVector(ScalarTyNumElements, CommonMask);
       transformScalarShuffleIndiciesToVector(ScalarTyNumElements, NewExtMask);
       ExtMask = NewExtMask;
+      transformScalarShuffleIndiciesToVector(ScalarTyNumElements,
+                                             NewSubVectorsMask);
+      SubVectorsMask = NewSubVectorsMask;
     }
     if (Action) {
       Value *Vec = InVectors.front();
diff --git a/llvm/test/Transforms/SLPVectorizer/RISCV/revec.ll b/llvm/test/Transforms/SLPVectorizer/RISCV/revec.ll
index 4267828d42e3d0..dd7a21198ac1f5 100644
--- a/llvm/test/Transforms/SLPVectorizer/RISCV/revec.ll
+++ b/llvm/test/Transforms/SLPVectorizer/RISCV/revec.ll
@@ -233,6 +233,25 @@ define ptr @test4() {
 }
 
 define i32 @test5() {
+; CHECK-LABEL: @test5(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = call <4 x double> @llvm.vector.insert.v4f64.v2f64(<4 x double> poison, <2 x double> zeroinitializer, i64 0)
+; CHECK-NEXT:    [[TMP1:%.*]] = call <4 x double> @llvm.vector.insert.v4f64.v2f64(<4 x double> [[TMP0]], <2 x double> zeroinitializer, i64 2)
+; CHECK-NEXT:    [[TMP2:%.*]] = fdiv <4 x double> [[TMP1]], [[TMP1]]
+; CHECK-NEXT:    [[TMP3:%.*]] = call <8 x double> @llvm.vector.insert.v8f64.v2f64(<8 x double> poison, <2 x double> zeroinitializer, i64 0)
+; CHECK-NEXT:    [[TMP4:%.*]] = call <8 x double> @llvm.vector.insert.v8f64.v2f64(<8 x double> [[TMP3]], <2 x double> zeroinitializer, i64 2)
+; CHECK-NEXT:    [[TMP5:%.*]] = call <8 x double> @llvm.vector.insert.v8f64.v2f64(<8 x double> [[TMP4]], <2 x double> zeroinitializer, i64 4)
+; CHECK-NEXT:    [[TMP6:%.*]] = call <8 x double> @llvm.vector.insert.v8f64.v2f64(<8 x double> [[TMP5]], <2 x double> zeroinitializer, i64 6)
+; CHECK-NEXT:    [[TMP7:%.*]] = call <8 x double> @llvm.vector.insert.v8f64.v2f64(<8 x double> poison, <2 x double> zeroinitializer, i64 2)
+; CHECK-NEXT:    [[TMP8:%.*]] = call <8 x double> @llvm.vector.insert.v8f64.v2f64(<8 x double> [[TMP7]], <2 x double> zeroinitializer, i64 6)
+; CHECK-NEXT:    [[TMP9:%.*]] = call <8 x double> @llvm.vector.insert.v8f64.v4f64(<8 x double> poison, <4 x double> [[TMP2]], i64 0)
+; CHECK-NEXT:    [[TMP10:%.*]] = shufflevector <8 x double> [[TMP9]], <8 x double> [[TMP8]], <8 x i32> <i32 0, i32 1, i32 10, i32 11, i32 2, i32 3, i32 14, i32 15>
+; CHECK-NEXT:    [[TMP11:%.*]] = fadd <8 x double> [[TMP6]], [[TMP10]]
+; CHECK-NEXT:    br label [[FOR_END47:%.*]]
+; CHECK:       for.end47:
+; CHECK-NEXT:    [[TMP12:%.*]] = phi <8 x double> [ [[TMP11]], [[ENTRY:%.*]] ]
+; CHECK-NEXT:    ret i32 0
+;
 entry:
   %div0 = fdiv <2 x double> zeroinitializer, zeroinitializer
   %div1 = fdiv <2 x double> zeroinitializer, zeroinitializer