[SDAG] allow more cast folding for vector sext-of-setcc

[spatel]

Summary

This is a follow-up to D103280 that eases the use restrictions, so we can handle the motivating case from:
https://llvm.org/PR50055

The loop code is adapted from similar use checks in ExtendUsesToFormExtLoad() and SliceUpLoad(). I did not see an easier way to filter out non-chain uses of load values.

[lebedev.ri]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10846

What if it is extending to the even wider type than VT?

[spatel]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10846

It seems neutral whether we have wider or narrower other type - trade a sext of the compare for a zext of the input. This might be because the transform to create zext-loads is conservative about uses.

Let me know if this matches what you have in mind (I can add a test either way):

define <8 x i32> @multi_use_wider_size(<8 x i8>* %src, <8 x i16>* %dst) nounwind {
  %load = load <8 x i8>, <8 x i8>* %src
  %zext = zext <8 x i8> %load to <8 x i32>
  %icmp = icmp eq <8 x i8> %load, zeroinitializer
  %sext = sext <8 x i1> %icmp to <8 x i16>
  store <8 x i16> %sext, <8 x i16>* %dst
  ret <8 x i32> %zext
}

Currently, we get:

vmovq	(%rdi), %xmm1
vpmovzxbd	%xmm1, %ymm0 ; 256-bit return value
vpxor	%xmm2, %xmm2, %xmm2
vpcmpeqb %xmm2, %xmm1, %xmm1
vpmovsxbw %xmm1, %xmm1
vmovdqa	%xmm1, (%rsi) ; 128-bit compare value

If we allow type mismatch:

vmovq	(%rdi), %xmm1
vpmovzxbd	%xmm1, %ymm0 
vpmovzxbw	%xmm1, %xmm1 
vpxor	%xmm2, %xmm2, %xmm2
vpcmpeqw	%xmm2, %xmm1, %xmm1
vmovdqa	%xmm1, (%rsi)

[lebedev.ri]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10846

My point being, iff the other ZEXT is to wider type, then i would expect that the ZEXT we create
would be absorbed by the load, and that other ZEXT would become ZEXR of ZEXT_LOAD,
but it seems like we end up with a plain load + 2 ZEXT's?

[spatel]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10846

Right - stepping through the debug spew for this example, if we allow the transform here, we have:

t7: v8i8,ch = load<(load 8 from %ir.src)> t0, t2, undef:i64
t21: v8i16 = zero_extend t7 -- feeds into the setcc
t8: v8i32 = zero_extend t7 -- return value

But the code in ExtendUsesToFormExtLoad() won't convert that load because it doesn't handle the extra use. It already has another FIXME comment about setcc patterns, so we could add one more there.

As a hack, I eased the restriction to see what we would end up with on this x86 example, and it's going to need more work. We truncate the zext load and then zext it back out, so we get this:

vpmovzxbw	(%rdi), %xmm1   
vpackuswb	%xmm1, %xmm1, %xmm0
vpmovzxbd	%xmm0, %ymm0

Ideally, we'd have something like:

vpmovzxbw	(%rdi), %xmm1   
vpmovzxwd	%xmm1, %ymm0

So it's going to be a multi-patch process if we want to allow more patterns here, and it's not clear to me that we actually gain from that transform (got rid of a sext, but added a zext?).

[lebedev.ri]

Seems fine to me unless others have other comments.
Thanks.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10844

// The only extra users that we are okay with is the exact same cast we are about to create.

10846

Thanks for checking.

[RKSimon]

LGTM - cheers!

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (!ISD::isSignedIntSetCC(CC) && N0.hasOneUse() &&
       // if extended to the destination type. See if the compare operands
       // can be freely extended to the destination type.
       auto IsFreeToZext = [&](SDValue V) {
         if (isConstantOrConstantVector(V, /*NoOpaques*/ true))
           return true;
 
         // Match a simple, non-extended load that can be converted to a
         // legal zext-load.
-        // TODO: Handle more than one use if the other uses are free to zext.
         // TODO: Allow widening of an existing zext-load?
-        return ISD::isNON_EXTLoad(V.getNode()) &&
-               ISD::isUNINDEXEDLoad(V.getNode()) &&
-               cast<LoadSDNode>(V)->isSimple() &&
-               TLI.isLoadExtLegal(ISD::ZEXTLOAD, VT, V.getValueType()) &&
-               V.hasOneUse();
+        if (!(ISD::isNON_EXTLoad(V.getNode()) &&
+              TLI.isLoadExtLegal(ISD::ZEXTLOAD, VT, V.getValueType())))
+          return false;
+
+        // Non-chain users of this value must either be the setcc in this
+        // sequence or zexts that can be folded into the new zext-load.
+        for (SDNode::use_iterator UI = V->use_begin(), UE = V->use_end();
+             UI != UE; ++UI) {
+          // Skip uses of the chain and the setcc.
+          SDNode *User = *UI;
+          if (UI.getUse().getResNo() != 0 || User == N0.getNode())
+            continue;
+          // Other users should get folded into the load.
+          if (User->getOpcode() != ISD::ZERO_EXTEND ||
+              User->getValueType(0) != VT)
+            return false;
+        }
+        return true;
       };
 
       if (IsFreeToZext(N00) && IsFreeToZext(N01)) {
         SDValue Ext0 = DAG.getZExtOrTrunc(N00, DL, VT);
         SDValue Ext1 = DAG.getZExtOrTrunc(N01, DL, VT);
         return DAG.getSetCC(DL, VT, Ext0, Ext1, CC);
       }
     }

llvm/test/CodeGen/X86/sext-vsetcc.ll

 ; AVX512-NEXT:    retq
   %loadx = load <8 x i16>, <8 x i16>* %x
   %loady = load <8 x i16>, <8 x i16>* %y
   %icmp = icmp ne <8 x i16> %loadx, %loady
   %sext = sext <8 x i1> %icmp to <8 x i32>
   ret <8 x i32> %sext
 }
 
