[llvm] [AArch64] Keep v16i8 -> v2i64 partial_reduce fixed-length for VL > 128 (PR #204938)

[Benjamin Maxwell via llvm-commits]

================
@@ -33517,29 +33469,42 @@ AArch64TargetLowering::LowerPARTIAL_REDUCE_MLA(SDValue Op,
   SDValue DotNode = DAG.getNode(Op.getOpcode(), DL, DotVT,
                                 DAG.getConstant(0, DL, DotVT), LHS, RHS);
 
-  SDValue Res;
   bool IsUnsigned = Op.getOpcode() == ISD::PARTIAL_REDUCE_UMLA;
-  if (Subtarget->hasSVE2() || Subtarget->isStreamingSVEAvailable()) {
+
+  // UADDW{B,T}/SADDW{B,T} fold the dot in the scalable domain, spreading the
+  // sums across all VL/64 lanes. That is only valid for a genuinely scalable
+  // result; a fixed-length result must convert from the scalable container
+  // before splitting (below), else the trailing extract drops the high lanes
+  // for any VL > the fixed width.
+  if (OrigResultVT.isScalableVector() &&
+      (Subtarget->hasSVE2() || Subtarget->isStreamingSVEAvailable())) {
     unsigned LoOpcode = IsUnsigned ? AArch64ISD::UADDWB : AArch64ISD::SADDWB;
     unsigned HiOpcode = IsUnsigned ? AArch64ISD::UADDWT : AArch64ISD::SADDWT;
     SDValue Lo = DAG.getNode(LoOpcode, DL, ResultVT, Acc, DotNode);
-    Res = DAG.getNode(HiOpcode, DL, ResultVT, Lo, DotNode);
-  } else {
-    // Fold (nx)v4i32 into (nx)v2i64
-    auto [DotNodeLo, DotNodeHi] = DAG.SplitVector(DotNode, DL);
-    if (IsUnsigned) {
-      DotNodeLo = DAG.getZExtOrTrunc(DotNodeLo, DL, ResultVT);
-      DotNodeHi = DAG.getZExtOrTrunc(DotNodeHi, DL, ResultVT);
-    } else {
-      DotNodeLo = DAG.getSExtOrTrunc(DotNodeLo, DL, ResultVT);
-      DotNodeHi = DAG.getSExtOrTrunc(DotNodeHi, DL, ResultVT);
-    }
-    auto Lo = DAG.getNode(ISD::ADD, DL, ResultVT, Acc, DotNodeLo);
-    Res = DAG.getNode(ISD::ADD, DL, ResultVT, Lo, DotNodeHi);
+    return DAG.getNode(HiOpcode, DL, ResultVT, Lo, DotNode);
   }
 
-  return ConvertToScalable ? convertFromScalableVector(DAG, OrigResultVT, Res)
-                           : Res;
+  // Fold the (nx)v4i32 dot into the (nx)v2i64 result. For a fixed-length
+  // result, convert from the scalable container before splitting: the sums sit
+  // in the low i32 lanes regardless of VL, so splitting after the extract would
+  // drop a 128-bit result's high lanes for VL > 128. See PR #177119 / issue
+  // #176954.
+  SDValue FoldDot = DotNode;
+  if (ConvertToScalable) {
+    EVT FixedDotVT = EVT::getVectorVT(*DAG.getContext(), MVT::i32,
+                                      OrigResultVT.getVectorNumElements() * 2);
----------------
MacDue wrote:

Ah, this is handling the v4i64 case (for fixed length 256-bit SVE). I think what you've done here is okay.

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