[X86] Detect SAD patterns and emit psadbw instructions on X86 redux

[mkuper]

Summary

This fixes PR27539 in the SAD detection code originally commited in http://reviews.llvm.org/rL267649
The patch is intentionally against r267722 and not against trunk (in which r267649 has been reverted), to highlight the diff vs. the original patch.

The original patch wasn't happy with reductions where the reduction vector was smaller than the result of the psadbw.
This can be handled by taking the low part of the result of the psadbw - this should be safe, since when the the reduction type is narrower than the psadbw result, the input vectors are also narrow, so the non-zero part of the psadbw result is even narrower.

[wmi]

lib/Target/X86/X86ISelLowering.cpp

29480

The condition can be represented using NumConcat == 0

29480

Can be represented as NumConcat == 0?

29481

The EXTRACT_SUBVECTOR generated a useless pshufd for the sad_2i8 test.

test/CodeGen/X86/sad.ll

990

The pshufd is generated because of the EXTRACT_SUBVECTOR to take the lower part of the sad result. however, since we know the sad result except the lower 16 bits are all 0, the pshufd is useless.

[mkuper]

Thanks, Wei!

lib/Target/X86/X86ISelLowering.cpp

29480

Right, I thought it was clearer this way. I can change it to NumConcat == 0 if you prefer. But, to be honest, I'd rather actually get rid of NumConcat in the condition above, and make that an explicit comparison as well.
What do you think?

test/CodeGen/X86/sad.ll

990

Right, because of the way we legalize v2i32, the extract becomes an anyext instead of a nop.
I'll see if I can get rid of it. Thanks!

[wmi]

lib/Target/X86/X86ISelLowering.cpp

29480

Get rid of NumConcat is better.

[mkuper]

Got rid of NumConcat.

Regarding the shuffle - I don't see a good way to get rid of it locally. If you have an idea on how to do that, I'll be happy to hear it.
It can probably be cleaned up with a DAGCombine that recognizes the shuffle we get from the extract is a nop (because the only lanes it shuffles around are zero). Are you ok with leaving a fixme, and following up on that separately if we think it's worthwhile?

[wmi]

Regarding the shuffle - I don't see a good way to get rid of it locally. If you have an idea on how to do that, I'll be happy to hear it.

Can we use TRUNCATE instead? SadVT is 2xi64, truncate it to 2xi32.

if (VT.getSizeInBits() < ResVT.getSizeInBits()) {
  MVT ResVT = MVT::getVectorVT(MVT::i32, VT.getSizeInBits() / 32);
  Sad = DAG.getNode(ISD::TRUNCATE, DL, ResVT, SadVT);
}

[mkuper]

That works (after a bit of massaging the types). In fact, it makes more sense than using extract_subvector. Thanks a lot!

[mkuper]

Eliminated redundant pshufd.

[wmi]

LGTM. But I find that the code generated for sad_16i8/sad_32i8 are not very good because of INSERT_SUBVECTOR. Legalization of INSERT_SUBVECTOR introduces many useless stack saves/restores. We may need to fix them in a separate patch before recommit r267649 and this patch.

[mkuper]

I agree the spill code is complete nonsense, but I don't think it should block this patch, for two reasons:

1. We seem to have a real problem with inserts into oversize vectors (e.g. <16 x i32> on SSE2), but I'm not entirely sure it's a high priority. I'm not sure how much of those the vectorizer actually generates - although, for reductions, that may be more common, I haven't checked.

2. More importantly, even with the spills, it still looks like a large net improvement. E.g. the SSE2 code for the sad16_i8 loop before this patch is:

