[SelectionDAG] Improve type legalisation for PARTIAL_REDUCE_MLA (#130935)

Implement proper splitting functions for PARTIAL_REDUCE_MLA ISD nodes.
This makes the udot_8to64 and sdot_8to64 tests generate dot product
instructions for when the new ISD nodes are used.

---------

Co-authored-by: James Chesterman <[email protected]>

Author: NickGuy-Arm

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

@@ -3220,8 +3220,30 @@ void DAGTypeLegalizer::SplitVecRes_VP_REVERSE(SDNode *N, SDValue &Lo,
 void DAGTypeLegalizer::SplitVecRes_PARTIAL_REDUCE_MLA(SDNode *N, SDValue &Lo,
                                                       SDValue &Hi) {
   SDLoc DL(N);
-  SDValue Expanded = TLI.expandPartialReduceMLA(N, DAG);
-  std::tie(Lo, Hi) = DAG.SplitVector(Expanded, DL);
+  SDValue Acc = N->getOperand(0);
+  SDValue Input1 = N->getOperand(1);
+  SDValue Input2 = N->getOperand(2);
+
+  SDValue AccLo, AccHi;
+  std::tie(AccLo, AccHi) = DAG.SplitVector(Acc, DL);
+  unsigned Opcode = N->getOpcode();
+
+  // If the input types don't need splitting, just accumulate into the
+  // low part of the accumulator.
+  if (getTypeAction(Input1.getValueType()) != TargetLowering::TypeSplitVector) {
+    Lo = DAG.getNode(Opcode, DL, AccLo.getValueType(), AccLo, Input1, Input2);
+    Hi = AccHi;
+    return;
+  }
+
+  SDValue Input1Lo, Input1Hi;
+  SDValue Input2Lo, Input2Hi;
+  std::tie(Input1Lo, Input1Hi) = DAG.SplitVector(Input1, DL);
+  std::tie(Input2Lo, Input2Hi) = DAG.SplitVector(Input2, DL);
+  EVT ResultVT = AccLo.getValueType();
+
+  Lo = DAG.getNode(Opcode, DL, ResultVT, AccLo, Input1Lo, Input2Lo);
+  Hi = DAG.getNode(Opcode, DL, ResultVT, AccHi, Input1Hi, Input2Hi);
 }
 
 void DAGTypeLegalizer::SplitVecRes_VECTOR_DEINTERLEAVE(SDNode *N) {
@@ -4501,7 +4523,20 @@ SDValue DAGTypeLegalizer::SplitVecOp_VECTOR_HISTOGRAM(SDNode *N) {
 }
 
 SDValue DAGTypeLegalizer::SplitVecOp_PARTIAL_REDUCE_MLA(SDNode *N) {
-  return TLI.expandPartialReduceMLA(N, DAG);
+  SDValue Acc = N->getOperand(0);
+  assert(getTypeAction(Acc.getValueType()) != TargetLowering::TypeSplitVector &&
+         "Accumulator should already be a legal type, and shouldn't need "
+         "further splitting");
+
+  SDLoc DL(N);
+  SDValue Input1Lo, Input1Hi, Input2Lo, Input2Hi;
+  std::tie(Input1Lo, Input1Hi) = DAG.SplitVector(N->getOperand(1), DL);
+  std::tie(Input2Lo, Input2Hi) = DAG.SplitVector(N->getOperand(2), DL);
+  unsigned Opcode = N->getOpcode();
+  EVT ResultVT = Acc.getValueType();
+
+  SDValue Lo = DAG.getNode(Opcode, DL, ResultVT, Acc, Input1Lo, Input2Lo);
+  return DAG.getNode(Opcode, DL, ResultVT, Lo, Input1Hi, Input2Hi);
 }
 
 //===----------------------------------------------------------------------===//