[X86][AVX] Combine shuffles extracted from a common vector

A small step towards combining shuffles across vector sizes - this recognizes when a shuffle's operands are all extracted from the same larger source and tries to combine to an unary shuffle of that source instead. Fixes one of the test cases from PR34380.

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

llvm-svn: 359292

1 files changed, 45 insertions, 0 deletions

diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index 5a63d00b405..e91724c3f32 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -31772,6 +31772,51 @@ static SDValue combineX86ShuffleChain(ArrayRef<SDValue> Inputs, SDValue Root,
     return DAG.getBitcast(RootVT, Res);
   }
 
+  // If that failed and both inputs are extracted from the same source then
+  // try to combine as an unary shuffle with the larger type.
+  if (!UnaryShuffle && V1.getOpcode() == ISD::EXTRACT_SUBVECTOR &&
+      V2.getOpcode() == ISD::EXTRACT_SUBVECTOR &&
+      isa<ConstantSDNode>(V1.getOperand(1)) &&
+      isa<ConstantSDNode>(V2.getOperand(1))) {
+    SDValue Src1 = V1.getOperand(0);
+    SDValue Src2 = V2.getOperand(0);
+    if (Src1 == Src2) {
+      unsigned Offset1 = V1.getConstantOperandVal(1);
+      unsigned Offset2 = V2.getConstantOperandVal(1);
+      assert(((Offset1 % VT1.getVectorNumElements()) == 0 ||
+              (Offset2 % VT2.getVectorNumElements()) == 0 ||
+              (Src1.getValueSizeInBits() % RootSizeInBits) == 0) &&
+             "Unexpected subvector extraction");
+      // Convert extraction indices to mask size.
+      Offset1 /= VT1.getVectorNumElements();
+      Offset2 /= VT2.getVectorNumElements();
+      Offset1 *= NumMaskElts;
+      Offset2 *= NumMaskElts;
+
+      // Create new mask for larger type.
+      SmallVector<int, 64> NewMask(Mask);
+      for (int &M : NewMask) {
+        if (M < 0)
+          continue;
+        if (M < (int)NumMaskElts)
+          M += Offset1;
+        else
+          M = (M - NumMaskElts) + Offset2;
+      }
+      unsigned Scale = Src1.getValueSizeInBits() / RootSizeInBits;
+      NewMask.append((Scale - 1) * NumMaskElts, SM_SentinelUndef);
+
+      SDValue NewInputs[] = {Src1};
+      if (SDValue Res = combineX86ShuffleChain(
+              NewInputs, Src1, NewMask, Depth, HasVariableMask,
+              AllowVariableMask, DAG, Subtarget)) {
+        Res = DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, VT1, Res,
+                          DAG.getIntPtrConstant(0, DL));
+        return DAG.getBitcast(RootVT, Res);
+      }
+    }
+  }
+
   // Failed to find any combines.
   return SDValue();
 }