Reassoc FMF should not optimize FMA(a, 0, b) to (b)

[svenvh]

Summary

Optimizing (a * 0 + b) to (b) requires assuming that a is finite and not
NaN. DAGCombiner will do this optimization when the reassoc fast math
flag is set, which is not correct. Change DAGCombiner to only consider
UnsafeMath for this optimization.

Co-authored-by: Andrea Faulds <[e-mail]>

[spatel]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15029–15030

We are slowly trying to get away from using the global settings by replacing those checks with fast-math-flags. Can we do that here?

This should match the conditions that we use to simplify X * 0.0. That's done via DAG.simplifyFPBinop():

// X * 0.0 --> 0.0
if (Opcode == ISD::FMUL && Flags.hasNoNaNs() && Flags.hasNoSignedZeros())
  if (YC->getValueAPF().isZero())
    return getConstantFP(0.0, SDLoc(Y), Y.getValueType());

[svenvh]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15029–15030

Sure, but is it okay to do the modernization in a separate followup commit? This change just intends to fix the immediate bug (i.e. restore the status quo to before 8e570c3390d20, which is where this bug was introduced).

I guess the followup patch will be like this then?

@@ -15026,7 +15026,7 @@ SDValue DAGCombiner::visitFMA(SDNode *N) {
        CostN1 == TargetLowering::NegatibleCost::Cheaper))
     return DAG.getNode(ISD::FMA, DL, VT, NegN0, NegN1, N2);
 
