[llvm] [AArch64][SVE] Add dot product lowering for PARTIAL_REDUCE_MLA node (PR #130933)

[Benjamin Maxwell via llvm-commits]

================
@@ -1585,6 +1585,44 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
       setOperationAction(ISD::MSTORE, VT, Custom);
     }
 
+    if (EnablePartialReduceNodes) {
+      // Mark known legal pairs as 'Legal' (these will expand to UDOT or SDOT).
+      // Other pairs will default to 'Expand'.
+      setPartialReduceMLAAction(MVT::nxv2i64, MVT::nxv8i16, Legal);
+      setPartialReduceMLAAction(MVT::nxv4i32, MVT::nxv16i8, Legal);
+
+      auto CanSplitToLegalPartialReduce = [&](MVT AccTy, MVT InTy) {
+        while (true) {
+          switch (getTypeAction(AccTy)) {
+          case TargetLoweringBase::TypeLegal:
+            return isPartialReduceMLALegalOrCustom(AccTy, InTy);
+          case TargetLoweringBase::TypeSplitVector:
+            if (!InTy.getVectorElementCount().isKnownEven())
+              return false;
+            // Currently, we only implement spillting for partial reductions,
+            // which splits the result and both operands.
+            AccTy = AccTy.getHalfNumVectorElementsVT();
+            InTy = InTy.getHalfNumVectorElementsVT();
+            break;
+          default:
+            // Assume all other type pairs are expanded.
+            return false;
+          }
+        }
+      };
+
+      for (MVT VT : MVT::integer_scalable_vector_valuetypes()) {
+        for (MVT InnerVT : MVT::integer_scalable_vector_valuetypes()) {
+          // Mark illegal type pairs that split to a legal pair as legal.
+          // This is needed as otherwise we may not apply useful combines.
+          if (!isTypeLegal(VT) || !isTypeLegal(InnerVT)) {
+            if (CanSplitToLegalPartialReduce(VT, InnerVT))
+              setPartialReduceMLAAction(VT, InnerVT, Legal);
+          }
+        }
+      }
----------------
MacDue wrote:

It's probably best to move this to the patch that implements splitting. I think that's a few patches down the line. 

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