[llvm] [InstCombine] canonicalize operands of fcmp ord/uno if they are known NaN (PR #97763)

[Alex MacLean via llvm-commits]

https://github.com/AlexMaclean updated https://github.com/llvm/llvm-project/pull/97763

>From 0778ce4484ba2e5845509e56c1186488eb2b1a1a Mon Sep 17 00:00:00 2001
From: Alex MacLean <amaclean at nvidia.com>
Date: Mon, 1 Jul 2024 19:26:28 +0000
Subject: [PATCH 1/2] [InstCombine] canonicalize operands of fcmp ord/uno if
 they are known NaN

---
 llvm/include/llvm/Analysis/ValueTracking.h    |  5 +++++
 .../InstCombine/InstCombineCompares.cpp       | 21 ++++++++++++-------
 .../Transforms/InstCombine/fcmp-special.ll    | 16 ++++++++++++++
 3 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/llvm/include/llvm/Analysis/ValueTracking.h b/llvm/include/llvm/Analysis/ValueTracking.h
index a67ad501982d22..51db54aa9b61d5 100644
--- a/llvm/include/llvm/Analysis/ValueTracking.h
+++ b/llvm/include/llvm/Analysis/ValueTracking.h
@@ -276,6 +276,8 @@ struct KnownFPClass {
     return (KnownFPClasses & Mask) == fcNone;
   }
 
+  bool isKnownAlways(FPClassTest Mask) const { return isKnownNever(~Mask); }
+
   bool isUnknown() const {
     return KnownFPClasses == fcAllFlags && !SignBit;
   }
@@ -285,6 +287,9 @@ struct KnownFPClass {
     return isKnownNever(fcNan);
   }
 
+  /// Return true if it's known this must always be a nan.
+  bool isKnownAlwaysNaN() const { return isKnownAlways(fcNan); }
+
   /// Return true if it's known this can never be an infinity.
   bool isKnownNeverInfinity() const {
     return isKnownNever(fcInf);
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
index abadf54a967678..040dcab9310686 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -8103,13 +8103,20 @@ Instruction *InstCombinerImpl::visitFCmpInst(FCmpInst &I) {
   // If we're just checking for a NaN (ORD/UNO) and have a non-NaN operand,
   // then canonicalize the operand to 0.0.
   if (Pred == CmpInst::FCMP_ORD || Pred == CmpInst::FCMP_UNO) {
-    if (!match(Op0, m_PosZeroFP()) &&
-        isKnownNeverNaN(Op0, 0, getSimplifyQuery().getWithInstruction(&I)))
-      return replaceOperand(I, 0, ConstantFP::getZero(OpType));
-
-    if (!match(Op1, m_PosZeroFP()) &&
-        isKnownNeverNaN(Op1, 0, getSimplifyQuery().getWithInstruction(&I)))
-      return replaceOperand(I, 1, ConstantFP::getZero(OpType));
+    if (!match(Op0, m_PosZeroFP())) {
+      KnownFPClass Known0 = computeKnownFPClass(Op0, fcAllFlags, &I);
+      if (Known0.isKnownNeverNaN())
+        return replaceOperand(I, 0, ConstantFP::getZero(OpType));
+      if (Known0.isKnownAlwaysNaN())
+        return replaceOperand(I, 0, ConstantFP::getNaN(OpType));
+    }
+    if (!match(Op1, m_PosZeroFP())) {
+      KnownFPClass Known1 = computeKnownFPClass(Op1, fcAllFlags, &I);
+      if (Known1.isKnownNeverNaN())
+        return replaceOperand(I, 1, ConstantFP::getZero(OpType));
+      if (Known1.isKnownAlwaysNaN())
+        return replaceOperand(I, 1, ConstantFP::getNaN(OpType));
+    }
   }
 
   // fcmp pred (fneg X), (fneg Y) -> fcmp swap(pred) X, Y
diff --git a/llvm/test/Transforms/InstCombine/fcmp-special.ll b/llvm/test/Transforms/InstCombine/fcmp-special.ll
index 64bc86f4266c78..6c196534931f5b 100644
--- a/llvm/test/Transforms/InstCombine/fcmp-special.ll
+++ b/llvm/test/Transforms/InstCombine/fcmp-special.ll
@@ -186,6 +186,22 @@ define i1 @nnan_ops_to_fcmp_uno(float %x, float %y) {
   ret i1 %cmp
 }
 
+define i1 @nan_ops_to_fcmp_ord(float nofpclass(inf zero sub norm) %x, float nofpclass(inf zero sub norm) %y) {
+; CHECK-LABEL: @nan_ops_to_fcmp_ord(
+; CHECK-NEXT:    ret i1 false
+;
+  %cmp = fcmp ord float %x, %y
+  ret i1 %cmp
+}
+
+define i1 @nan_ops_to_fcmp_uno(float nofpclass(inf zero sub norm) %x, float nofpclass(inf zero sub norm) %y) {
+; CHECK-LABEL: @nan_ops_to_fcmp_uno(
+; CHECK-NEXT:    ret i1 true
+;
+  %cmp = fcmp uno float %x, %y
+  ret i1 %cmp
+}
+
 ; TODO: For any predicate/type/FMF, comparison to -0.0 is the same as comparison to +0.0.
 
 define i1 @negative_zero_oeq(float %x) {

>From e7a5762c24e1ef66a97a7308c062426f8b4fdfc6 Mon Sep 17 00:00:00 2001
From: Alex MacLean <amaclean at nvidia.com>
Date: Fri, 5 Jul 2024 19:03:58 +0000
Subject: [PATCH 2/2] address comments

---
 llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
index 040dcab9310686..267cbd917bdf42 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -8108,14 +8108,16 @@ Instruction *InstCombinerImpl::visitFCmpInst(FCmpInst &I) {
       if (Known0.isKnownNeverNaN())
         return replaceOperand(I, 0, ConstantFP::getZero(OpType));
       if (Known0.isKnownAlwaysNaN())
-        return replaceOperand(I, 0, ConstantFP::getNaN(OpType));
+        return replaceInstUsesWith(I,
+                                   Builder.getInt1(Pred == CmpInst::FCMP_UNO));
     }
     if (!match(Op1, m_PosZeroFP())) {
       KnownFPClass Known1 = computeKnownFPClass(Op1, fcAllFlags, &I);
       if (Known1.isKnownNeverNaN())
         return replaceOperand(I, 1, ConstantFP::getZero(OpType));
       if (Known1.isKnownAlwaysNaN())
-        return replaceOperand(I, 1, ConstantFP::getNaN(OpType));
+        return replaceInstUsesWith(I,
+                                   Builder.getInt1(Pred == CmpInst::FCMP_UNO));
     }
   }