-  if (Options.UnsafeFPMath) {
+  if (N->getFlags().hasNoNaNs() && N->getFlags().hasNoSignedZeros()) {
     if (N0CFP && N0CFP->isZero())
       return N2;
     if (N1CFP && N1CFP->isZero())

[spatel]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15029–15030

Right - it should be a one-line patch. Just need to see if any regression tests fail with that. If you want to do that as another patch, please put a FIXME note on it in this patch, so we don't lose track that it needs to be fixed.

llvm/test/CodeGen/AArch64/neon-fma-FMF.ll

64

Prefer to add positive CHECK lines for correctness rather than NOT lines for the absence of one particular kind of wrong.
It would be great to update this whole file using the script at "utils/update_llc_test_checks.py" - that way, it's complete and easier to update.

[svenvh]

Added FIXME and updated CHECK lines to avoid CHECK-NOT.

llvm/test/CodeGen/AArch64/neon-fma-FMF.ll

64

It would be great to update this whole file using the script at "utils/update_llc_test_checks.py"

I've tried that, but it seems to remove some CHECK patterns presumably because the IR is not exactly the same for the two llc RUN invocations (yet we still want the test to verify that fma fusing happened).

Either way, I've updated the CHECK lines for this function with what update_llc_test_checks.py generates.

[spatel]

LGTM

llvm/test/CodeGen/AArch64/neon-fma-FMF.ll

64

OK - thanks for trying.

Probably need to add unique prefixes for each RUN line (presumably there is some potential difference between the runs that someone was trying to test for). There are many test file examples of that in this directory if you want to do that as a follow-up.

I assume we'll add another test like the one you're adding in this patch -- except with nsz + nnan FMF to match the code when it is updated again.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitFMA(SDNode *N) {
   ConstantFPSDNode *N0CFP = dyn_cast<ConstantFPSDNode>(N0);
   ConstantFPSDNode *N1CFP = dyn_cast<ConstantFPSDNode>(N1);
   EVT VT = N->getValueType(0);
   SDLoc DL(N);
   const TargetOptions &Options = DAG.getTarget().Options;
   // FMA nodes have flags that propagate to the created nodes.
   SelectionDAG::FlagInserter FlagsInserter(DAG, N);
 
-  bool UnsafeFPMath =
+  bool CanReassociate =
       Options.UnsafeFPMath || N->getFlags().hasAllowReassociation();
 
   // Constant fold FMA.
   if (isa<ConstantFPSDNode>(N0) &&
       isa<ConstantFPSDNode>(N1) &&
       isa<ConstantFPSDNode>(N2)) {
     return DAG.getNode(ISD::FMA, DL, VT, N0, N1, N2);
   }
 
   // (-N0 * -N1) + N2 --> (N0 * N1) + N2
   TargetLowering::NegatibleCost CostN0 =
       TargetLowering::NegatibleCost::Expensive;
   TargetLowering::NegatibleCost CostN1 =
       TargetLowering::NegatibleCost::Expensive;
   SDValue NegN0 =
       TLI.getNegatedExpression(N0, DAG, LegalOperations, ForCodeSize, CostN0);
   SDValue NegN1 =
       TLI.getNegatedExpression(N1, DAG, LegalOperations, ForCodeSize, CostN1);
   if (NegN0 && NegN1 &&
       (CostN0 == TargetLowering::NegatibleCost::Cheaper ||
        CostN1 == TargetLowering::NegatibleCost::Cheaper))
     return DAG.getNode(ISD::FMA, DL, VT, NegN0, NegN1, N2);
 
-  if (UnsafeFPMath) {
+  // FIXME: use fast math flags instead of Options.UnsafeFPMath
+  if (Options.UnsafeFPMath) {
     if (N0CFP && N0CFP->isZero())
       return N2;
     if (N1CFP && N1CFP->isZero())
       return N2;
   }
 
   if (N0CFP && N0CFP->isExactlyValue(1.0))
     return DAG.getNode(ISD::FADD, SDLoc(N), VT, N1, N2);
   if (N1CFP && N1CFP->isExactlyValue(1.0))
     return DAG.getNode(ISD::FADD, SDLoc(N), VT, N0, N2);
 
   // Canonicalize (fma c, x, y) -> (fma x, c, y)
   if (DAG.isConstantFPBuildVectorOrConstantFP(N0) &&
      !DAG.isConstantFPBuildVectorOrConstantFP(N1))
     return DAG.getNode(ISD::FMA, SDLoc(N), VT, N1, N0, N2);
 
-  if (UnsafeFPMath) {
+  if (CanReassociate) {
     // (fma x, c1, (fmul x, c2)) -> (fmul x, c1+c2)
     if (N2.getOpcode() == ISD::FMUL && N0 == N2.getOperand(0) &&
         DAG.isConstantFPBuildVectorOrConstantFP(N1) &&
         DAG.isConstantFPBuildVectorOrConstantFP(N2.getOperand(1))) {
       return DAG.getNode(ISD::FMUL, DL, VT, N0,
                          DAG.getNode(ISD::FADD, DL, VT, N1, N2.getOperand(1)));
     }
 
if (N0.getOpcode() == ISD::FNEG &&
         (TLI.isOperationLegal(ISD::ConstantFP, VT) ||
          (N1.hasOneUse() && !TLI.isFPImmLegal(N1CFP->getValueAPF(), VT,
                                               ForCodeSize)))) {
       return DAG.getNode(ISD::FMA, DL, VT, N0.getOperand(0),
                          DAG.getNode(ISD::FNEG, DL, VT, N1), N2);
     }
   }
 
-  if (UnsafeFPMath) {
+  if (CanReassociate) {
     // (fma x, c, x) -> (fmul x, (c+1))
     if (N1CFP && N0 == N2) {
       return DAG.getNode(
           ISD::FMUL, DL, VT, N0,
           DAG.getNode(ISD::FADD, DL, VT, N1, DAG.getConstantFP(1.0, DL, VT)));
     }
 
     // (fma x, c, (fneg x)) -> (fmul x, (c-1))

llvm/test/CodeGen/AArch64/neon-fma-FMF.ll

 define <2 x float> @no_fma_sub_1(<2 x float> %A, <2 x float> %B, <2 x float> %C) {
 ; CHECK-LABEL: no_fma_sub_1:
 ; CHECK: fmul
 ; CHECK: fsub
   %tmp1 = fmul contract <2 x float> %A, %B;
   %tmp2 = fsub <2 x float> %C, %tmp1;
   ret <2 x float> %tmp2
 }
+
+; Regression test: contract FMF allows folding (A * 0 + B) to FMA(A, 0, B), but
+; reassoc FMF must not allow further folding to just (B) without additional
+; FMFs (ninf, nnan)
+define float @fma_zero(float %A, float %B) {
+; CHECK-LABEL: fma_zero:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    movi d2, #0000000000000000
+; CHECK-NEXT:    fmadd s0, s0, s2, s1
+; CHECK-NEXT:    ret
+  %tmp1 = fmul contract reassoc float %A, 0.0e+0;
+  %tmp2 = fadd contract reassoc float %B, %tmp1;
+  ret float %tmp2
+}