[X86][SSE] MatchVectorAllZeroTest - handle OR vector reductions

[RKSimon]

Summary

This patch extends to handle OR vector reduction patterns where the result is compared against zero.

Fixes PR45378

[spatel]

llvm/lib/Target/X86/X86ISelLowering.cpp

45534

Code comments like this and the function name don't apply if we're matching a reduction, right?
I think we need to refactor and/or update comments to reduce confusion.

llvm/test/CodeGen/X86/pr45378.ll

10–11

Add 'nounwind' attribute to remove .cfi noise?
Would it make sense to minimize the test to return the i1 icmp rather than branch and make an external call?

[RKSimon]

rebased tests

[RKSimon]

Update comments to describe reductions as well

[RKSimon]

I think I can put this in LowerVectorAllZeroTest instead

[RKSimon]

Refreshed the patch to be based off LowerVectorAllZeroTest instead of combineVectorSizedSetCCEquality, which already handles scalar reductions and with a suitable refactor can handle vector reductions as well.

[RKSimon]

More changes to come - still finding regressions

[RKSimon]

Rebased now that the LowerVectorAllZero refactor is complete

[spatel]

llvm/lib/Target/X86/X86ISelLowering.cpp

21395

Can we assert that CC is ISD::SETEQ or ISD::SETNE here?

[RKSimon]

Add condition code assertions

[spatel]

LGTM

[hans]

This seems to be causing Chromium test failures. Diffing the asm looks suspicious:

 Disassembly of section .text._ZN5glcts18ComputeShaderTestsC2ERN4deqp7ContextE:
@@ -29795,7 +29795,7 @@
     11c4:      66 0f 63 c3             packsswb %xmm3,%xmm0
     11c8:      66 0f d7 c0             pmovmskb %xmm0,%eax
     11cc:      66 85 c0                test   %ax,%ax
-    11cf:      0f 84 65 01 00 00       je     133a <_ZN5glcts12_GLOBAL__N_128AdvancedPipelineComputeChain3RunEv+0x133a>
+    11cf:      0f 85 65 01 00 00       jne    133a <_ZN5glcts12_GLOBAL__N_128AdvancedPipelineComputeChain3RunEv+0x133a>
     11d5:      49 8b 45 20             mov    0x20(%r13),%rax
     11d9:      48 8b 00                mov    (%rax),%rax
     11dc:      48 8b 40 10             mov    0x10(%rax),%rax
@@ -29894,7 +29894,7 @@
     137d:      66 0f 63 c2             packsswb %xmm2,%xmm0
     1381:      66 0f d7 c0             pmovmskb %xmm0,%eax
     1385:      66 85 c0                test   %ax,%ax
-    1388:      0f 84 64 01 00 00       je     14f2 <_ZN5glcts12_GLOBAL__N_128AdvancedPipelineComputeChain3RunEv+0x14f2>
+    1388:      0f 85 64 01 00 00       jne    14f2 <_ZN5glcts12_GLOBAL__N_128AdvancedPipelineComputeChain3RunEv+0x14f2>
     138e:      49 8b 45 20             mov    0x20(%r13),%rax
     1392:      48 8b 00                mov    (%rax),%rax
     1395:      48 8b 40 10             mov    0x10(%rax),%rax

It seems the conditions on the jumps there are getting flipped, which really looks like a bug. I've posted a repro here: https://bugs.chromium.org/p/chromium/issues/detail?id=1097758#c9

I've reverted in 1357c065783e8d66c9db4be59aca389d1dc6c05f in the meantime.

