[llvm-branch-commits] [llvm] ValueTracking: Improve nan tracking for fma square special case (PR #175999)

[Matt Arsenault via llvm-branch-commits]

https://github.com/arsenm updated https://github.com/llvm/llvm-project/pull/175999

>From 550d190e3342f5029a184169db3cbf31a5008509 Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault at amd.com>
Date: Mon, 12 Jan 2026 14:28:25 +0100
Subject: [PATCH 1/3] ValueTracking: Improve nan tracking for fma square
 special case

In the square multiply case, we can infer if the add of opposite
sign infinities can occur.
---
 llvm/lib/Analysis/ValueTracking.cpp              |  4 ++++
 llvm/lib/Support/KnownFPClass.cpp                | 12 +++++++++++-
 llvm/test/Transforms/Attributor/nofpclass-fma.ll |  4 ++--
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index d1f620ac9eb3d..6b0abc99632bc 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -5661,6 +5661,10 @@ void computeKnownFPClass(const Value *V, const APInt &DemandedElts,
                           Q, Depth + 1);
     }
 
+    // TODO: Improve accuracy in unfused FMA pattern. We can prove an additional
+    // not-nan if the addend is known-not negative infinity if the multiply is
+    // known-not infinity.
+
     computeKnownFPClass(Op->getOperand(0), DemandedElts, fcAllFlags, KnownLHS,
                         Q, Depth + 1);
 
