[llvm] InstCombine: Preserve flags when simplifying exp (PR #174078)

[Matt Arsenault via llvm-commits]

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

>From 4ef3ef8f59ec31700d7d34b625e906023fe3493d Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault at amd.com>
Date: Wed, 31 Dec 2025 10:42:31 +0100
Subject: [PATCH] InstCombine: Preserve flags when simplifying exp

---
 .../InstCombineSimplifyDemanded.cpp           | 19 ++++++++++-------
 .../simplify-demanded-fpclass-exp.ll          | 21 +++++++++++++++++++
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
index 85672cdfe1377..00e42843215ca 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
@@ -2056,12 +2056,16 @@ Value *InstCombinerImpl::SimplifyDemandedUseFPClass(Value *V,
   if (!I->hasOneUse())
     return nullptr;
 
+  FastMathFlags FMF;
   if (auto *FPOp = dyn_cast<FPMathOperator>(I)) {
-    if (FPOp->hasNoNaNs())
+    FMF = FPOp->getFastMathFlags();
+    if (FMF.noNaNs())
       DemandedMask &= ~fcNan;
-    if (FPOp->hasNoInfs())
+
+    if (FMF.noInfs())
       DemandedMask &= ~fcInf;
   }
+
   switch (I->getOpcode()) {
   case Instruction::FNeg: {
     if (SimplifyDemandedFPClass(I, 0, llvm::fneg(DemandedMask), Known,
@@ -2172,8 +2176,8 @@ Value *InstCombinerImpl::SimplifyDemandedUseFPClass(Value *V,
 
         // fadd +/-0, 1.0 => 1.0
         // fadd nan, 1.0 => nan
-        return Builder.CreateFAdd(CI->getArgOperand(0),
-                                  ConstantFP::get(VTy, 1.0));
+        return Builder.CreateFAddFMF(CI->getArgOperand(0),
+                                     ConstantFP::get(VTy, 1.0), FMF);
       }
 
       if (KnownSrc.isKnownAlways(fcInf | fcNan)) {
@@ -2184,9 +2188,10 @@ Value *InstCombinerImpl::SimplifyDemandedUseFPClass(Value *V,
 
         // Note: Dropping canonicalize / quiet of signaling nan.
         Value *X = CI->getArgOperand(0);
-        Value *IsPosInfOrNan =
-            Builder.CreateFCmpUEQ(X, ConstantFP::getInfinity(VTy));
-        return Builder.CreateSelect(IsPosInfOrNan, X, ConstantFP::getZero(VTy));
+        Value *IsPosInfOrNan = Builder.CreateFCmpFMF(
+            FCmpInst::FCMP_UEQ, X, ConstantFP::getInfinity(VTy), FMF);
+        return Builder.CreateSelectFMF(IsPosInfOrNan, X,
+                                       ConstantFP::getZero(VTy), FMF);
       }
 
       // Only perform nan propagation.
diff --git a/llvm/test/Transforms/InstCombine/simplify-demanded-fpclass-exp.ll b/llvm/test/Transforms/InstCombine/simplify-demanded-fpclass-exp.ll
index 8311243e45580..737f7aee0e0f6 100644
--- a/llvm/test/Transforms/InstCombine/simplify-demanded-fpclass-exp.ll
+++ b/llvm/test/Transforms/InstCombine/simplify-demanded-fpclass-exp.ll
@@ -417,6 +417,16 @@ define nofpclass(nzero) float @source_is_known_nan(float nofpclass(inf norm sub
   ret float %exp
 }
 
+define nofpclass(nzero) float @source_is_known_nan_preserve_flags(float nofpclass(inf norm sub zero) %must.be.nan) {
+; CHECK-LABEL: define nofpclass(nzero) float @source_is_known_nan_preserve_flags(
+; CHECK-SAME: float nofpclass(inf zero sub norm) [[MUST_BE_NAN:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = fadd nsz contract float [[MUST_BE_NAN]], 1.000000e+00
+; CHECK-NEXT:    ret float [[TMP1]]
+;
+  %exp = call contract nsz float @llvm.exp2.f32(float %must.be.nan)
+  ret float %exp
+}
+
 define nofpclass(nzero) float @source_is_known_inf_or_nan(float nofpclass(norm sub zero) %must.be.inf.or.nan) {
 ; CHECK-LABEL: define nofpclass(nzero) float @source_is_known_inf_or_nan(
 ; CHECK-SAME: float nofpclass(zero sub norm) [[MUST_BE_INF_OR_NAN:%.*]]) {
@@ -428,6 +438,17 @@ define nofpclass(nzero) float @source_is_known_inf_or_nan(float nofpclass(norm s
   ret float %exp
 }
 
+define nofpclass(nzero) float @source_is_known_inf_or_nan_preserve_flags(float nofpclass(norm sub zero) %must.be.inf.or.nan) {
+; CHECK-LABEL: define nofpclass(nzero) float @source_is_known_inf_or_nan_preserve_flags(
+; CHECK-SAME: float nofpclass(zero sub norm) [[MUST_BE_INF_OR_NAN:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = fcmp nsz contract ueq float [[MUST_BE_INF_OR_NAN]], 0x7FF0000000000000
+; CHECK-NEXT:    [[TMP2:%.*]] = select nsz contract i1 [[TMP1]], float [[MUST_BE_INF_OR_NAN]], float 0.000000e+00
+; CHECK-NEXT:    ret float [[TMP2]]
+;
+  %exp = call contract nsz float @llvm.exp2.f32(float %must.be.inf.or.nan)
+  ret float %exp
+}
+
 define nofpclass(nzero nan) float @source_is_known_inf_or_nan__nnan_result(float nofpclass(norm sub zero) %must.be.inf.or.nan) {
 ; CHECK-LABEL: define nofpclass(nan nzero) float @source_is_known_inf_or_nan__nnan_result(
 ; CHECK-SAME: float nofpclass(zero sub norm) [[MUST_BE_INF_OR_NAN:%.*]]) {