[NFC][X86] `canonicalizeShuffleWithBinOps()`: refactor and generalize, NFC

[lebedev.ri]

Summary

This is completely NFC.
Follow-ups will enhance this xform to handle binary ops where one operand must be ignored,
and to make those changes easier to review, this is the refactoring-only change.

[RKSimon]

llvm/lib/Target/X86/X86ISelLowering.cpp

40498

Given we know that N0 is a binop - all this generic zip/enumate iteration code seems unnecessary tbh.

[lebedev.ri]

@RKSimon if you prefer, you may write these changes by yourself. I can abandon these patches.

llvm/lib/Target/X86/X86ISelLowering.cpp

40498

This is a generalization. Currently, the code does exactly the same,
but the implicit loop is manually unrolled. With unrolled loops,
you suffer the cost of duplicate LOC, having double the LOC to verify for correctness,
and verify that they do the right thing for right "iteration" (operand).
Doing that is counter-productive and anti-general.

In fact, that is the whole reason why LLVM is seeing such an explosion of code size,
because everything is special-coded, instead of having a "few" general utils,
that are both easy to prove to be correct, and don't incur bloat,
and are easy to use, and make easier to write new code.

[lebedev.ri]

llvm/lib/Target/X86/X86ISelLowering.cpp

40498

(to be most specific: the code to deal unops is still a specialization of the code
to deal with binops, they can be deduplicated further.)

[RKSimon]

In D141877#4059374, @lebedev.ri wrote:

@RKSimon if you prefer, you may write these changes by yourself. I can abandon these patches.

If you can wait a couple of days - I'm happy to get this done.

[lebedev.ri]

In D141877#4060254, @RKSimon wrote:

In D141877#4059374, @lebedev.ri wrote:

@RKSimon if you prefer, you may write these changes by yourself. I can abandon these patches.

If you can wait a couple of days - I'm happy to get this done.

I just don't understand why would anyone write this any other way.
Manual loop unrolling results in more code, and worse code.

[lebedev.ri]

llvm/lib/Target/X86/X86ISelLowering.cpp

40504

(Note how we immediately make use of the .index() part in D141806)

[lebedev.ri]

@RKSimon just so we can save everyone's time, can you please explain how you request this to look?

[lebedev.ri]

Use less smart functions.
@RKSimon is this what you had in mind?