diff --git a/llvm/lib/Support/KnownFPClass.cpp b/llvm/lib/Support/KnownFPClass.cpp
index dc7f6d3ca237d..a192f19e129f4 100644
--- a/llvm/lib/Support/KnownFPClass.cpp
+++ b/llvm/lib/Support/KnownFPClass.cpp
@@ -357,7 +357,17 @@ KnownFPClass KnownFPClass::fma(const KnownFPClass &KnownLHS,
 KnownFPClass KnownFPClass::fma_square(const KnownFPClass &KnownSquared,
                                       const KnownFPClass &KnownAddend,
                                       DenormalMode Mode) {
-  return fadd_impl(square(KnownSquared, Mode), KnownAddend, Mode);
+  KnownFPClass Squared = square(KnownSquared, Mode);
+  KnownFPClass Known = fadd_impl(Squared, KnownAddend, Mode);
+
+  // Since we know the squared input must be positive, the add of opposite sign
+  // infinities nan hazard only applies for negative nan.
+  if (KnownAddend.isKnownNever(fcNegInf | fcNan) &&
+      Known.isKnownNever(fcPosInf | fcNan) &&
+      KnownSquared.isKnownNeverLogicalZero(Mode))
+    Known.knownNot(fcNan);
+
+  return Known;
 }
 
 KnownFPClass KnownFPClass::exp(const KnownFPClass &KnownSrc) {
diff --git a/llvm/test/Transforms/Attributor/nofpclass-fma.ll b/llvm/test/Transforms/Attributor/nofpclass-fma.ll
index c4cdfd5e2b5a7..df42d9fb30058 100644
--- a/llvm/test/Transforms/Attributor/nofpclass-fma.ll
+++ b/llvm/test/Transforms/Attributor/nofpclass-fma.ll
@@ -454,9 +454,9 @@ define half @ret_fma_square__no_nan_no_inf__no_nan_no_ninf(half noundef nofpclas
 }
 
 define half @ret_fma_square__no_nan_no_inf_no_zero__no_nan_no_ninf(half noundef nofpclass(nan inf zero) %arg0, half nofpclass(nan ninf) %arg1) {
-; CHECK-LABEL: define half @ret_fma_square__no_nan_no_inf_no_zero__no_nan_no_ninf
+; CHECK-LABEL: define nofpclass(nan) half @ret_fma_square__no_nan_no_inf_no_zero__no_nan_no_ninf
 ; CHECK-SAME: (half noundef nofpclass(nan inf zero) [[ARG0:%.*]], half nofpclass(nan ninf) [[ARG1:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[CALL:%.*]] = call half @llvm.fma.f16(half noundef nofpclass(nan inf zero) [[ARG0]], half noundef nofpclass(nan inf zero) [[ARG0]], half nofpclass(nan ninf) [[ARG1]]) #[[ATTR2]]
+; CHECK-NEXT:    [[CALL:%.*]] = call nofpclass(nan) half @llvm.fma.f16(half noundef nofpclass(nan inf zero) [[ARG0]], half noundef nofpclass(nan inf zero) [[ARG0]], half nofpclass(nan ninf) [[ARG1]]) #[[ATTR2]]
 ; CHECK-NEXT:    ret half [[CALL]]
 ;
   %call = call half @llvm.fma.f16(half %arg0, half %arg0, half %arg1)

>From 55835d0354adbd9f756ea626772728bee6e4f233 Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault at amd.com>
Date: Thu, 15 Jan 2026 19:40:26 +0100
Subject: [PATCH 2/3] fix too conservative

---
 llvm/lib/Support/KnownFPClass.cpp                | 4 +---
 llvm/test/Transforms/Attributor/nofpclass-fma.ll | 4 ++--
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Support/KnownFPClass.cpp b/llvm/lib/Support/KnownFPClass.cpp
index a192f19e129f4..8a13c91780d50 100644
--- a/llvm/lib/Support/KnownFPClass.cpp
+++ b/llvm/lib/Support/KnownFPClass.cpp
@@ -362,9 +362,7 @@ KnownFPClass KnownFPClass::fma_square(const KnownFPClass &KnownSquared,
 
   // Since we know the squared input must be positive, the add of opposite sign
   // infinities nan hazard only applies for negative nan.
-  if (KnownAddend.isKnownNever(fcNegInf | fcNan) &&
-      Known.isKnownNever(fcPosInf | fcNan) &&
-      KnownSquared.isKnownNeverLogicalZero(Mode))
+  if (KnownAddend.isKnownNever(fcNegInf | fcNan) && Squared.isKnownNever(fcNan))
     Known.knownNot(fcNan);
 
   return Known;
diff --git a/llvm/test/Transforms/Attributor/nofpclass-fma.ll b/llvm/test/Transforms/Attributor/nofpclass-fma.ll
index df42d9fb30058..6d4ce9de7f90d 100644
--- a/llvm/test/Transforms/Attributor/nofpclass-fma.ll
+++ b/llvm/test/Transforms/Attributor/nofpclass-fma.ll
@@ -444,9 +444,9 @@ define half @ret_fma_square__no_nan_no_inf__no_nan_no_pinf(half noundef nofpclas
 }
 
 define half @ret_fma_square__no_nan_no_inf__no_nan_no_ninf(half noundef nofpclass(nan inf) %arg0, half nofpclass(nan ninf) %arg1) {
-; CHECK-LABEL: define half @ret_fma_square__no_nan_no_inf__no_nan_no_ninf
+; CHECK-LABEL: define nofpclass(nan) half @ret_fma_square__no_nan_no_inf__no_nan_no_ninf
 ; CHECK-SAME: (half noundef nofpclass(nan inf) [[ARG0:%.*]], half nofpclass(nan ninf) [[ARG1:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[CALL:%.*]] = call half @llvm.fma.f16(half noundef nofpclass(nan inf) [[ARG0]], half noundef nofpclass(nan inf) [[ARG0]], half nofpclass(nan ninf) [[ARG1]]) #[[ATTR2]]
+; CHECK-NEXT:    [[CALL:%.*]] = call nofpclass(nan) half @llvm.fma.f16(half noundef nofpclass(nan inf) [[ARG0]], half noundef nofpclass(nan inf) [[ARG0]], half nofpclass(nan ninf) [[ARG1]]) #[[ATTR2]]
 ; CHECK-NEXT:    ret half [[CALL]]
 ;
   %call = call half @llvm.fma.f16(half %arg0, half %arg0, half %arg1)

>From 3071adc01f43c1234fbaadf3f8e9a19040858a02 Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault at amd.com>
Date: Thu, 15 Jan 2026 19:52:06 +0100
Subject: [PATCH 3/3] Add comment

---
 llvm/lib/Support/KnownFPClass.cpp | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Support/KnownFPClass.cpp b/llvm/lib/Support/KnownFPClass.cpp
index 8a13c91780d50..614235df39782 100644
--- a/llvm/lib/Support/KnownFPClass.cpp
+++ b/llvm/lib/Support/KnownFPClass.cpp
@@ -361,7 +361,11 @@ KnownFPClass KnownFPClass::fma_square(const KnownFPClass &KnownSquared,
   KnownFPClass Known = fadd_impl(Squared, KnownAddend, Mode);
 
   // Since we know the squared input must be positive, the add of opposite sign
-  // infinities nan hazard only applies for negative nan.
+  // infinities nan hazard only applies for negative inf.
+  //
+  // TODO: Alternatively to proving addend is not -inf, we could know Squared is
+  // not pinf. Other than the degenerate always-subnormal input case, we can't
+  // prove that without a known range.
   if (KnownAddend.isKnownNever(fcNegInf | fcNan) && Squared.isKnownNever(fcNan))
     Known.knownNot(fcNan);