[llvm] [AArch64] Fold four and eight way partial reductions with [SU]ADDLP (PR #214636)

[Adam Scott via llvm-commits]

================
@@ -1755,11 +1755,41 @@ class LLVM_ABI TargetLoweringBase {
   }
 
   /// Return true if a PARTIAL_REDUCE_U/SMLA node with the specified types is
-  /// legal or custom for this target.
+  /// legal or custom for this target. A target that does not take the fold
+  /// directly can still build it from a ladder of narrower steps, each halving
+  /// the element count and doubling the width.
   bool isPartialReduceMLALegalOrCustom(unsigned Opc, EVT AccVT,
                                        EVT InputVT) const {
-    LegalizeAction Action = getPartialReduceMLAAction(Opc, AccVT, InputVT);
-    return Action == Legal || Action == Custom;
+    auto Supported = [&](EVT Acc, EVT In) {
+      LegalizeAction Action = getPartialReduceMLAAction(Opc, Acc, In);
+      return Action == Legal || Action == Custom;
+    };
+    if (Supported(AccVT, InputVT))
+      return true;
+    // The rungs below widen an integer element type.
+    if (Opc == ISD::PARTIAL_REDUCE_FMLA)
+      return false;
+
+    ElementCount AccEC = AccVT.getVectorElementCount();
+    ElementCount InEC = InputVT.getVectorElementCount();
+    // Scalable and fixed lengths are not comparable below.
+    if (AccEC.isScalable() != InEC.isScalable())
+      return false;
+
+    // Step down a rung at a time, taking each step the target supports, until
+    // one of them lands on the accumulator.
+    EVT In = InputVT;
+    while (InEC.getKnownMinValue() > 2 * AccEC.getKnownMinValue()) {
+      InEC = InEC.divideCoefficientBy(2);
+      MVT Next = MVT::getVectorVT(
+          MVT::getIntegerVT(In.getScalarSizeInBits() * 2), InEC);
+      if (!Next.isValid() || !Supported(Next, In))
----------------
as4230 wrote:

Yes, spot on with the problem I ran into. I'll leave isPartialReduceMLALegalOrCustom alone and make a separate query that walks the rungs and calls isPartialReduceMLALegalOrCustom for each one.

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