.LBB0_1:                                # %vector.body
                                        # =>This Inner Loop Header: Depth=1
	movdqu	a+1024(%rax), %xmm7
	pshufd	$78, %xmm7, %xmm5       # xmm5 = xmm7[2,3,0,1]
	punpcklbw	%xmm10, %xmm5   # xmm5 = xmm5[0],xmm10[0],xmm5[1],xmm10[1],xmm5[2],xmm10[2],xmm5[3],xmm10[3],xmm5[4],xmm10[4],xmm5[5],xmm10[5],xmm5[6],xmm10[6],xmm5[7],xmm10[7]
	movdqa	%xmm5, %xmm6
	punpcklwd	%xmm10, %xmm6   # xmm6 = xmm6[0],xmm10[0],xmm6[1],xmm10[1],xmm6[2],xmm10[2],xmm6[3],xmm10[3]
	punpckhwd	%xmm10, %xmm5   # xmm5 = xmm5[4],xmm10[4],xmm5[5],xmm10[5],xmm5[6],xmm10[6],xmm5[7],xmm10[7]
	punpcklbw	%xmm10, %xmm7   # xmm7 = xmm7[0],xmm10[0],xmm7[1],xmm10[1],xmm7[2],xmm10[2],xmm7[3],xmm10[3],xmm7[4],xmm10[4],xmm7[5],xmm10[5],xmm7[6],xmm10[6],xmm7[7],xmm10[7]
	movdqa	%xmm7, %xmm0
	punpcklwd	%xmm10, %xmm0   # xmm0 = xmm0[0],xmm10[0],xmm0[1],xmm10[1],xmm0[2],xmm10[2],xmm0[3],xmm10[3]
	punpckhwd	%xmm10, %xmm7   # xmm7 = xmm7[4],xmm10[4],xmm7[5],xmm10[5],xmm7[6],xmm10[6],xmm7[7],xmm10[7]
	movdqu	b+1024(%rax), %xmm1
	pshufd	$78, %xmm1, %xmm2       # xmm2 = xmm1[2,3,0,1]
	punpcklbw	%xmm10, %xmm2   # xmm2 = xmm2[0],xmm10[0],xmm2[1],xmm10[1],xmm2[2],xmm10[2],xmm2[3],xmm10[3],xmm2[4],xmm10[4],xmm2[5],xmm10[5],xmm2[6],xmm10[6],xmm2[7],xmm10[7]
	movdqa	%xmm2, %xmm3
	punpcklwd	%xmm10, %xmm3   # xmm3 = xmm3[0],xmm10[0],xmm3[1],xmm10[1],xmm3[2],xmm10[2],xmm3[3],xmm10[3]
	punpckhwd	%xmm10, %xmm2   # xmm2 = xmm2[4],xmm10[4],xmm2[5],xmm10[5],xmm2[6],xmm10[6],xmm2[7],xmm10[7]
	punpcklbw	%xmm10, %xmm1   # xmm1 = xmm1[0],xmm10[0],xmm1[1],xmm10[1],xmm1[2],xmm10[2],xmm1[3],xmm10[3],xmm1[4],xmm10[4],xmm1[5],xmm10[5],xmm1[6],xmm10[6],xmm1[7],xmm10[7]
	movdqa	%xmm1, %xmm4
	punpcklwd	%xmm10, %xmm4   # xmm4 = xmm4[0],xmm10[0],xmm4[1],xmm10[1],xmm4[2],xmm10[2],xmm4[3],xmm10[3]
	punpckhwd	%xmm10, %xmm1   # xmm1 = xmm1[4],xmm10[4],xmm1[5],xmm10[5],xmm1[6],xmm10[6],xmm1[7],xmm10[7]
	psubd	%xmm1, %xmm7
	psubd	%xmm4, %xmm0
	psubd	%xmm2, %xmm5
	psubd	%xmm3, %xmm6
	movdqa	%xmm6, %xmm1
	psrad	$31, %xmm1
	paddd	%xmm1, %xmm6
	pxor	%xmm1, %xmm6
	movdqa	%xmm5, %xmm1
	psrad	$31, %xmm1
	paddd	%xmm1, %xmm5
	pxor	%xmm1, %xmm5
	movdqa	%xmm0, %xmm1
	psrad	$31, %xmm1
	paddd	%xmm1, %xmm0
	pxor	%xmm1, %xmm0
	movdqa	%xmm7, %xmm1
	psrad	$31, %xmm1
	paddd	%xmm1, %xmm7
	pxor	%xmm1, %xmm7
	paddd	%xmm7, %xmm8
	paddd	%xmm0, %xmm9
	paddd	%xmm5, %xmm11
	paddd	%xmm6, %xmm12
	addq	$4, %rax
	jne	.LBB0_1

And with this patch:

