[x86] use PMOVMSK for vector-sized equality comparisons

[spatel]

Summary

We could do better by splitting any oversized type into whatever vector size the target supports, but I left that for future work if it ever comes up. The motivating case is memcmp() calls on 16-byte structs, so I think we can wire that up with a TLI hook that feeds into this.

[zvi]

LGTM.

lib/Target/X86/X86ISelLowering.cpp

34041 ↗

(On Diff #92392)

Sure!
And a TODO for AVX-512

[RKSimon]

LGTM

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

if (SDValue NewAdd = promoteExtBeforeAdd(N, DAG, Subtarget))
     return NewAdd;
 
   if (SDValue R = combineOrCmpEqZeroToCtlzSrl(N, DAG, DCI, Subtarget))
     return R;
 
   return SDValue();
 }
 
+/// Try to map a 128-bit or larger integer comparison to vector instructions
+/// before type legalization splits it up into chunks.
+static SDValue combineVectorSizedSetCCEquality(SDNode *SetCC, SelectionDAG &DAG,
+                                               const X86Subtarget &Subtarget) {
+  ISD::CondCode CC = cast<CondCodeSDNode>(SetCC->getOperand(2))->get();
+  assert((CC == ISD::SETNE || CC == ISD::SETEQ) && "Bad comparison predicate");
+
+  // We're looking for an oversized integer equality comparison, but ignore a
+  // comparison with zero because that gets special treatment in EmitTest().
+  SDValue X = SetCC->getOperand(0);
+  SDValue Y = SetCC->getOperand(1);
+  EVT OpVT = X.getValueType();
+  unsigned OpSize = OpVT.getSizeInBits();
+  if (!OpVT.isScalarInteger() || OpSize < 128 || isNullConstant(Y))
+    return SDValue();
+
+  // TODO: Use PXOR + PTEST for SSE4.1 or later?
+  // TODO: Add support for AVX-512.
+  EVT VT = SetCC->getValueType(0);
+  SDLoc DL(SetCC);
+  if ((OpSize == 128 && Subtarget.hasSSE2()) ||
+      (OpSize == 256 && Subtarget.hasAVX2())) {
+    EVT VecVT = OpSize == 128 ? MVT::v16i8 : MVT::v32i8;
+    SDValue VecX = DAG.getBitcast(VecVT, X);
+    SDValue VecY = DAG.getBitcast(VecVT, Y);
+
+    // If all bytes match (bitmask is 0x(FFFF)FFFF), that's equality.
+    // setcc i128 X, Y, eq --> setcc (pmovmskb (pcmpeqb X, Y)), 0xFFFF, eq
+    // setcc i128 X, Y, ne --> setcc (pmovmskb (pcmpeqb X, Y)), 0xFFFF, ne
+    // setcc i256 X, Y, eq --> setcc (vpmovmskb (vpcmpeqb X, Y)), 0xFFFFFFFF, eq
+    // setcc i256 X, Y, ne --> setcc (vpmovmskb (vpcmpeqb X, Y)), 0xFFFFFFFF, ne
+    SDValue Cmp = DAG.getNode(X86ISD::PCMPEQ, DL, VecVT, VecX, VecY);
+    SDValue MovMsk = DAG.getNode(X86ISD::MOVMSK, DL, MVT::i32, Cmp);
+    SDValue FFFFs = DAG.getConstant(OpSize == 128 ? 0xFFFF : 0xFFFFFFFF, DL,
+                                    MVT::i32);
+    return DAG.getSetCC(DL, VT, MovMsk, FFFFs, CC);
+  }
+
+  return SDValue();
+}
+
 static SDValue combineSetCC(SDNode *N, SelectionDAG &DAG,
                             const X86Subtarget &Subtarget) {
   ISD::CondCode CC = cast<CondCodeSDNode>(N->getOperand(2))->get();
   SDValue LHS = N->getOperand(0);
   SDValue RHS = N->getOperand(1);
   EVT VT = N->getValueType(0);
   SDLoc DL(N);
 
   if (CC == ISD::SETNE || CC == ISD::SETEQ) {
     EVT OpVT = LHS.getValueType();
     // 0-x == y --> x+y == 0
     // 0-x != y --> x+y != 0
     if (LHS.getOpcode() == ISD::SUB && isNullConstant(LHS.getOperand(0)) &&
         LHS.hasOneUse()) {
       SDValue Add = DAG.getNode(ISD::ADD, DL, OpVT, RHS, LHS.getOperand(1));
       return DAG.getSetCC(DL, VT, Add, DAG.getConstant(0, DL, OpVT), CC);
     }
     // x == 0-y --> x+y == 0
     // x != 0-y --> x+y != 0
     if (RHS.getOpcode() == ISD::SUB && isNullConstant(RHS.getOperand(0)) &&
         RHS.hasOneUse()) {
       SDValue Add = DAG.getNode(ISD::ADD, DL, OpVT, LHS, RHS.getOperand(1));
       return DAG.getSetCC(DL, VT, Add, DAG.getConstant(0, DL, OpVT), CC);
     }
+
+    if (SDValue V = combineVectorSizedSetCCEquality(N, DAG, Subtarget))
+      return V;
   }
 
   if (VT.getScalarType() == MVT::i1 &&
       (CC == ISD::SETNE || CC == ISD::SETEQ || ISD::isSignedIntSetCC(CC))) {
     bool IsSEXT0 =
         (LHS.getOpcode() == ISD::SIGN_EXTEND) &&
         (LHS.getOperand(0).getValueType().getScalarType() == MVT::i1);
     bool IsVZero1 = ISD::isBuildVectorAllZeros(RHS.getNode());

llvm/trunk/test/CodeGen/X86/setcc-wide-types.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=sse2 | FileCheck %s --check-prefix=SSE2
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=avx2 | FileCheck %s --check-prefix=AVX2
 
-; FIXME: Equality checks of 128/256-bit values can use PMOVMSK or PTEST to avoid scalarization.
+; Equality checks of 128/256-bit values can use PMOVMSK or PTEST to avoid scalarization.
 
 define i32 @ne_i128(<2 x i64> %x, <2 x i64> %y) {
 ; SSE2-LABEL: ne_i128:
 ; SSE2:       # BB#0:
-; SSE2-NEXT:    pshufd {{.*#+}} xmm2 = xmm0[2,3,0,1]
-; SSE2-NEXT:    movd %xmm2, %rax
-; SSE2-NEXT:    movd %xmm0, %rcx
-; SSE2-NEXT:    pshufd {{.*#+}} xmm0 = xmm1[2,3,0,1]
-; SSE2-NEXT:    movd %xmm0, %rdx
-; SSE2-NEXT:    movd %xmm1, %rsi
-; SSE2-NEXT:    xorq %rcx, %rsi
-; SSE2-NEXT:    xorq %rax, %rdx
+; SSE2-NEXT:    pcmpeqb %xmm1, %xmm0
+; SSE2-NEXT:    pmovmskb %xmm0, %ecx
 ; SSE2-NEXT:    xorl %eax, %eax
-; SSE2-NEXT:    orq %rsi, %rdx
+; SSE2-NEXT:    cmpl $65535, %ecx # imm = 0xFFFF
 ; SSE2-NEXT:    setne %al
 ; SSE2-NEXT:    retq
 ;
 ; AVX2-LABEL: ne_i128:
 ; AVX2:       # BB#0:
-; AVX2-NEXT:    vmovq %xmm0, %rax
-; AVX2-NEXT:    vpextrq $1, %xmm0, %rcx
-; AVX2-NEXT:    vmovq %xmm1, %rdx
-; AVX2-NEXT:    vpextrq $1, %xmm1, %rsi
-; AVX2-NEXT:    xorq %rcx, %rsi
-; AVX2-NEXT:    xorq %rax, %rdx
+; AVX2-NEXT:    vpcmpeqb %xmm1, %xmm0, %xmm0
+; AVX2-NEXT:    vpmovmskb %xmm0, %ecx
 ; AVX2-NEXT:    xorl %eax, %eax
-; AVX2-NEXT:    orq %rsi, %rdx
+; AVX2-NEXT:    cmpl $65535, %ecx # imm = 0xFFFF
 ; AVX2-NEXT:    setne %al
 ; AVX2-NEXT:    retq
   %bcx = bitcast <2 x i64> %x to i128
   %bcy = bitcast <2 x i64> %y to i128
   %cmp = icmp ne i128 %bcx, %bcy
   %zext = zext i1 %cmp to i32
   ret i32 %zext
 }
 
 define i32 @eq_i128(<2 x i64> %x, <2 x i64> %y) {
 ; SSE2-LABEL: eq_i128:
 ; SSE2:       # BB#0:
-; SSE2-NEXT:    pshufd {{.*#+}} xmm2 = xmm0[2,3,0,1]
-; SSE2-NEXT:    movd %xmm2, %rax
-; SSE2-NEXT:    movd %xmm0, %rcx
-; SSE2-NEXT:    pshufd {{.*#+}} xmm0 = xmm1[2,3,0,1]
-; SSE2-NEXT:    movd %xmm0, %rdx
-; SSE2-NEXT:    movd %xmm1, %rsi
-; SSE2-NEXT:    xorq %rcx, %rsi
-; SSE2-NEXT:    xorq %rax, %rdx
+; SSE2-NEXT:    pcmpeqb %xmm1, %xmm0
+; SSE2-NEXT:    pmovmskb %xmm0, %ecx
 ; SSE2-NEXT:    xorl %eax, %eax
-; SSE2-NEXT:    orq %rsi, %rdx
+; SSE2-NEXT:    cmpl $65535, %ecx # imm = 0xFFFF
 ; SSE2-NEXT:    sete %al
 ; SSE2-NEXT:    retq
 ;
 ; AVX2-LABEL: eq_i128:
 ; AVX2:       # BB#0:
-; AVX2-NEXT:    vmovq %xmm0, %rax
-; AVX2-NEXT:    vpextrq $1, %xmm0, %rcx
-; AVX2-NEXT:    vmovq %xmm1, %rdx
-; AVX2-NEXT:    vpextrq $1, %xmm1, %rsi
-; AVX2-NEXT:    xorq %rcx, %rsi
-; AVX2-NEXT:    xorq %rax, %rdx
+; AVX2-NEXT:    vpcmpeqb %xmm1, %xmm0, %xmm0
+; AVX2-NEXT:    vpmovmskb %xmm0, %ecx
 ; AVX2-NEXT:    xorl %eax, %eax
-; AVX2-NEXT:    orq %rsi, %rdx
+; AVX2-NEXT:    cmpl $65535, %ecx # imm = 0xFFFF
 ; AVX2-NEXT:    sete %al
 ; AVX2-NEXT:    retq
   %bcx = bitcast <2 x i64> %x to i128
   %bcy = bitcast <2 x i64> %y to i128
   %cmp = icmp eq i128 %bcx, %bcy
   %zext = zext i1 %cmp to i32
   ret i32 %zext
 }

 ; SSE2-NEXT:    orq %rax, %rdi
 ; SSE2-NEXT:    xorl %eax, %eax
 ; SSE2-NEXT:    orq %rcx, %rdi
 ; SSE2-NEXT:    setne %al
 ; SSE2-NEXT:    retq
 ;
 ; AVX2-LABEL: ne_i256:
 ; AVX2:       # BB#0:
-; AVX2-NEXT:    vmovq %xmm0, %r8
-; AVX2-NEXT:    vextracti128 $1, %ymm0, %xmm2
-; AVX2-NEXT:    vmovq %xmm2, %r9
-; AVX2-NEXT:    vpextrq $1, %xmm0, %r10
-; AVX2-NEXT:    vpextrq $1, %xmm2, %rsi
-; AVX2-NEXT:    vmovq %xmm1, %rdi
-; AVX2-NEXT:    vextracti128 $1, %ymm1, %xmm0
-; AVX2-NEXT:    vmovq %xmm0, %rax
-; AVX2-NEXT:    vpextrq $1, %xmm1, %rcx
-; AVX2-NEXT:    vpextrq $1, %xmm0, %rdx
-; AVX2-NEXT:    xorq %rsi, %rdx
-; AVX2-NEXT:    xorq %r10, %rcx
-; AVX2-NEXT:    orq %rdx, %rcx
-; AVX2-NEXT:    xorq %r9, %rax
-; AVX2-NEXT:    xorq %r8, %rdi
-; AVX2-NEXT:    orq %rax, %rdi
+; AVX2-NEXT:    vpcmpeqb %ymm1, %ymm0, %ymm0
+; AVX2-NEXT:    vpmovmskb %ymm0, %ecx
 ; AVX2-NEXT:    xorl %eax, %eax
-; AVX2-NEXT:    orq %rcx, %rdi
+; AVX2-NEXT:    cmpl $-1, %ecx
 ; AVX2-NEXT:    setne %al
 ; AVX2-NEXT:    vzeroupper
 ; AVX2-NEXT:    retq
   %bcx = bitcast <4 x i64> %x to i256
   %bcy = bitcast <4 x i64> %y to i256
   %cmp = icmp ne i256 %bcx, %bcy
   %zext = zext i1 %cmp to i32
   ret i32 %zext

 ; SSE2-NEXT:    orq %rax, %rdi
 ; SSE2-NEXT:    xorl %eax, %eax
 ; SSE2-NEXT:    orq %rcx, %rdi
 ; SSE2-NEXT:    sete %al
 ; SSE2-NEXT:    retq
 ;
 ; AVX2-LABEL: eq_i256:
 ; AVX2:       # BB#0:
-; AVX2-NEXT:    vmovq %xmm0, %r8
-; AVX2-NEXT:    vextracti128 $1, %ymm0, %xmm2
-; AVX2-NEXT:    vmovq %xmm2, %r9
-; AVX2-NEXT:    vpextrq $1, %xmm0, %r10
-; AVX2-NEXT:    vpextrq $1, %xmm2, %rsi
-; AVX2-NEXT:    vmovq %xmm1, %rdi
-; AVX2-NEXT:    vextracti128 $1, %ymm1, %xmm0
-; AVX2-NEXT:    vmovq %xmm0, %rax
-; AVX2-NEXT:    vpextrq $1, %xmm1, %rcx
-; AVX2-NEXT:    vpextrq $1, %xmm0, %rdx
-; AVX2-NEXT:    xorq %rsi, %rdx
-; AVX2-NEXT:    xorq %r10, %rcx
-; AVX2-NEXT:    orq %rdx, %rcx
-; AVX2-NEXT:    xorq %r9, %rax
-; AVX2-NEXT:    xorq %r8, %rdi
-; AVX2-NEXT:    orq %rax, %rdi
+; AVX2-NEXT:    vpcmpeqb %ymm1, %ymm0, %ymm0
+; AVX2-NEXT:    vpmovmskb %ymm0, %ecx
 ; AVX2-NEXT:    xorl %eax, %eax
-; AVX2-NEXT:    orq %rcx, %rdi
+; AVX2-NEXT:    cmpl $-1, %ecx
 ; AVX2-NEXT:    sete %al
 ; AVX2-NEXT:    vzeroupper
 ; AVX2-NEXT:    retq
   %bcx = bitcast <4 x i64> %x to i256
   %bcy = bitcast <4 x i64> %y to i256
   %cmp = icmp eq i256 %bcx, %bcy
   %zext = zext i1 %cmp to i32
   ret i32 %zext
 }