Adjust how usdot cases are lowered

Author: NickGuy-Arm

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

@@ -925,18 +925,8 @@ 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();
-  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;
-  }
+  if (TLI.getOperationAction(Opcode, VT) != TargetLowering::Custom)
+    return false;
 
   SmallVector<SDValue, 8> Results;
   if (LegalizeResult)
@@ -957,7 +947,6 @@ bool DAGTypeLegalizer::CustomLowerNode(SDNode *N, EVT VT, bool LegalizeResult) {
   return true;
 }
 
-
 /// Widen the node's results with custom code provided by the target and return
 /// "true", or do nothing and return "false".
 bool DAGTypeLegalizer::CustomWidenLowerNode(SDNode *N, EVT VT) {

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -1880,6 +1880,8 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
       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);
     }
   }