[llvm] r368753 - [GlobalISel]: Fix lowering of G_SHUFFLE_VECTOR with scalar sources

[Aditya Nandakumar via llvm-commits]

Author: aditya_nandakumar
Date: Tue Aug 13 14:49:11 2019
New Revision: 368753

URL: http://llvm.org/viewvc/llvm-project?rev=368753&view=rev
Log:
[GlobalISel]: Fix lowering of G_SHUFFLE_VECTOR with scalar sources

https://reviews.llvm.org/D66171

Modified:
    llvm/trunk/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
    llvm/trunk/test/CodeGen/AMDGPU/GlobalISel/legalize-shuffle-vector.mir
    llvm/trunk/test/MachineVerifier/test_g_shuffle_vector.mir

Modified: llvm/trunk/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/lib/CodeGen/GlobalISel/LegalizerHelper.cpp?rev=368753&r1=368752&r2=368753&view=diff
==============================================================================
--- llvm/trunk/lib/CodeGen/GlobalISel/LegalizerHelper.cpp (original)
+++ llvm/trunk/lib/CodeGen/GlobalISel/LegalizerHelper.cpp Tue Aug 13 14:49:11 2019
@@ -3821,6 +3821,7 @@ LegalizerHelper::lowerShuffleVector(Mach
   Register DstReg = MI.getOperand(0).getReg();
   Register Src0Reg = MI.getOperand(1).getReg();
   Register Src1Reg = MI.getOperand(2).getReg();
+  LLT Src0Ty = MRI.getType(Src0Reg);
   LLT DstTy = MRI.getType(DstReg);
   LLT EltTy = DstTy.getElementType();
   int NumElts = DstTy.getNumElements();
@@ -3842,11 +3843,15 @@ LegalizerHelper::lowerShuffleVector(Mach
       continue;
     }
 
-    Register SrcVec = Idx < NumElts ? Src0Reg : Src1Reg;
-    int ExtractIdx = Idx < NumElts ? Idx : Idx - NumElts;
-    auto IdxK = MIRBuilder.buildConstant(IdxTy, ExtractIdx);
-    auto Extract = MIRBuilder.buildExtractVectorElement(EltTy, SrcVec, IdxK);
-    BuildVec.push_back(Extract.getReg(0));
+    if (Src0Ty.isScalar()) {
+      BuildVec.push_back(Idx == 0 ? Src0Reg : Src1Reg);
+    } else {
+      Register SrcVec = Idx < NumElts ? Src0Reg : Src1Reg;
+      int ExtractIdx = Idx < NumElts ? Idx : Idx - NumElts;
+      auto IdxK = MIRBuilder.buildConstant(IdxTy, ExtractIdx);
+      auto Extract = MIRBuilder.buildExtractVectorElement(EltTy, SrcVec, IdxK);
+      BuildVec.push_back(Extract.getReg(0));
+    }
   }
 
   MIRBuilder.buildBuildVector(DstReg, BuildVec);

Modified: llvm/trunk/test/CodeGen/AMDGPU/GlobalISel/legalize-shuffle-vector.mir
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/test/CodeGen/AMDGPU/GlobalISel/legalize-shuffle-vector.mir?rev=368753&r1=368752&r2=368753&view=diff
==============================================================================
--- llvm/trunk/test/CodeGen/AMDGPU/GlobalISel/legalize-shuffle-vector.mir (original)
+++ llvm/trunk/test/CodeGen/AMDGPU/GlobalISel/legalize-shuffle-vector.mir Tue Aug 13 14:49:11 2019
@@ -2,6 +2,26 @@
 # RUN: llc -mtriple=amdgcn-mesa-mesa3d -mcpu=fiji -O0 -run-pass=legalizer %s -o - | FileCheck %s
 
 ---
+name: shufflevector_scalar_src
+tracksRegLiveness: true
+
+body: |
+  bb.0:
+    liveins: $vgpr0, $vgpr1
+
+    ; CHECK-LABEL: name: shufflevector_scalar_src
+    ; CHECK: liveins: $vgpr0, $vgpr1
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $vgpr0
+    ; CHECK: [[COPY1:%[0-9]+]]:_(s32) = COPY $vgpr1
+    ; CHECK: [[BUILD_VECTOR:%[0-9]+]]:_(<2 x s32>) = G_BUILD_VECTOR [[COPY]](s32), [[COPY1]](s32)
+    ; CHECK: $vgpr0_vgpr1 = COPY [[BUILD_VECTOR]](<2 x s32>)
+    %0:_(s32) = COPY $vgpr0
+    %1:_(s32) = COPY $vgpr1
+    %2:_(<2 x s32>) = G_SHUFFLE_VECTOR %0, %1, shufflemask(0, 1)
+    $vgpr0_vgpr1 = COPY %2
+
+...
+---
 name: shufflevector_v2s32_0_1
 tracksRegLiveness: true
 

Modified: llvm/trunk/test/MachineVerifier/test_g_shuffle_vector.mir
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/test/MachineVerifier/test_g_shuffle_vector.mir?rev=368753&r1=368752&r2=368753&view=diff
==============================================================================
--- llvm/trunk/test/MachineVerifier/test_g_shuffle_vector.mir (original)
+++ llvm/trunk/test/MachineVerifier/test_g_shuffle_vector.mir Tue Aug 13 14:49:11 2019
@@ -52,4 +52,8 @@ body:             |
     ; CHECK: Bad machine code: G_SHUFFLE_VECTOR cannot change element type
     %21:_(s16) = G_SHUFFLE_VECTOR %3, %4, shufflemask(0)
 
+    ; CHECK: Bad machine code: Out of bounds shuffle index
+    %22:_(s32) = G_IMPLICIT_DEF
+    %20:_(<2 x s32>) = G_SHUFFLE_VECTOR %22, %22, shufflemask(0, 2)
+
 ...