.LBB0_1:                                # %vector.body
                                        # =>This Inner Loop Header: Depth=1
	movdqa	%xmm0, %xmm4
	movdqu	a+1024(%rax), %xmm5
	movdqu	b+1024(%rax), %xmm0
	movdqa	%xmm4, (%rsp)
	movdqa	%xmm1, 16(%rsp)
	movdqa	%xmm3, 32(%rsp)
	movdqa	%xmm2, 48(%rsp)
	psadbw	%xmm5, %xmm0
	paddd	%xmm4, %xmm0
	movdqa	%xmm0, (%rsp)
	movdqa	16(%rsp), %xmm1
	movdqa	32(%rsp), %xmm3
	movdqa	48(%rsp), %xmm2
	addq	$4, %rax

Are you ok with me committing as is, and filing a PR on the insert_subvector issue?

lib/Target/X86/X86ISelLowering.cpp

static SDValue detectSADPattern(SDNode *N, SelectionDAG &DAG,
   Ops[0] = Op1.getOperand(0);
   Op1 = DAG.getNode(ISD::CONCAT_VECTORS, DL, ExtendedVT, Ops);
 
   // The output of PSADBW is a vector of i64.
   MVT SadVT = MVT::getVectorVT(MVT::i64, RegSize / 64);
   SDValue Sad = DAG.getNode(X86ISD::PSADBW, DL, SadVT, Op0, Op1);
 
   // We need to turn the vector of i64 into a vector of i32.
-  MVT ResVT = MVT::getVectorVT(MVT::i32, RegSize / 32);
+  // If the reduction vector is at least as wide as the psadbw result, just
+  // bitcast. If it's narrower, truncate - the high i32 of each i64 is zero
+  // anyway.
+  MVT ResVT = MVT::getVectorVT(MVT::i32, RegSize / 32);  
+  if (VT.getSizeInBits() >= ResVT.getSizeInBits())
     Sad = DAG.getNode(ISD::BITCAST, DL, ResVT, Sad);
+  else 
+    Sad = DAG.getNode(ISD::TRUNCATE, DL, VT, Sad);
 
-  NumConcat = VT.getSizeInBits() / ResVT.getSizeInBits();
-  if (NumConcat > 1) {
+  if (VT.getSizeInBits() > ResVT.getSizeInBits()) {
     // Update part of elements of the reduction vector. This is done by first
     // extracting a sub-vector from it, updating this sub-vector, and inserting
     // it back.
     SDValue SubPhi = DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, ResVT, Phi,
                                  DAG.getIntPtrConstant(0, DL));
     SDValue Res = DAG.getNode(ISD::ADD, DL, ResVT, Sad, SubPhi);
     return DAG.getNode(ISD::INSERT_SUBVECTOR, DL, VT, Phi, Res,
                        DAG.getIntPtrConstant(0, DL));
   } else
     return DAG.getNode(ISD::ADD, DL, VT, Sad, Phi);
 }
 
 static SDValue combineAdd(SDNode *N, SelectionDAG &DAG,
                           const X86Subtarget &Subtarget) {
   const SDNodeFlags *Flags = &cast<BinaryWithFlagsSDNode>(N)->Flags;
   if (Flags->hasVectorReduction()) {
     if (SDValue Sad = detectSADPattern(N, DAG, Subtarget))
       return Sad;
   }
 
   EVT VT = N->getValueType(0);

test/CodeGen/X86/sad.ll

middle.block:
   %bin.rdx4 = add <64 x i32> %bin.rdx3, %rdx.shuf4
   %rdx.shuf5 = shufflevector <64 x i32> %bin.rdx4, <64 x i32> undef, <64 x i32> <i32 2, i32 3, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef>
   %bin.rdx5 = add <64 x i32> %bin.rdx4, %rdx.shuf5
   %rdx.shuf6 = shufflevector <64 x i32> %bin.rdx5, <64 x i32> undef, <64 x i32> <i32 1, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef>
   %bin.rdx6 = add <64 x i32> %bin.rdx5, %rdx.shuf6
   %12 = extractelement <64 x i32> %bin.rdx6, i32 0
   ret i32 %12
 }
