[DAG] Update `isKnownNeverNaN` for `FMA/FMAD`

[Pierre-vh]

Summary

We can still get a NaN even if none of the operands are NaN,
e.g. from +inf/-inf. D50804 didn't catch that.

[foad]

Code change looks good. Is it possible to add a test as described in https://reviews.llvm.org/D50804#1202429?

[Pierre-vh]

In D134854#3823061, @foad wrote:

Code change looks good. Is it possible to add a test as described in https://reviews.llvm.org/D50804#1202429?

Do you mean we need the same test but in AMDGPU, or does the test need to be changed as well?

[foad]

In D134854#3823109, @Pierre-vh wrote:

In D134854#3823061, @foad wrote:

Code change looks good. Is it possible to add a test as described in https://reviews.llvm.org/D50804#1202429?

Do you mean we need the same test but in AMDGPU, or does the test need to be changed as well?

The test added in D50804 was for fadd and fmul. I'm suggesting you could extend it now to cover fma and fmad as well.

[Pierre-vh]

I did some tries but I can't get a DAG combine to involve FMA/FMAD in DAGISel that's affected by this change (I tried both on AArch64 and AMDGPU). Any idea of how I can write such a test and what the expected output should be?
The only combine I found involving this function only works if -enable-no-signed-zeros-fp-math is used

[arsenm]

In D134854#3836573, @Pierre-vh wrote:

I did some tries but I can't get a DAG combine to involve FMA/FMAD in DAGISel that's affected by this change (I tried both on AArch64 and AMDGPU). Any idea of how I can write such a test and what the expected output should be?
The only combine I found involving this function only works if -enable-no-signed-zeros-fp-math is used

Then use "no-signed-zeros-fp-math"="true" in the test. That combine should also be updated to use nsz flags

[Pierre-vh]

Add fma test

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

