[llvm] 6ccbf1d - [X86] combineSelect - use SelectableOp helper to match the zero operand as well as the target shuffle

[Simon Pilgrim via llvm-commits]

Author: Simon Pilgrim
Date: 2024-11-06T18:35:03Z
New Revision: 6ccbf1da6c9225fddaf6911e7bb49ee011e845a6

URL: https://github.com/llvm/llvm-project/commit/6ccbf1da6c9225fddaf6911e7bb49ee011e845a6
DIFF: https://github.com/llvm/llvm-project/commit/6ccbf1da6c9225fddaf6911e7bb49ee011e845a6.diff

LOG: [X86] combineSelect - use SelectableOp helper to match the zero operand as well as the target shuffle

For the "select(mask, extract_subvector(shuffle(x)), zero) --> extract_subvector(select(insert_subvector(mask), shuffle(x), zero))" fold, match the zero operand inside the SelectableOp helper.

Prep work for #113400 - we will be able to relax the zero operand requirement for some target shuffles.

Added: 
    

Modified: 
    llvm/lib/Target/X86/X86ISelLowering.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index d0d082020b3d24..c15517249eb819 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -46838,23 +46838,21 @@ static SDValue combineSelect(SDNode *N, SelectionDAG &DAG,
   // AVX512 - Extend select with zero to merge with target shuffle.
   // select(mask, extract_subvector(shuffle(x)), zero) -->
   // extract_subvector(select(insert_subvector(mask), shuffle(x), zero))
-  // TODO - support non target shuffles as well.
+  // TODO - support non target shuffles as well with canCombineAsMaskOperation.
   if (Subtarget.hasAVX512() && CondVT.isVector() &&
       CondVT.getVectorElementType() == MVT::i1) {
-    auto SelectableOp = [&TLI](SDValue Op) {
+    auto SelectableOp = [&TLI](SDValue Op, SDValue Alt) {
       return Op.getOpcode() == ISD::EXTRACT_SUBVECTOR &&
              isTargetShuffle(Op.getOperand(0).getOpcode()) &&
              isNullConstant(Op.getOperand(1)) &&
              TLI.isTypeLegal(Op.getOperand(0).getValueType()) &&
-             Op.hasOneUse() && Op.getOperand(0).hasOneUse();
+             Op.hasOneUse() && Op.getOperand(0).hasOneUse() &&
+             ISD::isBuildVectorAllZeros(Alt.getNode());
     };
 
-    bool SelectableLHS = SelectableOp(LHS);
-    bool SelectableRHS = SelectableOp(RHS);
-    bool ZeroLHS = ISD::isBuildVectorAllZeros(LHS.getNode());
-    bool ZeroRHS = ISD::isBuildVectorAllZeros(RHS.getNode());
-
-    if ((SelectableLHS && ZeroRHS) || (SelectableRHS && ZeroLHS)) {
+    bool SelectableLHS = SelectableOp(LHS, RHS);
+    bool SelectableRHS = SelectableOp(RHS, LHS);
+    if (SelectableLHS || SelectableRHS) {
       EVT SrcVT = SelectableLHS ? LHS.getOperand(0).getValueType()
                                 : RHS.getOperand(0).getValueType();
       EVT SrcCondVT = SrcVT.changeVectorElementType(MVT::i1);