[GlobalISel]: Fix lowering of G_Shuffle_vector where we pick up the wrong source index

[aditya_nandakumar]

Summary

<3 x s32> a = {1, 2, 3}
<3 x s32> b = {4, 5,6}
<2 x s32> c = G_SHUFFLE_VECTOR a, b, shufflemask(2,1)
produces
BUILD_VECTOR {4, 2}
but should produce
BUILD_VECTOR {3, 2}

[arsenm]

LGTM, although I copied this from the DAG lowering, which does seem to use the result number of elements?

[aditya_nandakumar]

Committed in 368781.
I haven't had a chance to look into the DAG's implementation yet but our out of tree bots caught this.

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

 LegalizerHelper::LegalizeResult
 LegalizerHelper::lowerShuffleVector(MachineInstr &MI) {
   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();
   LLT IdxTy = LLT::scalar(32);
 
   const Constant *ShufMask = MI.getOperand(3).getShuffleMask();
 
   SmallVector<int, 32> Mask;
   ShuffleVectorInst::getShuffleMask(ShufMask, Mask);
 
   Register Undef;
   SmallVector<Register, 32> BuildVec;
 
   for (int Idx : Mask) {
     if (Idx < 0) {
       if (!Undef.isValid())
         Undef = MIRBuilder.buildUndef(EltTy).getReg(0);
       BuildVec.push_back(Undef);
       continue;
     }
 
     if (Src0Ty.isScalar()) {
       BuildVec.push_back(Idx == 0 ? Src0Reg : Src1Reg);
     } else {
+      int NumElts = Src0Ty.getNumElements();
       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);
   MI.eraseFromParent();
   return Legalized;
 }

llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-shuffle-vector.mir

bb.0:
     %0:_(<3 x s32>) = COPY $vgpr0_vgpr1_vgpr2
     %1:_(<3 x s32>) = COPY $vgpr3_vgpr4_vgpr5
     %2:_(<3 x s32>) = G_SHUFFLE_VECTOR %0, %1, shufflemask(3, 2, 1)
     $vgpr0_vgpr1_vgpr2 = COPY %2
 
 ...
 
 ---
+name: shufflevector_v3s32_3_2_1_smaller
+tracksRegLiveness: true
+
+body: |
+  bb.0:
+    liveins: $vgpr0_vgpr1_vgpr2, $vgpr3_vgpr4_vgpr5
+
+    ; CHECK-LABEL: name: shufflevector_v3s32_3_2_1_smaller
+    ; CHECK: liveins: $vgpr0_vgpr1_vgpr2, $vgpr3_vgpr4_vgpr5
+    ; CHECK: [[COPY:%[0-9]+]]:_(<3 x s32>) = COPY $vgpr0_vgpr1_vgpr2
+    ; CHECK: [[COPY1:%[0-9]+]]:_(<3 x s32>) = COPY $vgpr3_vgpr4_vgpr5
+    ; CHECK: [[EXTRACT:%[0-9]+]]:_(s32) = G_EXTRACT [[COPY1]](<3 x s32>), 0
+    ; CHECK: [[EXTRACT1:%[0-9]+]]:_(s32) = G_EXTRACT [[COPY]](<3 x s32>), 32
+    ; CHECK: [[BUILD_VECTOR:%[0-9]+]]:_(<2 x s32>) = G_BUILD_VECTOR [[EXTRACT]](s32), [[EXTRACT1]](s32)
+    ; CHECK: $vgpr0_vgpr1 = COPY [[BUILD_VECTOR]](<2 x s32>)
+    %0:_(<3 x s32>) = COPY $vgpr0_vgpr1_vgpr2
+    %1:_(<3 x s32>) = COPY $vgpr3_vgpr4_vgpr5
+    %2:_(<2 x s32>) = G_SHUFFLE_VECTOR %0, %1, shufflemask(3, 1)
+    $vgpr0_vgpr1 = COPY %2
+
+...
+---
 name: shufflevector_v2s16_0_1
 tracksRegLiveness: true
 
 body: |
   bb.0:
     liveins: $vgpr0, $vgpr1
 
     ; CHECK-LABEL: name: shufflevector_v2s16_0_1
bb.0:
     ; CHECK: [[COPY1:%[0-9]+]]:_(<4 x s16>) = COPY $vgpr2_vgpr3
     ; CHECK: [[EXTRACT:%[0-9]+]]:_(<3 x s16>) = G_EXTRACT [[COPY]](<4 x s16>), 0
     ; CHECK: [[EXTRACT1:%[0-9]+]]:_(<3 x s16>) = G_EXTRACT [[COPY1]](<4 x s16>), 0
     ; CHECK: [[UV:%[0-9]+]]:_(s16), [[UV1:%[0-9]+]]:_(s16), [[UV2:%[0-9]+]]:_(s16) = G_UNMERGE_VALUES [[EXTRACT1]](<3 x s16>)
     ; CHECK: [[SEXT:%[0-9]+]]:_(s32) = G_SEXT [[UV]](s16)
     ; CHECK: [[SEXT1:%[0-9]+]]:_(s32) = G_SEXT [[UV1]](s16)
     ; CHECK: [[SEXT2:%[0-9]+]]:_(s32) = G_SEXT [[UV2]](s16)
     ; CHECK: [[BUILD_VECTOR:%[0-9]+]]:_(<3 x s32>) = G_BUILD_VECTOR [[SEXT]](s32), [[SEXT1]](s32), [[SEXT2]](s32)
-    ; CHECK: [[EXTRACT2:%[0-9]+]]:_(s32) = G_EXTRACT [[BUILD_VECTOR]](<3 x s32>), 32
+    ; CHECK: [[EXTRACT2:%[0-9]+]]:_(s32) = G_EXTRACT [[BUILD_VECTOR]](<3 x s32>), 64
     ; CHECK: [[TRUNC:%[0-9]+]]:_(s16) = G_TRUNC [[EXTRACT2]](s32)
     ; CHECK: [[UV3:%[0-9]+]]:_(s16), [[UV4:%[0-9]+]]:_(s16), [[UV5:%[0-9]+]]:_(s16) = G_UNMERGE_VALUES [[EXTRACT]](<3 x s16>)
     ; CHECK: [[SEXT3:%[0-9]+]]:_(s32) = G_SEXT [[UV3]](s16)
     ; CHECK: [[SEXT4:%[0-9]+]]:_(s32) = G_SEXT [[UV4]](s16)
     ; CHECK: [[SEXT5:%[0-9]+]]:_(s32) = G_SEXT [[UV5]](s16)
     ; CHECK: [[BUILD_VECTOR1:%[0-9]+]]:_(<3 x s32>) = G_BUILD_VECTOR [[SEXT3]](s32), [[SEXT4]](s32), [[SEXT5]](s32)
     ; CHECK: [[EXTRACT3:%[0-9]+]]:_(s32) = G_EXTRACT [[BUILD_VECTOR1]](<3 x s32>), 32
     ; CHECK: [[TRUNC1:%[0-9]+]]:_(s16) = G_TRUNC [[EXTRACT3]](s32)
-    ; CHECK: [[DEF:%[0-9]+]]:_(s32) = G_IMPLICIT_DEF
-    ; CHECK: [[TRUNC2:%[0-9]+]]:_(s16) = G_TRUNC [[DEF]](s32)
-    ; CHECK: [[UV6:%[0-9]+]]:_(s16), [[UV7:%[0-9]+]]:_(s16), [[UV8:%[0-9]+]]:_(s16) = G_UNMERGE_VALUES [[EXTRACT]](<3 x s16>)
+    ; CHECK: [[UV6:%[0-9]+]]:_(s16), [[UV7:%[0-9]+]]:_(s16), [[UV8:%[0-9]+]]:_(s16) = G_UNMERGE_VALUES [[EXTRACT1]](<3 x s16>)
     ; CHECK: [[SEXT6:%[0-9]+]]:_(s32) = G_SEXT [[UV6]](s16)
     ; CHECK: [[SEXT7:%[0-9]+]]:_(s32) = G_SEXT [[UV7]](s16)
     ; CHECK: [[SEXT8:%[0-9]+]]:_(s32) = G_SEXT [[UV8]](s16)
     ; CHECK: [[BUILD_VECTOR2:%[0-9]+]]:_(<3 x s32>) = G_BUILD_VECTOR [[SEXT6]](s32), [[SEXT7]](s32), [[SEXT8]](s32)
     ; CHECK: [[EXTRACT4:%[0-9]+]]:_(s32) = G_EXTRACT [[BUILD_VECTOR2]](<3 x s32>), 0
-    ; CHECK: [[TRUNC3:%[0-9]+]]:_(s16) = G_TRUNC [[EXTRACT4]](s32)
-    ; CHECK: [[BUILD_VECTOR3:%[0-9]+]]:_(<4 x s16>) = G_BUILD_VECTOR [[TRUNC]](s16), [[TRUNC1]](s16), [[TRUNC2]](s16), [[TRUNC3]](s16)
-    ; CHECK: $vgpr0_vgpr1 = COPY [[BUILD_VECTOR3]](<4 x s16>)
+    ; CHECK: [[TRUNC2:%[0-9]+]]:_(s16) = G_TRUNC [[EXTRACT4]](s32)
+    ; CHECK: [[UV9:%[0-9]+]]:_(s16), [[UV10:%[0-9]+]]:_(s16), [[UV11:%[0-9]+]]:_(s16) = G_UNMERGE_VALUES [[EXTRACT]](<3 x s16>)
+    ; CHECK: [[SEXT9:%[0-9]+]]:_(s32) = G_SEXT [[UV9]](s16)
+    ; CHECK: [[SEXT10:%[0-9]+]]:_(s32) = G_SEXT [[UV10]](s16)
+    ; CHECK: [[SEXT11:%[0-9]+]]:_(s32) = G_SEXT [[UV11]](s16)
+    ; CHECK: [[BUILD_VECTOR3:%[0-9]+]]:_(<3 x s32>) = G_BUILD_VECTOR [[SEXT9]](s32), [[SEXT10]](s32), [[SEXT11]](s32)
+    ; CHECK: [[EXTRACT5:%[0-9]+]]:_(s32) = G_EXTRACT [[BUILD_VECTOR3]](<3 x s32>), 0
+    ; CHECK: [[TRUNC3:%[0-9]+]]:_(s16) = G_TRUNC [[EXTRACT5]](s32)
+    ; CHECK: [[BUILD_VECTOR4:%[0-9]+]]:_(<4 x s16>) = G_BUILD_VECTOR [[TRUNC]](s16), [[TRUNC1]](s16), [[TRUNC2]](s16), [[TRUNC3]](s16)
+    ; CHECK: $vgpr0_vgpr1 = COPY [[BUILD_VECTOR4]](<4 x s16>)
     %0:_(<4 x s16>) = COPY $vgpr0_vgpr1
     %1:_(<4 x s16>) = COPY $vgpr2_vgpr3
     %2:_(<3 x s16>) = G_EXTRACT %0, 0
     %3:_(<3 x s16>) = G_EXTRACT %1, 0
     %4:_(<4 x s16>) = G_SHUFFLE_VECTOR %2, %3, shufflemask(5, 1, 3, 0)
     $vgpr0_vgpr1 = COPY %4
 
 ...