bool SelectionDAG::isKnownNeverNaN(SDValue Op, bool SNaN, unsigned Depth) const {
   unsigned Opcode = Op.getOpcode();
   switch (Opcode) {
   case ISD::FADD:
   case ISD::FSUB:
   case ISD::FMUL:
   case ISD::FDIV:
   case ISD::FREM:
   case ISD::FSIN:
-  case ISD::FCOS: {
+  case ISD::FCOS:
+  case ISD::FMA:
+  case ISD::FMAD: {
     if (SNaN)
       return true;
     // TODO: Need isKnownNeverInfinity
     return false;
   }
   case ISD::FCANONICALIZE:
   case ISD::FEXP:
   case ISD::FEXP2:
bool SelectionDAG::isKnownNeverNaN(SDValue Op, bool SNaN, unsigned Depth) const {
   case ISD::FP_ROUND: {
     if (SNaN)
       return true;
     return isKnownNeverNaN(Op.getOperand(0), SNaN, Depth + 1);
   }
   case ISD::SINT_TO_FP:
   case ISD::UINT_TO_FP:
     return true;
-  case ISD::FMA:
-  case ISD::FMAD: {
-    if (SNaN)
-      return true;
-    return isKnownNeverNaN(Op.getOperand(0), SNaN, Depth + 1) &&
-           isKnownNeverNaN(Op.getOperand(1), SNaN, Depth + 1) &&
-           isKnownNeverNaN(Op.getOperand(2), SNaN, Depth + 1);
-  }
   case ISD::FSQRT: // Need is known positive
   case ISD::FLOG:
   case ISD::FLOG2:
   case ISD::FLOG10:
   case ISD::FPOWI:
   case ISD::FPOW: {
     if (SNaN)
       return true;

llvm/test/CodeGen/AMDGPU/known-never-nan.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=fiji -verify-machineinstrs < %s | FileCheck -enable-var-scope -check-prefix=GCN %s
+; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx1010 -verify-machineinstrs < %s | FileCheck %s
 
 define half @known_nnan_extract_vector_elt(float %a, float %b, i32 %idx, half %c) {
-; GCN-LABEL: known_nnan_extract_vector_elt:
-; GCN:       ; %bb.0:
-; GCN-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
-; GCN-NEXT:    v_cvt_pkrtz_f16_f32 v0, v0, v1
-; GCN-NEXT:    v_lshlrev_b32_e32 v1, 4, v2
-; GCN-NEXT:    v_lshrrev_b32_e32 v0, v1, v0
-; GCN-NEXT:    v_add_f16_e32 v1, 1.0, v3
-; GCN-NEXT:    v_cmp_lt_f16_e32 vcc, v0, v1
-; GCN-NEXT:    v_cndmask_b32_e32 v0, v1, v0, vcc
-; GCN-NEXT:    s_setpc_b64 s[30:31]
+; CHECK-LABEL: known_nnan_extract_vector_elt:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; CHECK-NEXT:    s_waitcnt_vscnt null, 0x0
+; CHECK-NEXT:    v_cvt_pkrtz_f16_f32_e32 v0, v0, v1
+; CHECK-NEXT:    v_lshlrev_b32_e32 v1, 4, v2
+; CHECK-NEXT:    v_add_f16_e32 v2, 1.0, v3
+; CHECK-NEXT:    v_lshrrev_b32_e32 v0, v1, v0
+; CHECK-NEXT:    v_cmp_lt_f16_e32 vcc_lo, v0, v2
+; CHECK-NEXT:    v_cndmask_b32_e32 v0, v2, v0, vcc_lo
+; CHECK-NEXT:    s_setpc_b64 s[30:31]
   %cvt = call nnan <2 x half> @llvm.amdgcn.cvt.pkrtz(float %a, float %b)
   %extract = extractelement <2 x half> %cvt, i32 %idx
   %canon = call half @llvm.canonicalize.f16(half %extract)
   %nnan.c = fadd nnan nsz half %c, 1.0
   %cmp = fcmp olt half %extract, %nnan.c
   %select = select i1 %cmp, half %extract, half %nnan.c
   ret half %select
 }
 
-declare <2 x half> @llvm.amdgcn.cvt.pkrtz(float, float) #0
-declare half @llvm.canonicalize.f16(half) #0
+; should not emit v_max
+define float @fma_not_fmaxnm_maybe_nan(i32 %i1, i32 %i2, i32 %i3) #0 {
+; CHECK-LABEL: fma_not_fmaxnm_maybe_nan:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; CHECK-NEXT:    s_waitcnt_vscnt null, 0x0
+; CHECK-NEXT:    v_cvt_f32_u32_e32 v0, v0
+; CHECK-NEXT:    v_cvt_f32_u32_e32 v1, v1
+; CHECK-NEXT:    v_fmaak_f32 v0, v1, v0, 0xff800000
+; CHECK-NEXT:    v_cmp_nge_f32_e32 vcc_lo, 0xff800000, v0
+; CHECK-NEXT:    v_cndmask_b32_e32 v0, 0xff800000, v0, vcc_lo
+; CHECK-NEXT:    s_setpc_b64 s[30:31]
+  %f1 = uitofp i32 %i1 to float
+  %f2 = uitofp i32 %i2 to float
+  %f3 = uitofp i32 %i2 to float
+  %fma = tail call float @llvm.fma.f32(float %f2, float %f1, float 0xfff0000000000000)
+  %cmp = fcmp ugt float %fma, 0xfff0000000000000
+  %val = select i1 %cmp, float %fma, float 0xfff0000000000000
+  ret float %val
+}
+
+declare float @llvm.fma.f32(float, float, float) nounwind readnone
+declare float @llvm.amdgcn.fmad.ftz.f32(float, float, float)
+
+declare <2 x half> @llvm.amdgcn.cvt.pkrtz(float, float)
+declare half @llvm.canonicalize.f16(half)
+
+attributes #0 = { nounwind "no-signed-zeros-fp-math"="true" "target-features"="+mad-mac-f32-insts" }