[DAG] Allow build-to-shuffle combine to combine builds from two wide vectors.

This allows us to, in some cases, create a vector_shuffle out of a build_vector, when
the inputs to the build are extract_elements from two different vectors, at least one
of which is wider than the output. (E.g. a <8 x i16> being constructed out of
elements from a <16 x i16> and a <8 x i16>).

Differential Revision: https://reviews.llvm.org/D24491

llvm-svn: 281402

1 files changed, 53 insertions, 27 deletions

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 269a55357aa..9b2bf7f88ad 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -12974,9 +12974,15 @@ SDValue DAGCombiner::reduceBuildVecToShuffle(SDNode *N) {
   EVT InVT1 = VecIn1.getValueType();
   EVT InVT2 = VecIn2.getNode() ? VecIn2.getValueType() : InVT1;
   unsigned Vec2Offset = InVT1.getVectorNumElements();
+  unsigned ShuffleNumElems = NumElems;
+
+  MVT IdxTy = TLI.getVectorIdxTy(DAG.getDataLayout());
+  SDValue ZeroIdx = DAG.getConstant(0, dl, IdxTy);
 
   // We can't generate a shuffle node with mismatched input and output types.
   // Try to make the types match.
+  // TODO: Should this fire if InVT1/InVT2 are not legal types, or should
+  // we let legalization run its course first?
   if (InVT1 != VT || InVT2 != VT) {
     // Both inputs and the output must have the same base element type.
     EVT ElemType = VT.getVectorElementType();
@@ -12984,6 +12990,9 @@ SDValue DAGCombiner::reduceBuildVecToShuffle(SDNode *N) {
         ElemType != InVT2.getVectorElementType())
       return SDValue();
 
+    // TODO: Canonicalize this so that if the vectors have different lengths,
+    // VecIn1 is always longer.
+
     // The element types match, now figure out the lengths.
     if (InVT1.getSizeInBits() * 2 == VT.getSizeInBits() && InVT1 == InVT2) {
       // If both input vectors are exactly half the size of the output, concat
@@ -12997,26 +13006,36 @@ SDValue DAGCombiner::reduceBuildVecToShuffle(SDNode *N) {
       if (UsesZeroVector)
         Vec2Offset = NumElems;
     } else if (InVT1.getSizeInBits() == VT.getSizeInBits() * 2) {
-      // If we only have one input vector, and it's twice the size of the
-      // output, split it in two.
       if (!TLI.isExtractSubvectorCheap(VT, NumElems))
         return SDValue();
 
-      // TODO: Support the case where we have one input that's too wide, and
-      // another input which is wide/"correct"/narrow. We can do this by
-      // widening the narrow input, shuffling the wide vectors, and then
-      // extracting the low subvector.
-      if (UsesZeroVector || VecIn2.getNode())
+      if (UsesZeroVector)
         return SDValue();
 
-      MVT IdxTy = TLI.getVectorIdxTy(DAG.getDataLayout());
-      VecIn2 = DAG.getNode(ISD::EXTRACT_SUBVECTOR, dl, VT, VecIn1,
-                           DAG.getConstant(NumElems, dl, IdxTy));
-      VecIn1 = DAG.getNode(ISD::EXTRACT_SUBVECTOR, dl, VT, VecIn1,
-                           DAG.getConstant(0, dl, IdxTy));
-      // Since we now have shorter input vectors, adjust the offset of the
-      // second vector's start.
-      Vec2Offset = NumElems;
+      if (!VecIn2.getNode()) {
+        // If we only have one input vector, and it's twice the size of the
+        // output, split it in two.        
+        VecIn2 = DAG.getNode(ISD::EXTRACT_SUBVECTOR, dl, VT, VecIn1,
+                             DAG.getConstant(NumElems, dl, IdxTy));
+        VecIn1 = DAG.getNode(ISD::EXTRACT_SUBVECTOR, dl, VT, VecIn1, ZeroIdx);
+        // Since we now have shorter input vectors, adjust the offset of the
+        // second vector's start.
+        Vec2Offset = NumElems;
+      } else if (InVT2.getSizeInBits() <= InVT1.getSizeInBits()) {
+        // VecIn1 is wider than the output, and we have another, possibly
+        // smaller input. Pad the smaller input with undefs, shuffle at the
+        // input vector width, and extract the output.
+
+        // The shuffle type is different than VT, so check legality again.
+        if (LegalOperations &&
+            !TLI.isOperationLegal(ISD::VECTOR_SHUFFLE, InVT1))
+          return SDValue();
+
+        if (InVT1 != InVT2)
+          VecIn2 = DAG.getNode(ISD::INSERT_SUBVECTOR, dl, InVT1,
+                               DAG.getUNDEF(InVT1), VecIn2, ZeroIdx);
+        ShuffleNumElems = NumElems * 2;
+      }
     } else {
       // TODO: Support cases where the length mismatch isn't exactly by a
       // factor of 2.
@@ -13024,18 +13043,20 @@ SDValue DAGCombiner::reduceBuildVecToShuffle(SDNode *N) {
     }
   }
 
-  SmallVector<int, 8> Mask;
+  // Initialize mask to undef.
+  SmallVector<int, 8> Mask(ShuffleNumElems, -1);
 
+  // Only need to run up to the number of elements actually used, not the
+  // total number of elements in the shuffle - if we are shuffling a wider
+  // vector, the high lanes should be set to undef.
   for (unsigned i = 0; i != NumElems; ++i) {
-    if (VectorMask[i] == -1) {
-      Mask.push_back(-1);
+    if (VectorMask[i] == -1)
       continue;
-    }
 
     // If we are trying to blend with zero, we need to take a zero from the
     // correct position in the second input.
     if (VectorMask[i] == 0) {
-      Mask.push_back(Vec2Offset + i);
+      Mask[i] = Vec2Offset + i;
       continue;
     }
 
@@ -13044,12 +13065,12 @@ SDValue DAGCombiner::reduceBuildVecToShuffle(SDNode *N) {
         cast<ConstantSDNode>(Extract.getOperand(1))->getZExtValue();
 
     if (VectorMask[i] == 1) {
-      Mask.push_back(ExtIndex);
+      Mask[i] = ExtIndex;
       continue;
     }
 
     assert(VectorMask[i] == 2 && "Expected input to be from second vector");
-    Mask.push_back(Vec2Offset + ExtIndex);
+    Mask[i] = Vec2Offset + ExtIndex;
   }
 
   // Avoid introducing illegal shuffles with zero.
@@ -13059,18 +13080,23 @@ SDValue DAGCombiner::reduceBuildVecToShuffle(SDNode *N) {
   if (UsesZeroVector && !TLI.isVectorClearMaskLegal(Mask, VT))
     return SDValue();
 
+  // The type the input vectors may have changed above.
+  InVT1 = VecIn1.getValueType();
+
   // If we already have a VecIn2, it should have the same type as VecIn1.
   // If we don't, get an undef/zero vector of the appropriate type.
-  VecIn2 =
-      getRightHandValue(DAG, dl, VecIn2, VecIn1.getValueType(), UsesZeroVector);
-  assert(VecIn1.getValueType() == VecIn2.getValueType() &&
-         "Unexpected second input type.");
+  VecIn2 = getRightHandValue(DAG, dl, VecIn2, InVT1, UsesZeroVector);
+  assert(InVT1 == VecIn2.getValueType() && "Unexpected second input type.");
 
   // Return the new VECTOR_SHUFFLE node.
   SDValue Ops[2];
   Ops[0] = VecIn1;
   Ops[1] = VecIn2;
-  return DAG.getVectorShuffle(VT, dl, Ops[0], Ops[1], Mask);
+  SDValue Shuffle = DAG.getVectorShuffle(InVT1, dl, Ops[0], Ops[1], Mask);
+  if (ShuffleNumElems > NumElems)
+    Shuffle = DAG.getNode(ISD::EXTRACT_SUBVECTOR, dl, VT, Shuffle, ZeroIdx);
+
+  return Shuffle;
 }
 
 SDValue DAGCombiner::visitBUILD_VECTOR(SDNode *N) {