[X86][SSE] Keep 32-bit target i64 vector shifts on SSE unit.

[RKSimon]

Summary

This patch improves the 32-bit target i64 constant matching to detect the shuffle vector splats that are introduced by i64 vector shift vectorization (D8416).

[qcolombet]

lib/Target/X86/X86ISelLowering.cpp

17157

We should check that the SplatIndex is covered by the first operand of Amt, otherwise the following code is wrong.

[RKSimon]

Thanks Quentin - updated to get the correct operand based on splat index.

[qcolombet]

Hi Simon,

Thanks for the updated version. I have mixed feeling about it though. Indeed, I think the code is correct, but given how the lowering works, I believe we cannot reach it.
In other words, currently, I do not think we can produce a test case that covers this code. If you can add a test case to cover it, then awesome :).

How about not doing the transformation when the splat index is not covered by the first operand, so that if it happens we would have a test case to add with the code support it?

Other than that LGTM, no need to wait for a new review.

Thanks,
-Quentin

[RKSimon]

In D11327#214333, @qcolombet wrote:

Thanks for the updated version. I have mixed feeling about it though. Indeed, I think the code is correct, but given how the lowering works, I believe we cannot reach it.
In other words, currently, I do not think we can produce a test case that covers this code. If you can add a test case to cover it, then awesome :).

How about not doing the transformation when the splat index is not covered by the first operand, so that if it happens we would have a test case to add with the code support it?

Yes the canonicalization in DAG.getVectorShuffle /should/ stop it happening. Would adding an assert to check the splat comes from the first operand be too harsh?

[qcolombet]

I’m fine with an assert.

Thanks,
-Quentin

[RKSimon]

In D11327#214353, @qcolombet wrote:

I’m fine with an assert.

All done. Thanks Quentin.

lib/Target/X86/X86ISelLowering.cpp

if (auto *ShiftConst = BVAmt->getConstantSplatNode()) {
         }
         llvm_unreachable("Unknown shift opcode.");
       }
     }
   }
 
   // Special case in 32-bit mode, where i64 is expanded into high and low parts.
   if (!Subtarget->is64Bit() &&
-      (VT == MVT::v2i64 || (Subtarget->hasInt256() && VT == MVT::v4i64)) &&
-      Amt.getOpcode() == ISD::BITCAST &&
-      Amt.getOperand(0).getOpcode() == ISD::BUILD_VECTOR) {
+      (VT == MVT::v2i64 || (Subtarget->hasInt256() && VT == MVT::v4i64))) {
+
+    // Peek through any splat that was introduced for i64 shift vectorization.
+    int SplatIndex = -1;
+    if (ShuffleVectorSDNode *SVN = dyn_cast<ShuffleVectorSDNode>(Amt.getNode()))
+      if (SVN->isSplat()) {
+        SplatIndex = SVN->getSplatIndex();
+        Amt = Amt.getOperand(0);
+      }
+
+    if (Amt.getOpcode() != ISD::BITCAST ||
+        Amt.getOperand(0).getOpcode() != ISD::BUILD_VECTOR)
+      return SDValue();
+
     Amt = Amt.getOperand(0);
     unsigned Ratio = Amt.getSimpleValueType().getVectorNumElements() /
                      VT.getVectorNumElements();
     unsigned RatioInLog2 = Log2_32_Ceil(Ratio);
     uint64_t ShiftAmt = 0;
+    unsigned BaseOp = (SplatIndex < 0 ? 0 : SplatIndex * Ratio);
     for (unsigned i = 0; i != Ratio; ++i) {
-      ConstantSDNode *C = dyn_cast<ConstantSDNode>(Amt.getOperand(i));
+      ConstantSDNode *C = dyn_cast<ConstantSDNode>(Amt.getOperand(i + BaseOp));
       if (!C)
         return SDValue();
       // 6 == Log2(64)
       ShiftAmt |= C->getZExtValue() << (i * (1 << (6 - RatioInLog2)));
     }
-    // Check remaining shift amounts.
+
+    // Check remaining shift amounts (if not a splat).
+    if (SplatIndex < 0) {
       for (unsigned i = Ratio; i != Amt.getNumOperands(); i += Ratio) {
        uint64_t ShAmt = 0;
         for (unsigned j = 0; j != Ratio; ++j) {
           ConstantSDNode *C =
             dyn_cast<ConstantSDNode>(Amt.getOperand(i + j));
           if (!C)
             return SDValue();
           // 6 == Log2(64)
          ShAmt |= C->getZExtValue() << (j * (1 << (6 - RatioInLog2)));
         }
         if (ShAmt != ShiftAmt)
           return SDValue();
       }
+    }
 
     if (SupportedVectorShiftWithImm(VT, Subtarget, Op.getOpcode()))
       return getTargetVShiftByConstNode(X86Opc, dl, VT, R, ShiftAmt, DAG);
 
     if (Op.getOpcode() == ISD::SRA)
       return ArithmeticShiftRight64(ShiftAmt);
   }
 

