From d5c292d8ef590f64d26c16d12afebb6ad7f50373 Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper@sifive.com>
Date: Thu, 29 Aug 2024 12:35:50 -0700
Subject: [PATCH] [GISel][RISCV] Correctly handle scalable vector shuffles of
 pointer vectors in IRTranslator. (#106580)

---
 llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp  | 11 ++++-----
 llvm/lib/CodeGen/MachineVerifier.cpp          |  4 ++--
 .../GlobalISel/irtranslator/shufflevector.ll  | 23 +++++++++++++++++--
 .../MachineVerifier/test_g_splat_vector.mir   |  4 ++--
 4 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
index 968d0a2a5c75..b290d7fb4ce4 100644
--- a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
@@ -3183,15 +3183,14 @@ bool IRTranslator::translateExtractElement(const User &U,
 
 bool IRTranslator::translateShuffleVector(const User &U,
                                           MachineIRBuilder &MIRBuilder) {
-  // A ShuffleVector that has operates on scalable vectors is a splat vector
-  // where the value of the splat vector is the 0th element of the first
-  // operand, since the index mask operand is the zeroinitializer (undef and
+  // A ShuffleVector that operates on scalable vectors is a splat vector where
+  // the value of the splat vector is the 0th element of the first operand,
+  // since the index mask operand is the zeroinitializer (undef and
   // poison are treated as zeroinitializer here).
   if (U.getOperand(0)->getType()->isScalableTy()) {
-    Value *Op0 = U.getOperand(0);
+    Register Val = getOrCreateVReg(*U.getOperand(0));
     auto SplatVal = MIRBuilder.buildExtractVectorElementConstant(
-        LLT::scalar(Op0->getType()->getScalarSizeInBits()),
-        getOrCreateVReg(*Op0), 0);
+        MRI->getType(Val).getElementType(), Val, 0);
     MIRBuilder.buildSplatVector(getOrCreateVReg(U), SplatVal);
     return true;
   }
diff --git a/llvm/lib/CodeGen/MachineVerifier.cpp b/llvm/lib/CodeGen/MachineVerifier.cpp
index 5e9bb4c27ffb..759201ed9dad 100644
--- a/llvm/lib/CodeGen/MachineVerifier.cpp
+++ b/llvm/lib/CodeGen/MachineVerifier.cpp
@@ -1835,8 +1835,8 @@ void MachineVerifier::verifyPreISelGenericInstruction(const MachineInstr *MI) {
       break;
     }
 
-    if (!SrcTy.isScalar()) {
-      report("Source type must be a scalar", MI);
+    if (!SrcTy.isScalar() && !SrcTy.isPointer()) {
+      report("Source type must be a scalar or pointer", MI);
       break;
     }
 
diff --git a/llvm/test/CodeGen/RISCV/GlobalISel/irtranslator/shufflevector.ll b/llvm/test/CodeGen/RISCV/GlobalISel/irtranslator/shufflevector.ll
index 7ea67073bc28..89c7bfe81d5f 100644
--- a/llvm/test/CodeGen/RISCV/GlobalISel/irtranslator/shufflevector.ll
+++ b/llvm/test/CodeGen/RISCV/GlobalISel/irtranslator/shufflevector.ll
@@ -1770,5 +1770,24 @@ define <vscale x 16 x i64> @shufflevector_nxv16i64_2(<vscale x 16 x i64> %a) {
   ret <vscale x 16 x i64> %b
 }
 
-
-
+define <vscale x 1 x ptr> @shufflevector_nxv1p0_0() {
+  ; RV32-LABEL: name: shufflevector_nxv1p0_0
+  ; RV32: bb.1 (%ir-block.0):
+  ; RV32-NEXT:   [[DEF:%[0-9]+]]:_(<vscale x 1 x p0>) = G_IMPLICIT_DEF
+  ; RV32-NEXT:   [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
+  ; RV32-NEXT:   [[EVEC:%[0-9]+]]:_(p0) = G_EXTRACT_VECTOR_ELT [[DEF]](<vscale x 1 x p0>), [[C]](s32)
+  ; RV32-NEXT:   [[SPLAT_VECTOR:%[0-9]+]]:_(<vscale x 1 x p0>) = G_SPLAT_VECTOR [[EVEC]](p0)
+  ; RV32-NEXT:   $v8 = COPY [[SPLAT_VECTOR]](<vscale x 1 x p0>)
+  ; RV32-NEXT:   PseudoRET implicit $v8
+  ;
+  ; RV64-LABEL: name: shufflevector_nxv1p0_0
+  ; RV64: bb.1 (%ir-block.0):
+  ; RV64-NEXT:   [[DEF:%[0-9]+]]:_(<vscale x 1 x p0>) = G_IMPLICIT_DEF
+  ; RV64-NEXT:   [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 0
+  ; RV64-NEXT:   [[EVEC:%[0-9]+]]:_(p0) = G_EXTRACT_VECTOR_ELT [[DEF]](<vscale x 1 x p0>), [[C]](s64)
+  ; RV64-NEXT:   [[SPLAT_VECTOR:%[0-9]+]]:_(<vscale x 1 x p0>) = G_SPLAT_VECTOR [[EVEC]](p0)
+  ; RV64-NEXT:   $v8 = COPY [[SPLAT_VECTOR]](<vscale x 1 x p0>)
+  ; RV64-NEXT:   PseudoRET implicit $v8
+  %a = shufflevector <vscale x 1 x ptr> poison, <vscale x 1 x ptr> poison, <vscale x 1 x i32> zeroinitializer
+  ret <vscale x 1 x ptr> %a
+}
diff --git a/llvm/test/MachineVerifier/test_g_splat_vector.mir b/llvm/test/MachineVerifier/test_g_splat_vector.mir
index 00074349776f..a5bde496a3f2 100644
--- a/llvm/test/MachineVerifier/test_g_splat_vector.mir
+++ b/llvm/test/MachineVerifier/test_g_splat_vector.mir
@@ -16,10 +16,10 @@ body:             |
     ; CHECK: Destination type must be a scalable vector
     %4:_(<2 x s32>) = G_SPLAT_VECTOR %0
 
-    ; CHECK: Source type must be a scalar
+    ; CHECK: Source type must be a scalar or pointer
     %5:_(<vscale x 2 x s32>) = G_SPLAT_VECTOR %1
 
-    ; CHECK: Source type must be a scalar
+    ; CHECK: Source type must be a scalar or pointer
     %6:_(<vscale x 2 x s32>) = G_SPLAT_VECTOR %2
 
     ; CHECK: Element type of the destination must be the same size or smaller than the source type