llvm/lib/Target/X86/X86ISelLowering.cpp

   SDValue C = N->getOperand(IsStrict ? 3 : 2);
 
   auto invertIfNegative = [&DAG, &TLI, &DCI](SDValue &V) {
     bool CodeSize = DAG.getMachineFunction().getFunction().hasOptSize();
     bool LegalOperations = !DCI.isBeforeLegalizeOps();
     if (SDValue NegV = TLI.getCheaperNegatedExpression(V, DAG, LegalOperations,
                                                        CodeSize)) {
       V = NegV;
       return true;
     }
     // Look through extract_vector_elts. If it comes from an FNEG, create a
     // new extract from the FNEG input.
     if (V.getOpcode() == ISD::EXTRACT_VECTOR_ELT &&
         isNullConstant(V.getOperand(1))) {
       SDValue Vec = V.getOperand(0);
       if (SDValue NegV = TLI.getCheaperNegatedExpression(
               Vec, DAG, LegalOperations, CodeSize)) {

       return DAG.getSetCC(DL, CmpVT, A, B, ISD::SETNE);
     if (HasPT)
       return DAG.getNode(ISD::XOR, DL, VecVT, A, B);
     return DAG.getSetCC(DL, CmpVT, A, B, ISD::SETEQ);
   }
   llvm_unreachable("Impossible");
 }
 
-/// Try to map a 128-bit or larger integer comparison to vector instructions
-/// before type legalization splits it up into chunks.
+/// Try to map a 128-bit or larger integer comparison/or-vector-reduction 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.
+  // We're looking for an oversized integer equality comparison or an
+  // or-vector-reduction compared agaionst zero.
   SDValue X = SetCC->getOperand(0);
   SDValue Y = SetCC->getOperand(1);
   EVT OpVT = X.getValueType();
+  EVT VT = SetCC->getValueType(0);
+  bool IsCmpWithZero = isNullConstant(Y);
   unsigned OpSize = OpVT.getSizeInBits();
-  if (!OpVT.isScalarInteger() || OpSize < 128)
+  if (!OpVT.isScalarInteger())
     return SDValue();
 
+  // If we're comparing a smaller integer - see if it came from a OR vector
+  // reduction, compared against zero, else bail out.
+  SDLoc DL(SetCC);
+  bool IsReduction = false;
+  if (OpSize < 128) {
+    if (!IsCmpWithZero || X.getOpcode() != ISD::EXTRACT_VECTOR_ELT)
+      return SDValue();
+    ISD::NodeType BinOp;
+    SDValue Match = DAG.matchBinOpReduction(X.getNode(), BinOp, {ISD::OR});
+    if (!Match || Match.getValueSizeInBits() < 128)
+      return SDValue();
+    // Adjust the comparison values to the reductions's vector sources.
+    IsReduction = true;
+    OpSize = Match.getValueSizeInBits();
+    X = Match;
+    Y = DAG.getConstant(0, DL, X.getValueType());
+  }
+
   // Ignore a comparison with zero because that gets special treatment in
   // EmitTest(). But make an exception for the special case of a pair of
   // logically-combined vector-sized operands compared to zero. This pattern may
   // be generated by the memcmp expansion pass with oversized integer compares
   // (see PR33325).
-  bool IsOrXorXorTreeCCZero = isNullConstant(Y) && isOrXorXorTree(X);
-  if (isNullConstant(Y) && !IsOrXorXorTreeCCZero)
+  bool IsOrXorXorTreeCCZero = IsCmpWithZero && isOrXorXorTree(X);
+  if (IsCmpWithZero && !IsOrXorXorTreeCCZero && !IsReduction)
     return SDValue();
 
   // Don't perform this combine if constructing the vector will be expensive.
   auto IsVectorBitCastCheap = [](SDValue X) {
     X = peekThroughBitcasts(X);
     return isa<ConstantSDNode>(X) || X.getValueType().isVector() ||
            X.getOpcode() == ISD::LOAD;
   };
-  if ((!IsVectorBitCastCheap(X) || !IsVectorBitCastCheap(Y)) &&
-      !IsOrXorXorTreeCCZero)
+  if (!IsOrXorXorTreeCCZero && !IsReduction &&
+      (!IsVectorBitCastCheap(X) || !IsVectorBitCastCheap(Y)))
     return SDValue();
 
-  EVT VT = SetCC->getValueType(0);
-  SDLoc DL(SetCC);
-  bool HasAVX = Subtarget.hasAVX();
-
   // Use XOR (plus OR) and PTEST after SSE4.1 for 128/256-bit operands.
   // Use PCMPNEQ (plus OR) and KORTEST for 512-bit operands.
   // Otherwise use PCMPEQ (plus AND) and mask testing.
   if ((OpSize == 128 && Subtarget.hasSSE2()) ||
-      (OpSize == 256 && HasAVX) ||
+      (OpSize == 256 && Subtarget.hasAVX()) ||
       (OpSize == 512 && Subtarget.useAVX512Regs())) {
     bool HasPT = Subtarget.hasSSE41();
 
     // PTEST and MOVMSK are slow on Knights Landing and Knights Mill and widened
     // vector registers are essentially free. (Technically, widening registers
     // prevents load folding, but the tradeoff is worth it.)
     bool PreferKOT = Subtarget.preferMaskRegisters();
     bool NeedZExt = PreferKOT && !Subtarget.hasVLX() && OpSize != 512;

llvm/test/CodeGen/X86/pr45378.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-prefixes=CHECK,SSE,SSE2
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=sse4.1   | FileCheck %s --check-prefixes=CHECK,SSE,SSE41
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=avx      | FileCheck %s --check-prefixes=CHECK,AVX,AVX1OR2,AVX1
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=avx2     | FileCheck %s --check-prefixes=CHECK,AVX,AVX1OR2,AVX2
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=avx512f  | FileCheck %s --check-prefixes=CHECK,AVX,AVX512,AVX512F
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=avx512bw | FileCheck %s --check-prefixes=CHECK,AVX,AVX512,AVX512BW
 
 declare i64 @llvm.experimental.vector.reduce.or.v2i64(<2 x i64>)
 
 define i1 @parseHeaders(i64 * %ptr) nounwind {
-; SSE-LABEL: parseHeaders:
-; SSE:       # %bb.0:
-; SSE-NEXT:    movdqu (%rdi), %xmm0
-; SSE-NEXT:    pshufd {{.*#+}} xmm1 = xmm0[2,3,0,1]
-; SSE-NEXT:    por %xmm0, %xmm1
-; SSE-NEXT:    movq %xmm1, %rax
-; SSE-NEXT:    testq %rax, %rax
-; SSE-NEXT:    sete %al
-; SSE-NEXT:    retq
-;
-; AVX1-LABEL: parseHeaders:
-; AVX1:       # %bb.0:
-; AVX1-NEXT:    vmovdqu (%rdi), %xmm0
-; AVX1-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[2,3,0,1]
-; AVX1-NEXT:    vpor %xmm1, %xmm0, %xmm0
-; AVX1-NEXT:    vmovq %xmm0, %rax
-; AVX1-NEXT:    testq %rax, %rax
-; AVX1-NEXT:    sete %al
-; AVX1-NEXT:    retq
-;
-; AVX2-LABEL: parseHeaders:
-; AVX2:       # %bb.0:
-; AVX2-NEXT:    vpbroadcastq 8(%rdi), %xmm0
-; AVX2-NEXT:    vpor (%rdi), %xmm0, %xmm0
-; AVX2-NEXT:    vmovq %xmm0, %rax
-; AVX2-NEXT:    testq %rax, %rax
-; AVX2-NEXT:    sete %al
-; AVX2-NEXT:    retq
-;
-; AVX512-LABEL: parseHeaders:
-; AVX512:       # %bb.0:
-; AVX512-NEXT:    vpbroadcastq 8(%rdi), %xmm0
-; AVX512-NEXT:    vpor (%rdi), %xmm0, %xmm0
-; AVX512-NEXT:    vmovq %xmm0, %rax
-; AVX512-NEXT:    testq %rax, %rax
-; AVX512-NEXT:    sete %al
-; AVX512-NEXT:    retq
+; SSE2-LABEL: parseHeaders:
+; SSE2:       # %bb.0:
+; SSE2-NEXT:    movdqu (%rdi), %xmm0
+; SSE2-NEXT:    pxor %xmm1, %xmm1
+; SSE2-NEXT:    pcmpeqb %xmm0, %xmm1
+; SSE2-NEXT:    pmovmskb %xmm1, %eax
+; SSE2-NEXT:    cmpl $65535, %eax # imm = 0xFFFF
+; SSE2-NEXT:    sete %al
+; SSE2-NEXT:    retq
+;
+; SSE41-LABEL: parseHeaders:
+; SSE41:       # %bb.0:
+; SSE41-NEXT:    movdqu (%rdi), %xmm0
+; SSE41-NEXT:    ptest %xmm0, %xmm0
+; SSE41-NEXT:    sete %al
+; SSE41-NEXT:    retq
+;
+; AVX-LABEL: parseHeaders:
+; AVX:       # %bb.0:
+; AVX-NEXT:    vmovdqu (%rdi), %xmm0
+; AVX-NEXT:    vptest %xmm0, %xmm0
+; AVX-NEXT:    sete %al
+; AVX-NEXT:    retq
   %vptr = bitcast i64 * %ptr to <2 x i64> *
   %vload = load <2 x i64>, <2 x i64> * %vptr, align 8
   %vreduce = call i64 @llvm.experimental.vector.reduce.or.v2i64(<2 x i64> %vload)
   %vcheck = icmp eq i64 %vreduce, 0
   ret i1 %vcheck
 }
 
 define i1 @parseHeaders2_scalar_or(i64 * %ptr) nounwind {