[AArch64] Combine UADDVs to generate vector add

[mivnay]

Summary

ADD(UADDV a, UADDV b) --> UADDV(ADD a, b)

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

[mivnay]

Add support for all integer types.

[mivnay]

ping!

[dmgreen]

The pre-commit check is not being helpful here... :-/

Seems like a useful optimisation to me though.

Can you add test for i16 and i8. As far as I understand they will not fold because we will have legalized types and the return type will not match the vector element type? It's still worth having the tests. We could think of doing this target independent instead. Folding add(vecreduce(x), vecreduce(y)) -> vecreduce(add(x, y)). It sounds generally useful to me. World that work in your case, or would it specifically need to work on UADDV?

Can you also run the update_llc_test_checks.py script on the file and pre-commit the tests, just showing the changes here.

lib/Target/AArch64/AArch64ISelLowering.cpp

12279 ↗

(On Diff #297403)

--> UADDV(ADD a, b)

12281 ↗

(On Diff #297403)

You can get DAG from DCI.DAG (and I don't think think method uses DCI, which might simplify things).

12376 ↗

(On Diff #297403)

I think this can be:

if (SDValue Val = performUADDVCombine(N, DCI, DAG))
  return Val;

[mivnay]

Can you also run the update_llc_test_checks.py script on the file and pre-commit the tests, just showing the changes here.

Sure. I have created new patch : https://reviews.llvm.org/D89365

[mivnay]

In D88731#2328076, @dmgreen wrote:

The pre-commit check is not being helpful here... :-/

Seems like a useful optimisation to me though.

Can you add test for i16 and i8. As far as I understand they will not fold because we will have legalized types and the return type will not match the vector element type?

Yes!

We could think of doing this target independent instead. Folding add(vecreduce(x), vecreduce(y)) -> vecreduce(add(x, y)). It sounds generally useful to me. World that work in your case, or would it specifically need to work on UADDV?

I don't need to use UADDV specifically. But, I am not sure what impact it has on other targets (different lowerings, patterns, etc..).

[dmgreen]

Yeah. I was going to say that this would replace 3 instructions with 2, so would generally be useful.

But thinking about it, that wouldn't actually be the case for MVE. The two reductions (VADDVA's) would likely overlap better than back to back loads.

The code here looks good to me. If we find we need to handle more cases in the future we can look at it there, and perhaps put some sort of "undo" in for MVE. But in the meantime this patch LGTM.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

Cmp = getAArch64Cmp(
         ISD::getSetCCInverse(InfoAndKind.Info.Generic.CC, CmpVT), CCVal, DAG,
         dl);
 
   EVT VT = Op->getValueType(0);
   LHS = DAG.getNode(ISD::ADD, dl, VT, RHS, DAG.getConstant(1, dl, VT));
   return DAG.getNode(AArch64ISD::CSEL, dl, VT, RHS, LHS, CCVal, Cmp);
 }
 
+// ADD(UADDV a, UADDV b) -->  UADDV((ADD a, b))
+static SDValue performUADDVCombine(SDNode *N,
+                                   TargetLowering::DAGCombinerInfo &DCI,
+                                   SelectionDAG &DAG) {
+  MVT VT = N->getSimpleValueType(0);
+  // TODO: Currently handled for 32 bit integer vectors.
+  if (N->getOpcode() != ISD::ADD || VT != MVT::i32)
+    return SDValue();
+
+  SDValue LHS = N->getOperand(0);
+  SDValue RHS = N->getOperand(1);
+  if (LHS.getOpcode() != ISD::EXTRACT_VECTOR_ELT ||
+      RHS.getOpcode() != ISD::EXTRACT_VECTOR_ELT || LHS.getValueType() != VT)
+    return SDValue();
+
+  auto LHSN1 = dyn_cast<ConstantSDNode>(LHS->getOperand(1));
+  auto RHSN1 = dyn_cast<ConstantSDNode>(RHS->getOperand(1));
+  if (!LHSN1 || !RHSN1 || !LHSN1->isNullValue() || !RHSN1->isNullValue())
+    return SDValue();
+
+  SDValue Op1 = LHS->getOperand(0);
+  SDValue Op2 = RHS->getOperand(0);
+  if (Op1.getOpcode() != AArch64ISD::UADDV ||
+      Op1.getValueType() != MVT::v4i32 ||
+      Op2.getOpcode() != AArch64ISD::UADDV || Op2.getValueType() != MVT::v4i32)
+    return SDValue();
+
+  SDValue Val1 = Op1.getOperand(0);
+  SDValue Val2 = Op2.getOperand(0);
+  EVT ValVT = Val1->getValueType(0);
+  SDLoc DL(N);
+  SDValue AddVal = DAG.getNode(ISD::ADD, DL, ValVT, Val1, Val2);
+  return DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, VT,
+                     DAG.getNode(AArch64ISD::UADDV, DL, ValVT, AddVal),
+                     DAG.getConstant(0, DL, MVT::i64));
+}
+
 // The basic add/sub long vector instructions have variants with "2" on the end
 // which act on the high-half of their inputs. They are normally matched by
 // patterns like:
 //
 // (add (zeroext (extract_high LHS)),
 //      (zeroext (extract_high RHS)))
 // -> uaddl2 vD, vN, vM
 //
if (!VT.is128BitVector()) {
     if (N->getOpcode() == ISD::ADD)
       return performSetccAddFolding(N, DAG);
     return SDValue();
   }
 
   // Make sure both branches are extended in the same way.
   SDValue LHS = N->getOperand(0);
   SDValue RHS = N->getOperand(1);
   if ((LHS.getOpcode() != ISD::ZERO_EXTEND &&
        LHS.getOpcode() != ISD::SIGN_EXTEND) ||
       LHS.getOpcode() != RHS.getOpcode())
     return SDValue();
 
   unsigned ExtType = LHS.getOpcode();
 
   // It's not worth doing if at least one of the inputs isn't already an
   // extract, but we don't know which it'll be so we have to try both.
   if (isEssentiallyExtractHighSubvector(LHS.getOperand(0))) {
     RHS = tryExtendDUPToExtractHigh(RHS.getOperand(0), DAG);
     if (!RHS.getNode())
       return SDValue();
 
     RHS = DAG.getNode(ExtType, SDLoc(N), VT, RHS);
   } else if (isEssentiallyExtractHighSubvector(RHS.getOperand(0))) {
     LHS = tryExtendDUPToExtractHigh(LHS.getOperand(0), DAG);
     if (!LHS.getNode())
       return SDValue();
 
     LHS = DAG.getNode(ExtType, SDLoc(N), VT, LHS);
   }
 
   return DAG.getNode(N->getOpcode(), SDLoc(N), VT, LHS, RHS);
 }
 
