[llvm] [SLP][REVEC] Expand getelementptr into vector form. (PR #103704)

[Han-Kuan Chen via llvm-commits]

https://github.com/HanKuanChen created https://github.com/llvm/llvm-project/pull/103704

None

>From bb85c3bc71c787b33c581534ca0f35b7e561e8dc Mon Sep 17 00:00:00 2001
From: Han-Kuan Chen <hankuan.chen at sifive.com>
Date: Fri, 26 Jul 2024 00:13:45 -0700
Subject: [PATCH 1/2] [SLP][REVEC] Pre-commit test.

---
 .../Transforms/SLPVectorizer/RISCV/revec.ll   | 25 +++++++++++++++++++
 1 file changed, 25 insertions(+)
 create mode 100644 llvm/test/Transforms/SLPVectorizer/RISCV/revec.ll

diff --git a/llvm/test/Transforms/SLPVectorizer/RISCV/revec.ll b/llvm/test/Transforms/SLPVectorizer/RISCV/revec.ll
new file mode 100644
index 00000000000000..00d16f857ca2d0
--- /dev/null
+++ b/llvm/test/Transforms/SLPVectorizer/RISCV/revec.ll
@@ -0,0 +1,25 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -mtriple=riscv64 -mcpu=sifive-x280 -passes=slp-vectorizer -S -slp-revec -slp-max-reg-size=1024 -slp-threshold=-100 %s | FileCheck %s
+
+define i32 @test() {
+entry:
+  %getelementptr0 = getelementptr i8, ptr null, i64 64036
+  %getelementptr1 = getelementptr i8, ptr null, i64 64064
+  br label %if.end.i87
+
+if.end.i87:                                       ; preds = %entry
+  %0 = load <2 x i32>, ptr %getelementptr0, align 4
+  %1 = load <2 x i32>, ptr %getelementptr1, align 8
+  switch i32 0, label %sw.bb509.i [
+  i32 1, label %sw.bb509.i
+  i32 0, label %if.then458.i
+  ]
+
+if.then458.i:                                     ; preds = %if.end.i87
+  br label %sw.bb509.i
+
+sw.bb509.i:                                       ; preds = %if.then458.i, %if.end.i87, %if.end.i87
+  %4 = phi <2 x i32> [ %0, %if.then458.i ], [ %0, %if.end.i87 ], [ %0, %if.end.i87 ]
+  %5 = phi <2 x i32> [ %1, %if.then458.i ], [ zeroinitializer, %if.end.i87 ], [ zeroinitializer, %if.end.i87 ]
+  ret i32 0
+}

>From d0b504ab7d215f59258396dbca6a240614da239a Mon Sep 17 00:00:00 2001
From: Han-Kuan Chen <hankuan.chen at sifive.com>
Date: Thu, 4 Jul 2024 00:00:16 -0700
Subject: [PATCH 2/2] [SLP][REVEC] Expand getelementptr into vector form.

---
 .../Transforms/Vectorize/SLPVectorizer.cpp    | 20 +++++++++++++++++++
 .../Transforms/SLPVectorizer/RISCV/revec.ll   | 18 +++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index feffd9ae3c99b7..66d3b6a0036f4c 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -13760,6 +13760,26 @@ Value *BoUpSLP::vectorizeTree(TreeEntry *E, bool PostponedPHIs) {
           LLVM_DEBUG(dbgs() << "SLP: Diamond merged for " << *VL0 << ".\n");
           return E->VectorizedValue;
         }
+        if (isa<FixedVectorType>(ScalarTy)) {
+          assert(SLPReVec && "FixedVectorType is not expected.");
+          // CreateMaskedGather expects VecTy and VecPtr have same size. We need
+          // to expand VecPtr if ScalarTy is a vector type.
+          unsigned ScalarTyNumElements =
+              cast<FixedVectorType>(ScalarTy)->getNumElements();
+          unsigned VecTyNumElements =
+              cast<FixedVectorType>(VecTy)->getNumElements();
+          SmallVector<Constant *> Indices(VecTyNumElements);
+          transform(seq(VecTyNumElements), Indices.begin(), [=](unsigned I) {
+            return Builder.getInt64(I % ScalarTyNumElements);
+          });
+          unsigned VF =
+              cast<FixedVectorType>(VecPtr->getType())->getNumElements();
+          VecPtr = Builder.CreateGEP(
+              VecTy->getElementType(),
+              Builder.CreateShuffleVector(
+                  VecPtr, createReplicatedMask(VecTyNumElements / VF, VF)),
+              ConstantVector::get(Indices));
+        }
         // Use the minimum alignment of the gathered loads.
         Align CommonAlignment = computeCommonAlignment<LoadInst>(E->Scalars);
         NewLI = Builder.CreateMaskedGather(VecTy, VecPtr, CommonAlignment);
diff --git a/llvm/test/Transforms/SLPVectorizer/RISCV/revec.ll b/llvm/test/Transforms/SLPVectorizer/RISCV/revec.ll
index 00d16f857ca2d0..d822a24220df2c 100644
--- a/llvm/test/Transforms/SLPVectorizer/RISCV/revec.ll
+++ b/llvm/test/Transforms/SLPVectorizer/RISCV/revec.ll
@@ -2,6 +2,24 @@
 ; RUN: opt -mtriple=riscv64 -mcpu=sifive-x280 -passes=slp-vectorizer -S -slp-revec -slp-max-reg-size=1024 -slp-threshold=-100 %s | FileCheck %s
 
 define i32 @test() {
+; CHECK-LABEL: @test(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[IF_END_I87:%.*]]
+; CHECK:       if.end.i87:
+; CHECK-NEXT:    [[TMP0:%.*]] = call <4 x i32> @llvm.masked.gather.v4i32.v4p0(<4 x ptr> getelementptr (i32, <4 x ptr> <ptr inttoptr (i64 64036 to ptr), ptr inttoptr (i64 64036 to ptr), ptr inttoptr (i64 64064 to ptr), ptr inttoptr (i64 64064 to ptr)>, <4 x i64> <i64 0, i64 1, i64 0, i64 1>), i32 4, <4 x i1> <i1 true, i1 true, i1 true, i1 true>, <4 x i32> poison)
+; CHECK-NEXT:    [[TMP1:%.*]] = call <4 x i32> @llvm.vector.insert.v4i32.v2i32(<4 x i32> poison, <2 x i32> poison, i64 0)
+; CHECK-NEXT:    [[TMP2:%.*]] = call <4 x i32> @llvm.vector.insert.v4i32.v2i32(<4 x i32> [[TMP1]], <2 x i32> zeroinitializer, i64 2)
+; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <4 x i32> [[TMP0]], <4 x i32> [[TMP2]], <4 x i32> <i32 0, i32 1, i32 6, i32 7>
+; CHECK-NEXT:    switch i32 0, label [[SW_BB509_I:%.*]] [
+; CHECK-NEXT:      i32 1, label [[SW_BB509_I]]
+; CHECK-NEXT:      i32 0, label [[IF_THEN458_I:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       if.then458.i:
+; CHECK-NEXT:    br label [[SW_BB509_I]]
+; CHECK:       sw.bb509.i:
+; CHECK-NEXT:    [[TMP4:%.*]] = phi <4 x i32> [ [[TMP0]], [[IF_THEN458_I]] ], [ [[TMP3]], [[IF_END_I87]] ], [ [[TMP3]], [[IF_END_I87]] ]
+; CHECK-NEXT:    ret i32 0
+;
 entry:
   %getelementptr0 = getelementptr i8, ptr null, i64 64036
   %getelementptr1 = getelementptr i8, ptr null, i64 64064