[llvm] [AArch64][SVE] Add lowering for PARTIAL_REDUCE_U/SMLA to USDOT (PR #131327)

[Nicholas Guy via llvm-commits]

================
@@ -1869,7 +1869,20 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
     setPartialReduceMLAAction(MVT::nxv2i64, MVT::nxv8i16, Legal);
     setPartialReduceMLAAction(MVT::nxv4i32, MVT::nxv16i8, Legal);
 
+    // 8to64
     setPartialReduceMLAAction(MVT::nxv2i64, MVT::nxv16i8, Custom);
+
+    if (Subtarget->hasMatMulInt8()) {
+      // USDOT
+      setPartialReduceMLAAction(MVT::nxv2i64, MVT::nxv8i64, Custom);
+      setPartialReduceMLAAction(MVT::nxv4i32, MVT::nxv16i32, Custom);
+      setPartialReduceMLAAction(MVT::nxv2i64, MVT::nxv4i64, Custom);
+      setPartialReduceMLAAction(MVT::nxv4i32, MVT::nxv8i32, Custom);
+      setPartialReduceMLAAction(MVT::nxv8i16, MVT::nxv16i16, Custom);
+      setPartialReduceMLAAction(MVT::nxv16i8, MVT::nxv32i8, Custom);
+
+      setOperationAction(ISD::PARTIAL_REDUCE_UMLA, MVT::nxv16i32, Custom);
----------------
NickGuy-Arm wrote:

The issue presents itself as the extend type being obscured by numerous (depending on the illegal types) pairs of AArch64ISD::UUNPKLO/HI. While it would technically be possible to traverse these to find the pre-extended type, we would have to traverse the entire tree of a potentially unknown size to verify that it's all valid.
By running this lowering before type legalization, we can make our adjustments and insert the node hiererchy representing USDOT, which itself can then benefit from type legalization.

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