[AArch64] Generate dot for v16i8 sum reduction to i32

[mivnay]

Summary

Convert VECREDUCE_ADD( EXTEND(v16i8_type) ) to VECREDUCE_ADD( DOTv16i8(v16i8_type) ) whenever the result type is i32. This gains in one of the SPECCPU 2017 benchmark.

This partially solves the bug: https://bugs.llvm.org/show_bug.cgi?id=46888
Meta ticket: https://bugs.llvm.org/show_bug.cgi?id=46929

[SjoerdMeijer]

I haven't looked in much detail at this patch, but this looks like some straightforward lowering of llvm.experimental.vector.reduce.add. Absolutely nothing wrong with that, but I am curious who's going to produce this intrinsic? The vectoriser, the matrix pass? In other words, any ideas on the bigger picture?

[mivnay]

In D88577#2305592, @SjoerdMeijer wrote:

I haven't looked in much detail at this patch, but this looks like some straightforward lowering of llvm.experimental.vector.reduce.add. Absolutely nothing wrong with that, but I am curious who's going to produce this intrinsic? The vectoriser, the matrix pass? In other words, any ideas on the bigger picture?

Thanks for looking into the patch. The pattern added in lit test gets generated from the below C code:

#include <stdint.h>
#include <stdlib.h>

int func(uint8_t *a) {
  int sum = 0;
  for (int i = 0; i < 16; i++) {
    sum += a[i];
  }

  return sum;
}

EDIT: SLP vectorizer generates the pattern in this case.

[dmgreen]

Hello

Can you update with full context? -U999999. It makes phabriactor reviews easier to follow.

I had thought about this somewhat in reference to inloop reductions. I had presumed that it would need some form of partial reduction though, as you would want part of the reduction would then happen outside the loop (I think)

Improving codegen on it's own is good, but I'm interested in seeing how this fits with the other patches.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

10970

We ideally shouldn't be just producing a machine node here. Can you add a AArch64ISD::UDOT node?

We should be doing the same for SDOT too.

[mivnay]

In D88577#2305658, @dmgreen wrote:

Hello

Can you update with full context? -U999999. It makes phabriactor reviews easier to follow.

I had thought about this somewhat in reference to inloop reductions. I had presumed that it would need some form of partial reduction though, as you would want part of the reduction would then happen outside the loop (I think)

Improving codegen on it's own is good, but I'm interested in seeing how this fits with the other patches.

Hi,

I am working on the performance related issues mentioned in the bug and meta ticket. I have three unrelated patches (i.e., patterns) for codegen improvements. This is the first one.

[mivnay]

Added support for SDOT

[mivnay]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

10970

I have added the support for SDOT.

Can you add a AArch64ISD::UDOT node?

Can this be done in a later patch?

[dmgreen]

EDIT: SLP vectorizer generates the pattern in this case.

Ah I see.

https://llvm.org/docs/Phabricator.html#phabricator-request-review-web has some details on creating diffs with extra context.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

10949–10951

Any single statement blocks do not need { } brackets.

10964–10965

This second isScalable check isn't needed, and maybe move the num elements == 16 check up to the VT check, by checking for v16i32 directly?

10970

It's best not create machine nodes directly in the DAG...

Maybe create a int_aarch64_neon_sdot intrinsic node instead? That will avoid the need for a new node, and go through the existing tablegen patterns.

It would be nice to be able to fold add(x, UDOT(0, y)) -> UDOT(x, y), which is where having a node will really become useful. That needn't be done here though.

10972

EVT::getVectorVT(*DAG.getContext(), MVT::i32, 4) -> MVT::v4i32

I wasn't aware that you could use getConstant directly on vector types like this. Neat.

10981

Why is this is calling ReplaceAllUsesOfValueWith? Is returning the value not enough?

llvm/test/CodeGen/AArch64/neon-dot-product.ll

261

Can you update the tests to show all the instructions? The 1's and 0's and addv's are important here too.

I would just use the update_llc_test_checks script.

[mivnay]

Fixed context and review comments

[dmgreen]

Thanks. LGTM with one minor modification.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

11015

EVT::getVectorVT(*DAG.getContext(), MVT::i32, 4) -> MVT::v4i32 :)

[mivnay]

