From 2ccbb3f1ceeb326309c0685305d1dfd38d5dbcc2 Mon Sep 17 00:00:00 2001
From: Geoffrey Martin-Noble <gcmn@google.com>
Date: Thu, 12 Sep 2019 15:43:26 -0700
Subject: [PATCH] Cmpf constant folding for nan and inf

PiperOrigin-RevId: 268783645
---
 mlir/lib/Dialect/StandardOps/Ops.cpp    |  8 +--
 mlir/test/Transforms/constant-fold.mlir | 84 ++++++++++++++++++++++---
 2 files changed, 80 insertions(+), 12 deletions(-)

diff --git a/mlir/lib/Dialect/StandardOps/Ops.cpp b/mlir/lib/Dialect/StandardOps/Ops.cpp
index 86966db8cb0b..ef7b795d5f0e 100644
--- a/mlir/lib/Dialect/StandardOps/Ops.cpp
+++ b/mlir/lib/Dialect/StandardOps/Ops.cpp
@@ -1011,10 +1011,10 @@ OpFoldResult CmpFOp::fold(ArrayRef<Attribute> operands) {
 
   auto lhs = operands.front().dyn_cast_or_null<FloatAttr>();
   auto rhs = operands.back().dyn_cast_or_null<FloatAttr>();
-  if (!lhs || !rhs ||
-      // TODO(b/122019992) Implement and test constant folding for nan/inf when
-      // it is possible to have constant nan/inf
-      !lhs.getValue().isFinite() || !rhs.getValue().isFinite())
+
+  // TODO(gcmn) We could actually do some intelligent things if we know only one
+  // of the operands, but it's inf or nan.
+  if (!lhs || !rhs)
     return {};
 
   auto val = applyCmpPredicate(getPredicate(), lhs.getValue(), rhs.getValue());
diff --git a/mlir/test/Transforms/constant-fold.mlir b/mlir/test/Transforms/constant-fold.mlir
index 7e2a19301038..e1433b5a72c0 100644
--- a/mlir/test/Transforms/constant-fold.mlir
+++ b/mlir/test/Transforms/constant-fold.mlir
@@ -375,19 +375,87 @@ func @cmpf_normal_numbers() -> (i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1,
 // -----
 
 // CHECK-LABEL: func @cmpf_nan
-func @cmpf_nans() {
-  // TODO(b/122019992) Add tests for nan constant folding when it's possible to
-  // have nan constants
-  return
+func @cmpf_nan() -> (i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1) {
+  %c42 = constant 42. : f32
+  %cqnan = constant 0xFFFFFFFF : f32
+  // CHECK-DAG: [[F:%.+]] = constant 0 : i1
+  // CHECK-DAG: [[T:%.+]] = constant 1 : i1
+  // CHECK-NEXT: return [[F]],
+  %0 = cmpf "false", %c42, %cqnan : f32
+  // CHECK-SAME: [[F]]
+  %1 = cmpf "oeq", %c42, %cqnan : f32
+  // CHECK-SAME: [[F]],
+  %2 = cmpf "ogt", %c42, %cqnan : f32
+  // CHECK-SAME: [[F]],
+  %3 = cmpf "oge", %c42, %cqnan : f32
+  // CHECK-SAME: [[F]],
+  %4 = cmpf "olt", %c42, %cqnan : f32
+  // CHECK-SAME: [[F]],
+  %5 = cmpf "ole", %c42, %cqnan : f32
+  // CHECK-SAME: [[F]],
+  %6 = cmpf "one", %c42, %cqnan : f32
+  // CHECK-SAME: [[F]],
+  %7 = cmpf "ord", %c42, %cqnan : f32
+  // CHECK-SAME: [[T]],
+  %8 = cmpf "ueq", %c42, %cqnan : f32
+  // CHECK-SAME: [[T]],
+  %9 = cmpf "ugt", %c42, %cqnan : f32
+  // CHECK-SAME: [[T]],
+  %10 = cmpf "uge", %c42, %cqnan : f32
+  // CHECK-SAME: [[T]],
+  %11 = cmpf "ult", %c42, %cqnan : f32
+  // CHECK-SAME: [[T]],
+  %12 = cmpf "ule", %c42, %cqnan : f32
+  // CHECK-SAME: [[T]],
+  %13 = cmpf "une", %c42, %cqnan : f32
+  // CHECK-SAME: [[T]],
+  %14 = cmpf "uno", %c42, %cqnan : f32
+  // CHECK-SAME: [[T]]
+  %15 = cmpf "true", %c42, %cqnan : f32
+  return %0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15 : i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1
 }
 
 // -----
 
 // CHECK-LABEL: func @cmpf_inf
-func @cmpf_inf() {
-  // TODO(b/122019992) Add tests for inf constant folding when it's possible to
-  // have inf constants
-  return
+func @cmpf_inf() -> (i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1) {
+  %c42 = constant 42. : f32
+  %cpinf = constant 0x7F800000 : f32
+  // CHECK-DAG: [[F:%.+]] = constant 0 : i1
+  // CHECK-DAG: [[T:%.+]] = constant 1 : i1
+  // CHECK-NEXT: return [[F]],
+  %0 = cmpf "false", %c42, %cpinf: f32
+  // CHECK-SAME: [[F]]
+  %1 = cmpf "oeq", %c42, %cpinf: f32
+  // CHECK-SAME: [[F]],
+  %2 = cmpf "ogt", %c42, %cpinf: f32
+  // CHECK-SAME: [[F]],
+  %3 = cmpf "oge", %c42, %cpinf: f32
+  // CHECK-SAME: [[T]],
+  %4 = cmpf "olt", %c42, %cpinf: f32
+  // CHECK-SAME: [[T]],
+  %5 = cmpf "ole", %c42, %cpinf: f32
+  // CHECK-SAME: [[T]],
+  %6 = cmpf "one", %c42, %cpinf: f32
+  // CHECK-SAME: [[T]],
+  %7 = cmpf "ord", %c42, %cpinf: f32
+  // CHECK-SAME: [[F]],
+  %8 = cmpf "ueq", %c42, %cpinf: f32
+  // CHECK-SAME: [[F]],
+  %9 = cmpf "ugt", %c42, %cpinf: f32
+  // CHECK-SAME: [[F]],
+  %10 = cmpf "uge", %c42, %cpinf: f32
+  // CHECK-SAME: [[T]],
+  %11 = cmpf "ult", %c42, %cpinf: f32
+  // CHECK-SAME: [[T]],
+  %12 = cmpf "ule", %c42, %cpinf: f32
+  // CHECK-SAME: [[T]],
+  %13 = cmpf "une", %c42, %cpinf: f32
+  // CHECK-SAME: [[F]],
+  %14 = cmpf "uno", %c42, %cpinf: f32
+  // CHECK-SAME: [[T]]
+  %15 = cmpf "true", %c42, %cpinf: f32
+  return %0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15 : i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1, i1
 }
 
 // -----