test/CodeGen/X86/vector-shift-lshr-128.ll

 ;
 ; AVX2-LABEL: constant_shift_v2i64:
 ; AVX2:       # BB#0:
 ; AVX2-NEXT:    vpsrlvq {{.*}}(%rip), %xmm0, %xmm0
 ; AVX2-NEXT:    retq
 ;
 ; X32-SSE-LABEL: constant_shift_v2i64:
 ; X32-SSE:       # BB#0:
-; X32-SSE-NEXT:    movl $7, %eax
-; X32-SSE-NEXT:    movd %eax, %xmm2
 ; X32-SSE-NEXT:    movdqa %xmm0, %xmm1
-; X32-SSE-NEXT:    psrlq %xmm2, %xmm1
-; X32-SSE-NEXT:    movl $1, %eax
-; X32-SSE-NEXT:    movd %eax, %xmm2
-; X32-SSE-NEXT:    psrlq %xmm2, %xmm0
+; X32-SSE-NEXT:    psrlq $7, %xmm1
+; X32-SSE-NEXT:    psrlq $1, %xmm0
 ; X32-SSE-NEXT:    movsd {{.*#+}} xmm1 = xmm0[0],xmm1[1]
 ; X32-SSE-NEXT:    movapd %xmm1, %xmm0
 ; X32-SSE-NEXT:    retl
   %shift = lshr <2 x i64> %a, <i64 1, i64 7>
   ret <2 x i64> %shift
 }
 
 define <4 x i32> @constant_shift_v4i32(<4 x i32> %a) nounwind {

test/CodeGen/X86/vector-shift-shl-128.ll

 ;
 ; AVX2-LABEL: constant_shift_v2i64:
 ; AVX2:       # BB#0:
 ; AVX2-NEXT:    vpsllvq {{.*}}(%rip), %xmm0, %xmm0
 ; AVX2-NEXT:    retq
 ;
 ; X32-SSE-LABEL: constant_shift_v2i64:
 ; X32-SSE:       # BB#0:
-; X32-SSE-NEXT:    movl $7, %eax
-; X32-SSE-NEXT:    movd %eax, %xmm2
 ; X32-SSE-NEXT:    movdqa %xmm0, %xmm1
-; X32-SSE-NEXT:    psllq %xmm2, %xmm1
-; X32-SSE-NEXT:    movl $1, %eax
-; X32-SSE-NEXT:    movd %eax, %xmm2
-; X32-SSE-NEXT:    psllq %xmm2, %xmm0
+; X32-SSE-NEXT:    psllq $7, %xmm1
+; X32-SSE-NEXT:    psllq $1, %xmm0
 ; X32-SSE-NEXT:    movsd {{.*#+}} xmm1 = xmm0[0],xmm1[1]
 ; X32-SSE-NEXT:    movapd %xmm1, %xmm0
 ; X32-SSE-NEXT:    retl
   %shift = shl <2 x i64> %a, <i64 1, i64 7>
   ret <2 x i64> %shift
 }
 
 define <4 x i32> @constant_shift_v4i32(<4 x i32> %a) nounwind {