[X86] Also try to zero elts when lowering v32i8 shuffle with PSHUFB.

[ab]

Summary

Otherwise we fallback to pshufb+blend later on.

We could generalize lowerVectorShuffleAsPSHUFB to ymm, but when I tried that was less readable than doing it inline.

[ab]

• break early when the shuffle isn't single-input, even with zeroing.

[spatel]

lib/Target/X86/X86ISelLowering.cpp

11427–11445

Would it make sense to enhance isSingleInputShuffleMask() like:
bool isSingleInputShuffleMask(ArrayRef<int> Mask, bool OrZeroes = false) {}
and then optionally check for zeroes there?

It seems like vperm2* and maybe some other instructions would then be able to use the same helper.

11447–11455

Please add a comment here to describe the pshufb encoding. Something like:
If the most-significant-bit of any byte of the PSHUFB mask is set, the destination byte is zeroed.
I hate having to look back at the manual. :)

11448

Use SM_SentinelUndef?

[RKSimon]

We could generalize lowerVectorShuffleAsPSHUFB to ymm, but when I tried that was less readable than doing it inline.

How much did you have to do to alter lowerVectorShuffleAsPSHUFB? I'd have expected it to end up quite similar to the implementation in combineX86ShuffleChain which I thought was quite straightforward.

lib/Target/X86/X86ISelLowering.cpp

11448

Convention for lowering is to accept any negative value as undefined - its only later on that we make use of sentinel constants.

[ab]

In D19661#415401, @RKSimon wrote:

We could generalize lowerVectorShuffleAsPSHUFB to ymm, but when I tried that was less readable than doing it inline.

How much did you have to do to alter lowerVectorShuffleAsPSHUFB? I'd have expected it to end up quite similar to the implementation in combineX86ShuffleChain which I thought was quite straightforward.

Sorry: it's not ymm that's the problem, it's that lowerVectorShuffleAsPSHUFB should really be lowerVectorShuffleAsBlendOfPSHUFBs; if you don't want the final blend, there's not that much you can reuse.
It's also tricky to extract the PSHUFB creation out of this hypothetical lowerVectorShuffleAsBlendOfPSHUFB. It wants a PSHUFB that ignores the lanes from the other operand. But we want a PSHUFB if there are no lanes from the other operand.

I'll give it another think; maybe we can do better.

lib/Target/X86/X86ISelLowering.cpp

11427–11445

I wanted to do that at first, but you need to look at the operands, and at that point you basically duplicated computeZeroableShuffleElements.

An alternative I considered was to do some kind of computeZeroableShuffleMask(Mask, V1, V2, NewMask), which returns a mask with SM_SentinelZero and SM_SentinelUndef. Then, users can check that it isSingleInputShuffleMask and do their thing. WDYT?

11447–11455

Heh, fair enough, will do!

[spatel]

lib/Target/X86/X86ISelLowering.cpp

11427–11445

That would be good - slow down the code explosion for x86 vector lowering.
But I don't think we have to hold up this patch for that; a TODO comment should be ok.

LGTM, but I'll let Simon give the final word because he has a much better understanding of everything going on with shuffles.

[RKSimon]

LGTM with a couple of nits:

Please can you add a comment explaining why lowerVectorShuffleAsPSHUFB can't be used?

I think we're getting to the point where we should probably just generate Zeroable as standard and pass it around the same as Mask - we're generating it a lot these days. That would also make the addition of a 'isSingleInputShuffleMaskWithZeroables' helper more straightforward. We can look into this for a future patch.

lib/Target/X86/X86ISelLowering.cpp

11448

Sorry I was wrong here (I missed the Zeroable case afterward) - it should read Mask[i] == SM_SentinelUndef

[ab]

I did a final review before committing, but realized I missed the v16i16 case.

Here's an alternative patch, with a generalized PSHUFB lowering, and "zeroable mask" computation. WDYT?

I went ahead and renamed lowerVectorShuffleAsPSHUFB in r271024.

[RKSimon]

Minor comments but my previous LGTM still stands!

lib/Target/X86/X86ISelLowering.cpp

7150

const int NumEltBytes = VT.getScalarSizeInBits() / 8;

7157

Can't we just pass Subtarget as an argument like we do in other cases?

7166

Rephrase 'Sign bit set in i8 mask means zero element'?

