[GlobalISel]: Fix lowering of G_SHUFFLE_VECTOR with scalar sources

https://reviews.llvm.org/D66171

llvm-svn: 368753
diff --git a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
index db796e9..4d7db27 100644
--- a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
@@ -3821,6 +3821,7 @@
   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 @@
       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);