[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

22

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

 
 #include "X86ISelLowering.h"
 #include "Utils/X86ShuffleDecode.h"
 #include "X86CallingConv.h"
 #include "X86FrameLowering.h"
 #include "X86InstrBuilder.h"
 #include "X86MachineFunctionInfo.h"
 #include "X86TargetMachine.h"
 #include "X86TargetObjectFile.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/CodeGen/IntrinsicLowering.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
static SDValue LowerFABSorFNEG(SDValue Op, SelectionDAG &DAG) {
 
   // FIXME: Use function attribute "OptimizeForSize" and/or CodeGenOpt::Level to
   // decide if we should generate a 16-byte constant mask when we only need 4 or
   // 8 bytes for the scalar case.
 
   MVT LogicVT;
   MVT EltVT;
   unsigned NumElts;
-  
   if (VT.isVector()) {
     LogicVT = VT;
     EltVT = VT.getVectorElementType();
     NumElts = VT.getVectorNumElements();
   } else {
     // There are no scalar bitwise logical SSE/AVX instructions, so we
     // generate a 16-byte vector constant and logic op even for the scalar case.
     // Using a 16-byte mask allows folding the load of the mask with
case INTR_TYPE_4OP:
       return DAG.getNode(IntrData->Opc0, dl, Op.getValueType(), Op.getOperand(1),
         Op.getOperand(2), Op.getOperand(3), Op.getOperand(4));
     case INTR_TYPE_1OP_MASK_RM: {
       SDValue Src = Op.getOperand(1);
       SDValue PassThru = Op.getOperand(2);
       SDValue Mask = Op.getOperand(3);
       SDValue RoundingMode;
       // We allways add rounding mode to the Node.
-      // If the rounding mode is not specified, we add the 
+      // If the rounding mode is not specified, we add the
       // "current direction" mode.
       if (Op.getNumOperands() == 4)
         RoundingMode =
           DAG.getConstant(X86::STATIC_ROUNDING::CUR_DIRECTION, dl, MVT::i32);
       else
         RoundingMode = Op.getOperand(4);
       unsigned IntrWithRoundingModeOpcode = IntrData->Opc1;
       if (IntrWithRoundingModeOpcode != 0)
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 = SplatIndex > (int)VT.getVectorNumElements() ? Amt.getOperand(1)
+                                                          : 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));
+          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 {