[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

10966

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

   setTargetDAGCombine(ISD::INTRINSIC_W_CHAIN);
   setTargetDAGCombine(ISD::INSERT_VECTOR_ELT);
   setTargetDAGCombine(ISD::EXTRACT_VECTOR_ELT);
+  setTargetDAGCombine(ISD::VECREDUCE_ADD);
 
   setTargetDAGCombine(ISD::GlobalAddress);
 

   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();
+  }
+
+  EVT VT = Op0.getValueType();
+  if (VT.isScalableVector() || VT.getVectorElementType() != MVT::i32) {
+    return SDValue();
+  }
+
+  EVT Op0VT = Op0.getOperand(0).getValueType();
+  unsigned ExtOpcode = Op0.getOpcode();
+  if (ExtOpcode != ISD::ZERO_EXTEND && ExtOpcode != ISD::SIGN_EXTEND) {
+    return SDValue();
+  }
+
+  if (Op0VT.isScalableVector() || Op0VT.getVectorElementType() != MVT::i8 ||
+      Op0VT.getVectorNumElements() != 16) {
+    return SDValue();
+  }
+
+  SDValue Ones =
+      DAG.getConstant(1, SDLoc(Op0), Op0.getOperand(0).getValueType());
+  SDValue Zeros = DAG.getConstant(
+      0, SDLoc(Op0), EVT::getVectorVT(*DAG.getContext(), MVT::i32, 4));
+
+  auto DotOpcode =
+      (ExtOpcode == ISD::ZERO_EXTEND) ? AArch64::UDOTv16i8 : AArch64::SDOTv16i8;
+  MachineSDNode *ABD =
+      DAG.getMachineNode(DotOpcode, SDLoc(Op0), Zeros.getValueType(), Zeros,
+                         Ones, Op0.getOperand(0));
+  SDValue FinalABD = DAG.getNode(ISD::VECREDUCE_ADD, SDLoc(N),
+                                 N->getValueType(0), SDValue(ABD, 0));
+  DAG.ReplaceAllUsesOfValueWith(SDValue(N, 0), FinalABD);
+  return FinalABD;
+}
+
 static SDValue performXorCombine(SDNode *N, SelectionDAG &DAG,
                                  TargetLowering::DAGCombinerInfo &DCI,
                                  const AArch64Subtarget *Subtarget) {

     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()) {

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

   ret i32 %op.extra
 }
 
+define i32 @test_udot_v16i8_2(i8* nocapture readonly %a1) {
+entry:
+; CHECK-LABEL: test_udot_v16i8_2:
+; CHECK:  udot {{v[0-9]+}}.4s, {{v[0-9]+}}.16b, {{v[0-9]+}}.16b
+  %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:

   %op.extra = add nsw i32 %7, %sum
   ret i32 %op.extra
 }
+
+define i32 @test_sdot_v16i8_2(i8* nocapture readonly %a1) {
+entry:
+; CHECK-LABEL: test_sdot_v16i8_2:
+; CHECK:  sdot {{v[0-9]+}}.4s, {{v[0-9]+}}.16b, {{v[0-9]+}}.16b
+  %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
+}