Address comments

Move the getPartialReduceAdd function around.
Make the new codepath work for fixed length NEON vectors too.

Author: JamesChesterman

llvm/include/llvm/CodeGen/SelectionDAG.h

@@ -1607,11 +1607,6 @@ class SelectionDAG {
   /// the target's desired shift amount type.
   SDValue getShiftAmountOperand(EVT LHSTy, SDValue Op);
 
-  /// Create the DAG equivalent of vector_partial_reduce where Op1 and Op2 are
-  /// its operands and ReducedTY is the intrinsic's return type.
-  SDValue getPartialReduceAdd(SDLoc DL, EVT ReducedTy, SDValue Op1,
-                              SDValue Op2);
-
   /// Expands a node with multiple results to an FP or vector libcall. The
   /// libcall is expected to take all the operands of the \p Node followed by
   /// output pointers for each of the results. \p CallRetResNo can be optionally

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -5636,7 +5636,14 @@ class TargetLowering : public TargetLoweringBase {
 
   // Expands PARTIAL_REDUCE_S/UMLA nodes to a series of simpler operations,
   // consisting of zext/sext, extract_subvector, mul and add operations.
-  SDValue expandPartialReduceMLA(SDNode *N, SelectionDAG &DAG) const;
+  SDValue expandPartialReduceMLA(SDLoc DL, SDValue Acc, SDValue Input1,
+                                 SDValue Input2, SelectionDAG &DAG) const;
+
+  // Create the DAG equivalent of vector_partial_reduce where Op1 and Op2 are
+  // its operands and ReducedTY is the return type.
+  static SDValue getPartialReduceAdd(SDLoc DL, EVT ReducedTy, EVT FullTy,
+                                     SDValue Op1, SDValue Op2,
+                                     SelectionDAG &DAG);
 
 private:
   SDValue foldSetCCWithAnd(EVT VT, SDValue N0, SDValue N1, ISD::CondCode Cond,

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorOps.cpp

@@ -1198,9 +1198,13 @@ void VectorLegalizer::Expand(SDNode *Node, SmallVectorImpl<SDValue> &Results) {
     Results.push_back(TLI.expandVecReduce(Node, DAG));
     return;
   case ISD::PARTIAL_REDUCE_UMLA:
-  case ISD::PARTIAL_REDUCE_SMLA:
-    Results.push_back(TLI.expandPartialReduceMLA(Node, DAG));
+  case ISD::PARTIAL_REDUCE_SMLA: {
+    SDLoc DL(Node);
+    Results.push_back(TLI.expandPartialReduceMLA(DL, Node->getOperand(0),
+                                                 Node->getOperand(1),
+                                                 Node->getOperand(2), DAG));
     return;
+  }
   case ISD::VECREDUCE_SEQ_FADD:
   case ISD::VECREDUCE_SEQ_FMUL:
     Results.push_back(TLI.expandVecReduceSeq(Node, DAG));

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

@@ -3186,7 +3186,9 @@ void DAGTypeLegalizer::SplitVecRes_VP_REVERSE(SDNode *N, SDValue &Lo,
 }
 
 void DAGTypeLegalizer::SplitVecRes_PARTIAL_REDUCE_MLA(SDNode *N) {
-  SDValue Res = TLI.expandPartialReduceMLA(N, DAG);
+  SDLoc DL(N);
+  SDValue Res = TLI.expandPartialReduceMLA(
+      DL, N->getOperand(0), N->getOperand(1), N->getOperand(2), DAG);
   ReplaceValueWith(SDValue(N, 0), Res);
 }
 
@@ -4447,7 +4449,9 @@ SDValue DAGTypeLegalizer::SplitVecOp_VECTOR_HISTOGRAM(SDNode *N) {
 }
 
 SDValue DAGTypeLegalizer::SplitVecOp_PARTIAL_REDUCE_MLA(SDNode *N) {
-  SDValue Res = TLI.expandPartialReduceMLA(N, DAG);
+  SDLoc DL(N);
+  SDValue Res = TLI.expandPartialReduceMLA(
+      DL, N->getOperand(0), N->getOperand(1), N->getOperand(2), DAG);
   ReplaceValueWith(SDValue(N, 0), Res);
   return SDValue();
 }

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

@@ -2474,35 +2474,6 @@ SDValue SelectionDAG::getShiftAmountOperand(EVT LHSTy, SDValue Op) {
   return getZExtOrTrunc(Op, SDLoc(Op), ShTy);
 }
 
-SDValue SelectionDAG::getPartialReduceAdd(SDLoc DL, EVT ReducedTy, SDValue Op1,
-                                          SDValue Op2) {
-  EVT FullTy = Op2.getValueType();
-
-  unsigned Stride = ReducedTy.getVectorMinNumElements();
-  unsigned ScaleFactor = FullTy.getVectorMinNumElements() / Stride;
-
-  // Collect all of the subvectors
-  std::deque<SDValue> Subvectors = {Op1};
-  for (unsigned I = 0; I < ScaleFactor; I++) {
-    auto SourceIndex = getVectorIdxConstant(I * Stride, DL);
-    Subvectors.push_back(
-        getNode(ISD::EXTRACT_SUBVECTOR, DL, ReducedTy, {Op2, SourceIndex}));
-  }
-
-  // Flatten the subvector tree
-  while (Subvectors.size() > 1) {
-    Subvectors.push_back(
-        getNode(ISD::ADD, DL, ReducedTy, {Subvectors[0], Subvectors[1]}));
-    Subvectors.pop_front();
-    Subvectors.pop_front();
-  }
-
-  assert(Subvectors.size() == 1 &&
-         "There should only be one subvector after tree flattening");
-
-  return Subvectors[0];
-}
-
 /// Given a store node \p StoreNode, return true if it is safe to fold that node
 /// into \p FPNode, which expands to a library call with output pointers.
 static bool canFoldStoreIntoLibCallOutputPointers(StoreSDNode *StoreNode,

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -8125,21 +8125,19 @@ void SelectionDAGBuilder::visitIntrinsicCall(const CallInst &I,
     return;
   }
   case Intrinsic::experimental_vector_partial_reduce_add: {
+    SDValue Acc = getValue(I.getOperand(0));
+    EVT AccVT = Acc.getValueType();
+    SDValue Input = getValue(I.getOperand(1));
+    EVT InputVT = Input.getValueType();
+
+    assert(AccVT.getVectorElementType() == InputVT.getVectorElementType() &&
+           "Expected operands to have the same vector element type!");
+    assert(InputVT.getVectorElementCount().getKnownMinValue() %
+                   AccVT.getVectorElementCount().getKnownMinValue() ==
+               0 &&
+           "Expected the element count of the Input operand to be a positive "
+           "integer multiple of the element count of the Accumulator operand!");
     if (NewPartialReduceLowering) {
-      SDValue Acc = getValue(I.getOperand(0));
-      EVT AccVT = Acc.getValueType();
-      SDValue Input = getValue(I.getOperand(1));
-      EVT InputVT = Input.getValueType();
-
-      assert(AccVT.getVectorElementType() == InputVT.getVectorElementType() &&
-             "Expected operands to have the same vector element type!");
-      assert(
-          InputVT.getVectorElementCount().getKnownMinValue() %
-                  AccVT.getVectorElementCount().getKnownMinValue() ==
-              0 &&
-          "Expected the element count of the Input operand to be a positive "
-          "integer multiple of the element count of the Accumulator operand!");
-
       // ISD::PARTIAL_REDUCE_UMLA is chosen arbitrarily and would function the
       // same if ISD::PARTIAL_REDUCE_SMLA was chosen instead. It should be
       // changed to its correct signedness when combining or expanding,
@@ -8154,9 +8152,8 @@ void SelectionDAGBuilder::visitIntrinsicCall(const CallInst &I,
       return;
     }
 
-    setValue(&I, DAG.getPartialReduceAdd(sdl, EVT::getEVT(I.getType()),
-                                         getValue(I.getOperand(0)),
-                                         getValue(I.getOperand(1))));
+    setValue(&I, TLI.expandPartialReduceMLA(
+                     sdl, Acc, Input, DAG.getConstant(1, sdl, InputVT), DAG));
     return;
   }
   case Intrinsic::experimental_cttz_elts: {

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

@@ -34,6 +34,7 @@
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Target/TargetMachine.h"
 #include <cctype>
+#include <deque>
 using namespace llvm;
 
 /// NOTE: The TargetMachine owns TLOF.
@@ -12189,20 +12190,15 @@ SDValue TargetLowering::scalarizeExtractedVectorLoad(EVT ResultVT,
   return Load;
 }
 
-SDValue TargetLowering::expandPartialReduceMLA(SDNode *N,
+SDValue TargetLowering::expandPartialReduceMLA(SDLoc DL, SDValue Acc,
+                                               SDValue Input1, SDValue Input2,
                                                SelectionDAG &DAG) const {
-  SDLoc DL(N);
-  SDValue Acc = N->getOperand(0);
-  SDValue Input1 = N->getOperand(1);
-  SDValue Input2 = N->getOperand(2);
-
   EVT ReducedTy = Acc.getValueType();
   EVT FullTy = Input1.getValueType();
 
   auto ExtendToAccEltVT = [&](SDValue V) {
-    unsigned ExtOpc = N->getOpcode() == ISD::PARTIAL_REDUCE_UMLA
-                          ? ISD::ZERO_EXTEND
-                          : ISD::SIGN_EXTEND;
+    unsigned ExtOpc = V->getOpcode() == ISD::SIGN_EXTEND ? ISD::SIGN_EXTEND
+                                                         : ISD::ZERO_EXTEND;
     EVT ExtVT = V.getValueType().changeVectorElementType(
         Acc.getValueType().getVectorElementType());
     if (ExtVT != FullTy)
@@ -12224,5 +12220,34 @@ SDValue TargetLowering::expandPartialReduceMLA(SDNode *N,
     Input = ExtendToAccEltVT(Input1);
   }
 
-  return DAG.getPartialReduceAdd(DL, ReducedTy, Acc, Input);
+  return TargetLowering::getPartialReduceAdd(DL, ReducedTy, FullTy, Acc, Input,
+                                             DAG);
+}
+
+SDValue TargetLowering::getPartialReduceAdd(SDLoc DL, EVT ReducedTy, EVT FullTy,
+                                            SDValue Op1, SDValue Op2,
+                                            SelectionDAG &DAG) {
+  unsigned Stride = ReducedTy.getVectorMinNumElements();
+  unsigned ScaleFactor = FullTy.getVectorMinNumElements() / Stride;
+
+  // Collect all of the subvectors
+  std::deque<SDValue> Subvectors = {Op1};
+  for (unsigned I = 0; I < ScaleFactor; I++) {
+    auto SourceIndex = DAG.getVectorIdxConstant(I * Stride, DL);
+    Subvectors.push_back(
+        DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, ReducedTy, {Op2, SourceIndex}));
+  }
+
+  // Flatten the subvector tree
+  while (Subvectors.size() > 1) {
+    Subvectors.push_back(
+        DAG.getNode(ISD::ADD, DL, ReducedTy, {Subvectors[0], Subvectors[1]}));
+    Subvectors.pop_front();
+    Subvectors.pop_front();
+  }
+
+  assert(Subvectors.size() == 1 &&
+         "There should only be one subvector after tree flattening");
+
+  return Subvectors[0];
 }

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -1358,6 +1358,9 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
       setOperationAction(ISD::BSWAP, VT, Expand);
       setOperationAction(ISD::CTTZ, VT, Expand);
 
+      setOperationAction(ISD::PARTIAL_REDUCE_UMLA, VT, Expand);
+      setOperationAction(ISD::PARTIAL_REDUCE_SMLA, VT, Expand);
+
       for (MVT InnerVT : MVT::fixedlen_vector_valuetypes()) {
         setTruncStoreAction(VT, InnerVT, Expand);
         setLoadExtAction(ISD::SEXTLOAD, VT, InnerVT, Expand);
@@ -22015,8 +22018,10 @@ static SDValue performIntrinsicCombine(SDNode *N,
       return Dot;
     if (SDValue WideAdd = tryLowerPartialReductionToWideAdd(N, Subtarget, DAG))
       return WideAdd;
-    return DAG.getPartialReduceAdd(SDLoc(N), N->getValueType(0),
-                                   N->getOperand(1), N->getOperand(2));
+    SDValue Input = N->getOperand(2);
+    return TargetLowering::getPartialReduceAdd(SDLoc(N), N->getValueType(0),
+                                               Input.getValueType(),
+                                               N->getOperand(1), Input, DAG);
   }
   case Intrinsic::aarch64_neon_vcvtfxs2fp:
   case Intrinsic::aarch64_neon_vcvtfxu2fp: