[AVX-512] Add support for changing VSHUFF64x2 to VSHUFF32x4 when its feeding a vselect with 32-bit element size.

[craig.topper]

Summary

Shuffle lowering may have widened the element size of a i32 shuffle to i64 before selecting X86ISD::SHUF128. If this shuffle was used by a vselect this can prevent us from selecting masked operations.

This patch detects this and changes the element size to match the vselect.

I don't handle changing integer to floating point or vice versa as its not clear if its better to push such a bitcast to the inputs of the shuffle or to the user of the vselect. So I'm ignoring that case for now.

[zvi]

LGTM

[craig.topper]

Zvi, did you forget to accept the revision or did you want someone else to look at it?

[RKSimon]

LGTM

lib/Target/X86/X86ISelLowering.cpp

27779 ↗

(On Diff #79187)

Add a TODO to investigate this in the future?

llvm/trunk/lib/Target/X86/X86ISelLowering.cpp

static bool combineBitcastForMaskedOp(SDValue OrigOp, SelectionDAG &DAG,
   // do this combine as that would replicate the operation.
   if (!Op.hasOneUse())
     return false;
 
   MVT VT = OrigOp.getSimpleValueType();
   MVT EltVT = VT.getVectorElementType();
   SDLoc DL(Op.getNode());
 
-  switch (Op.getOpcode()) {
+  auto BitcastAndCombineShuffle = [&](unsigned Opcode, SDValue Op0, SDValue Op1,
+                                      SDValue Op2) {
+    Op0 = DAG.getBitcast(VT, Op0);
+    DCI.AddToWorklist(Op0.getNode());
+    Op1 = DAG.getBitcast(VT, Op1);
+    DCI.AddToWorklist(Op1.getNode());
+    DCI.CombineTo(OrigOp.getNode(),
+                  DAG.getNode(Opcode, DL, VT, Op0, Op1, Op2));
+    return true;
+  };
+
+  unsigned Opcode = Op.getOpcode();
+  switch (Opcode) {
   case X86ISD::PALIGNR:
     // PALIGNR can be converted to VALIGND/Q for 128-bit vectors.
     if (!VT.is128BitVector())
       return false;
+    Opcode = X86ISD::VALIGN;
     LLVM_FALLTHROUGH;
   case X86ISD::VALIGN: {
     if (EltVT != MVT::i32 && EltVT != MVT::i64)
       return false;
     uint64_t Imm = cast<ConstantSDNode>(Op.getOperand(2))->getZExtValue();
     MVT OpEltVT = Op.getSimpleValueType().getVectorElementType();
     unsigned ShiftAmt = Imm * OpEltVT.getSizeInBits();
     unsigned EltSize = EltVT.getSizeInBits();
     // Make sure we can represent the same shift with the new VT.
     if ((ShiftAmt % EltSize) != 0)
       return false;
     Imm = ShiftAmt / EltSize;
-    SDValue Op0 = DAG.getBitcast(VT, Op.getOperand(0));
-    DCI.AddToWorklist(Op0.getNode());
-    SDValue Op1 = DAG.getBitcast(VT, Op.getOperand(1));
-    DCI.AddToWorklist(Op1.getNode());
-    DCI.CombineTo(OrigOp.getNode(),
-                  DAG.getNode(X86ISD::VALIGN, DL, VT, Op0, Op1,
-                              DAG.getConstant(Imm, DL, MVT::i8)));
-    return true;
+    return BitcastAndCombineShuffle(Opcode, Op.getOperand(0), Op.getOperand(1),
+                                    DAG.getConstant(Imm, DL, MVT::i8));
+  }
+  case X86ISD::SHUF128: {
+    if (EltVT.getSizeInBits() != 32 && EltVT.getSizeInBits() != 64)
+      return false;
+    // Only change element size, not type.
+    if (VT.isInteger() != Op.getSimpleValueType().isInteger())
+      return false;
+    return BitcastAndCombineShuffle(Opcode, Op.getOperand(0), Op.getOperand(1),
+                                    Op.getOperand(2));
   }
   }
 
   return false;
 }
 
 /// Do target-specific dag combines on SELECT and VSELECT nodes.
 static SDValue combineSelect(SDNode *N, SelectionDAG &DAG,

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

 }
 
 define <16 x float> @test_vshuff32x4_512_mask(<16 x float> %x, <16 x float> %x1, <16 x float> %y, <16 x i1> %mask) nounwind {
 ; AVX512F-LABEL: test_vshuff32x4_512_mask:
 ; AVX512F:       # BB#0:
 ; AVX512F-NEXT:    vpmovsxbd %xmm3, %zmm3
 ; AVX512F-NEXT:    vpslld $31, %zmm3, %zmm3
 ; AVX512F-NEXT:    vptestmd %zmm3, %zmm3, %k1
-; AVX512F-NEXT:    vshuff64x2 {{.*#+}} zmm0 = zmm0[0,1,2,3],zmm1[2,3,0,1]
-; AVX512F-NEXT:    vblendmps %zmm0, %zmm2, %zmm0 {%k1}
+; AVX512F-NEXT:    vshuff32x4 {{.*#+}} zmm2 {%k1} = zmm0[0,1,2,3,4,5,6,7],zmm1[4,5,6,7,0,1,2,3]
+; AVX512F-NEXT:    vmovaps %zmm2, %zmm0
 ; AVX512F-NEXT:    retq
 ;
 ; AVX512BW-LABEL: test_vshuff32x4_512_mask:
 ; AVX512BW:       # BB#0:
 ; AVX512BW-NEXT:    vpsllw $7, %xmm3, %xmm3
 ; AVX512BW-NEXT:    vpmovb2m %zmm3, %k1
-; AVX512BW-NEXT:    vshuff64x2 {{.*#+}} zmm0 = zmm0[0,1,2,3],zmm1[2,3,0,1]
-; AVX512BW-NEXT:    vblendmps %zmm0, %zmm2, %zmm0 {%k1}
+; AVX512BW-NEXT:    vshuff32x4 {{.*#+}} zmm2 {%k1} = zmm0[0,1,2,3,4,5,6,7],zmm1[4,5,6,7,0,1,2,3]
+; AVX512BW-NEXT:    vmovaps %zmm2, %zmm0
 ; AVX512BW-NEXT:    retq
   %x2 = shufflevector <16 x float> %x, <16 x float> %x1, <16 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 20, i32 21, i32 22, i32 23, i32 16, i32 17, i32 18, i32 19>
   %res = select <16 x i1> %mask, <16 x float> %x2, <16 x float> %y
   ret <16 x float> %res
 }
 
 define <16 x i32> @test_vshufi32x4_512_mask(<16 x i32> %x, <16 x i32> %x1, <16 x i32> %y, <16 x i1> %mask) nounwind {
 ; AVX512F-LABEL: test_vshufi32x4_512_mask:
 ; AVX512F:       # BB#0:
 ; AVX512F-NEXT:    vpmovsxbd %xmm3, %zmm3
 ; AVX512F-NEXT:    vpslld $31, %zmm3, %zmm3
 ; AVX512F-NEXT:    vptestmd %zmm3, %zmm3, %k1
-; AVX512F-NEXT:    vshufi64x2 {{.*#+}} zmm0 = zmm0[0,1,2,3],zmm1[2,3,0,1]
-; AVX512F-NEXT:    vpblendmd %zmm0, %zmm2, %zmm0 {%k1}
+; AVX512F-NEXT:    vshufi32x4 {{.*#+}} zmm2 {%k1} = zmm0[0,1,2,3,4,5,6,7],zmm1[4,5,6,7,0,1,2,3]
+; AVX512F-NEXT:    vmovdqa64 %zmm2, %zmm0
 ; AVX512F-NEXT:    retq
 ;
 ; AVX512BW-LABEL: test_vshufi32x4_512_mask:
 ; AVX512BW:       # BB#0:
 ; AVX512BW-NEXT:    vpsllw $7, %xmm3, %xmm3
 ; AVX512BW-NEXT:    vpmovb2m %zmm3, %k1
-; AVX512BW-NEXT:    vshufi64x2 {{.*#+}} zmm0 = zmm0[0,1,2,3],zmm1[2,3,0,1]
-; AVX512BW-NEXT:    vpblendmd %zmm0, %zmm2, %zmm0 {%k1}
+; AVX512BW-NEXT:    vshufi32x4 {{.*#+}} zmm2 {%k1} = zmm0[0,1,2,3,4,5,6,7],zmm1[4,5,6,7,0,1,2,3]
+; AVX512BW-NEXT:    vmovdqa64 %zmm2, %zmm0
 ; AVX512BW-NEXT:    retq
   %x2 = shufflevector <16 x i32> %x, <16 x i32> %x1, <16 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 20, i32 21, i32 22, i32 23, i32 16, i32 17, i32 18, i32 19>
   %res = select <16 x i1> %mask, <16 x i32> %x2, <16 x i32> %y
   ret <16 x i32> %res
 }