Use PMADDWD to expand reduction in a loop

[danielcdh]

Summary

PMADDWD can help improve 8/16 bit integer mutliply-add operation performance for cases like:

for (int i = 0; i < count; i++)

a += x[i] * y[i];

[wmi]

lib/Target/X86/X86ISelLowering.cpp

34588–34599

Maybe use std::swap, so that Op0 and Op1 are unnecessary.
MulOp = N->getOperand(0);
Phi = N->getOperand(1);
if (MulOp.getOpcode() != ISD::MUL) {
std::swap(MulOp, Phi);
if (MulOp.getOpcode() != ISD::MUL)
return SDValue();
}

34605–34607

Just realize maybe we need to prove no wrap before use PMADDUBSW? It is possible that i8*i8+i8*i8 will overflow for i16 and PMADDUBSW will generate saturate result. However, the original i32 operations will not have overflow issue.

[danielcdh]

remove the support for PMADDUBSW as it cannot handle overflow case.

lib/Target/X86/X86ISelLowering.cpp

34605–34607

Good catch!

Looks like we should not use compiler to generate PMADDUBSW directly. Patch updated.

[wmi]

test/CodeGen/X86/madd.ll

25

The test can be simpler by removing the preheader and exit. just like what test/CodeGen/X86/sad.ll does. It will be easier if it is needed to extend the test and check different kinds of element types or different vector lengths.

[zvi]

Thanks for working on this patch. Regarding support for PMADDUBSW, can we match something like the following?

for (int i = 0; i < count; i++) {
  a = saturate(a + x[i] * y[i]);
}

[danielcdh]

In D31679#719786, @zvi wrote:

Thanks for working on this patch. Regarding support for PMADDUBSW, can we match something like the following?

for (int i = 0; i < count; i++) {
  a = saturate(a + x[i] * y[i]);
}

How does user specify "saturate"? Is it a general builtin in clang?

[danielcdh]

simplify test

[mkuper]

I suggest we leave the PMADDUBSW discussion for a separate patch.

Some minor comments inline.

lib/Target/X86/X86ISelLowering.cpp

34600

Since this version always uses PMADDWD, we no longer care about Mode at all, and this entire if can go away.

34630

Isn't Madd.getSimpleValueType() always MVT::i32?

34630

Just a sanity check - this is correct irrespective of whether the initial value for the reduction phi is 0 or not, right?

34708

Bikeshedding: maybe call this combineLoopMAddPattern() to match the other name?

34708

Sad -> Madd?

[danielcdh]

update

lib/Target/X86/X86ISelLowering.cpp

34630

I suppose yes.

34630

It should be vector type instead of scalar type.

[mkuper]

LGTM

lib/Target/X86/X86ISelLowering.cpp

34630

Argh, right, sorry.

[zvi]

In D31679#720193, @danielcdh wrote:

How does user specify "saturate"? Is it a general builtin in clang?

I'm not aware of such a builtin and my snippet above was more of pseudo-code.
In c++ you may implement saturation as:

int sat_sint16(int x) {
  return std::min(32767, std::max(-32768, x));
}

AFAIK, the loop vectorizer will not vectorize the reduction for PMADDUBSW, so i agree with @mkuper to do this in a different patch,

lib/Target/X86/X86ISelLowering.cpp

                        DAG.getConstant(-1ULL, DL, VT), NewCmp);
 
   // X - (Z == 0) --> sub X, (zext(sete  Z, 0)) --> sbb X, 0, (cmp Z, 1)
   // X + (Z == 0) --> add X, (zext(sete  Z, 0)) --> adc X, 0, (cmp Z, 1)
   return DAG.getNode(IsSub ? X86ISD::SBB : X86ISD::ADC, DL, VT, X,
                      DAG.getConstant(0, DL, VT), NewCmp);
 }
 
