[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

       if (!I->second.isAllOnesValue())
         return false;
     }
   }
 
   return true;
 }
 
-// Check whether an OR'd tree is PTEST-able, or if we can fallback to
+// Check whether an OR'd reduction tree is PTEST-able, or if we can fallback to
 // CMP(MOVMSK(PCMPEQB(X,0))).
 static SDValue LowerVectorAllZeroTest(SDValue Op, ISD::CondCode CC,
                                       const X86Subtarget &Subtarget,
                                       SelectionDAG &DAG, SDValue &X86CC) {
-  assert(Op.getOpcode() == ISD::OR && "Only check OR'd tree.");
-
-  if (!Subtarget.hasSSE2() || !Op->hasOneUse())
+  if (!Subtarget.hasSSE2())
     return SDValue();
 
+  bool UsePTEST = Subtarget.hasSSE41();
+
+  auto LowerCmpZero = [&](SDValue V, const SDLoc &DL) {
+    X86CC = DAG.getTargetConstant(CC == ISD::SETEQ ? X86::COND_E : X86::COND_NE,
+                                  DL, MVT::i8);
+
+    if (UsePTEST)
+      return DAG.getNode(X86ISD::PTEST, DL, MVT::i32, V, V);
+
+    SDValue Result = DAG.getNode(X86ISD::PCMPEQ, DL, MVT::v16i8, V,
+                                 getZeroVector(MVT::v16i8, Subtarget, DAG, DL));
+    Result = DAG.getNode(X86ISD::MOVMSK, DL, MVT::i32, Result);
+    return DAG.getNode(X86ISD::CMP, DL, MVT::i32, Result,
+                       DAG.getConstant(0xFFFF, DL, MVT::i32));
+  };
+
   SmallVector<SDValue, 8> VecIns;
-  if (!matchScalarReduction(Op, ISD::OR, VecIns))
-    return SDValue();
+  if (Op.getOpcode() == ISD::OR && matchScalarReduction(Op, ISD::OR, VecIns)) {
+    if (!Op->hasOneUse())
+      return SDValue();
 
     // Quit if not 128/256-bit vector.
     EVT VT = VecIns[0].getValueType();
     if (!VT.is128BitVector() && !VT.is256BitVector())
       return SDValue();
 
     SDLoc DL(Op);
-  bool UsePTEST = Subtarget.hasSSE41();
-  MVT TestVT =
-      VT.is128BitVector() ? (UsePTEST ? MVT::v2i64 : MVT::v16i8) : MVT::v4i64;
+    MVT TestVT =
+        VT.is128BitVector() ? (UsePTEST ? MVT::v2i64 : MVT::v16i8) : MVT::v4i64;
 
     // Cast all vectors into TestVT for PTEST/PCMPEQ.
     for (unsigned i = 0, e = VecIns.size(); i < e; ++i)
       VecIns[i] = DAG.getBitcast(TestVT, VecIns[i]);
 
     // If more than one full vector is evaluated, OR them first before PTEST.
-  for (unsigned Slot = 0, e = VecIns.size(); e - Slot > 1; Slot += 2, e += 1) {
-    // Each iteration will OR 2 nodes and append the result until there is only
-    // 1 node left, i.e. the final OR'd value of all vectors.
-    SDValue LHS = VecIns[Slot];
-    SDValue RHS = VecIns[Slot + 1];
-    VecIns.push_back(DAG.getNode(ISD::OR, DL, TestVT, LHS, RHS));
-  }
+    for (unsigned Slot = 0, e = VecIns.size(); e - Slot > 1;
+         Slot += 2, e += 1) {
+      // Each iteration will OR 2 nodes and append the result until there is
+      // only 1 node left, i.e. the final OR'd value of all vectors.
+      SDValue LHS = VecIns[Slot];
+      SDValue RHS = VecIns[Slot + 1];
+      VecIns.push_back(DAG.getNode(ISD::OR, DL, TestVT, LHS, RHS));
+    }
 
-  X86CC = DAG.getTargetConstant(CC == ISD::SETEQ ? X86::COND_E : X86::COND_NE,
-                                DL, MVT::i8);
+    return LowerCmpZero(VecIns.back(), DL);
+  }
 
-  if (UsePTEST)
-    return DAG.getNode(X86ISD::PTEST, DL, MVT::i32, VecIns.back(),
-                       VecIns.back());
+  ISD::NodeType BinOp;
+  if (Op.getOpcode() == ISD::EXTRACT_VECTOR_ELT)
+    if (SDValue Match = DAG.matchBinOpReduction(Op.getNode(), BinOp, {ISD::OR}))
+      return LowerCmpZero(Match, SDLoc(Op));
 
-  SDValue Result = DAG.getNode(X86ISD::PCMPEQ, DL, MVT::v16i8, VecIns.back(),
-                               getZeroVector(MVT::v16i8, Subtarget, DAG, DL));
-  Result = DAG.getNode(X86ISD::MOVMSK, DL, MVT::i32, Result);
-  return DAG.getNode(X86ISD::CMP, DL, MVT::i32, Result,
-                     DAG.getConstant(0xFFFF, DL, MVT::i32));
+  return SDValue();
 }
 
 /// return true if \c Op has a use that doesn't just read flags.
 static bool hasNonFlagsUse(SDValue Op) {
   for (SDNode::use_iterator UI = Op->use_begin(), UE = Op->use_end(); UI != UE;
        ++UI) {
     SDNode *User = *UI;
     unsigned UOpNo = UI.getOperandNo();

   if (Op0.getOpcode() == ISD::AND && Op0.hasOneUse() && isNullConstant(Op1) &&
       (CC == ISD::SETEQ || CC == ISD::SETNE)) {
     if (SDValue BT = LowerAndToBT(Op0, CC, dl, DAG, X86CC))
       return BT;
   }
 
   // Try to use PTEST/PMOVMSKB for a tree ORs equality compared with 0.
   // TODO: We could do AND tree with all 1s as well by using the C flag.
-  if (Op0.getOpcode() == ISD::OR && isNullConstant(Op1) &&
-      (CC == ISD::SETEQ || CC == ISD::SETNE)) {
+  if (isNullConstant(Op1) && (CC == ISD::SETEQ || CC == ISD::SETNE))
     if (SDValue CmpZ = LowerVectorAllZeroTest(Op0, CC, Subtarget, DAG, X86CC))
       return CmpZ;
-  }
 
   // Try to lower using KORTEST or KTEST.
   if (SDValue Test = EmitAVX512Test(Op0, Op1, CC, dl, DAG, Subtarget, X86CC))
     return Test;
 
   // Look for X == 0, X == 1, X != 0, or X != 1.  We can simplify some forms of
   // these.
   if ((isOneConstant(Op1) || isNullConstant(Op1)) &&

   // When CET and shadow stack is enabled, we need to fix the Shadow Stack.
   if (MF->getMMI().getModule()->getModuleFlag("cf-protection-return")) {
     thisMBB = emitLongJmpShadowStackFix(MI, thisMBB);
   }
 
   // Reload FP
   MIB = BuildMI(*thisMBB, MI, DL, TII->get(PtrLoadOpc), FP);
   for (unsigned i = 0; i < X86::AddrNumOperands; ++i) {
     const MachineOperand &MO = MI.getOperand(i);

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 {