[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