Address comments on patch. Remove shouldExpandPartialReductionIntrinsic().

Author: JamesChesterman

llvm/include/llvm/CodeGen/ISDOpcodes.h

@@ -1451,9 +1451,18 @@ enum NodeType {
   VECREDUCE_UMAX,
   VECREDUCE_UMIN,
 
-  // Nodes used to represent a partial reduction addition operation (signed and
-  // unsigned).
-  // Operands: Accumulator, Input
+  // Partial Reduction nodes. These represent multiply-add instructions because
+  // Input1 and Input2 are multiplied together first. This result is then
+  // reduced, by addition, to the number of elements that the Accumulator's type
+  // has.
+  // Input1 and Input2 must be the same type. Accumulator's element type must
+  // match that of Input1 and Input2. The number of elements in Input1 and
+  // Input2 must be a positive integer multiple of the number of elements in the
+  // Accumulator.
+  // The signedness of this node will dictate the signedness of nodes expanded
+  // from it. The signedness of the node is dictated by the signedness of
+  // Input1.
+  // Operands: Accumulator, Input1, Input2
   // Outputs: Output
   PARTIAL_REDUCE_SMLA,
   PARTIAL_REDUCE_UMLA,

llvm/include/llvm/CodeGen/SelectionDAG.h

@@ -1604,10 +1604,6 @@ class SelectionDAG {
   /// the target's desired shift amount type.
   SDValue getShiftAmountOperand(EVT LHSTy, SDValue Op);
 
-  /// Expands PARTIAL_REDUCE_S/UMLA nodes.
-  /// \p Op1 Accumulator for where the result is stored for the partial
-  /// reduction operation.
-  /// \p Op2 Input for the partial reduction operation.
   /// Expands PARTIAL_REDUCE_S/UMLA nodes.
   /// \p Acc Accumulator for where the result is stored for the partial
   /// reduction operation.

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -455,13 +455,6 @@ class TargetLoweringBase {
     return true;
   }
 
-  /// Return true if the @llvm.experimental.vector.partial.reduce.* intrinsic
-  /// should be expanded using generic code in SelectionDAGBuilder.
-  virtual bool
-  shouldExpandPartialReductionIntrinsic(const IntrinsicInst *I) const {
-    return true;
-  }
-
   /// Return true if the @llvm.get.active.lane.mask intrinsic should be expanded
   /// using generic code in SelectionDAGBuilder.
   virtual bool shouldExpandGetActiveLaneMask(EVT VT, EVT OpVT) const {

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -8122,15 +8122,22 @@ void SelectionDAGBuilder::visitIntrinsicCall(const CallInst &I,
     SDValue Acc = getValue(I.getOperand(0));
     EVT AccVT = Acc.getValueType();
     SDValue Input = getValue(I.getOperand(1));
-
-    if (!TLI.shouldExpandPartialReductionIntrinsic(cast<IntrinsicInst>(&I))) {
-      setValue(&I, DAG.getNode(ISD::PARTIAL_REDUCE_UMLA, dl, AccVT, Acc, Input,
-                               DAG.getConstant(1, dl, Input.getValueType())));
-      return;
-    }
-    setValue(&I,
-             DAG.expandPartialReduceAdd(
-                 dl, Acc, Input, DAG.getConstant(1, dl, Input.getValueType())));
+    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 used instead. It should be changed
+    // to its correct signedness when combining or expanding, according to
+    // extends being performed on Input.
+    setValue(&I, DAG.getNode(ISD::PARTIAL_REDUCE_UMLA, dl, AccVT, Acc, Input,
+                             DAG.getConstant(1, dl, InputVT)));
     return;
   }
   case Intrinsic::experimental_cttz_elts: {

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -2049,28 +2049,6 @@ bool AArch64TargetLowering::shouldExpandGetActiveLaneMask(EVT ResVT,
   return false;
 }
 
-bool AArch64TargetLowering::shouldExpandPartialReductionIntrinsic(
-    const IntrinsicInst *I) const {
-  if (I->getIntrinsicID() != Intrinsic::experimental_vector_partial_reduce_add)
-    return true;
-
-  EVT VT = EVT::getEVT(I->getType());
-  auto Input = I->getOperand(1);
-  EVT InputVT = EVT::getEVT(Input->getType());
-
-  if ((InputVT == MVT::nxv4i64 && VT == MVT::nxv2i64) ||
-      (InputVT == MVT::nxv8i32 && VT == MVT::nxv4i32) ||
-      (InputVT == MVT::nxv16i16 && VT == MVT::nxv8i16) ||
-      (InputVT == MVT::nxv16i64 && VT == MVT::nxv4i64) ||
-      (InputVT == MVT::nxv16i32 && VT == MVT::nxv4i32) ||
-      (InputVT == MVT::nxv8i64 && VT == MVT::nxv2i64) ||
-      (InputVT == MVT::v16i64 && VT == MVT::v4i64) ||
-      (InputVT == MVT::v16i32 && VT == MVT::v4i32) ||
-      (InputVT == MVT::v8i32 && VT == MVT::v2i32))
-    return false;
-  return true;
-}
-
 bool AArch64TargetLowering::shouldExpandCttzElements(EVT VT) const {
   if (!Subtarget->isSVEorStreamingSVEAvailable())
     return true;
@@ -22037,9 +22015,9 @@ SDValue tryCombineToDotProduct(SDValue &Op0, SDValue &Op1, SDValue &Op2,
                                const AArch64Subtarget *Subtarget, SDLoc &DL) {
   bool Scalable = Op0->getValueType(0).isScalableVector();
   if (Scalable && !Subtarget->isSVEorStreamingSVEAvailable())
-    return DAG.expandPartialReduceAdd(DL, Op0, Op1, Op2);
+    return SDValue();
   if (!Scalable && (!Subtarget->isNeonAvailable() || !Subtarget->hasDotProd()))
-    return DAG.expandPartialReduceAdd(DL, Op0, Op1, Op2);
+    return SDValue();
 
   unsigned Op1Opcode = Op1->getOpcode();
   SDValue MulOpLHS, MulOpRHS;
@@ -22056,7 +22034,7 @@ SDValue tryCombineToDotProduct(SDValue &Op0, SDValue &Op1, SDValue &Op2,
     unsigned ExtMulOpRHSOpcode = ExtMulOpRHS->getOpcode();
     if (!ISD::isExtOpcode(ExtMulOpLHSOpcode) ||
         !ISD::isExtOpcode(ExtMulOpRHSOpcode))
-      return DAG.expandPartialReduceAdd(DL, Op0, Op1, Op2);
+      return SDValue();
 
     MulOpLHSIsSigned = ExtMulOpLHSOpcode == ISD::SIGN_EXTEND;
     MulOpRHSIsSigned = ExtMulOpRHSOpcode == ISD::SIGN_EXTEND;
@@ -22066,7 +22044,7 @@ SDValue tryCombineToDotProduct(SDValue &Op0, SDValue &Op1, SDValue &Op2,
     EVT MulOpLHSVT = MulOpLHS.getValueType();
 
     if (MulOpLHSVT != MulOpRHS.getValueType())
-      return DAG.expandPartialReduceAdd(DL, Op0, Op1, Op2);
+      return SDValue();
 
     Op2 = DAG.getAnyExtOrTrunc(Op2, DL, MulOpLHSVT);
     MulOpLHS = DAG.getNode(ISD::MUL, DL, MulOpLHSVT, MulOpLHS, Op2);
@@ -22092,12 +22070,12 @@ SDValue tryCombineToDotProduct(SDValue &Op0, SDValue &Op1, SDValue &Op2,
   unsigned DotOpcode = MulOpLHSIsSigned ? AArch64ISD::SDOT : AArch64ISD::UDOT;
   if (MulOpLHSIsSigned != MulOpRHSIsSigned) {
     if (!Subtarget->hasMatMulInt8())
-      return DAG.expandPartialReduceAdd(DL, Op0, Op1, Op2);
+      return SDValue();
 
     bool Scalable = ReducedVT.isScalableVT();
     // There's no nxv2i64 version of usdot
     if (Scalable && ReducedVT != MVT::nxv4i32 && ReducedVT != MVT::nxv4i64)
-      return DAG.expandPartialReduceAdd(DL, Op0, Op1, Op2);
+      return SDValue();
 
     if (!MulOpRHSIsSigned)
       std::swap(MulOpLHS, MulOpRHS);
@@ -22134,10 +22112,10 @@ SDValue tryCombineToWideAdd(SDValue &Op0, SDValue &Op1, SDValue &Op2,
                             SelectionDAG &DAG,
                             const AArch64Subtarget *Subtarget, SDLoc &DL) {
   if (!Subtarget->hasSVE2() && !Subtarget->isStreamingSVEAvailable())
-    return DAG.expandPartialReduceAdd(DL, Op0, Op1, Op2);
+    return SDValue();
   unsigned Op1Opcode = Op1->getOpcode();
   if (!ISD::isExtOpcode(Op1Opcode))
-    return DAG.expandPartialReduceAdd(DL, Op0, Op1, Op2);
+    return SDValue();
 
   EVT AccVT = Op0->getValueType(0);
   Op1 = Op1->getOperand(0);
@@ -22146,7 +22124,7 @@ SDValue tryCombineToWideAdd(SDValue &Op0, SDValue &Op1, SDValue &Op2,
   SDValue Input = DAG.getNode(ISD::MUL, DL, Op1VT, Op1, Op2);
 
   if (!AccVT.isScalableVector())
-    return DAG.expandPartialReduceAdd(DL, Op0, Op1, Op2);
+    return SDValue();
 
   if (!(Op1VT == MVT::nxv4i32 && AccVT == MVT::nxv2i64) &&
       !(Op1VT == MVT::nxv8i16 && AccVT == MVT::nxv4i32) &&
@@ -22177,7 +22155,10 @@ SDValue performPartialReduceAddCombine(SDNode *N, SelectionDAG &DAG,
     return Dot;
   if (auto WideAdd = tryCombineToWideAdd(Op0, Op1, Op2, DAG, Subtarget, DL))
     return WideAdd;
-  return SDValue();
+  // N->getOperand needs calling again because the Op variables may have been
+  // changed by the functions above
+  return DAG.expandPartialReduceAdd(DL, N->getOperand(0), N->getOperand(1),
+                                    N->getOperand(2));
 }
 
 static SDValue performIntrinsicCombine(SDNode *N,

llvm/lib/Target/AArch64/AArch64ISelLowering.h

@@ -993,9 +993,6 @@ class AArch64TargetLowering : public TargetLowering {
 
   bool shouldExpandGetActiveLaneMask(EVT VT, EVT OpVT) const override;
 
-  bool
-  shouldExpandPartialReductionIntrinsic(const IntrinsicInst *I) const override;
-
   bool shouldExpandCttzElements(EVT VT) const override;
 
   bool shouldExpandVectorMatch(EVT VT, unsigned SearchSize) const override;