+
+define i32 @sad_2i8() nounwind {
+; SSE2-LABEL: sad_2i8:
+; SSE2:       # BB#0: # %entry
+; SSE2-NEXT:    pxor %xmm0, %xmm0
+; SSE2-NEXT:    movq $-1024, %rax # imm = 0xFFFFFFFFFFFFFC00
+; SSE2-NEXT:    movl $65535, %ecx # imm = 0xFFFF
+; SSE2-NEXT:    movd %ecx, %xmm1
+; SSE2-NEXT:    .p2align 4, 0x90
+; SSE2-NEXT:  .LBB3_1: # %vector.body
+; SSE2-NEXT:    # =>This Inner Loop Header: Depth=1
+; SSE2-NEXT:    movd {{.*#+}} xmm2 = mem[0],zero,zero,zero
+; SSE2-NEXT:    movd {{.*#+}} xmm3 = mem[0],zero,zero,zero
+; SSE2-NEXT:    pand %xmm1, %xmm3
+; SSE2-NEXT:    pand %xmm1, %xmm2
+; SSE2-NEXT:    psadbw %xmm3, %xmm2
+; SSE2-NEXT:    paddq %xmm2, %xmm0
+; SSE2-NEXT:    addq $4, %rax
+; SSE2-NEXT:    jne .LBB3_1
+; SSE2-NEXT:  # BB#2: # %middle.block
+; SSE2-NEXT:    pshufd {{.*#+}} xmm1 = xmm0[2,3,0,1]
+; SSE2-NEXT:    paddq %xmm0, %xmm1
+; SSE2-NEXT:    movd %xmm1, %eax
+; SSE2-NEXT:    retq
+;
+; AVX2-LABEL: sad_2i8:
+; AVX2:       # BB#0: # %entry
+; AVX2-NEXT:    vpxor %xmm0, %xmm0, %xmm0
+; AVX2-NEXT:    movq $-1024, %rax # imm = 0xFFFFFFFFFFFFFC00
+; AVX2-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; AVX2-NEXT:    .p2align 4, 0x90
+; AVX2-NEXT:  .LBB3_1: # %vector.body
+; AVX2-NEXT:    # =>This Inner Loop Header: Depth=1
+; AVX2-NEXT:    vmovd {{.*#+}} xmm2 = mem[0],zero,zero,zero
+; AVX2-NEXT:    vmovd {{.*#+}} xmm3 = mem[0],zero,zero,zero
+; AVX2-NEXT:    vpblendw {{.*#+}} xmm3 = xmm3[0],xmm0[1,2,3,4,5,6,7]
+; AVX2-NEXT:    vpblendw {{.*#+}} xmm2 = xmm2[0],xmm0[1,2,3,4,5,6,7]
+; AVX2-NEXT:    vpsadbw %xmm3, %xmm2, %xmm2
+; AVX2-NEXT:    vpaddq %xmm1, %xmm2, %xmm1
+; AVX2-NEXT:    addq $4, %rax
+; AVX2-NEXT:    jne .LBB3_1
+; AVX2-NEXT:  # BB#2: # %middle.block
+; AVX2-NEXT:    vpshufd {{.*#+}} xmm0 = xmm1[2,3,0,1]
+; AVX2-NEXT:    vpaddq %xmm0, %xmm1, %xmm0
+; AVX2-NEXT:    vmovd %xmm0, %eax
+; AVX2-NEXT:    retq
+;
+; AVX512F-LABEL: sad_2i8:
+; AVX512F:       # BB#0: # %entry
+; AVX512F-NEXT:    vpxor %xmm0, %xmm0, %xmm0
+; AVX512F-NEXT:    movq $-1024, %rax # imm = 0xFFFFFFFFFFFFFC00
+; AVX512F-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; AVX512F-NEXT:    .p2align 4, 0x90
+; AVX512F-NEXT:  .LBB3_1: # %vector.body
+; AVX512F-NEXT:    # =>This Inner Loop Header: Depth=1
+; AVX512F-NEXT:    vmovd {{.*#+}} xmm2 = mem[0],zero,zero,zero
+; AVX512F-NEXT:    vmovd {{.*#+}} xmm3 = mem[0],zero,zero,zero
+; AVX512F-NEXT:    vpblendw {{.*#+}} xmm3 = xmm3[0],xmm0[1,2,3,4,5,6,7]
+; AVX512F-NEXT:    vpblendw {{.*#+}} xmm2 = xmm2[0],xmm0[1,2,3,4,5,6,7]
+; AVX512F-NEXT:    vpsadbw %xmm3, %xmm2, %xmm2
+; AVX512F-NEXT:    vpaddq %xmm1, %xmm2, %xmm1
+; AVX512F-NEXT:    addq $4, %rax
+; AVX512F-NEXT:    jne .LBB3_1
+; AVX512F-NEXT:  # BB#2: # %middle.block
+; AVX512F-NEXT:    vpshufd {{.*#+}} xmm0 = xmm1[2,3,0,1]
+; AVX512F-NEXT:    vpaddq %xmm0, %xmm1, %xmm0
+; AVX512F-NEXT:    vmovd %xmm0, %eax
+; AVX512F-NEXT:    retq
+;
+; AVX512BW-LABEL: sad_2i8:
+; AVX512BW:       # BB#0: # %entry
+; AVX512BW-NEXT:    vpxor %xmm0, %xmm0, %xmm0
+; AVX512BW-NEXT:    movq $-1024, %rax # imm = 0xFFFFFFFFFFFFFC00
+; AVX512BW-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; AVX512BW-NEXT:    .p2align 4, 0x90
+; AVX512BW-NEXT:  .LBB3_1: # %vector.body
+; AVX512BW-NEXT:    # =>This Inner Loop Header: Depth=1
+; AVX512BW-NEXT:    vmovd {{.*#+}} xmm2 = mem[0],zero,zero,zero
+; AVX512BW-NEXT:    vmovd {{.*#+}} xmm3 = mem[0],zero,zero,zero
+; AVX512BW-NEXT:    vpblendw {{.*#+}} xmm3 = xmm3[0],xmm0[1,2,3,4,5,6,7]
+; AVX512BW-NEXT:    vpblendw {{.*#+}} xmm2 = xmm2[0],xmm0[1,2,3,4,5,6,7]
+; AVX512BW-NEXT:    vpsadbw %xmm3, %xmm2, %xmm2
+; AVX512BW-NEXT:    vpaddq %xmm1, %xmm2, %xmm1
+; AVX512BW-NEXT:    addq $4, %rax
+; AVX512BW-NEXT:    jne .LBB3_1
+; AVX512BW-NEXT:  # BB#2: # %middle.block
+; AVX512BW-NEXT:    vpshufd {{.*#+}} xmm0 = xmm1[2,3,0,1]
+; AVX512BW-NEXT:    vpaddq %xmm0, %xmm1, %xmm0
+; AVX512BW-NEXT:    vmovd %xmm0, %eax
+; AVX512BW-NEXT:    retq
+entry:
+  br label %vector.body
+
+vector.body:
+  %index = phi i64 [ 0, %entry ], [ %index.next, %vector.body ]
+  %vec.phi = phi <2 x i32> [ zeroinitializer, %entry ], [ %10, %vector.body ]
+  %0 = getelementptr inbounds [1024 x i8], [1024 x i8]* @a, i64 0, i64 %index
+  %1 = bitcast i8* %0 to <2 x i8>*
+  %wide.load = load <2 x i8>, <2 x i8>* %1, align 4
+  %2 = zext <2 x i8> %wide.load to <2 x i32>
+  %3 = getelementptr inbounds [1024 x i8], [1024 x i8]* @b, i64 0, i64 %index
+  %4 = bitcast i8* %3 to <2 x i8>*
+  %wide.load1 = load <2 x i8>, <2 x i8>* %4, align 4
+  %5 = zext <2 x i8> %wide.load1 to <2 x i32>
+  %6 = sub nsw <2 x i32> %2, %5
+  %7 = icmp sgt <2 x i32> %6, <i32 -1, i32 -1>
+  %8 = sub nsw <2 x i32> zeroinitializer, %6
+  %9 = select <2 x i1> %7, <2 x i32> %6, <2 x i32> %8
+  %10 = add nsw <2 x i32> %9, %vec.phi
+  %index.next = add i64 %index, 4
+  %11 = icmp eq i64 %index.next, 1024
+  br i1 %11, label %middle.block, label %vector.body
+
+middle.block:
+  %.lcssa = phi <2 x i32> [ %10, %vector.body ]
+  %rdx.shuf = shufflevector <2 x i32> %.lcssa, <2 x i32> undef, <2 x i32> <i32 1, i32 undef>
+  %bin.rdx = add <2 x i32> %.lcssa, %rdx.shuf
+  %12 = extractelement <2 x i32> %bin.rdx, i32 0
+  ret i32 %12
+}
+