[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 61acc5d969e6d88013bfc88249ecac56de12bce1 Mon Sep 17 00:00:00 2001
From: Alex MacLean <amaclean at nvidia.com>
Date: Fri, 5 Jul 2024 19:52:55 +0000
Subject: [PATCH 1/2] pre-commit tests

---
 .../Transforms/InstCombine/fcmp-special.ll     | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/llvm/test/Transforms/InstCombine/fcmp-special.ll b/llvm/test/Transforms/InstCombine/fcmp-special.ll
index 64bc86f4266c7..4a82e6d7b265c 100644
--- a/llvm/test/Transforms/InstCombine/fcmp-special.ll
+++ b/llvm/test/Transforms/InstCombine/fcmp-special.ll
@@ -186,6 +186,24 @@ 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:    [[CMP:%.*]] = fcmp ord float [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %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:    [[CMP:%.*]] = fcmp uno float [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %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 add8ceef6c9d45b0e9cef65b38f76b129fa32584 Mon Sep 17 00:00:00 2001
From: Alex MacLean <amaclean at nvidia.com>
Date: Fri, 5 Jul 2024 19:55:17 +0000
Subject: [PATCH 2/2] [InstCombine] canonicalize operands of fcmp ord/uno if
 they are known NaN

---
 llvm/include/llvm/Analysis/ValueTracking.h    |  5 ++++
 .../InstCombine/InstCombineCompares.cpp       | 23 +++++++++++++------
 .../Transforms/InstCombine/fcmp-special.ll    |  6 ++---
 3 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/llvm/include/llvm/Analysis/ValueTracking.h b/llvm/include/llvm/Analysis/ValueTracking.h
index a67ad501982d2..51db54aa9b61d 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 abadf54a96767..267cbd917bdf4 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -8103,13 +8103,22 @@ 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 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 replaceInstUsesWith(I,
+                                   Builder.getInt1(Pred == CmpInst::FCMP_UNO));
+    }
   }
 
   // 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 4a82e6d7b265c..6c196534931f5 100644
--- a/llvm/test/Transforms/InstCombine/fcmp-special.ll
+++ b/llvm/test/Transforms/InstCombine/fcmp-special.ll
@@ -188,8 +188,7 @@ define i1 @nnan_ops_to_fcmp_uno(float %x, float %y) {
 
 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:    [[CMP:%.*]] = fcmp ord float [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
   %cmp = fcmp ord float %x, %y
   ret i1 %cmp
@@ -197,8 +196,7 @@ define i1 @nan_ops_to_fcmp_ord(float nofpclass(inf zero sub norm) %x, float nofp
 
 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:    [[CMP:%.*]] = fcmp uno float [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 true
 ;
   %cmp = fcmp uno float %x, %y
   ret i1 %cmp