[SLPVectorizer] Tidyup isShuffle helper

[RKSimon]

Summary

Ensure we keep track of the input vectors in all cases instead of just for SK_Alternate (this will be more useful once D47985 has landed).

Ideally we'd reuse the shuffle mask pattern matching in TargetTransformInfo::getInstructionThroughput here to easily add support for all TargetTransformInfo::ShuffleKind without mass code duplication, I've added a TODO for now.

[RKSimon]

Simplified SK_Select pattern matching

[RKSimon]

@spatel Should this wait for D48236 or can I commit it now?

[spatel]

In D48023#1135195, @RKSimon wrote:

@spatel Should this wait for D48236 or can I commit it now?

IIUC both are NFC and independent, so it's just a matter of preference in which order we want to handle the cleanup?

LGTM if you want to commit this 1st.

lib/Transforms/Vectorize/SLPVectorizer.cpp

 /// %7 = extractelement <4 x i8> %5, i32 1
 /// %ins2 = insertelement <4 x i8> %ins1, i8 %7, i32 1
 /// %8 = extractelement <4 x i8> %5, i32 2
 /// %ins3 = insertelement <4 x i8> %ins2, i8 %8, i32 2
 /// %9 = extractelement <4 x i8> %5, i32 3
 /// %ins4 = insertelement <4 x i8> %ins3, i8 %9, i32 3
 /// ret <4 x i8> %ins4
 /// InstCombiner transforms this into a shuffle and vector mul
+/// TODO: Can we split off and reuse the shuffle mask detection from
+/// TargetTransformInfo::getInstructionThroughput?
 static Optional<TargetTransformInfo::ShuffleKind>
 isShuffle(ArrayRef<Value *> VL) {
   auto *EI0 = cast<ExtractElementInst>(VL[0]);
   unsigned Size = EI0->getVectorOperandType()->getVectorNumElements();
   Value *Vec1 = nullptr;
   Value *Vec2 = nullptr;
   enum ShuffleMode {Unknown, FirstAlternate, SecondAlternate, Permute};
   ShuffleMode CommonShuffleMode = Unknown;
for (unsigned I = 0, E = VL.size(); I < E; ++I) {
     if (Idx->getValue().uge(Size))
       continue;
     unsigned IntIdx = Idx->getValue().getZExtValue();
     // We can extractelement from undef vector.
     if (isa<UndefValue>(Vec))
       continue;
     // For correct shuffling we have to have at most 2 different vector operands
     // in all extractelement instructions.
-    if (Vec1 && Vec2 && Vec != Vec1 && Vec != Vec2)
+    if (!Vec1 || Vec1 == Vec)
+      Vec1 = Vec;
+    else if (!Vec2 || Vec2 == Vec)
+      Vec2 = Vec;
+    else
       return None;
     if (CommonShuffleMode == Permute)
       continue;
     // If the extract index is not the same as the operation number, it is a
     // permutation.
     if (IntIdx != I) {
       CommonShuffleMode = Permute;
       continue;
     }
     // Check the shuffle mode for the current operation.
-    if (!Vec1)
-      Vec1 = Vec;
-    else if (Vec != Vec1)
-      Vec2 = Vec;
     // Example: shufflevector A, B, <0,5,2,7>
     // I is odd and IntIdx for A == I - FirstAlternate shuffle.
     // I is even and IntIdx for B == I - FirstAlternate shuffle.
     // Example: shufflevector A, B, <4,1,6,3>
     // I is even and IntIdx for A == I - SecondAlternate shuffle.
     // I is odd and IntIdx for B == I - SecondAlternate shuffle.
     const bool IIsEven = I & 1;
     const bool CurrVecIsA = Vec == Vec1;