llvm/lib/Target/X86/X86ISelLowering.cpp

         default:
           return SDValue(); // Nothing to combine.
 
         case X86ISD::PSHUFLW:
         case X86ISD::PSHUFHW:
           if (V.getOpcode() == CombineOp)
             break;
 
           Chain.push_back(V);
 
           [[fallthrough]];
         case ISD::BITCAST:
           V = V.getOperand(0);
           continue;
         }
         break;
       } while (V.hasOneUse());
       break;
     }
     // Break out of the loop if we break out of the switch.
     break;
   }
 

 }
 
 // Canonicalize SHUFFLE(BINOP(X,Y)) -> BINOP(SHUFFLE(X),SHUFFLE(Y)).
 static SDValue canonicalizeShuffleWithBinOps(SDValue N, SelectionDAG &DAG,
                                              const SDLoc &DL) {
   const TargetLowering &TLI = DAG.getTargetLoweringInfo();
   EVT ShuffleVT = N.getValueType();
 
-  auto IsMergeableWithShuffle = [&DAG](SDValue Op, bool FoldLoad = false) {
+  auto IsMergeableWithShuffle = [&DAG](unsigned Opcode, unsigned OpNo,
+                                       SDValue Op, bool FoldLoad = false) {
     // AllZeros/AllOnes constants are freely shuffled and will peek through
     // bitcasts. Other constant build vectors do not peek through bitcasts. Only
     // merge with target shuffles if it has one use so shuffle combining is
     // likely to kick in. Shuffles of splats are expected to be removed.
     return ISD::isBuildVectorAllOnes(Op.getNode()) ||
            ISD::isBuildVectorAllZeros(Op.getNode()) ||
            ISD::isBuildVectorOfConstantSDNodes(Op.getNode()) ||
            ISD::isBuildVectorOfConstantFPSDNodes(Op.getNode()) ||
            (isTargetShuffle(Op.getOpcode()) && Op->hasOneUse()) ||
            (FoldLoad && isShuffleFoldableLoad(Op)) ||
            DAG.isSplatValue(Op, /*AllowUndefs*/ false);
   };
   auto IsSafeToMoveShuffle = [ShuffleVT](SDValue Op, unsigned BinOp) {
     // Ensure we only shuffle whole vector src elements, unless its a logical
     // binops where we can more aggressively move shuffles from dst to src.
     return BinOp == ISD::AND || BinOp == ISD::OR || BinOp == ISD::XOR ||
            BinOp == X86ISD::ANDNP ||
            (Op.getScalarValueSizeInBits() <= ShuffleVT.getScalarSizeInBits());
   };
 
+  auto AreAllOperandMergeableWithShuffle =
+      [IsMergeableWithShuffle](unsigned Opcode, unsigned OpNo,
+                               ArrayRef<SDValue> Ops, bool FoldLoad = false) {
+        return all_of(
+            Ops, [IsMergeableWithShuffle, Opcode, OpNo, FoldLoad](SDValue Op) {
+              return IsMergeableWithShuffle(Opcode, OpNo, Op, FoldLoad);
+            });
+      };
+  auto AreAnyOperandMergeableWithShuffle =
+      [IsMergeableWithShuffle](unsigned Opcode, unsigned OpNo,
+                               ArrayRef<SDValue> Ops, bool FoldLoad = false) {
+        return any_of(
+            Ops, [IsMergeableWithShuffle, Opcode, OpNo, FoldLoad](SDValue Op) {
+              return IsMergeableWithShuffle(Opcode, OpNo, Op, FoldLoad);
+            });
+      };
+
+  auto IsAnyOperandMergeableWithShuffle =
+      [IsMergeableWithShuffle](unsigned Opcode, ArrayRef<SDValue> Ops,
+                               bool FoldLoad = false) {
+        return any_of(enumerate(Ops), [IsMergeableWithShuffle, Opcode,
+                                       FoldLoad](auto I) {
+          return IsMergeableWithShuffle(Opcode, I.index(), I.value(), FoldLoad);
+        });
+      };
+
   unsigned Opc = N.getOpcode();
   switch (Opc) {
   // Unary and Unary+Permute Shuffles.
   case X86ISD::PSHUFB: {
     // Don't merge PSHUFB if it contains zero'd elements.
     SmallVector<int> Mask;
     SmallVector<SDValue> Ops;
     if (!getTargetShuffleMask(N.getNode(), ShuffleVT.getSimpleVT(), false, Ops,
                               Mask))
       break;
     [[fallthrough]];
   }
   case X86ISD::VBROADCAST:
   case X86ISD::MOVDDUP:
   case X86ISD::PSHUFD:
   case X86ISD::PSHUFHW:
   case X86ISD::PSHUFLW:
   case X86ISD::VPERMI:
   case X86ISD::VPERMILPI: {
     if (N.getOperand(0).getValueType() == ShuffleVT &&
         N->isOnlyUserOf(N.getOperand(0).getNode())) {
       SDValue N0 = peekThroughOneUseBitcasts(N.getOperand(0));
       unsigned SrcOpcode = N0.getOpcode();
       if (TLI.isBinOp(SrcOpcode) && IsSafeToMoveShuffle(N0, SrcOpcode)) {
-        SDValue Op00 = peekThroughOneUseBitcasts(N0.getOperand(0));
-        SDValue Op01 = peekThroughOneUseBitcasts(N0.getOperand(1));
-        if (IsMergeableWithShuffle(Op00, Opc != X86ISD::PSHUFB) ||
-            IsMergeableWithShuffle(Op01, Opc != X86ISD::PSHUFB)) {
-          SDValue LHS, RHS;
-          Op00 = DAG.getBitcast(ShuffleVT, Op00);
-          Op01 = DAG.getBitcast(ShuffleVT, Op01);
-          if (N.getNumOperands() == 2) {
-            LHS = DAG.getNode(Opc, DL, ShuffleVT, Op00, N.getOperand(1));
-            RHS = DAG.getNode(Opc, DL, ShuffleVT, Op01, N.getOperand(1));
-          } else {
-            LHS = DAG.getNode(Opc, DL, ShuffleVT, Op00);
-            RHS = DAG.getNode(Opc, DL, ShuffleVT, Op01);
-          }
+        assert(N0.getNumOperands() == 2 && "Not a binop?");
+        std::array<SDValue, 2> N0Ops;
+        for (unsigned OpIdx : seq(0U, 2U))
+          N0Ops[OpIdx] = peekThroughOneUseBitcasts(N0.getOperand(OpIdx));
+        if (IsAnyOperandMergeableWithShuffle(SrcOpcode, N0Ops,
+                                             /*FoldLoad=*/Opc !=
+                                                 X86ISD::PSHUFB)) {
           EVT OpVT = N0.getValueType();
+          for (unsigned OpIdx : seq(0U, 2U)) {
+            SDValue &Op = N0Ops[OpIdx];
+            Op = DAG.getBitcast(ShuffleVT, Op);
+            if (N.getNumOperands() == 2)
+              Op = DAG.getNode(Opc, DL, ShuffleVT, Op, N.getOperand(1));
+            else
+              Op = DAG.getNode(Opc, DL, ShuffleVT, Op);
+            Op = DAG.getBitcast(OpVT, Op);
+          }
           return DAG.getBitcast(ShuffleVT,
-                                DAG.getNode(SrcOpcode, DL, OpVT,
-                                            DAG.getBitcast(OpVT, LHS),
-                                            DAG.getBitcast(OpVT, RHS)));
+                                DAG.getNode(SrcOpcode, DL, OpVT, N0Ops));
         }
       }
     }
     break;
   }
   // Binary and Binary+Permute Shuffles.
   case X86ISD::INSERTPS: {
     // Don't merge INSERTPS if it contains zero'd elements.
     unsigned InsertPSMask = N.getConstantOperandVal(2);
     unsigned ZeroMask = InsertPSMask & 0xF;
     if (ZeroMask != 0)
       break;
     [[fallthrough]];
   }
   case X86ISD::MOVSD:
   case X86ISD::MOVSS:
   case X86ISD::BLENDI:
   case X86ISD::SHUFP:
   case X86ISD::UNPCKH:
   case X86ISD::UNPCKL: {
     if (N->isOnlyUserOf(N.getOperand(0).getNode()) &&
         N->isOnlyUserOf(N.getOperand(1).getNode())) {
-      SDValue N0 = peekThroughOneUseBitcasts(N.getOperand(0));
-      SDValue N1 = peekThroughOneUseBitcasts(N.getOperand(1));
-      unsigned SrcOpcode = N0.getOpcode();
-      if (TLI.isBinOp(SrcOpcode) && N1.getOpcode() == SrcOpcode &&
-          IsSafeToMoveShuffle(N0, SrcOpcode) &&
-          IsSafeToMoveShuffle(N1, SrcOpcode)) {
-        SDValue Op00 = peekThroughOneUseBitcasts(N0.getOperand(0));
-        SDValue Op10 = peekThroughOneUseBitcasts(N1.getOperand(0));
-        SDValue Op01 = peekThroughOneUseBitcasts(N0.getOperand(1));
-        SDValue Op11 = peekThroughOneUseBitcasts(N1.getOperand(1));
+      std::array<SDValue, 2> OpsOfShuf;
+      for (unsigned OpIdx : seq(0U, 2U))
+        OpsOfShuf[OpIdx] = peekThroughOneUseBitcasts(N.getOperand(OpIdx));
+      unsigned SrcOpcode = OpsOfShuf[0].getOpcode();
+      if (TLI.isBinOp(SrcOpcode) && OpsOfShuf[1].getOpcode() == SrcOpcode &&
+          IsSafeToMoveShuffle(OpsOfShuf[0], SrcOpcode) &&
+          IsSafeToMoveShuffle(OpsOfShuf[1], SrcOpcode)) {
+        assert(OpsOfShuf[0].getNumOperands() == 2 &&
+               OpsOfShuf[1].getNumOperands() == 2 && "Not binops?");
+        std::array<std::array<SDValue, 2>, 2> NthOpsOfShufOps;
+        for (unsigned OpIdx : seq(0U, 2U)) {
+          for (unsigned ShufOpIdx : seq(0U, 2U)) {
+            NthOpsOfShufOps[OpIdx][ShufOpIdx] = peekThroughOneUseBitcasts(
+                OpsOfShuf[ShufOpIdx].getOperand(OpIdx));
+          }
+        }
         // Ensure the total number of shuffles doesn't increase by folding this
         // shuffle through to the source ops.
-        if (((IsMergeableWithShuffle(Op00) && IsMergeableWithShuffle(Op10)) ||
-             (IsMergeableWithShuffle(Op01) && IsMergeableWithShuffle(Op11))) ||
-            ((IsMergeableWithShuffle(Op00) || IsMergeableWithShuffle(Op10)) &&
-             (IsMergeableWithShuffle(Op01) || IsMergeableWithShuffle(Op11)))) {
-          SDValue LHS, RHS;
-          Op00 = DAG.getBitcast(ShuffleVT, Op00);
-          Op10 = DAG.getBitcast(ShuffleVT, Op10);
-          Op01 = DAG.getBitcast(ShuffleVT, Op01);
-          Op11 = DAG.getBitcast(ShuffleVT, Op11);
-          if (N.getNumOperands() == 3) {
-            LHS = DAG.getNode(Opc, DL, ShuffleVT, Op00, Op10, N.getOperand(2));
-            RHS = DAG.getNode(Opc, DL, ShuffleVT, Op01, Op11, N.getOperand(2));
-          } else {
-            LHS = DAG.getNode(Opc, DL, ShuffleVT, Op00, Op10);
-            RHS = DAG.getNode(Opc, DL, ShuffleVT, Op01, Op11);
+        if (any_of(enumerate(NthOpsOfShufOps),
+                   [&](auto I) {
+                     return AreAllOperandMergeableWithShuffle(
+                         SrcOpcode, I.index(), I.value());
+                   }) ||
+            all_of(enumerate(NthOpsOfShufOps), [&](auto I) {
+              return AreAnyOperandMergeableWithShuffle(SrcOpcode, I.index(),
+                                                       I.value());
+            })) {
+          EVT OpVT = OpsOfShuf[0].getValueType();
+          std::array<SDValue, 2> NewOps;
+          for (unsigned OpIdx : seq(0U, 2U)) {
+            MutableArrayRef<SDValue> NthOpsOfOpsOfShufOp =
+                NthOpsOfShufOps[OpIdx];
+            SDValue &NewOp = NewOps[OpIdx];
+            for (SDValue &V : NthOpsOfOpsOfShufOp)
+              V = DAG.getBitcast(ShuffleVT, V);
+            if (N.getNumOperands() == 3)
+              NewOp = DAG.getNode(Opc, DL, ShuffleVT, NthOpsOfOpsOfShufOp[0],
+                                  NthOpsOfOpsOfShufOp[1], N.getOperand(2));
+            else
+              NewOp = DAG.getNode(Opc, DL, ShuffleVT, NthOpsOfOpsOfShufOp[0],
+                                  NthOpsOfOpsOfShufOp[1]);
+            NewOp = DAG.getBitcast(OpVT, NewOp);
           }
-          EVT OpVT = N0.getValueType();
           return DAG.getBitcast(ShuffleVT,
-                                DAG.getNode(SrcOpcode, DL, OpVT,
-                                            DAG.getBitcast(OpVT, LHS),
-                                            DAG.getBitcast(OpVT, RHS)));
+                                DAG.getNode(SrcOpcode, DL, OpVT, NewOps));
         }
       }
     }
     break;
   }
   }
   return SDValue();
 }