lib/Target/X86/X86ISelLowering.cpp

static SDValue lowerVectorShuffleWithUNPCK(SDLoc DL, MVT VT, ArrayRef<int> Mask,
   SmallVector<int, 8> Unpckh;
 
   for (int i = 0; i < NumElts; ++i) {
     unsigned LaneStart = (i / NumEltsInLane) * NumEltsInLane;
     int LoPos = (i % NumEltsInLane) / 2 + LaneStart + NumElts * (i % 2);
     int HiPos = LoPos + NumEltsInLane / 2;
     Unpckl.push_back(LoPos);
     Unpckh.push_back(HiPos);
   }
 
   if (isShuffleEquivalent(V1, V2, Mask, Unpckl))
     return DAG.getNode(X86ISD::UNPCKL, DL, VT, V1, V2);
   if (isShuffleEquivalent(V1, V2, Mask, Unpckh))
     return DAG.getNode(X86ISD::UNPCKH, DL, VT, V1, V2);
 
   // Commute and try again.
   ShuffleVectorSDNode::commuteMask(Unpckl);
   if (isShuffleEquivalent(V1, V2, Mask, Unpckl))
     return DAG.getNode(X86ISD::UNPCKL, DL, VT, V2, V1);
 
   ShuffleVectorSDNode::commuteMask(Unpckh);
   if (isShuffleEquivalent(V1, V2, Mask, Unpckh))
     return DAG.getNode(X86ISD::UNPCKH, DL, VT, V2, V1);
 
   return SDValue();
 }
 
 /// \brief Try to emit a bitmask instruction for a shuffle.
 ///
 /// This handles cases where we can model a blend exactly as a bitmask due to
 /// one of the inputs being zeroable.
 static SDValue lowerVectorShuffleAsBitMask(SDLoc DL, MVT VT, SDValue V1,
                                            SDValue V2, ArrayRef<int> Mask,
if (SDValue Rotate = lowerVectorShuffleAsByteRotate(
     return Rotate;
 
   // Try to create an in-lane repeating shuffle mask and then shuffle the
   // the results into the target lanes.
   if (SDValue V = lowerShuffleAsRepeatedMaskAndLanePermute(
           DL, MVT::v32i8, V1, V2, Mask, Subtarget, DAG))
     return V;
 
-  if (isSingleInputShuffleMask(Mask)) {
-    // There are no generalized cross-lane shuffle operations available on i8
-    // element types.
-    if (is128BitLaneCrossingShuffleMask(MVT::v32i8, Mask))
-      return lowerVectorShuffleAsLanePermuteAndBlend(DL, MVT::v32i8, V1, V2,
-                                                     Mask, DAG);
+  SmallBitVector Zeroable = computeZeroableShuffleElements(Mask, V1, V2);
+  bool SingleInputMask = true;
+  bool SingleInputAndZeroesMask = true;
+  for (int i = 0, Size = Mask.size(); i < Size; ++i) {
+    if (Mask[i] >= Size) {
+      SingleInputMask = false;
+      if (!Zeroable[i]) {
+        SingleInputAndZeroesMask = false;
+        break;
+      }
+    }
+  }
+  if (SingleInputMask && is128BitLaneCrossingShuffleMask(MVT::v32i8, Mask))
+    return lowerVectorShuffleAsLanePermuteAndBlend(DL, MVT::v32i8, V1, V2, Mask,
+                                                   DAG);
 
+  if (SingleInputAndZeroesMask) {
     SDValue PSHUFBMask[32];
-    for (int i = 0; i < 32; ++i)
+    for (int i = 0; i < 32; ++i) {
+      if (Mask[i] < 0)
+        PSHUFBMask[i] = DAG.getUNDEF(MVT::i8);
+      else if (Zeroable[i])
+        PSHUFBMask[i] = DAG.getConstant(0x80, DL, MVT::i8);
+      else
         PSHUFBMask[i] =
-          Mask[i] < 0
-              ? DAG.getUNDEF(MVT::i8)
-              : DAG.getConstant(Mask[i] < 16 ? Mask[i] : Mask[i] - 16, DL,
-                                MVT::i8);
+            DAG.getConstant(Mask[i] < 16 ? Mask[i] : Mask[i] - 16, DL, MVT::i8);
+    }
 
     return DAG.getNode(X86ISD::PSHUFB, DL, MVT::v32i8, V1,
                        DAG.getBuildVector(MVT::v32i8, DL, PSHUFBMask));
   }
 
   // Try to simplify this by merging 128-bit lanes to enable a lane-based
   // shuffle.
   if (SDValue Result = lowerVectorShuffleByMerging128BitLanes(

test/CodeGen/X86/vector-shuffle-256-v32.ll

 ; ALL-LABEL: shuffle_v32i8_zz_01_zz_03_zz_05_zz_07_zz_09_zz_11_zz_13_zz_15_zz_17_zz_19_zz_21_zz_23_zz_25_zz_27_zz_29_zz_31:
 ; ALL:       # BB#0:
 ; ALL-NEXT:    vandps {{.*}}(%rip), %ymm0, %ymm0
 ; ALL-NEXT:    retq
   %shuffle = shufflevector <32 x i8> %a, <32 x i8> zeroinitializer, <32 x i32> <i32 32, i32 1, i32 34, i32 3, i32 36, i32 5, i32 38, i32 7, i32 40, i32 9, i32 42, i32 11, i32 44, i32 13, i32 46, i32 15, i32 48, i32 17, i32 50, i32 19, i32 52, i32 21, i32 54, i32 23, i32 56, i32 25, i32 58, i32 27, i32 60, i32 29, i32 62, i32 31>
   ret <32 x i8> %shuffle
 }
 
+define <32 x i8> @shuffle_v32i8_01_zz_02_zz_04_05_06_07_08_09_10_11_12_13_14_15_16_17_18_19_20_21_22_23_24_25_26_27_28_29_30_31(<32 x i8> %a) {
+; AVX1-LABEL: shuffle_v32i8_01_zz_02_zz_04_05_06_07_08_09_10_11_12_13_14_15_16_17_18_19_20_21_22_23_24_25_26_27_28_29_30_31:
+; AVX1:       # BB#0:
+; AVX1-NEXT:    vpshufb {{.*#+}} xmm1 = xmm0[1],zero,xmm0[2],zero,xmm0[4,5,6,7,8,9,10,11,12,13,14,15]
+; AVX1-NEXT:    vextractf128 $1, %ymm0, %xmm0
+; AVX1-NEXT:    vinsertf128 $1, %xmm0, %ymm1, %ymm0
+; AVX1-NEXT:    retq
+;
+; AVX2-LABEL: shuffle_v32i8_01_zz_02_zz_04_05_06_07_08_09_10_11_12_13_14_15_16_17_18_19_20_21_22_23_24_25_26_27_28_29_30_31:
+; AVX2:       # BB#0:
+; AVX2-NEXT:    vpshufb {{.*#+}} ymm0 = ymm0[1],zero,ymm0[2],zero,ymm0[4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31]
+; AVX2-NEXT:    retq
+  %shuffle = shufflevector <32 x i8> %a, <32 x i8> zeroinitializer, <32 x i32> <i32 1, i32 32, i32 2, i32 32, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15, i32 16, i32 17, i32 18, i32 19, i32 20, i32 21, i32 22, i32 23, i32 24, i32 25, i32 26, i32 27, i32 28, i32 29, i32 30, i32 31>
+  ret <32 x i8> %shuffle
+}
+
 define <32 x i8> @shuffle_v32i8_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32(<32 x i8> %a, <32 x i8> %b) {
 ; AVX1-LABEL: shuffle_v32i8_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32_00_32:
 ; AVX1:       # BB#0:
 ; AVX1-NEXT:    vpunpcklbw {{.*#+}} xmm0 = xmm0[0],xmm1[0],xmm0[1],xmm1[1],xmm0[2],xmm1[2],xmm0[3],xmm1[3],xmm0[4],xmm1[4],xmm0[5],xmm1[5],xmm0[6],xmm1[6],xmm0[7],xmm1[7]
 ; AVX1-NEXT:    vpshuflw {{.*#+}} xmm0 = xmm0[0,0,0,0,4,5,6,7]
 ; AVX1-NEXT:    vpshufd {{.*#+}} xmm0 = xmm0[0,0,1,1]
 ; AVX1-NEXT:    vinsertf128 $1, %xmm0, %ymm0, %ymm0
 ; AVX1-NEXT:    retq