[X86][DAGCombiner] Teach narrowShuffle to use concat_vectors instead of inserting into undef

[craig.topper]

Summary

Concat_vectors is more canonical during early DAG combine. For example, its what's used by SelectionDAGBuilder when converting IR shuffles into SelectionDAG shuffles when element counts between inputs and mask don't match. We also have combines in DAGCombiner than can pull concat_vectors through a shuffle. See partitionShuffleOfConcats. So it seems like concat_vectors is a better operation to use here. I had to teach DAGCombiner's SimplifyVBinOp to also handle concat_vectors with undef. I haven't checked yet if we can remove the INSERT_SUBVECTOR version in there or not.

I didn't want to mess with the other caller of getShuffleHalfVectors that's used during shuffle lowering where insert_subvector probably is what we want to produce so I've enabled this via a boolean passed to the function.

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

19482

Not sure if its worth it but you could put the check for getOpcode() == ISD::CONCAT_VECTORS in here as well.

[craig.topper]

Address review comment

[RKSimon]

LGTM with one optional change

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

19482

I meant

return Concat.getOpcode() == ISD::CONCAT_VECTORS && std::all_of();

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (NarrowVT == Y.getValueType() &&
       SDLoc DL(N);
       SDValue VecC =
           DAG.getNode(Opcode, DL, VT, DAG.getUNDEF(VT), DAG.getUNDEF(VT));
       SDValue NarrowBO = DAG.getNode(Opcode, DL, NarrowVT, X, Y);
       return DAG.getNode(ISD::INSERT_SUBVECTOR, DL, VT, VecC, NarrowBO, Z);
     }
   }
 
+  // Make sure all but the first op are undef.
+  auto ConcatWithUndef = [](SDValue Concat) {
+    return std::all_of(std::next(Concat->op_begin()), Concat->op_end(),
+                       [](const SDValue &Op) {
+                         return Op.isUndef();
+                       });
+  };
+
+  // The following pattern is likely to emerge with vector reduction ops. Moving
+  // the binary operation ahead of the concat may allow using a narrower vector
+  // instruction that has better performance than the wide version of the op:
+  // VBinOp (concat X, undef), (concat Y, undef) --> concat (VBinOp X, Y), VecC
+  if (LHS.getOpcode() == ISD::CONCAT_VECTORS && ConcatWithUndef(LHS) &&
+      RHS.getOpcode() == ISD::CONCAT_VECTORS && ConcatWithUndef(RHS) &&
+      (LHS.hasOneUse() || RHS.hasOneUse())) {
+    SDValue X = LHS.getOperand(0);
+    SDValue Y = RHS.getOperand(0);
+    EVT NarrowVT = X.getValueType();
+    if (NarrowVT == Y.getValueType() &&
+        TLI.isOperationLegalOrCustomOrPromote(Opcode, NarrowVT)) {
+      // (binop undef, undef) may not return undef, so compute that result.
+      SDLoc DL(N);
+      SDValue VecC =
+          DAG.getNode(Opcode, DL, NarrowVT, DAG.getUNDEF(NarrowVT),
+                      DAG.getUNDEF(NarrowVT));
+      SmallVector<SDValue, 4> Ops(LHS.getNumOperands(), VecC);
+      Ops[0] = DAG.getNode(Opcode, DL, NarrowVT, X, Y);
+      return DAG.getNode(ISD::CONCAT_VECTORS, DL, VT, Ops);
+    }
+  }
+
   if (SDValue V = scalarizeBinOpOfSplats(N, DAG))
     return V;
 
   return SDValue();
 }
 
 SDValue DAGCombiner::SimplifySelect(const SDLoc &DL, SDValue N0, SDValue N1,
                                     SDValue N2) {

llvm/lib/Target/X86/X86ISelLowering.cpp

getHalfShuffleMask(ArrayRef<int> Mask, MutableArrayRef<int> HalfMask,
   return true;
 }
 
 /// Given the output values from getHalfShuffleMask(), create a half width
 /// shuffle of extracted vectors followed by an insert back to full width.
 static SDValue getShuffleHalfVectors(const SDLoc &DL, SDValue V1, SDValue V2,
                                      ArrayRef<int> HalfMask, int HalfIdx1,
                                      int HalfIdx2, bool UndefLower,
-                                     SelectionDAG &DAG) {
+                                     SelectionDAG &DAG, bool UseConcat = false) {
   assert(V1.getValueType() == V2.getValueType() && "Different sized vectors?");
   assert(V1.getValueType().isSimple() && "Expecting only simple types");
 
   MVT VT = V1.getSimpleValueType();
   unsigned NumElts = VT.getVectorNumElements();
   unsigned HalfNumElts = NumElts / 2;
   MVT HalfVT = MVT::getVectorVT(VT.getVectorElementType(), HalfNumElts);
 
   auto getHalfVector = [&](int HalfIdx) {
     if (HalfIdx < 0)
       return DAG.getUNDEF(HalfVT);
     SDValue V = (HalfIdx < 2 ? V1 : V2);
     HalfIdx = (HalfIdx % 2) * HalfNumElts;
     return DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, HalfVT, V,
                        DAG.getIntPtrConstant(HalfIdx, DL));
   };
 
   // ins undef, (shuf (ext V1, HalfIdx1), (ext V2, HalfIdx2), HalfMask), Offset
   SDValue Half1 = getHalfVector(HalfIdx1);
   SDValue Half2 = getHalfVector(HalfIdx2);
   SDValue V = DAG.getVectorShuffle(HalfVT, DL, Half1, Half2, HalfMask);
+  if (UseConcat) {
+    SDValue Op0 = V;
+    SDValue Op1 = DAG.getUNDEF(HalfVT);
+    if (UndefLower)
+      std::swap(Op0, Op1);
+    return DAG.getNode(ISD::CONCAT_VECTORS, DL, VT, Op0, Op1);
+  }
+
   unsigned Offset = UndefLower ? HalfNumElts : 0;
   return DAG.getNode(ISD::INSERT_SUBVECTOR, DL, VT, DAG.getUNDEF(VT), V,
                      DAG.getIntPtrConstant(Offset, DL));
 }
 
 /// Lower shuffles where an entire half of a 256 or 512-bit vector is UNDEF.
 /// This allows for fast cases such as subvector extraction/insertion
 /// or shuffling smaller vector types which can lower more efficiently.
static SDValue narrowShuffle(ShuffleVectorSDNode *Shuf, SelectionDAG &DAG) {
 
   // Create a half-width shuffle to replace the unnecessarily wide shuffle.
   // The trick is knowing that all of the insert/extract are actually free
   // subregister (zmm<->ymm or ymm<->xmm) ops. That leaves us with a shuffle
   // of narrow inputs into a narrow output, and that is always cheaper than
   // the wide shuffle that we started with.
   return getShuffleHalfVectors(SDLoc(Shuf), Shuf->getOperand(0),
                                Shuf->getOperand(1), HalfMask, HalfIdx1,
-                               HalfIdx2, false, DAG);
+                               HalfIdx2, false, DAG, /*UseConcat*/true);
 }
 
 static SDValue combineShuffle(SDNode *N, SelectionDAG &DAG,
                               TargetLowering::DAGCombinerInfo &DCI,
                               const X86Subtarget &Subtarget) {
   if (auto *Shuf = dyn_cast<ShuffleVectorSDNode>(N))
     if (SDValue V = narrowShuffle(Shuf, DAG))
       return V;