[llvm-branch-commits] [llvm] ValueTracking: Improve nan tracking for fma square special case (PR #175999)
Matt Arsenault via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Thu Jan 15 13:57:33 PST 2026
https://github.com/arsenm updated https://github.com/llvm/llvm-project/pull/175999
>From 6d01a14d16e596c6b4a626e970abedcd3a309561 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/4] 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 2fd0dd81622f6..96d9a63fafe2f 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 9f5259b5f079fd432ea8d8b6dcad5b1b322c8329 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/4] 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 96d9a63fafe2f..5a4c582bfd56d 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 9a13f3e079b087a24eb68665f5c5502cce28879c 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/4] 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);
>From fcd62a012a474aaf2c5f1cd07d0e3a8b1b39abea Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault at amd.com>
Date: Thu, 15 Jan 2026 21:49:28 +0100
Subject: [PATCH 4/4] propagate flags into fma queries
---
llvm/lib/Analysis/ValueTracking.cpp | 15 +++++++
.../Transforms/Attributor/nofpclass-fma.ll | 41 +++++++++++++++++++
2 files changed, 56 insertions(+)
diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 6b0abc99632bc..d450d77ece057 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -5143,6 +5143,16 @@ void computeKnownFPClass(const Value *V, const APInt &DemandedElts,
DenormalMode Mode =
F ? F->getDenormalMode(FltSem) : DenormalMode::getDynamic();
+ if (KnownNotFromFlags & fcNan) {
+ KnownSrc.knownNot(fcNan);
+ KnownAddend.knownNot(fcNan);
+ }
+
+ if (KnownNotFromFlags & fcInf) {
+ KnownSrc.knownNot(fcInf);
+ KnownAddend.knownNot(fcInf);
+ }
+
Known = KnownFPClass::fma_square(KnownSrc, KnownAddend, Mode);
break;
}
@@ -5153,6 +5163,11 @@ void computeKnownFPClass(const Value *V, const APInt &DemandedElts,
InterestedClasses, KnownSrc[I], Q, Depth + 1);
if (KnownSrc[I].isUnknown())
return;
+
+ if (KnownNotFromFlags & fcNan)
+ KnownSrc[I].knownNot(fcNan);
+ if (KnownNotFromFlags & fcInf)
+ KnownSrc[I].knownNot(fcInf);
}
const Function *F = II->getFunction();
diff --git a/llvm/test/Transforms/Attributor/nofpclass-fma.ll b/llvm/test/Transforms/Attributor/nofpclass-fma.ll
index 5a4c582bfd56d..0f376b9834496 100644
--- a/llvm/test/Transforms/Attributor/nofpclass-fma.ll
+++ b/llvm/test/Transforms/Attributor/nofpclass-fma.ll
@@ -462,5 +462,46 @@ define half @ret_fma_square__no_nan_no_inf_no_zero__no_nan_no_ninf(half noundef
%call = call half @llvm.fma.f16(half %arg0, half %arg0, half %arg1)
ret half %call
}
+
+define half @ret_fma_ninf_square__no_nan__no_nan(half noundef nofpclass(nan) %arg0, half nofpclass(nan) %arg1) {
+; CHECK-LABEL: define nofpclass(nan inf) half @ret_fma_ninf_square__no_nan__no_nan
+; CHECK-SAME: (half noundef nofpclass(nan) [[ARG0:%.*]], half nofpclass(nan) [[ARG1:%.*]]) #[[ATTR1]] {
+; CHECK-NEXT: [[CALL:%.*]] = call ninf nofpclass(nan inf) half @llvm.fma.f16(half noundef nofpclass(nan) [[ARG0]], half noundef nofpclass(nan) [[ARG0]], half nofpclass(nan) [[ARG1]]) #[[ATTR2]]
+; CHECK-NEXT: ret half [[CALL]]
+;
+ %call = call ninf half @llvm.fma.f16(half %arg0, half %arg0, half %arg1)
+ ret half %call
+}
+
+define nofpclass(inf) half @ret_noinf_fma_square__no_nan__no_nan(half noundef nofpclass(nan) %arg0, half nofpclass(nan) %arg1) {
+; CHECK-LABEL: define nofpclass(inf) half @ret_noinf_fma_square__no_nan__no_nan
+; CHECK-SAME: (half noundef nofpclass(nan) [[ARG0:%.*]], half nofpclass(nan) [[ARG1:%.*]]) #[[ATTR1]] {
+; CHECK-NEXT: [[CALL:%.*]] = call half @llvm.fma.f16(half noundef nofpclass(nan) [[ARG0]], half noundef nofpclass(nan) [[ARG0]], half nofpclass(nan) [[ARG1]]) #[[ATTR2]]
+; CHECK-NEXT: ret half [[CALL]]
+;
+ %call = call half @llvm.fma.f16(half %arg0, half %arg0, half %arg1)
+ ret half %call
+}
+
+define nofpclass(nan) half @ret_nonan_fma_square__no_nan__no_nan(half noundef nofpclass(nan) %arg0, half nofpclass(nan) %arg1) {
+; CHECK-LABEL: define nofpclass(nan) half @ret_nonan_fma_square__no_nan__no_nan
+; CHECK-SAME: (half noundef nofpclass(nan) [[ARG0:%.*]], half nofpclass(nan) [[ARG1:%.*]]) #[[ATTR1]] {
+; CHECK-NEXT: [[CALL:%.*]] = call half @llvm.fma.f16(half noundef nofpclass(nan) [[ARG0]], half noundef nofpclass(nan) [[ARG0]], half nofpclass(nan) [[ARG1]]) #[[ATTR2]]
+; CHECK-NEXT: ret half [[CALL]]
+;
+ %call = call half @llvm.fma.f16(half %arg0, half %arg0, half %arg1)
+ ret half %call
+}
+
+define half @ret_fma_ninf__no_nan_inputs(half nofpclass(nan) %arg0, half nofpclass(nan) %arg1, half nofpclass(nan) %arg2) {
+; CHECK-LABEL: define nofpclass(nan inf) half @ret_fma_ninf__no_nan_inputs
+; CHECK-SAME: (half nofpclass(nan) [[ARG0:%.*]], half nofpclass(nan) [[ARG1:%.*]], half nofpclass(nan) [[ARG2:%.*]]) #[[ATTR1]] {
+; CHECK-NEXT: [[CALL:%.*]] = call ninf nofpclass(nan inf) half @llvm.fma.f16(half nofpclass(nan) [[ARG0]], half nofpclass(nan) [[ARG1]], half nofpclass(nan) [[ARG2]]) #[[ATTR2]]
+; CHECK-NEXT: ret half [[CALL]]
+;
+ %call = call ninf half @llvm.fma.f16(half %arg0, half %arg1, half %arg2)
+ ret half %call
+}
+
;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
; TUNIT: {{.*}}
More information about the llvm-branch-commits
mailing list