Move fma expansions to default expand rule

Author: peterbell10

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -5325,6 +5325,23 @@ class TargetLowering : public TargetLoweringBase {
                            SDNodeFlags Flags, const SDLoc &DL,
                            SelectionDAG &DAG) const;
 
+  /// Expand floating point add
+  /// \param N Node to expand
+  /// \returns The expansion result or SDValue() if it fails.
+  SDValue expandFADD(SDNode *N, SelectionDAG &DAG) const;
+
+  /// Expand floating point multiply
+  /// \param N Node to expand
+  /// \param Result output after conversion
+  /// \returns The expansion result or SDValue() if it fails.
+  SDValue expandFMUL(SDNode *N, SelectionDAG &DAG) const;
+
+  /// Expand floating point subtract
+  /// \param N Node to expand
+  /// \param Result output after conversion
+  /// \returns The expansion result or SDValue() if it fails.
+  SDValue expandFSUB(SDNode *N, SelectionDAG &DAG) const;
+
   /// Expand CTPOP nodes. Expands vector/scalar CTPOP nodes,
   /// vector nodes can only succeed if all operations are legal/custom.
   /// \param N Node to expand

llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

@@ -3671,14 +3671,21 @@ bool SelectionDAGLegalize::ExpandNode(SDNode *Node) {
     Results.push_back(ExpandConstant(CP));
     break;
   }
