Adjust how usdot cases are lowered

llvm · NickGuy-Arm · May 6, 2025 · May 6, 2025 · May 7, 2025 · May 8, 2025
commit 22636ac1a9bd8e73943f9b52dc5ff063dc5ed258
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp
@@ -925,8 +925,18 @@ SDValue DAGTypeLegalizer::CreateStackStoreLoad(SDValue Op,
 bool DAGTypeLegalizer::CustomLowerNode(SDNode *N, EVT VT, bool LegalizeResult) {
  // See if the target wants to custom lower this node.
  unsigned Opcode = N->getOpcode();
-  if (TLI.getOperationAction(Opcode, VT) != TargetLowering::Custom)
-    return false;
+  bool IsPRMLAOpcode =
+      Opcode == ISD::PARTIAL_REDUCE_UMLA || Opcode == ISD::PARTIAL_REDUCE_SMLA;
+
+  if (IsPRMLAOpcode) {
+    if (TLI.getPartialReduceMLAAction(N->getValueType(0),
+                                      N->getOperand(1).getValueType()) !=
+        TargetLowering::Custom)
+      return false;
+  } else {
+    if (TLI.getOperationAction(Opcode, VT) != TargetLowering::Custom)
+      return false;
+  }

  SmallVector<SDValue, 8> Results;
  if (LegalizeResult)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -1872,17 +1872,9 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
    // 8to64
    setPartialReduceMLAAction(MVT::nxv2i64, MVT::nxv16i8, Custom);

-    if (Subtarget->hasMatMulInt8()) {
-      // USDOT
-      setPartialReduceMLAAction(MVT::nxv2i64, MVT::nxv8i64, Custom);
+    // USDOT
+    if (Subtarget->hasMatMulInt8())
      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);
-    }
  }

  // Handle operations that are only available in non-streaming SVE mode.
@@ -7755,8 +7747,8 @@ SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
    return LowerFLDEXP(Op, DAG);
  case ISD::EXPERIMENTAL_VECTOR_HISTOGRAM:
    return LowerVECTOR_HISTOGRAM(Op, DAG);
-  case ISD::PARTIAL_REDUCE_UMLA:
  case ISD::PARTIAL_REDUCE_SMLA:
+  case ISD::PARTIAL_REDUCE_UMLA:
    return LowerPARTIAL_REDUCE_MLA(Op, DAG);
  }
 }