[X86][AVX512] Better lowering for _mm512_maskz_shuffle_epi32

llvm/lib/Target/X86/X86ISelLowering.cpp

@@ -17350,6 +17350,54 @@ static SDValue lowerV8I64Shuffle(const SDLoc &DL, ArrayRef<int> Mask,
   return lowerShuffleWithPERMV(DL, MVT::v8i64, Mask, V1, V2, Subtarget, DAG);
 }
 
+static SDValue lowerShuffleAsVSELECT(const SDLoc &DL, ArrayRef<int> Mask,
+                                     SDValue V1, SDValue V2,
+                                     SelectionDAG &DAG) {
+  if (V1.getOpcode() != ISD::BUILD_VECTOR &&
+      V2.getOpcode() != ISD::BUILD_VECTOR)
+    return SDValue();
+
+  bool IsV1BuildVector = V1.getOpcode() == ISD::BUILD_VECTOR;
+  SDValue BuildVector = IsV1BuildVector ? V1 : V2;
+
+  if (!ISD::isBuildVectorAllZeros(BuildVector.getNode()))
+    return SDValue();
+
+  // This relates to the lowering of `_mm512_maskz_shuffle_epi32` intrinsic.
+  // The `BUILD_VECTOR` contains the zeroing mask. If the corresponding
+  // element is UNDEF, then the bit in mask is set. If it is zero, the
+  // corresponding bit in mask is zero.
+  APInt DestMask(16, 0);
+  for (unsigned i = 0; i < 16; ++i) {
+    SDValue Op = BuildVector->getOperand(i);
+    if (Op.isUndef())
+      DestMask.setBit(i);
+  }
+  if (DestMask.isZero())
+    return SDValue();
+
+  SDValue Bitcast = DAG.getNode(ISD::BITCAST, DL, MVT::v16i1,
+                                DAG.getConstant(DestMask, DL, MVT::i16));
+
+  SmallVector<SDValue, 16> ZeroElements(16, DAG.getConstant(0, DL, MVT::i32));
+  SDValue Zeros = DAG.getBuildVector(MVT::v16i32, DL, ZeroElements);
+
+  SmallVector<int, 16> NewMask(16);
+  for (int I = 0; I < 16; ++I) {
+    if (IsV1BuildVector) {
+      NewMask[I] = Mask[I] >= 16 ? Mask[I] - 16 : Mask[I] + 16;
+    } else {
+      NewMask[I] = Mask[I];
+    }
+  }
+
+  return DAG.getNode(ISD::VSELECT, DL, MVT::v16i32, Bitcast,
+                     DAG.getVectorShuffle(MVT::v16i32, DL,
+                                          IsV1BuildVector ? V2 : V1,
+                                          DAG.getUNDEF(MVT::v16i32), NewMask),
+                     Zeros);
+}
+
 /// Handle lowering of 16-lane 32-bit integer shuffles.
 static SDValue lowerV16I32Shuffle(const SDLoc &DL, ArrayRef<int> Mask,
                                   const APInt &Zeroable, SDValue V1, SDValue V2,
@@ -17395,6 +17443,9 @@ static SDValue lowerV16I32Shuffle(const SDLoc &DL, ArrayRef<int> Mask,
     // Use dedicated unpack instructions for masks that match their pattern.
     if (SDValue V = lowerShuffleWithUNPCK(DL, MVT::v16i32, V1, V2, Mask, DAG))
       return V;
+
+    if (SDValue V = lowerShuffleAsVSELECT(DL, Mask, V1, V2, DAG))
+      return V;
   }
 
   // Try to use shift instructions.

llvm/test/CodeGen/X86/vector-shuffle-512-v16.ll

@@ -990,3 +990,38 @@ bb:
   ret void
 }
 
+define <16 x i32> @shuffle_v16f32_00_18_02_18_4_22_06_22_08_26_10_26_12_30_14_30(<16 x i32> %a) {
+; AVX512F-LABEL: shuffle_v16f32_00_18_02_18_4_22_06_22_08_26_10_26_12_30_14_30:
+; AVX512F:       # %bb.0:
+; AVX512F-NEXT:    movw $-21846, %ax # imm = 0xAAAA
+; AVX512F-NEXT:    kmovw %eax, %k1
+; AVX512F-NEXT:    vpshufd {{.*#+}} zmm0 {%k1} {z} = zmm0[2,2,2,2,6,6,6,6,10,10,10,10,14,14,14,14]
+; AVX512F-NEXT:    retq
+;
+; AVX512BW-LABEL: shuffle_v16f32_00_18_02_18_4_22_06_22_08_26_10_26_12_30_14_30:
+; AVX512BW:       # %bb.0:
+; AVX512BW-NEXT:    movw $-21846, %ax # imm = 0xAAAA
+; AVX512BW-NEXT:    kmovd %eax, %k1
+; AVX512BW-NEXT:    vpshufd {{.*#+}} zmm0 {%k1} {z} = zmm0[2,2,2,2,6,6,6,6,10,10,10,10,14,14,14,14]
+; AVX512BW-NEXT:    retq
+  %res = shufflevector <16 x i32> <i32 0, i32 poison, i32 0, i32 poison, i32 0, i32 poison, i32 0, i32 poison, i32 0, i32 poison, i32 0, i32 poison, i32 0, i32 poison, i32 0, i32 poison>, <16 x i32> %a, <16 x i32> <i32 0, i32 18, i32 2, i32 18, i32 4, i32 22, i32 6, i32 22, i32 8, i32 26, i32 10, i32 26, i32 12, i32 30, i32 14, i32 30>
+  ret <16 x i32> %res
+}
+
+define <16 x i32> @shuffle_v16f32_16_02_18_02_20_06_22_06_24_10_26_10_28_14_30_14(<16 x i32> %a) {
+; AVX512F-LABEL: shuffle_v16f32_16_02_18_02_20_06_22_06_24_10_26_10_28_14_30_14:
+; AVX512F:       # %bb.0:
+; AVX512F-NEXT:    movw $-21846, %ax # imm = 0xAAAA
+; AVX512F-NEXT:    kmovw %eax, %k1
+; AVX512F-NEXT:    vpshufd {{.*#+}} zmm0 {%k1} {z} = zmm0[2,2,2,2,6,6,6,6,10,10,10,10,14,14,14,14]
+; AVX512F-NEXT:    retq
+;
+; AVX512BW-LABEL: shuffle_v16f32_16_02_18_02_20_06_22_06_24_10_26_10_28_14_30_14:
+; AVX512BW:       # %bb.0:
+; AVX512BW-NEXT:    movw $-21846, %ax # imm = 0xAAAA
+; AVX512BW-NEXT:    kmovd %eax, %k1
+; AVX512BW-NEXT:    vpshufd {{.*#+}} zmm0 {%k1} {z} = zmm0[2,2,2,2,6,6,6,6,10,10,10,10,14,14,14,14]
+; AVX512BW-NEXT:    retq
+  %res = shufflevector <16 x i32> %a , <16 x i32> <i32 0, i32 poison, i32 0, i32 poison, i32 0, i32 poison, i32 0, i32 poison, i32 0, i32 poison, i32 0, i32 poison, i32 0, i32 poison, i32 0, i32 poison>, <16 x i32> <i32 16, i32 2, i32 18, i32 2, i32 20, i32 6, i32 22, i32 6, i32 24, i32 10, i32 26, i32 10, i32 28, i32 14, i32 30, i32 14>
+  ret <16 x i32> %res
+}