ValueTracking: Address todo for nan fmul handling in computeKnownFPClass

If both operands can't be zero or nan, the result can't be nan.

4 files changed, 74 insertions, 19 deletions

diff --git a/llvm/include/llvm/Analysis/ValueTracking.h b/llvm/include/llvm/Analysis/ValueTracking.h
index 849d8a7..c3093d5 100644
--- a/llvm/include/llvm/Analysis/ValueTracking.h
+++ b/llvm/include/llvm/Analysis/ValueTracking.h
@@ -249,6 +249,22 @@ struct KnownFPClass {
     return (KnownFPClasses & fcInf) == fcNone;
   }
 
+  /// Return true if it's known this can never be a subnormal
+  bool isKnownNeverSubnormal() const {
+    return (KnownFPClasses & fcSubnormal) == fcNone;
+  }
+
+  /// Return true if it's known this can never be a zero. This means a literal
+  /// [+-]0, and does not include denormal inputs implicitly treated as [+-]0.
+  bool isKnownNeverZero() const {
+    return (KnownFPClasses & fcZero) == fcNone;
+  }
+
+  /// Return true if it's know this can never be interpreted as a zero. This
+  /// extends isKnownNeverZero to cover the case where the assumed
+  /// floating-point mode for the function interprets denormals as zero.
+  bool isKnownNeverLogicalZero(const Function &F, Type *Ty) const;
+
   KnownFPClass &operator|=(const KnownFPClass &RHS) {
     KnownFPClasses = KnownFPClasses | RHS.KnownFPClasses;
 
diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 8ff71de..8a23100 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -4125,11 +4125,16 @@ bool llvm::isKnownNeverNaN(const Value *V, const TargetLibraryInfo *TLI,
 
 /// Return true if it's possible to assume IEEE treatment of input denormals in
 /// \p F for \p Val.
-static bool inputDenormalIsIEEE(const Function &F, const Value *Val) {
-  Type *Ty = Val->getType()->getScalarType();
+static bool inputDenormalIsIEEE(const Function &F, const Type *Ty) {
+  Ty = Ty->getScalarType();
   return F.getDenormalMode(Ty->getFltSemantics()).Input == DenormalMode::IEEE;
 }
 
+bool KnownFPClass::isKnownNeverLogicalZero(const Function &F, Type *Ty) const {
+  return isKnownNeverZero() &&
+         (isKnownNeverSubnormal() || inputDenormalIsIEEE(F, Ty));
+}
+
 /// Returns a pair of values, which if passed to llvm.is.fpclass, returns the
 /// same result as an fcmp with the given operands.
 std::pair<Value *, FPClassTest> llvm::fcmpToClassTest(FCmpInst::Predicate Pred,
@@ -4143,7 +4148,7 @@ std::pair<Value *, FPClassTest> llvm::fcmpToClassTest(FCmpInst::Predicate Pred,
   if (ConstRHS->isZero()) {
     // Compares with fcNone are only exactly equal to fcZero if input denormals are
     // not flushed.
-    if (FCmpInst::isEquality(Pred) && !inputDenormalIsIEEE(F, LHS))
+    if (FCmpInst::isEquality(Pred) && !inputDenormalIsIEEE(F, LHS->getType()))
       return {nullptr, fcNone};
 
     switch (Pred) {
@@ -4460,15 +4465,25 @@ void computeKnownFPClass(const Value *V, const APInt &DemandedElts,
   }
   case Instruction::FMul: {
     KnownFPClass KnownLHS, KnownRHS;
-    computeKnownFPClass(Op->getOperand(1), DemandedElts, fcNan | fcInf,
-                        KnownRHS, Depth + 1, Q, TLI);
-    if (KnownRHS.isKnownNeverNaN() && KnownRHS.isKnownNeverInfinity()) {
-      computeKnownFPClass(Op->getOperand(0), DemandedElts, fcNan | fcInf,
-                          KnownLHS, Depth + 1, Q, TLI);
+    computeKnownFPClass(Op->getOperand(1), DemandedElts,
+                        fcNan | fcInf | fcZero | fcSubnormal, KnownRHS,
+                        Depth + 1, Q, TLI);
+    if (KnownRHS.isKnownNeverNaN() &&
+        (KnownRHS.isKnownNeverInfinity() || KnownRHS.isKnownNeverZero())) {
+      computeKnownFPClass(Op->getOperand(0), DemandedElts,
+                          fcNan | fcInf | fcZero, KnownLHS, Depth + 1, Q, TLI);
+      if (!KnownLHS.isKnownNeverNaN())
+        break;
 
-      // Zero multiplied with infinity produces NaN.
-      // FIXME: If neither side can be zero fmul never produces NaN.
-      if (KnownLHS.isKnownNeverNaN() && KnownLHS.isKnownNeverInfinity())
+      const Function *F = cast<Instruction>(Op)->getFunction();
+
+      // If neither side can be zero (or nan) fmul never produces NaN.
+      // TODO: Check operand combinations.
+      // e.g. fmul nofpclass(inf nan zero), nofpclass(nan) -> nofpclass(nan)
+      if ((KnownLHS.isKnownNeverInfinity() ||
+           KnownLHS.isKnownNeverLogicalZero(*F, Op->getType())) &&
+          (KnownRHS.isKnownNeverInfinity() ||
+           KnownRHS.isKnownNeverLogicalZero(*F, Op->getType())))
         Known.knownNot(fcNan);
     }
 
diff --git a/llvm/test/Transforms/Attributor/nofpclass-nan-fmul.ll b/llvm/test/Transforms/Attributor/nofpclass-nan-fmul.ll
index 3b433dd..d4ea815 100644
--- a/llvm/test/Transforms/Attributor/nofpclass-nan-fmul.ll
+++ b/llvm/test/Transforms/Attributor/nofpclass-nan-fmul.ll
@@ -12,7 +12,7 @@ define float @ret_fmul_ieee_nonan__nonan(float nofpclass(nan) %arg0, float nofpc
 }
 
 define float @ret_fmul_ieee_nonan_nozero__nonan_nozero(float nofpclass(nan zero) %arg0, float nofpclass(nan zero) %arg1) #0 {
-; CHECK-LABEL: define float @ret_fmul_ieee_nonan_nozero__nonan_nozero
+; CHECK-LABEL: define nofpclass(nan) float @ret_fmul_ieee_nonan_nozero__nonan_nozero
 ; CHECK-SAME: (float nofpclass(nan zero) [[ARG0:%.*]], float nofpclass(nan zero) [[ARG1:%.*]]) #[[ATTR0]] {
 ; CHECK-NEXT:    [[FMUL:%.*]] = fmul float [[ARG0]], [[ARG1]]
 ; CHECK-NEXT:    ret float [[FMUL]]
@@ -62,7 +62,7 @@ define float @ret_fmul_ieee_nonan_noinf__nonan_noinf(float nofpclass(nan inf) %a
 }
 
 define float @ret_fmul_ieee_nonan_nozero__nonan_noinf(float nofpclass(nan zero) %arg0, float nofpclass(nan inf) %arg1) #0 {
-; CHECK-LABEL: define float @ret_fmul_ieee_nonan_nozero__nonan_noinf
+; CHECK-LABEL: define nofpclass(nan) float @ret_fmul_ieee_nonan_nozero__nonan_noinf
 ; CHECK-SAME: (float nofpclass(nan zero) [[ARG0:%.*]], float nofpclass(nan inf) [[ARG1:%.*]]) #[[ATTR0]] {
 ; CHECK-NEXT:    [[FMUL:%.*]] = fmul float [[ARG0]], [[ARG1]]
 ; CHECK-NEXT:    ret float [[FMUL]]
@@ -72,7 +72,7 @@ define float @ret_fmul_ieee_nonan_nozero__nonan_noinf(float nofpclass(nan zero)
 }
 
 define float @ret_fmul_ieee_nonan_noinf__nonan_nozero(float nofpclass(nan inf) %arg0, float nofpclass(nan zero) %arg1) #0 {
-; CHECK-LABEL: define float @ret_fmul_ieee_nonan_noinf__nonan_nozero
+; CHECK-LABEL: define nofpclass(nan) float @ret_fmul_ieee_nonan_noinf__nonan_nozero
 ; CHECK-SAME: (float nofpclass(nan inf) [[ARG0:%.*]], float nofpclass(nan zero) [[ARG1:%.*]]) #[[ATTR0]] {
 ; CHECK-NEXT:    [[FMUL:%.*]] = fmul float [[ARG0]], [[ARG1]]
 ; CHECK-NEXT:    ret float [[FMUL]]
@@ -102,7 +102,7 @@ define float @ret_fmul_ieee_nonan_noposzero__nonan_noinf(float nofpclass(nan pze
 }
 
 define float @ret_fmul_ieee_nonan_nozero_nosub__nonan_nozero_nosub(float nofpclass(nan zero sub) %arg0, float nofpclass(nan zero sub) %arg1) #0 {
-; CHECK-LABEL: define float @ret_fmul_ieee_nonan_nozero_nosub__nonan_nozero_nosub
+; CHECK-LABEL: define nofpclass(nan) float @ret_fmul_ieee_nonan_nozero_nosub__nonan_nozero_nosub
 ; CHECK-SAME: (float nofpclass(nan zero sub) [[ARG0:%.*]], float nofpclass(nan zero sub) [[ARG1:%.*]]) #[[ATTR0]] {
 ; CHECK-NEXT:    [[FMUL:%.*]] = fmul float [[ARG0]], [[ARG1]]
 ; CHECK-NEXT:    ret float [[FMUL]]
@@ -113,7 +113,7 @@ define float @ret_fmul_ieee_nonan_nozero_nosub__nonan_nozero_nosub(float nofpcla
 
 ; Denormal mode doesn't matter because sources are nofpclass(sub)
 define float @ret_fmul_daz_nonan_nozero_nosub__nonan_nozero_nosub(float nofpclass(nan zero sub) %arg0, float nofpclass(nan zero sub) %arg1) #1 {
-; CHECK-LABEL: define float @ret_fmul_daz_nonan_nozero_nosub__nonan_nozero_nosub
+; CHECK-LABEL: define nofpclass(nan) float @ret_fmul_daz_nonan_nozero_nosub__nonan_nozero_nosub
 ; CHECK-SAME: (float nofpclass(nan zero sub) [[ARG0:%.*]], float nofpclass(nan zero sub) [[ARG1:%.*]]) #[[ATTR1]] {
 ; CHECK-NEXT:    [[FMUL:%.*]] = fmul float [[ARG0]], [[ARG1]]
 ; CHECK-NEXT:    ret float [[FMUL]]
@@ -123,7 +123,7 @@ define float @ret_fmul_daz_nonan_nozero_nosub__nonan_nozero_nosub(float nofpclas
 }
 
 define float @ret_fmul_dapz_nonan_nozero_nosub__nonan_nozero_nosub(float nofpclass(nan zero sub) %arg0, float nofpclass(nan zero sub) %arg1) #2 {
-; CHECK-LABEL: define float @ret_fmul_dapz_nonan_nozero_nosub__nonan_nozero_nosub
+; CHECK-LABEL: define nofpclass(nan) float @ret_fmul_dapz_nonan_nozero_nosub__nonan_nozero_nosub
 ; CHECK-SAME: (float nofpclass(nan zero sub) [[ARG0:%.*]], float nofpclass(nan zero sub) [[ARG1:%.*]]) #[[ATTR2]] {
 ; CHECK-NEXT:    [[FMUL:%.*]] = fmul float [[ARG0]], [[ARG1]]
 ; CHECK-NEXT:    ret float [[FMUL]]
@@ -133,7 +133,7 @@ define float @ret_fmul_dapz_nonan_nozero_nosub__nonan_nozero_nosub(float nofpcla
 }
 
 define float @ret_fmul_dynamic_nonan_nozero_nosub__nonan_nozero_nosub(float nofpclass(nan zero sub) %arg0, float nofpclass(nan zero sub) %arg1) #3 {
-; CHECK-LABEL: define float @ret_fmul_dynamic_nonan_nozero_nosub__nonan_nozero_nosub
+; CHECK-LABEL: define nofpclass(nan) float @ret_fmul_dynamic_nonan_nozero_nosub__nonan_nozero_nosub
 ; CHECK-SAME: (float nofpclass(nan zero sub) [[ARG0:%.*]], float nofpclass(nan zero sub) [[ARG1:%.*]]) #[[ATTR3]] {
 ; CHECK-NEXT:    [[FMUL:%.*]] = fmul float [[ARG0]], [[ARG1]]
 ; CHECK-NEXT:    ret float [[FMUL]]
diff --git a/llvm/unittests/Analysis/ValueTrackingTest.cpp b/llvm/unittests/Analysis/ValueTrackingTest.cpp
index ba983d4..211959f 100644
--- a/llvm/unittests/Analysis/ValueTrackingTest.cpp
+++ b/llvm/unittests/Analysis/ValueTrackingTest.cpp
@@ -72,6 +72,8 @@ protected:
     A3 = findInstructionByNameOrNull(F, "A3");
     A4 = findInstructionByNameOrNull(F, "A4");
     A5 = findInstructionByNameOrNull(F, "A5");
+    A6 = findInstructionByNameOrNull(F, "A6");
+    A7 = findInstructionByNameOrNull(F, "A7");
 
     CxtI = findInstructionByNameOrNull(F, "CxtI");
     CxtI2 = findInstructionByNameOrNull(F, "CxtI2");
@@ -83,7 +85,8 @@ protected:
   Function *F = nullptr;
   Instruction *A = nullptr;
   // Instructions (optional)
-  Instruction *A2 = nullptr, *A3 = nullptr, *A4 = nullptr, *A5 = nullptr;
+  Instruction *A2 = nullptr, *A3 = nullptr, *A4 = nullptr, *A5 = nullptr,
+              *A6 = nullptr, *A7 = nullptr;
 
   // Context instructions (optional)
   Instruction *CxtI = nullptr, *CxtI2 = nullptr, *CxtI3 = nullptr;
@@ -1592,6 +1595,27 @@ TEST_F(ComputeKnownFPClassTest, FMul) {
   expectKnownFPClass(fcAllFlags, std::nullopt, A5);
 }
 
+TEST_F(ComputeKnownFPClassTest, FMulNoZero) {
+  parseAssembly(
+      "define float @test(float nofpclass(zero) %no.zero, float nofpclass(zero nan) %no.zero.nan, float nofpclass(nzero nan) %no.negzero.nan, float nofpclass(pzero nan) %no.poszero.nan, float nofpclass(inf nan) %no.inf.nan, float nofpclass(inf) %no.inf, float nofpclass(nan) %no.nan) {\n"
+      "  %A = fmul float %no.zero.nan, %no.zero.nan"
+      "  %A2 = fmul float %no.zero, %no.zero"
+      "  %A3 = fmul float %no.poszero.nan, %no.zero.nan"
+      "  %A4 = fmul float %no.nan, %no.zero"
+      "  %A5 = fmul float %no.zero, %no.inf"
+      "  %A6 = fmul float %no.zero.nan, %no.nan"
+      "  %A7 = fmul float %no.nan, %no.zero.nan"
+      "  ret float %A\n"
+      "}\n");
+  expectKnownFPClass(fcFinite | fcInf, std::nullopt, A);
+  expectKnownFPClass(fcAllFlags, std::nullopt, A2);
+  expectKnownFPClass(fcAllFlags, std::nullopt, A3);
+  expectKnownFPClass(fcAllFlags, std::nullopt, A4);
+  expectKnownFPClass(fcAllFlags, std::nullopt, A5);
+  expectKnownFPClass(fcAllFlags, std::nullopt, A6);
+  expectKnownFPClass(fcAllFlags, std::nullopt, A7);
+}
+
 TEST_F(ValueTrackingTest, isNonZeroRecurrence) {
   parseAssembly(R"(
     define i1 @test(i8 %n, i8 %r) {