+static SDValue performAddSubCombine(SDNode *N,
+                                    TargetLowering::DAGCombinerInfo &DCI,
+                                    SelectionDAG &DAG) {
+  // Try to change sum of two reductions.
+  SDValue Val = performUADDVCombine(N, DCI, DAG);
+  if (Val.getNode()) {
+    return Val;
+  }
+
+  return performAddSubLongCombine(N, DCI, DAG);
+}
+
 // Massage DAGs which we can use the high-half "long" operations on into
 // something isel will recognize better. E.g.
 //
 // (aarch64_neon_umull (extract_high vec) (dupv64 scalar)) -->
 //   (aarch64_neon_umull (extract_high (v2i64 vec)))
 //                     (extract_high (v2i64 (dup128 scalar)))))
 //
 static SDValue tryCombineLongOpWithDup(unsigned IID, SDNode *N,
SDValue AArch64TargetLowering::PerformDAGCombine(SDNode *N,
                                                  DAGCombinerInfo &DCI) const {
   SelectionDAG &DAG = DCI.DAG;
   switch (N->getOpcode()) {
   default:
     LLVM_DEBUG(dbgs() << "Custom combining: skipping\n");
     break;
   case ISD::ADD:
   case ISD::SUB:
-    return performAddSubLongCombine(N, DCI, DAG);
+    return performAddSubCombine(N, DCI, DAG);
   case ISD::XOR:
     return performXorCombine(N, DAG, DCI, Subtarget);
   case ISD::MUL:
     return performMulCombine(N, DAG, DCI, Subtarget);
   case ISD::SINT_TO_FP:
   case ISD::UINT_TO_FP:
     return performIntToFpCombine(N, DAG, Subtarget);
   case ISD::FP_TO_SINT:

llvm/test/CodeGen/AArch64/aarch64-addv.ll

 
 define i32 @oversized_ADDV_512(<16 x i32>* %arr)  {
 ; CHECK-LABEL: oversized_ADDV_512
 ; CHECK: addv {{s[0-9]+}}, {{v[0-9]+}}.4s
   %bin.rdx = load <16 x i32>, <16 x i32>* %arr
   %r = call i32 @llvm.experimental.vector.reduce.add.v16i32(<16 x i32> %bin.rdx)
   ret i32 %r
 }
+
+define i32 @addv_combine(i32* nocapture readonly %a1, i32* nocapture readonly %a2) {
+; CHECK-LABEL: addv_combine
+; CHECK: add {{v[0-9]+}}.4s, {{v[0-9]+}}.4s, {{v[0-9]+}}.4s
+; CHECK: addv {{s[0-9]+}}, {{v[0-9]+}}.4s
+entry:
+  %0 = bitcast i32* %a1 to <4 x i32>*
+  %1 = load <4 x i32>, <4 x i32>* %0
+  %2 = bitcast i32* %a2 to <4 x i32>*
+  %3 = load <4 x i32>, <4 x i32>* %2
+  %4 = call i32 @llvm.experimental.vector.reduce.add.v4i32(<4 x i32> %1)
+  %5 = call i32 @llvm.experimental.vector.reduce.add.v4i32(<4 x i32> %3)
+  %6 = add i32 %4, %5
+  ret i32 %6
+}