+  case ISD::FADD: {
+    if (SDValue Expand = TLI.expandFADD(Node, DAG)) {
+      Results.push_back(Expand);
+    }
+    break;
+  }
+  case ISD::FMUL: {
+    if (SDValue Expand = TLI.expandFMUL(Node, DAG)) {
+      Results.push_back(Expand);
+    }
+    break;
+  }
   case ISD::FSUB: {
-    EVT VT = Node->getValueType(0);
-    if (TLI.isOperationLegalOrCustom(ISD::FADD, VT) &&
-        TLI.isOperationLegalOrCustom(ISD::FNEG, VT)) {
-      const SDNodeFlags Flags = Node->getFlags();
-      Tmp1 = DAG.getNode(ISD::FNEG, dl, VT, Node->getOperand(1));
-      Tmp1 = DAG.getNode(ISD::FADD, dl, VT, Node->getOperand(0), Tmp1, Flags);
-      Results.push_back(Tmp1);
+    if (SDValue Expand = TLI.expandFSUB(Node, DAG)) {
+      Results.push_back(Expand);
     }
     break;
   }

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

@@ -9070,6 +9070,60 @@ SDValue TargetLowering::expandIS_FPCLASS(EVT ResultVT, SDValue Op,
   return Res;
 }
 
+SDValue TargetLowering::expandFADD(SDNode *Node, SelectionDAG &DAG) const {
+  auto VT = Node->getValueType(0);
+  if (!isOperationLegalOrCustom(ISD::FMA, VT)) {
+    return {};
+  }
+
+  // FADD(a, b) -> FMA(a, 1.0, b)
+  SDLoc DL(Node);
+  auto One = DAG.getConstantFP(1.0, DL, VT);
+  SmallVector<SDValue, 3> Operands{Node->getOperand(0), One,
+                                   Node->getOperand(1)};
+  return DAG.getNode(ISD::FMA, DL, VT, Operands, Node->getFlags());
+}
+
+SDValue TargetLowering::expandFMUL(SDNode *Node, SelectionDAG &DAG) const {
+  auto VT = Node->getValueType(0);
+  if (!isOperationLegalOrCustom(ISD::FMA, VT)) {
+    return {};
+  }
+
+  // FMUL(a, b) -> FMA(a, b, -0.0)
+  // NOTE: The identity is -0, not 0, because -0 + 0 == 0 for floats
+  SDLoc DL(Node);
+  auto NegZero = DAG.getConstantFP(-0.0, DL, VT);
+  SmallVector<SDValue, 3> Operands{Node->getOperand(0), Node->getOperand(1),
+                                   NegZero};
+  return DAG.getNode(ISD::FMA, DL, VT, Operands, Node->getFlags());
+}
+
+SDValue TargetLowering::expandFSUB(SDNode *Node, SelectionDAG &DAG) const {
+  SDLoc DL(Node);
+  SDNodeFlags SDFlags = Node->getFlags();
+  auto VT = Node->getValueType(0);
+
+  bool CanUseFMA = isOperationLegalOrCustom(ISD::FMA, VT);
+  bool CanUseAddSub = (isOperationLegalOrCustom(ISD::FADD, VT) &&
+                       isOperationLegalOrCustom(ISD::FNEG, VT));
+  bool PreferAddSub = CanUseAddSub && isFNegFree(VT);
+
+  // FSUB(a, b) -> FMA(b, -1.0, a)
+  if (CanUseFMA && !PreferAddSub) {
+    auto NegOne = DAG.getConstantFP(-1.0, DL, VT);
+    SmallVector<SDValue, 3> Operands{Node->getOperand(1), NegOne,
+                                     Node->getOperand(0)};
+    return DAG.getNode(ISD::FMA, DL, VT, Operands, SDFlags);
+  }
+  // FSUB(a, b) -> FADD(a, FNEG(b))
+  if (CanUseAddSub) {
+    auto Neg = DAG.getNode(ISD::FNEG, DL, VT, Node->getOperand(1));
+    return DAG.getNode(ISD::FADD, DL, VT, Node->getOperand(0), Neg, SDFlags);
+  }
+  return {};
+}
+
 // Only expand vector types if we have the appropriate vector bit operations.
 static bool canExpandVectorCTPOP(const TargetLowering &TLI, EVT VT) {
   assert(VT.isVector() && "Expected vector type");

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

@@ -2520,11 +2520,7 @@ SDValue NVPTXTargetLowering::LowerFADD(SDValue Op, SelectionDAG &DAG) const {
   }
 
   // FADD(a, b) -> FMA(a, 1.0, b)
-  SDLoc DL(Op);
-  auto VT = Op.getValueType();
-  auto One = DAG.getConstantFP(1.0, DL, VT);
-  SmallVector<SDValue, 3> Operands{Op->getOperand(0), One, Op->getOperand(1)};
-  return DAG.getNode(ISD::FMA, DL, VT, Operands);
+  return expandFADD(Op.getNode(), DAG);
 }
 
 SDValue NVPTXTargetLowering::LowerFSUB(SDValue Op, SelectionDAG &DAG) const {
@@ -2534,12 +2530,7 @@ SDValue NVPTXTargetLowering::LowerFSUB(SDValue Op, SelectionDAG &DAG) const {
   }
 
   // FSUB(a, b) -> FMA(b, -1.0, a)
-  SDLoc DL(Op);
-  auto VT = Op.getValueType();
-  auto NegOne = DAG.getConstantFP(-1.0, DL, VT);
-  SmallVector<SDValue, 3> Operands{Op->getOperand(1), NegOne,
-                                   Op->getOperand(0)};
-  return DAG.getNode(ISD::FMA, DL, VT, Operands);
+  return expandFSUB(Op.getNode(), DAG);
 }
 
 SDValue NVPTXTargetLowering::LowerFMUL(SDValue Op, SelectionDAG &DAG) const {
@@ -2549,13 +2540,7 @@ SDValue NVPTXTargetLowering::LowerFMUL(SDValue Op, SelectionDAG &DAG) const {
   }
 
   // FMUL(a, b) -> FMA(a, b, -0.0)
-  // NOTE: The identity is -0, not 0, because -0 + 0 == 0 for floats
-  SDLoc DL(Op);
-  auto VT = Op.getValueType();
-  auto NegZero = DAG.getConstantFP(-0.0, DL, VT);
-  SmallVector<SDValue, 3> Operands{Op->getOperand(0), Op->getOperand(1),
-                                   NegZero};
-  return DAG.getNode(ISD::FMA, DL, VT, Operands);
+  return expandFMUL(Op.getNode(), DAG);
 }
 
 SDValue NVPTXTargetLowering::LowerINT_TO_FP(SDValue Op,