The build failure is unrelated to this patch and happening in other patches as well (Example: https://reviews.llvm.org/D88471). What should I do? Also, I do not have commit access. Can you please help in committing this patch?

[dmgreen]

In D88577#2308416, @mivnay wrote:

The build failure is unrelated to this patch and happening in other patches as well (Example: https://reviews.llvm.org/D88471). What should I do? Also, I do not have commit access. Can you please help in committing this patch?

Yeah I tend to ignore those failures. I can check it though.

I can certainly commit this. I just need an author string to attribute it correctly. Is "Vinay Madhusudan <[e-mail]>" OK for that?

[mivnay]

I can certainly commit this. I just need an author string to attribute it correctly. Is "Vinay Madhusudan <[e-mail]>" OK for that?

Yes, thanks a lot.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
 
   setTargetDAGCombine(ISD::SELECT);
   setTargetDAGCombine(ISD::VSELECT);
 
   setTargetDAGCombine(ISD::INTRINSIC_VOID);
   setTargetDAGCombine(ISD::INTRINSIC_W_CHAIN);
   setTargetDAGCombine(ISD::INSERT_VECTOR_ELT);
   setTargetDAGCombine(ISD::EXTRACT_VECTOR_ELT);
+  setTargetDAGCombine(ISD::VECREDUCE_ADD);
 
   setTargetDAGCombine(ISD::GlobalAddress);
 
   // In case of strict alignment, avoid an excessive number of byte wide stores.
   MaxStoresPerMemsetOptSize = 8;
   MaxStoresPerMemset = Subtarget->requiresStrictAlign()
                        ? MaxStoresPerMemsetOptSize : 32;
 

 ///   cmge X, X, #0
 static SDValue foldVectorXorShiftIntoCmp(SDNode *N, SelectionDAG &DAG,
                                          const AArch64Subtarget *Subtarget) {
   EVT VT = N->getValueType(0);
   if (!Subtarget->hasNEON() || !VT.isVector())
     return SDValue();
 
   // There must be a shift right algebraic before the xor, and the xor must be a
   // 'not' operation.
   SDValue Shift = N->getOperand(0);
   SDValue Ones = N->getOperand(1);
   if (Shift.getOpcode() != AArch64ISD::VASHR || !Shift.hasOneUse() ||
       !ISD::isBuildVectorAllOnes(Ones.getNode()))
     return SDValue();
 
   // The shift should be smearing the sign bit across each vector element.
   auto *ShiftAmt = dyn_cast<ConstantSDNode>(Shift.getOperand(1));
   EVT ShiftEltTy = Shift.getValueType().getVectorElementType();
   if (!ShiftAmt || ShiftAmt->getZExtValue() != ShiftEltTy.getSizeInBits() - 1)
     return SDValue();
 
   return DAG.getNode(AArch64ISD::CMGEz, SDLoc(N), VT, Shift.getOperand(0));
 }
 
 // Generate SUBS and CSEL for integer abs.
 static SDValue performIntegerAbsCombine(SDNode *N, SelectionDAG &DAG) {
   EVT VT = N->getValueType(0);
 
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
   SDLoc DL(N);
 
   // Check pattern of XOR(ADD(X,Y), Y) where Y is SRA(X, size(X)-1)
   // and change it to SUB and CSEL.
   if (VT.isInteger() && N->getOpcode() == ISD::XOR &&
       N0.getOpcode() == ISD::ADD && N0.getOperand(1) == N1 &&
       N1.getOpcode() == ISD::SRA && N1.getOperand(0) == N0.getOperand(0))
     if (ConstantSDNode *Y1C = dyn_cast<ConstantSDNode>(N1.getOperand(1)))
       if (Y1C->getAPIntValue() == VT.getSizeInBits() - 1) {
         SDValue Neg = DAG.getNode(ISD::SUB, DL, VT, DAG.getConstant(0, DL, VT),
                                   N0.getOperand(0));
         // Generate SUBS & CSEL.
         SDValue Cmp =
             DAG.getNode(AArch64ISD::SUBS, DL, DAG.getVTList(VT, MVT::i32),
                         N0.getOperand(0), DAG.getConstant(0, DL, VT));
         return DAG.getNode(AArch64ISD::CSEL, DL, VT, N0.getOperand(0), Neg,
                            DAG.getConstant(AArch64CC::PL, DL, MVT::i32),
                            SDValue(Cmp.getNode(), 1));
       }
   return SDValue();
 }
 
+// VECREDUCE_ADD( EXTEND(v16i8_type) ) to
+// VECREDUCE_ADD( DOTv16i8(v16i8_type) )
+static SDValue performVecReduceAddCombine(SDNode *N, SelectionDAG &DAG,
+                                          const AArch64Subtarget *ST) {
+  SDValue Op0 = N->getOperand(0);
+  if (!ST->hasDotProd() || N->getValueType(0) != MVT::i32)
+    return SDValue();
+
+  if (Op0.getValueType().getVectorElementType() != MVT::i32)
+    return SDValue();
+
+  unsigned ExtOpcode = Op0.getOpcode();
+  if (ExtOpcode != ISD::ZERO_EXTEND && ExtOpcode != ISD::SIGN_EXTEND)
+    return SDValue();
+
+  EVT Op0VT = Op0.getOperand(0).getValueType();
+  if (Op0VT != MVT::v16i8)
+    return SDValue();
+
+  SDLoc DL(Op0);
+  SDValue Ones = DAG.getConstant(1, DL, Op0VT);
+  SDValue Zeros = DAG.getConstant(0, DL, MVT::v4i32);
+  auto DotIntrisic = (ExtOpcode == ISD::ZERO_EXTEND)
+                         ? Intrinsic::aarch64_neon_udot
+                         : Intrinsic::aarch64_neon_sdot;
+  SDValue Dot = DAG.getNode(ISD::INTRINSIC_WO_CHAIN, DL, Zeros.getValueType(),
+                            DAG.getConstant(DotIntrisic, DL, MVT::i32), Zeros,
+                            Ones, Op0.getOperand(0));
+  return DAG.getNode(ISD::VECREDUCE_ADD, DL, N->getValueType(0), Dot);
+}
+
 static SDValue performXorCombine(SDNode *N, SelectionDAG &DAG,
                                  TargetLowering::DAGCombinerInfo &DCI,
                                  const AArch64Subtarget *Subtarget) {
   if (DCI.isBeforeLegalizeOps())
     return SDValue();
 
   if (SDValue Cmp = foldVectorXorShiftIntoCmp(N, DAG, Subtarget))
     return Cmp;
SDValue AArch64TargetLowering::PerformDAGCombine(SDNode *N,
   case AArch64ISD::NVCAST:
     return performNVCASTCombine(N);
   case AArch64ISD::UZP1:
     return performUzpCombine(N, DAG);
   case ISD::INSERT_VECTOR_ELT:
     return performPostLD1Combine(N, DCI, true);
   case ISD::EXTRACT_VECTOR_ELT:
     return performExtractVectorEltCombine(N, DAG);
+  case ISD::VECREDUCE_ADD:
+    return performVecReduceAddCombine(N, DCI.DAG, Subtarget);
   case ISD::INTRINSIC_VOID:
   case ISD::INTRINSIC_W_CHAIN:
     switch (cast<ConstantSDNode>(N->getOperand(1))->getZExtValue()) {
     case Intrinsic::aarch64_sve_prfb_gather_scalar_offset:
       return combineSVEPrefetchVecBaseImmOff(N, DAG, 1 /*=ScalarSizeInBytes*/);
     case Intrinsic::aarch64_sve_prfh_gather_scalar_offset:
       return combineSVEPrefetchVecBaseImmOff(N, DAG, 2 /*=ScalarSizeInBytes*/);
     case Intrinsic::aarch64_sve_prfw_gather_scalar_offset:

llvm/test/CodeGen/AArch64/neon-dot-product.ll

; CHECK:  udot {{v[0-9]+}}.4s, {{v[0-9]+}}.16b, {{v[0-9]+}}.16b
   %4 = load <16 x i8>, <16 x i8>* %3
   %5 = zext <16 x i8> %4 to <16 x i32>
   %6 = mul nuw nsw <16 x i32> %5, %2
   %7 = call i32 @llvm.experimental.vector.reduce.add.v16i32(<16 x i32> %6)
   %op.extra = add i32 %7, %sum
   ret i32 %op.extra
 }
 
+define i32 @test_udot_v16i8_2(i8* nocapture readonly %a1) {
+; CHECK-LABEL: test_udot_v16i8_2:
+; CHECK:    movi {{v[0-9]+}}.16b, #1
+; CHECK:    movi {{v[0-9]+}}.2d, #0000000000000000
+; CHECK:    udot {{v[0-9]+}}.4s, {{v[0-9]+}}.16b, {{v[0-9]+}}.16b
+; CHECK:    addv s0, {{v[0-9]+}}.4s
+entry:
+  %0 = bitcast i8* %a1 to <16 x i8>*
+  %1 = load <16 x i8>, <16 x i8>* %0
+  %2 = zext <16 x i8> %1 to <16 x i32>
+  %3 = call i32 @llvm.experimental.vector.reduce.add.v16i32(<16 x i32> %2)
+  ret i32 %3
+}
+
 define i32 @test_sdot_v16i8(i8* nocapture readonly %a, i8* nocapture readonly %b, i32 %sum) {
 entry:
 ; CHECK-LABEL: test_sdot_v16i8:
 ; CHECK:  sdot {{v[0-9]+}}.4s, {{v[0-9]+}}.16b, {{v[0-9]+}}.16b
   %0 = bitcast i8* %a to <16 x i8>*
   %1 = load <16 x i8>, <16 x i8>* %0
   %2 = sext <16 x i8> %1 to <16 x i32>
   %3 = bitcast i8* %b to <16 x i8>*
   %4 = load <16 x i8>, <16 x i8>* %3
   %5 = sext <16 x i8> %4 to <16 x i32>
   %6 = mul nsw <16 x i32> %5, %2
   %7 = call i32 @llvm.experimental.vector.reduce.add.v16i32(<16 x i32> %6)
   %op.extra = add nsw i32 %7, %sum
   ret i32 %op.extra
 }
+
+define i32 @test_sdot_v16i8_2(i8* nocapture readonly %a1) {
+; CHECK-LABEL: test_sdot_v16i8_2:
+; CHECK:    movi {{v[0-9]+}}.16b, #1
+; CHECK:    movi {{v[0-9]+}}.2d, #0000000000000000
+; CHECK:    sdot {{v[0-9]+}}.4s, {{v[0-9]+}}.16b, {{v[0-9]+}}.16b
+; CHECK:    addv s0, {{v[0-9]+}}.4s
+entry:
+  %0 = bitcast i8* %a1 to <16 x i8>*
+  %1 = load <16 x i8>, <16 x i8>* %0
+  %2 = sext <16 x i8> %1 to <16 x i32>
+  %3 = call i32 @llvm.experimental.vector.reduce.add.v16i32(<16 x i32> %2)
+  ret i32 %3
+}