[llvm] [VectorCombine] Handle shuffle of selects (PR #128032)

[Simon Pilgrim via llvm-commits]

================
@@ -1899,6 +1900,55 @@ bool VectorCombine::foldShuffleOfBinops(Instruction &I) {
   return true;
 }
 
+/// Try to convert,
+/// (shuffle(select(c1,t1,f1)), (select(c2,t2,f2)), m) into
+/// (select (shuffle c1,c2,m), (shuffle t1,t2,m), (shuffle f1,f2,m))
+bool VectorCombine::foldShuffleOfSelects(Instruction &I) {
+  ArrayRef<int> Mask;
+  Value *C1, *T1, *F1, *C2, *T2, *F2;
+  if (!match(&I, m_Shuffle(
+                     m_OneUse(m_Select(m_Value(C1), m_Value(T1), m_Value(F1))),
+                     m_OneUse(m_Select(m_Value(C2), m_Value(T2), m_Value(F2))),
+                     m_Mask(Mask))))
+    return false;
+
+  auto SelectOp = Instruction::Select;
+
+  auto *DstVecTy = dyn_cast<FixedVectorType>(I.getType());
+  auto *C1VecTy = dyn_cast<FixedVectorType>(C1->getType());
+  auto *C2VecTy = dyn_cast<FixedVectorType>(C2->getType());
+  auto *T1VecTy = dyn_cast<FixedVectorType>(T1->getType());
+  auto *F1VecTy = dyn_cast<FixedVectorType>(F1->getType());
+
+  if (!C1VecTy || !C2VecTy)
+    return false;
+
+  InstructionCost OldCost = TTI.getCmpSelInstrCost(
+      SelectOp, T1->getType(), C1VecTy, CmpInst::BAD_ICMP_PREDICATE, CostKind);
+  OldCost += TTI.getCmpSelInstrCost(SelectOp, T2->getType(), C2VecTy,
+                                    CmpInst::BAD_ICMP_PREDICATE, CostKind);
+  OldCost += TTI.getShuffleCost(TargetTransformInfo::SK_PermuteTwoSrc, DstVecTy,
+                                Mask, CostKind);
+
+  InstructionCost NewCost =
+      TTI.getShuffleCost(TargetTransformInfo::SK_PermuteTwoSrc, C1VecTy, Mask,
+                         CostKind, 0, nullptr, {C1, C2}, &I);
+  NewCost += TTI.getShuffleCost(TargetTransformInfo::SK_PermuteTwoSrc, T1VecTy,
+                                Mask, CostKind, 0, nullptr, {T1, T2}, &I);
----------------
RKSimon wrote:

You can't use &I here - this is a new instruction - it should be on OldCost shuffle (and you can add similar instructions to to the OldCost getCmpSelInstrCost calls as well

https://github.com/llvm/llvm-project/pull/128032