+static SDValue combineLoopMAddPattern(SDNode *N, SelectionDAG &DAG,
+                                      const X86Subtarget &Subtarget) {
+  SDValue MulOp = N->getOperand(0);
+  SDValue Phi = N->getOperand(1);
+
+  if (MulOp.getOpcode() != ISD::MUL)
+    std::swap(MulOp, Phi);
+  if (MulOp.getOpcode() != ISD::MUL)
+    return SDValue();
+
+  ShrinkMode Mode;
+  if (!canReduceVMulWidth(MulOp.getNode(), DAG, Mode))
+    return SDValue();
+
+  EVT VT = N->getValueType(0);
+
+  unsigned RegSize = 128;
+  if (Subtarget.hasBWI())
+    RegSize = 512;
+  else if (Subtarget.hasAVX2())
+    RegSize = 256;
+  unsigned VectorSize = VT.getVectorNumElements() * 16;
+  // If the vector size is less than 128, or greater than the supported RegSize,
+  // do not use PMADD.
+  if (VectorSize < 128 || VectorSize > RegSize)
+    return SDValue();
+
+  SDLoc DL(N);
+  EVT ReducedVT = EVT::getVectorVT(*DAG.getContext(), MVT::i16,
+                                   VT.getVectorNumElements());
+  EVT MAddVT = EVT::getVectorVT(*DAG.getContext(), MVT::i32,
+                                VT.getVectorNumElements() / 2);
+
+  // Shrink the operands of mul.
+  SDValue N0 = DAG.getNode(ISD::TRUNCATE, DL, ReducedVT, MulOp->getOperand(0));
+  SDValue N1 = DAG.getNode(ISD::TRUNCATE, DL, ReducedVT, MulOp->getOperand(1));
+
+  // Madd vector size is half of the original vector size
+  SDValue Madd = DAG.getNode(X86ISD::VPMADDWD, DL, MAddVT, N0, N1);
+  // Fill the rest of the output with 0
+  SDValue Zero = getZeroVector(Madd.getSimpleValueType(), Subtarget, DAG, DL);
+  SDValue Concat = DAG.getNode(ISD::CONCAT_VECTORS, DL, VT, Madd, Zero);
+  return DAG.getNode(ISD::ADD, DL, VT, Concat, Phi);
+}
+
 static SDValue combineLoopSADPattern(SDNode *N, SelectionDAG &DAG,
                                      const X86Subtarget &Subtarget) {
   SDLoc DL(N);
   EVT VT = N->getValueType(0);
   SDValue Op0 = N->getOperand(0);
   SDValue Op1 = N->getOperand(1);
 
   // TODO: There's nothing special about i32, any integer type above i16 should

 }
 
 static SDValue combineAdd(SDNode *N, SelectionDAG &DAG,
                           const X86Subtarget &Subtarget) {
   const SDNodeFlags *Flags = &cast<BinaryWithFlagsSDNode>(N)->Flags;
   if (Flags->hasVectorReduction()) {
     if (SDValue Sad = combineLoopSADPattern(N, DAG, Subtarget))
       return Sad;
+    if (SDValue MAdd = combineLoopMAddPattern(N, DAG, Subtarget))
+      return MAdd;
   }
   EVT VT = N->getValueType(0);
   SDValue Op0 = N->getOperand(0);
   SDValue Op1 = N->getOperand(1);
 
   // Try to synthesize horizontal adds from adds of shuffles.
   if (((Subtarget.hasSSSE3() && (VT == MVT::v8i16 || VT == MVT::v4i32)) ||
        (Subtarget.hasInt256() && (VT == MVT::v16i16 || VT == MVT::v8i32))) &&

test/CodeGen/X86/madd.ll

+; 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
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+avx512f | FileCheck %s --check-prefix=AVX512
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+avx512bw | FileCheck %s --check-prefix=AVX512
+
+;SSE2-label: @_Z10test_shortPsS_i
+;SSE2:        movdqu
+;SSE2-NEXT:   movdqu
+;SSE2-NEXT:   pmaddwd
+;SSE2-NEXT:   paddd
+
+;AVX2-label: @_Z10test_shortPsS_i
+;AVX2:        vmovdqu
+;AVX2-NEXT:   vpmaddwd
+;AVX2-NEXT:   vinserti128
+;AVX2-NEXT:   vpaddd
+
+;AVX512-label: @_Z10test_shortPsS_i
+;AVX512:        vmovdqu
+;AVX512-NEXT:   vpmaddwd
+;AVX512-NEXT:   vinserti128
+;AVX512-NEXT:   vpaddd
+
+define i32 @_Z10test_shortPsS_i(i16* nocapture readonly, i16* nocapture readonly, i32) local_unnamed_addr #0 {
+entry:
+  %3 = zext i32 %2 to i64
+  br label %vector.body
+
+vector.body:
+  %index = phi i64 [ %index.next, %vector.body ], [ 0, %entry ]
+  %vec.phi = phi <8 x i32> [ %11, %vector.body ], [ zeroinitializer, %entry ]
+  %4 = getelementptr inbounds i16, i16* %0, i64 %index
+  %5 = bitcast i16* %4 to <8 x i16>*
+  %wide.load = load <8 x i16>, <8 x i16>* %5, align 2
+  %6 = sext <8 x i16> %wide.load to <8 x i32>
+  %7 = getelementptr inbounds i16, i16* %1, i64 %index
+  %8 = bitcast i16* %7 to <8 x i16>*
+  %wide.load14 = load <8 x i16>, <8 x i16>* %8, align 2
+  %9 = sext <8 x i16> %wide.load14 to <8 x i32>
+  %10 = mul nsw <8 x i32> %9, %6
+  %11 = add nsw <8 x i32> %10, %vec.phi
+  %index.next = add i64 %index, 8
+  %12 = icmp eq i64 %index.next, %3
+  br i1 %12, label %middle.block, label %vector.body
+
+middle.block:
+  %rdx.shuf = shufflevector <8 x i32> %11, <8 x i32> undef, <8 x i32> <i32 4, i32 5, i32 6, i32 7, i32 undef, i32 undef, i32 undef, i32 undef>
+  %bin.rdx = add <8 x i32> %11, %rdx.shuf
+  %rdx.shuf15 = shufflevector <8 x i32> %bin.rdx, <8 x i32> undef, <8 x i32> <i32 2, i32 3, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef>
+  %bin.rdx16 = add <8 x i32> %bin.rdx, %rdx.shuf15
+  %rdx.shuf17 = shufflevector <8 x i32> %bin.rdx16, <8 x i32> undef, <8 x i32> <i32 1, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef>
+  %bin.rdx18 = add <8 x i32> %bin.rdx16, %rdx.shuf17
+  %13 = extractelement <8 x i32> %bin.rdx18, i32 0
+  ret i32 %13
+}
+
+;AVX2-label: @_Z9test_charPcS_i
+;AVX2:       vpmovsxbw
+;AVX2-NEXT:  vpmovsxbw
+;AVX2-NEXT:  vpmaddwd
+;AVX2-NEXT:  vpaddd
+
+;AVX512-label: @_Z9test_charPcS_i
+;AVX512:       vpmovsxbw
+;AVX512-NEXT:  vpmovsxbw
+;AVX512-NEXT:  vpmaddwd
+;AVX512-NEXT:  vinserti64x4
+;AVX512-NEXT:  vpaddd
+
+define i32 @_Z9test_charPcS_i(i8* nocapture readonly, i8* nocapture readonly, i32) local_unnamed_addr #0 {
+entry:
+  %3 = zext i32 %2 to i64
+  br label %vector.body
+
+vector.body:
+  %index = phi i64 [ %index.next, %vector.body ], [ 0, %entry ]
+  %vec.phi = phi <16 x i32> [ %11, %vector.body ], [ zeroinitializer, %entry ]
+  %4 = getelementptr inbounds i8, i8* %0, i64 %index
+  %5 = bitcast i8* %4 to <16 x i8>*
+  %wide.load = load <16 x i8>, <16 x i8>* %5, align 1
+  %6 = sext <16 x i8> %wide.load to <16 x i32>
+  %7 = getelementptr inbounds i8, i8* %1, i64 %index
+  %8 = bitcast i8* %7 to <16 x i8>*
+  %wide.load14 = load <16 x i8>, <16 x i8>* %8, align 1
+  %9 = sext <16 x i8> %wide.load14 to <16 x i32>
+  %10 = mul nsw <16 x i32> %9, %6
+  %11 = add nsw <16 x i32> %10, %vec.phi
+  %index.next = add i64 %index, 16
+  %12 = icmp eq i64 %index.next, %3
+  br i1 %12, label %middle.block, label %vector.body
+
+middle.block:
+  %rdx.shuf = shufflevector <16 x i32> %11, <16 x i32> undef, <16 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef>
+  %bin.rdx = add <16 x i32> %11, %rdx.shuf
+  %rdx.shuf15 = shufflevector <16 x i32> %bin.rdx, <16 x i32> undef, <16 x i32> <i32 4, i32 5, i32 6, i32 7, 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.rdx16 = add <16 x i32> %bin.rdx, %rdx.shuf15
+  %rdx.shuf17 = shufflevector <16 x i32> %bin.rdx16, <16 x i32> undef, <16 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>
+  %bin.rdx18 = add <16 x i32> %bin.rdx16, %rdx.shuf17
+  %rdx.shuf19 = shufflevector <16 x i32> %bin.rdx18, <16 x i32> undef, <16 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>
+  %bin.rdx20 = add <16 x i32> %bin.rdx18, %rdx.shuf19
+  %13 = extractelement <16 x i32> %bin.rdx20, i32 0
+  ret i32 %13
+}