[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
Thu Mar 20 07:26:03 PDT 2025


================
@@ -1585,6 +1585,21 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
       setOperationAction(ISD::MSTORE, VT, Custom);
     }
 
+    for (MVT VT : MVT::integer_scalable_vector_valuetypes()) {
+      if (!EnablePartialReduceNodes)
+        break;
+      for (MVT InnerVT : MVT::integer_scalable_vector_valuetypes()) {
+        ElementCount VTElemCount = VT.getVectorElementCount();
+        if (VTElemCount.getKnownMinValue() == 1)
+          continue;
+        if (VTElemCount * 4 == InnerVT.getVectorElementCount())
+          setPartialReduceMLAAction(VT, InnerVT, Custom);
+        if (InnerVT.getVectorElementType().getSizeInBits() * 4 ==
+            VT.getVectorElementType().getSizeInBits())
+          setPartialReduceMLAAction(VT, InnerVT, Legal);
+      }
+    }
----------------
MacDue wrote:

Having played around with this a bit, I'd suggest setting these like so:
```suggestion
    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);
        }
      }
      // 3. Mark known legal pairs as 'Legal' (these will expand to USDOT).
      setPartialReduceMLAAction(MVT::nxv2i64, MVT::nxv8i16, Legal);
      setPartialReduceMLAAction(MVT::nxv4i32, MVT::nxv16i8, Legal);
    }
```
I think this makes it clearer what's actually a legal or custom combination (which is a fairly small number of combinations). This passes the existing tests too.

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


More information about the llvm-commits mailing list