Fix reviews

Author: abhishek-kaushik22

llvm/lib/Target/X86/X86ISelLowering.cpp

@@ -17350,24 +17350,23 @@ 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> RepeatedMask, SDValue V1,
-                                     SDValue V2, SelectionDAG &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();
-  SDValue BuildVector;
-  if (V1.getOpcode() == ISD::BUILD_VECTOR) {
-    BuildVector = V1;
-    if (V2.getOpcode() != ISD::BITCAST)
-      return SDValue();
-  } else {
-    BuildVector = V2;
-    if (V1.getOpcode() != ISD::BITCAST)
-      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);
@@ -17377,28 +17376,25 @@ static SDValue lowerShuffleAsVSELECT(const SDLoc &DL,
   if (DestMask.isZero())
     return SDValue();
 
-  unsigned Imm8 = 0;
-  for (unsigned i = 0; i < 4; ++i) {
-    if (V1.getOpcode() != ISD::BUILD_VECTOR) {
-      if (RepeatedMask[i] >= 4) {
-        continue;
-      }
-    } else if (RepeatedMask[i] < 4) {
-      continue;
-    }
-    Imm8 += (RepeatedMask[i] % 4) << (2 * i);
-  }
-
   SDValue Bitcast = DAG.getNode(ISD::BITCAST, DL, MVT::v16i1,
                                 DAG.getConstant(DestMask, DL, MVT::i16));
 
-  std::vector<SDValue> ZeroElements(16, DAG.getConstant(0, DL, MVT::i32));
+  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.getNode(X86ISD::PSHUFD, DL, MVT::v16i32,
-                                 V1.getOpcode() != ISD::BUILD_VECTOR ? V1 : V2,
-                                 DAG.getTargetConstant(Imm8, DL, MVT::i8)),
+                     DAG.getVectorShuffle(MVT::v16i32, DL,
+                                          IsV1BuildVector ? V2 : V1,
+                                          DAG.getUNDEF(MVT::v16i32), NewMask),
                      Zeros);
 }
 
@@ -17448,7 +17444,7 @@ static SDValue lowerV16I32Shuffle(const SDLoc &DL, ArrayRef<int> Mask,
     if (SDValue V = lowerShuffleWithUNPCK(DL, MVT::v16i32, V1, V2, Mask, DAG))
       return V;
 
-    if (SDValue V = lowerShuffleAsVSELECT(DL, RepeatedMask, V1, V2, DAG))
+    if (SDValue V = lowerShuffleAsVSELECT(DL, Mask, V1, V2, DAG))
       return V;
   }
 

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

@@ -990,23 +990,38 @@ bb:
   ret void
 }
 
-define <8 x i64> @pr121147(<8 x i64> %a) {
-; AVX512F-LABEL: pr121147:
-; AVX512F:       # %bb.0: # %entry
+define <16 x i32> @gen_VPSHUFD_AVX512_0(<16 x i32> %a) {
+; AVX512F-LABEL: gen_VPSHUFD_AVX512_0:
+; AVX512F:       # %bb.0:
 ; AVX512F-NEXT:    movw $-21846, %ax # imm = 0xAAAA
 ; AVX512F-NEXT:    kmovw %eax, %k1
-; AVX512F-NEXT:    vpshufd {{.*#+}} zmm0 {%k1} {z} = zmm0[0,2,0,2,4,6,4,6,8,10,8,10,12,14,12,14]
+; 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: pr121147:
-; AVX512BW:       # %bb.0: # %entry
+; AVX512BW-LABEL: gen_VPSHUFD_AVX512_0:
+; AVX512BW:       # %bb.0:
 ; AVX512BW-NEXT:    movw $-21846, %ax # imm = 0xAAAA
 ; AVX512BW-NEXT:    kmovd %eax, %k1
-; AVX512BW-NEXT:    vpshufd {{.*#+}} zmm0 {%k1} {z} = zmm0[0,2,0,2,4,6,4,6,8,10,8,10,12,14,12,14]
+; 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
-entry:
-  %0 = bitcast <8 x i64> %a to <16 x i32>
-  %1 = 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> %0, <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>
-  %2 = bitcast <16 x i32> %1 to <8 x i64>
-  ret <8 x i64> %2
+  %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> @gen_VPSHUFD_AVX512_1(<16 x i32> %a) {
+; AVX512F-LABEL: gen_VPSHUFD_AVX512_1:
+; 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: gen_VPSHUFD_AVX512_1:
+; 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
 }