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

Benjamin Maxwell via llvm-commits llvm-commits at lists.llvm.org
Tue Apr 8 07:40:55 PDT 2025


================
@@ -1585,6 +1585,26 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
       setOperationAction(ISD::MSTORE, VT, Custom);
     }
 
+    if (EnablePartialReduceNodes) {
+      for (MVT VT : MVT::integer_scalable_vector_valuetypes()) {
+        for (MVT InnerVT : MVT::integer_scalable_vector_valuetypes()) {
+          // 1. Set all combinations where a type is illegal to "Legal"
+          // - These will be legalized to a legal type pair
+          // - Avoid expanding them too early (or preventing folds)
+          if (!isTypeLegal(VT) || !isTypeLegal(InnerVT)) {
+            setPartialReduceMLAAction(VT, InnerVT, Legal);
+            continue;
+          }
+          //  2. Set all legal combinations to "Expand"
+          // - Not all of these can be lowered (via a Legal or Custom lowering).
+          setPartialReduceMLAAction(VT, InnerVT, Expand);
----------------
MacDue wrote:

Maybe something like this then? (which for now assumes we only split partial reductions):

https://gist.github.com/MacDue/cbb6383c452319725fc9abf8a6d81b90

This does worsen a `nxv16i8` to `nxv4i64` case, but looking at some later patches that case is poorly handled even with the full legalization. I think that's a case where we'd just want a DAG combine that reduces the accumulator type to `nxv4i32`. 

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


More information about the llvm-commits mailing list