MUL instructions now included in DAG combines.

JamesChesterman · JamesChesterman · commit 60924265c42d · 2025-01-20T15:28:54.000Z
diff --git a/llvm/include/llvm/CodeGen/SelectionDAG.h b/llvm/include/llvm/CodeGen/SelectionDAG.h
@@ -1608,7 +1608,13 @@ class SelectionDAG {
   /// \p Op1 Accumulator for where the result is stored for the partial
   /// reduction operation.
   /// \p Op2 Input for the partial reduction operation.
-  SDValue expandPartialReduceAdd(SDLoc DL, SDValue Op1, SDValue Op2);
+  /// Expands PARTIAL_REDUCE_S/UMLA nodes.
+  /// \p Acc Accumulator for where the result is stored for the partial
+  /// reduction operation.
+  /// \p Input1 First input for the partial reduction operation.
+  /// \p Input2 Second input for the partial reduction operation.
+  SDValue expandPartialReduceAdd(SDLoc DL, SDValue Acc, SDValue Input1,
+                                 SDValue Input2);
 
   /// 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
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -2467,20 +2467,24 @@ SDValue SelectionDAG::getShiftAmountOperand(EVT LHSTy, SDValue Op) {
   return getZExtOrTrunc(Op, SDLoc(Op), ShTy);
 }
 
-SDValue SelectionDAG::expandPartialReduceAdd(SDLoc DL, SDValue Op1,
-                                             SDValue Op2) {
-  EVT ReducedTy = Op1.getValueType();
-  EVT FullTy = Op2.getValueType();
+SDValue SelectionDAG::expandPartialReduceAdd(SDLoc DL, SDValue Acc,
+                                             SDValue Input1, SDValue Input2) {
+
+  EVT FullTy = Input1.getValueType();
+  Input2 = getAnyExtOrTrunc(Input2, DL, FullTy);
+  SDValue Input = getNode(ISD::MUL, DL, FullTy, Input1, Input2);
+
+  EVT ReducedTy = Acc.getValueType();
 
   unsigned Stride = ReducedTy.getVectorMinNumElements();
   unsigned ScaleFactor = FullTy.getVectorMinNumElements() / Stride;
 
   // Collect all of the subvectors
-  std::deque<SDValue> Subvectors = {Op1};
+  std::deque<SDValue> Subvectors = {Acc};
   for (unsigned I = 0; I < ScaleFactor; I++) {
     auto SourceIndex = getVectorIdxConstant(I * Stride, DL);
     Subvectors.push_back(
-        getNode(ISD::EXTRACT_SUBVECTOR, DL, ReducedTy, {Op2, SourceIndex}));
+        getNode(ISD::EXTRACT_SUBVECTOR, DL, ReducedTy, {Input, SourceIndex}));
   }
 
   // Flatten the subvector tree
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -8124,11 +8124,13 @@ void SelectionDAGBuilder::visitIntrinsicCall(const CallInst &I,
     SDValue Input = getValue(I.getOperand(1));
 
     if (!TLI.shouldExpandPartialReductionIntrinsic(cast<IntrinsicInst>(&I))) {
-      setValue(&I,
-               DAG.getNode(ISD::PARTIAL_REDUCE_UMLA, dl, AccVT, Acc, Input));
+      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));
+    setValue(&I,
+             DAG.expandPartialReduceAdd(
+                 dl, Acc, Input, DAG.getConstant(1, dl, Input.getValueType())));
     return;
   }
   case Intrinsic::experimental_cttz_elts: {
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -22032,30 +22032,31 @@ static SDValue tryCombineWhileLo(SDNode *N,
   return SDValue(N, 0);
 }
 
-SDValue tryCombineToDotProduct(SDValue &Acc, SDValue &Input, SelectionDAG &DAG,
+SDValue tryCombineToDotProduct(SDValue &Acc, SDValue &Input1, SDValue &Input2,
+                               SelectionDAG &DAG,
                                const AArch64Subtarget *Subtarget, SDLoc &DL) {
   bool Scalable = Acc.getValueType().isScalableVector();
   if (Scalable && !Subtarget->isSVEorStreamingSVEAvailable())
     return SDValue();
   if (!Scalable && (!Subtarget->isNeonAvailable() || !Subtarget->hasDotProd()))
     return SDValue();
 
-  unsigned InputOpcode = Input->getOpcode();
+  unsigned Input1Opcode = Input1->getOpcode();
   EVT AccVT = Acc->getValueType(0);
   if (AccVT.getVectorElementCount() * 4 ==
-          Input->getValueType(0).getVectorElementCount() &&
-      InputOpcode != ISD::MUL)
-    return DAG.expandPartialReduceAdd(DL, Acc, Input);
-  if (InputOpcode != ISD::MUL)
+          Input1->getValueType(0).getVectorElementCount() &&
+      Input1Opcode != ISD::MUL)
+    return DAG.expandPartialReduceAdd(DL, Acc, Input1, Input2);
+  if (Input1Opcode != ISD::MUL)
     return SDValue();
 
-  auto A = Input->getOperand(0);
-  auto B = Input->getOperand(1);
+  auto A = Input1->getOperand(0);
+  auto B = Input1->getOperand(1);
   unsigned AOpcode = A->getOpcode();
   unsigned BOpcode = B->getOpcode();
 
   if (!ISD::isExtOpcode(AOpcode) || !ISD::isExtOpcode(BOpcode))
-    return DAG.expandPartialReduceAdd(DL, Acc, Input);
+    return DAG.expandPartialReduceAdd(DL, Acc, Input1, Input2);
 
   bool AIsSigned = AOpcode == ISD::SIGN_EXTEND;
   bool BIsSigned = BOpcode == ISD::SIGN_EXTEND;
@@ -22064,6 +22065,10 @@ SDValue tryCombineToDotProduct(SDValue &Acc, SDValue &Input, SelectionDAG &DAG,
   B = B->getOperand(0);
   EVT MulSrcVT = A.getValueType();
 
+  Input2 = DAG.getAnyExtOrTrunc(Input2, DL, MulSrcVT);
+  A = DAG.getNode(ISD::MUL, DL, MulSrcVT, A, Input2);
+  B = DAG.getNode(ISD::MUL, DL, MulSrcVT, B, Input2);
+
   // Dot products operate on chunks of four elements so there must be four times
   // as many elements in the wide type
   if (!(AccVT == MVT::nxv4i64 && MulSrcVT == MVT::nxv16i8) &&
@@ -22072,17 +22077,17 @@ SDValue tryCombineToDotProduct(SDValue &Acc, SDValue &Input, SelectionDAG &DAG,
       !(AccVT == MVT::v4i64 && MulSrcVT == MVT::v16i8) &&
       !(AccVT == MVT::v4i32 && MulSrcVT == MVT::v16i8) &&
       !(AccVT == MVT::v2i32 && MulSrcVT == MVT::v8i8))
-    return DAG.expandPartialReduceAdd(DL, Acc, Input);
+    return DAG.expandPartialReduceAdd(DL, Acc, Input1, Input2);
 
   unsigned DotOpcode = AIsSigned ? AArch64ISD::SDOT : AArch64ISD::UDOT;
   if (AIsSigned != BIsSigned) {
     if (!Subtarget->hasMatMulInt8())
-      return DAG.expandPartialReduceAdd(DL, Acc, Input);
+      return DAG.expandPartialReduceAdd(DL, Acc, Input1, Input2);
 
     bool Scalable = AccVT.isScalableVT();
     // There's no nxv2i64 version of usdot
     if (Scalable && AccVT != MVT::nxv4i32 && AccVT != MVT::nxv4i64)
-      return DAG.expandPartialReduceAdd(DL, Acc, Input);
+      return DAG.expandPartialReduceAdd(DL, Acc, Input1, Input2);
 
     if (!BIsSigned)
       std::swap(A, B);
@@ -22109,32 +22114,37 @@ SDValue tryCombineToDotProduct(SDValue &Acc, SDValue &Input, SelectionDAG &DAG,
   }
 
   if (A.getValueType() != B.getValueType())
-    return DAG.expandPartialReduceAdd(DL, Acc, Input);
+    return DAG.expandPartialReduceAdd(DL, Acc, Input1, Input2);
 
   unsigned NewOpcode =
       AIsSigned ? ISD::PARTIAL_REDUCE_SMLA : ISD::PARTIAL_REDUCE_UMLA;
   return DAG.getNode(NewOpcode, DL, AccVT, Acc, A, B);
 }
 
-SDValue tryCombineToWideAdd(SDValue &Acc, SDValue &Input, SelectionDAG &DAG,
+SDValue tryCombineToWideAdd(SDValue &Acc, SDValue &Input1, SDValue &Input2,
+                            SelectionDAG &DAG,
                             const AArch64Subtarget *Subtarget, SDLoc &DL) {
   if (!Subtarget->hasSVE2() && !Subtarget->isStreamingSVEAvailable())
-    return DAG.expandPartialReduceAdd(DL, Acc, Input);
-  unsigned InputOpcode = Input->getOpcode();
-  if (!ISD::isExtOpcode(InputOpcode))
-    return DAG.expandPartialReduceAdd(DL, Acc, Input);
-  Input = Input->getOperand(0);
-  EVT InputVT = Input.getValueType();
+    return DAG.expandPartialReduceAdd(DL, Acc, Input1, Input2);
+  unsigned Input1Opcode = Input1->getOpcode();
+  if (!ISD::isExtOpcode(Input1Opcode))
+    return DAG.expandPartialReduceAdd(DL, Acc, Input1, Input2);
+
   EVT AccVT = Acc->getValueType(0);
+  Input1 = Input1->getOperand(0);
+  EVT InputVT = Input1.getValueType();
+  Input2 = DAG.getAnyExtOrTrunc(Input2, DL, InputVT);
+  SDValue Input = DAG.getNode(ISD::MUL, DL, InputVT, Input1, Input2);
+
   if (!AccVT.isScalableVector())
-    return DAG.expandPartialReduceAdd(DL, Acc, Input);
+    return DAG.expandPartialReduceAdd(DL, Acc, Input1, Input2);
 
   if (!(InputVT == MVT::nxv4i32 && AccVT == MVT::nxv2i64) &&
       !(InputVT == MVT::nxv8i16 && AccVT == MVT::nxv4i32) &&
       !(InputVT == MVT::nxv16i8 && AccVT == MVT::nxv8i16))
     return SDValue();
 
-  unsigned NewOpcode = InputOpcode == ISD::SIGN_EXTEND
+  unsigned NewOpcode = Input1Opcode == ISD::SIGN_EXTEND
                            ? ISD::PARTIAL_REDUCE_SMLA
                            : ISD::PARTIAL_REDUCE_UMLA;
   return DAG.getNode(NewOpcode, DL, AccVT, Acc, Input,
@@ -22145,18 +22155,21 @@ SDValue performPartialReduceAddCombine(SDNode *N, SelectionDAG &DAG,
                                        const AArch64Subtarget *Subtarget) {
   SDLoc DL(N);
   auto Acc = N->getOperand(0);
-  auto Input = N->getOperand(1);
+  auto Input1 = N->getOperand(1);
+  auto Input2 = N->getOperand(2);
   EVT AccElemVT = Acc.getValueType().getVectorElementType();
-  EVT InputElemVT = Input.getValueType().getVectorElementType();
+  EVT InputElemVT = Input1.getValueType().getVectorElementType();
 
   // If the exts have already been removed or it has already been lowered to an
   // usdot instruction, then the element types will not be equal
-  if (InputElemVT != AccElemVT || Input.getOpcode() == AArch64ISD::USDOT)
+  if (InputElemVT != AccElemVT || Input1.getOpcode() == AArch64ISD::USDOT)
     return SDValue(N, 0);
 
-  if (auto Dot = tryCombineToDotProduct(Acc, Input, DAG, Subtarget, DL))
+  if (auto Dot =
+          tryCombineToDotProduct(Acc, Input1, Input2, DAG, Subtarget, DL))
     return Dot;
-  if (auto WideAdd = tryCombineToWideAdd(Acc, Input, DAG, Subtarget, DL))
+  if (auto WideAdd =
+          tryCombineToWideAdd(Acc, Input1, Input2, DAG, Subtarget, DL))
     return WideAdd;
   return SDValue();
 }
