[llvm] [AArch64] Improve lowering of pure fixed-length partial add reductions. (PR #214691)

[Paul Walker via llvm-commits]

================
@@ -34748,6 +34710,33 @@ AArch64TargetLowering::LowerPARTIAL_REDUCE_MLA(SDValue Op,
   EVT OrigResultVT = ResultVT;
   EVT OpVT = LHS.getValueType();
 
+  // Two-way fixed-length integer add reductions.
+  if ((Op.getOpcode() == ISD::PARTIAL_REDUCE_UMLA ||
+       Op.getOpcode() == ISD::PARTIAL_REDUCE_SMLA) &&
+      Subtarget->isNeonAvailable() && ResultVT.isFixedLengthVector() &&
+      ResultVT.getScalarSizeInBits() == OpVT.getScalarSizeInBits() * 2 &&
+      ResultVT.getVectorNumElements() * 2 == OpVT.getVectorNumElements() &&
+      (OpVT.getSizeInBits() == 64 || OpVT.getSizeInBits() == 128)) {
+    // A pure partial reduction can lower to a [SU]ADALP.
+    if (isOneVector(RHS)) {
+      bool IsUnsigned = Op.getOpcode() == ISD::PARTIAL_REDUCE_UMLA;
+      unsigned Opc = IsUnsigned ? AArch64ISD::UADDLP : AArch64ISD::SADDLP;
+      return DAG.getNode(ISD::ADD, DL, ResultVT, Acc,
+                         DAG.getNode(Opc, DL, ResultVT, LHS));
----------------
paulwalker-arm wrote:

Can you return `Op` and handle this in isel like we do for SVE? 

https://github.com/llvm/llvm-project/pull/214691