ValueTracking: Handle no-nan check for computeKnownFPClass for fadd/fsub

Copy the logic from isKnownNeverNaN for fadd/fsub.

Author: arsenm

llvm/include/llvm/Analysis/ValueTracking.h

@@ -238,6 +238,17 @@ struct KnownFPClass {
   /// definitely set or false if the sign bit is definitely unset.
   std::optional<bool> SignBit;
 
+
+  /// Return true if it's known this can never be a nan.
+  bool isKnownNeverNaN() const {
+    return (KnownFPClasses & fcNan) == fcNone;
+  }
+
+  /// Return true if it's known this can never be an infinity.
+  bool isKnownNeverInfinity() const {
+    return (KnownFPClasses & fcInf) == fcNone;
+  }
+
   KnownFPClass &operator|=(const KnownFPClass &RHS) {
     KnownFPClasses = KnownFPClasses | RHS.KnownFPClasses;
 

llvm/lib/Analysis/ValueTracking.cpp

@@ -4439,6 +4439,25 @@ void computeKnownFPClass(const Value *V, const APInt &DemandedElts,
 
     break;
   }
+  case Instruction::FAdd:
+  case Instruction::FSub: {
+    KnownFPClass KnownLHS, KnownRHS;
+    computeKnownFPClass(Op->getOperand(1), DemandedElts, fcNan | fcInf,
+                        KnownRHS, Depth + 1, Q, TLI);
+    if (KnownRHS.isKnownNeverNaN()) {
+      // RHS is canonically cheaper to compute. Skip inspecting the LHS if
+      // there's no point.
+      computeKnownFPClass(Op->getOperand(0), DemandedElts, fcNan | fcInf,
+                          KnownLHS, Depth + 1, Q, TLI);
+      // Adding positive and negative infinity produces NaN.
+      // TODO: Check sign of infinities.
+      if (KnownLHS.isKnownNeverNaN() &&
+          (KnownLHS.isKnownNeverInfinity() || KnownRHS.isKnownNeverInfinity()))
+        Known.knownNot(fcNan);
+    }
+
+    break;
+  }
   case Instruction::SIToFP:
   case Instruction::UIToFP: {
     // Cannot produce nan

llvm/unittests/Analysis/ValueTrackingTest.cpp

@@ -71,6 +71,7 @@ class ValueTrackingTest : public testing::Test {
     A2 = findInstructionByNameOrNull(F, "A2");
     A3 = findInstructionByNameOrNull(F, "A3");
     A4 = findInstructionByNameOrNull(F, "A4");
+    A5 = findInstructionByNameOrNull(F, "A5");
 
     CxtI = findInstructionByNameOrNull(F, "CxtI");
     CxtI2 = findInstructionByNameOrNull(F, "CxtI2");
@@ -82,7 +83,7 @@ class ValueTrackingTest : public testing::Test {
   Function *F = nullptr;
   Instruction *A = nullptr;
   // Instructions (optional)
-  Instruction *A2 = nullptr, *A3 = nullptr, *A4 = nullptr;
+  Instruction *A2 = nullptr, *A3 = nullptr, *A4 = nullptr, *A5 = nullptr;
 
   // Context instructions (optional)
   Instruction *CxtI = nullptr, *CxtI2 = nullptr, *CxtI3 = nullptr;
@@ -1540,6 +1541,40 @@ TEST_F(ComputeKnownFPClassTest, SIToFP) {
   expectKnownFPClass(~fcNan, std::nullopt, A3);
 }
 
+TEST_F(ComputeKnownFPClassTest, FAdd) {
+  parseAssembly(
+      "define float @test(float nofpclass(nan inf) %nnan.ninf, float nofpclass(nan) %nnan, float nofpclass(qnan) %no.qnan, float %unknown) {\n"
+      "  %A = fadd float %nnan, %nnan.ninf"
+      "  %A2 = fadd float %nnan.ninf, %nnan"
+      "  %A3 = fadd float %nnan.ninf, %unknown"
+      "  %A4 = fadd float %nnan.ninf, %no.qnan"
+      "  %A5 = fadd float %nnan, %nnan"
+      "  ret float %A\n"
+      "}\n");
+  expectKnownFPClass(fcFinite | fcInf, std::nullopt, A);
+  expectKnownFPClass(fcFinite | fcInf, std::nullopt, A2);
+  expectKnownFPClass(fcAllFlags, std::nullopt, A3);
+  expectKnownFPClass(fcAllFlags, std::nullopt, A4);
+  expectKnownFPClass(fcAllFlags, std::nullopt, A5);
+}
+
+TEST_F(ComputeKnownFPClassTest, FSub) {
+  parseAssembly(
+      "define float @test(float nofpclass(nan inf) %nnan.ninf, float nofpclass(nan) %nnan, float nofpclass(qnan) %no.qnan, float %unknown) {\n"
+      "  %A = fsub float %nnan, %nnan.ninf"
+      "  %A2 = fsub float %nnan.ninf, %nnan"
+      "  %A3 = fsub float %nnan.ninf, %unknown"
+      "  %A4 = fsub float %nnan.ninf, %no.qnan"
+      "  %A5 = fsub float %nnan, %nnan"
+      "  ret float %A\n"
+      "}\n");
+  expectKnownFPClass(fcFinite | fcInf, std::nullopt, A);
+  expectKnownFPClass(fcFinite | fcInf, std::nullopt, A2);
+  expectKnownFPClass(fcAllFlags, std::nullopt, A3);
+  expectKnownFPClass(fcAllFlags, std::nullopt, A4);
+  expectKnownFPClass(fcAllFlags, std::nullopt, A5);
+}
+
 TEST_F(ValueTrackingTest, isNonZeroRecurrence) {
   parseAssembly(R"(
     define i1 @test(i8 %n, i8 %r) {