[VectorCombine] Guard against the lane zero select predicate being scalar

All but the first lane was being checked, but this could leave the first lane
with a scalar select predicate. This just extends the check to make sure the
types are all the same

Author: davemgreen

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

@@ -1902,7 +1902,9 @@ bool VectorCombine::foldShuffleToIdentity(Instruction &I) {
             if (CI->getPredicate() != cast<CmpInst>(FrontV)->getPredicate())
               return false;
           if (auto *SI = dyn_cast<SelectInst>(V))
-            if (!isa<VectorType>(SI->getOperand(0)->getType()))
+            if (!isa<VectorType>(SI->getOperand(0)->getType()) ||
+                SI->getOperand(0)->getType() !=
+                    cast<SelectInst>(FrontV)->getOperand(0)->getType())
               return false;
           if (isa<CallInst>(V) && !isa<IntrinsicInst>(V))
             return false;

llvm/test/Transforms/VectorCombine/AArch64/shuffletoidentity.ll

@@ -993,4 +993,23 @@ define void @maximal_legal_fpmath(ptr %addr1, ptr %addr2, ptr %result, float %va
   ret void
 }
 
+define <2 x float> @first_scalar_select(<2 x float> %0, <2 x float> %1, float %x) {
+; CHECK-LABEL: @first_scalar_select(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP_I903:%.*]] = fcmp ogt float [[X:%.*]], 0.000000e+00
+; CHECK-NEXT:    [[SEL1639:%.*]] = select i1 [[CMP_I903]], <2 x float> [[TMP0:%.*]], <2 x float> [[TMP1:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = fcmp ogt <2 x float> [[TMP0]], zeroinitializer
+; CHECK-NEXT:    [[SEL48_I913:%.*]] = select <2 x i1> [[TMP2]], <2 x float> [[TMP0]], <2 x float> [[TMP1]]
+; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <2 x float> [[SEL1639]], <2 x float> [[SEL48_I913]], <2 x i32> <i32 0, i32 3>
+; CHECK-NEXT:    ret <2 x float> [[TMP3]]
+;
+entry:
+  %cmp.i903 = fcmp ogt float %x, 0.000000e+00
+  %sel1639 = select i1 %cmp.i903, <2 x float> %0, <2 x float> %1
+  %3 = fcmp ogt <2 x float> %0, zeroinitializer
+  %sel48.i913 = select <2 x i1> %3, <2 x float> %0, <2 x float> %1
+  %4 = shufflevector <2 x float> %sel1639, <2 x float> %sel48.i913, <2 x i32> <i32 0, i32 3>
+  ret <2 x float> %4
+}
+
 declare void @use(<4 x i8>)