[X86][SSE] Improve PSHUFB shuffle mask decoding.

In cases where the PSHUFB shuffle mask is shared it might not be bitcasted to a vXi8 byte vector. This patch adds support for decoding these wider shuffle masks from the ConstantPool.

The test case in question makes use of this to recognise the shuffle mask is an unary UNPCKL pattern and simplifies accordingly.

llvm-svn: 261201

llvm/lib/Target/X86/X86ShuffleDecodeConstantPool.cpp

@@ -40,24 +40,43 @@ void DecodePSHUFBMask(const Constant *C, SmallVectorImpl<int> &ShuffleMask) {
   assert(MaskTySize == 128 || MaskTySize == 256 || MaskTySize == 512);
 #endif
 
-  // This is a straightforward byte vector.
+  if (!MaskTy->isVectorTy())
-  if (MaskTy->isVectorTy() && MaskTy->getVectorElementType()->isIntegerTy(8)) {
+    return;
-    int NumElements = MaskTy->getVectorNumElements();
+  int NumElts = MaskTy->getVectorNumElements();
-    ShuffleMask.reserve(NumElements);
 
-    for (int i = 0; i < NumElements; ++i) {
+  Type *EltTy = MaskTy->getVectorElementType();
-      // For AVX vectors with 32 bytes the base of the shuffle is the 16-byte
+  if (!EltTy->isIntegerTy())
-      // lane of the vector we're inside.
+    return;
-      int Base = i & ~0xf;
+
+  // The shuffle mask requires a byte vector - decode cases with
+  // wider elements as well.
+  unsigned BitWidth = cast<IntegerType>(EltTy)->getBitWidth();
+  if ((BitWidth % 8) != 0)
+    return;
+
+  int Scale = BitWidth / 8;
+  int NumBytes = NumElts * Scale;
+  ShuffleMask.reserve(NumBytes);
+
+  for (int i = 0; i != NumElts; ++i) {
     Constant *COp = C->getAggregateElement(i);
     if (!COp) {
       ShuffleMask.clear();
       return;
     } else if (isa<UndefValue>(COp)) {
-        ShuffleMask.push_back(SM_SentinelUndef);
+      ShuffleMask.append(Scale, SM_SentinelUndef);
       continue;
     }
-      uint64_t Element = cast<ConstantInt>(COp)->getZExtValue();
+
+    APInt APElt = cast<ConstantInt>(COp)->getValue();
+    for (int j = 0; j != Scale; ++j) {
+      // For AVX vectors with 32 bytes the base of the shuffle is the 16-byte
+      // lane of the vector we're inside.
+      int Base = ((i * Scale) + j) & ~0xf;
+
+      uint64_t Element = APElt.getLoBits(8).getZExtValue();
+      APElt = APElt.lshr(8);
+
       // If the high bit (7) of the byte is set, the element is zeroed.
       if (Element & (1 << 7))
         ShuffleMask.push_back(SM_SentinelZero);
@@ -68,7 +87,8 @@ void DecodePSHUFBMask(const Constant *C, SmallVectorImpl<int> &ShuffleMask) {
       }
     }
   }
-  // TODO: Handle funny-looking vectors too.
+
+  assert(NumBytes == (int)ShuffleMask.size() && "Unexpected shuffle mask size");
 }
 
 void DecodeVPERMILPMask(const Constant *C, unsigned ElSize,

llvm/test/CodeGen/X86/pshufb-mask-comments.ll

@@ -75,9 +75,9 @@ define <16 x i8> @test5(<16 x i8> %V) {
 define <16 x i8> @test6(<16 x i8> %V, <2 x i64>* %P) {
 ; CHECK-LABEL: test6:
 ; CHECK:       # BB#0:
-; CHECK-NEXT:    movdqa {{.*#+}} xmm1 = [217019414673948672,506380106026255364]
+; CHECK-NEXT:    movaps {{.*#+}} xmm1 = [217019414673948672,506380106026255364]
-; CHECK-NEXT:    movdqa %xmm1, (%rdi)
+; CHECK-NEXT:    movaps %xmm1, (%rdi)
-; CHECK-NEXT:    pshufb %xmm1, %xmm0
+; CHECK-NEXT:    punpcklbw {{.*#+}} xmm0 = xmm0[0,0,1,1,2,2,3,3,4,4,5,5,6,6,7,7]
 ; CHECK-NEXT:    retq
   %1 = insertelement <2 x i64> undef, i64 217019414673948672, i32 0
   %2 = insertelement <2 x i64>    %1, i64 506380106026255364, i32 1