-; negative test - extra use (TODO)
+; Both uses of the load can be absorbed by the zext-load, so we eliminate the explicit casts.
 
 define <8 x i32> @PR50055(<8 x i8>* %src, <8 x i32>* %dst) nounwind {
 ; SSE-LABEL: PR50055:
 ; SSE:       # %bb.0:
 ; SSE-NEXT:    movq {{.*#+}} xmm2 = mem[0],zero
 ; SSE-NEXT:    pxor %xmm3, %xmm3
 ; SSE-NEXT:    movdqa %xmm2, %xmm1
 ; SSE-NEXT:    punpcklbw {{.*#+}} xmm1 = xmm1[0],xmm3[0],xmm1[1],xmm3[1],xmm1[2],xmm3[2],xmm1[3],xmm3[3],xmm1[4],xmm3[4],xmm1[5],xmm3[5],xmm1[6],xmm3[6],xmm1[7],xmm3[7]
 ; SSE-NEXT:    movdqa %xmm1, %xmm0
 ; SSE-NEXT:    punpcklwd {{.*#+}} xmm0 = xmm0[0],xmm3[0],xmm0[1],xmm3[1],xmm0[2],xmm3[2],xmm0[3],xmm3[3]
 ; SSE-NEXT:    punpckhwd {{.*#+}} xmm1 = xmm1[4],xmm3[4],xmm1[5],xmm3[5],xmm1[6],xmm3[6],xmm1[7],xmm3[7]
 ; SSE-NEXT:    pcmpeqb %xmm3, %xmm2
 ; SSE-NEXT:    pcmpeqd %xmm3, %xmm3
 ; SSE-NEXT:    pxor %xmm2, %xmm3
 ; SSE-NEXT:    punpcklbw {{.*#+}} xmm2 = xmm2[0],xmm3[0],xmm2[1],xmm3[1],xmm2[2],xmm3[2],xmm2[3],xmm3[3],xmm2[4],xmm3[4],xmm2[5],xmm3[5],xmm2[6],xmm3[6],xmm2[7],xmm3[7]
 ; SSE-NEXT:    punpcklwd {{.*#+}} xmm3 = xmm3[0],xmm2[0],xmm3[1],xmm2[1],xmm3[2],xmm2[2],xmm3[3],xmm2[3]
 ; SSE-NEXT:    psrad $24, %xmm3
 ; SSE-NEXT:    punpckhwd {{.*#+}} xmm2 = xmm2[4,4,5,5,6,6,7,7]
 ; SSE-NEXT:    psrad $24, %xmm2
 ; SSE-NEXT:    movdqa %xmm2, 16(%rsi)
 ; SSE-NEXT:    movdqa %xmm3, (%rsi)
 ; SSE-NEXT:    retq
 ;
-; AVX2-LABEL: PR50055:
-; AVX2:       # %bb.0:
-; AVX2-NEXT:    vmovq {{.*#+}} xmm1 = mem[0],zero
-; AVX2-NEXT:    vpmovzxbd {{.*#+}} ymm0 = xmm1[0],zero,zero,zero,xmm1[1],zero,zero,zero,xmm1[2],zero,zero,zero,xmm1[3],zero,zero,zero,xmm1[4],zero,zero,zero,xmm1[5],zero,zero,zero,xmm1[6],zero,zero,zero,xmm1[7],zero,zero,zero
-; AVX2-NEXT:    vpxor %xmm2, %xmm2, %xmm2
-; AVX2-NEXT:    vpcmpeqb %xmm2, %xmm1, %xmm1
-; AVX2-NEXT:    vpcmpeqd %xmm2, %xmm2, %xmm2
-; AVX2-NEXT:    vpxor %xmm2, %xmm1, %xmm1
-; AVX2-NEXT:    vpmovsxbd %xmm1, %ymm1
-; AVX2-NEXT:    vmovdqa %ymm1, (%rsi)
-; AVX2-NEXT:    retq
-;
-; AVX512-LABEL: PR50055:
-; AVX512:       # %bb.0:
-; AVX512-NEXT:    vmovq {{.*#+}} xmm1 = mem[0],zero
-; AVX512-NEXT:    vpmovzxbd {{.*#+}} ymm0 = xmm1[0],zero,zero,zero,xmm1[1],zero,zero,zero,xmm1[2],zero,zero,zero,xmm1[3],zero,zero,zero,xmm1[4],zero,zero,zero,xmm1[5],zero,zero,zero,xmm1[6],zero,zero,zero,xmm1[7],zero,zero,zero
-; AVX512-NEXT:    vpxor %xmm2, %xmm2, %xmm2
-; AVX512-NEXT:    vpcmpeqb %xmm2, %xmm1, %xmm1
-; AVX512-NEXT:    vpternlogq $15, %zmm1, %zmm1, %zmm1
-; AVX512-NEXT:    vpmovsxbd %xmm1, %ymm1
-; AVX512-NEXT:    vmovdqa %ymm1, (%rsi)
-; AVX512-NEXT:    retq
+; AVX-LABEL: PR50055:
+; AVX:       # %bb.0:
+; AVX-NEXT:    vpmovzxbd {{.*#+}} ymm0 = mem[0],zero,zero,zero,mem[1],zero,zero,zero,mem[2],zero,zero,zero,mem[3],zero,zero,zero,mem[4],zero,zero,zero,mem[5],zero,zero,zero,mem[6],zero,zero,zero,mem[7],zero,zero,zero
+; AVX-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; AVX-NEXT:    vpcmpgtd %ymm1, %ymm0, %ymm1
+; AVX-NEXT:    vmovdqa %ymm1, (%rsi)
+; AVX-NEXT:    retq
   %load = load <8 x i8>, <8 x i8>* %src
   %zext = zext <8 x i8> %load to <8 x i32>
   %icmp = icmp ne <8 x i8> %load, zeroinitializer
   %sext = sext <8 x i1> %icmp to <8 x i32>
   store <8 x i32> %sext, <8 x i32>* %dst
   ret <8 x i32> %zext
 }
 
+; negative test - extra uses must be absorbable by a zext-load.
+
 define <8 x i16> @multi_use_different_sizes(<8 x i8>* %src, <8 x i32>* %dst) nounwind {
 ; SSE-LABEL: multi_use_different_sizes:
 ; SSE:       # %bb.0:
 ; SSE-NEXT:    movq {{.*#+}} xmm1 = mem[0],zero
 ; SSE-NEXT:    pxor %xmm2, %xmm2
 ; SSE-NEXT:    movdqa %xmm1, %xmm0
 ; SSE-NEXT:    punpcklbw {{.*#+}} xmm0 = xmm0[0],xmm2[0],xmm0[1],xmm2[1],xmm0[2],xmm2[2],xmm0[3],xmm2[3],xmm0[4],xmm2[4],xmm0[5],xmm2[5],xmm0[6],xmm2[6],xmm0[7],xmm2[7]
 ; SSE-NEXT:    pcmpeqb %xmm2, %xmm1