From cb9da9e776f9a1e336bd4b808f1a9888939306b6 Mon Sep 17 00:00:00 2001
From: Benjamin Kramer <kramerb@google.com>
Date: Tue, 9 Feb 2021 04:36:43 -0800
Subject: [PATCH] [mlir][kernelgen] Test ops that can return NaN with the
 default set of inputs

NaNs no longer break comparisons, so it's safe to test them now. I also want to
put a NaN into the default set of inputs eventually, but right now that
uncovers some bugs in both Eigen and generated kernels.

PiperOrigin-RevId: 356473399
Change-Id: I2d39f0906ccbcbabb8d0bef48e0d23f6f58da620
---
 .../mlir_generated/gpu_binary_ops_test.cc     | 40 ++++------
 .../mlir_generated/gpu_unary_ops_test.cc      | 80 +++++++------------
 2 files changed, 42 insertions(+), 78 deletions(-)

diff --git a/tensorflow/core/kernels/mlir_generated/gpu_binary_ops_test.cc b/tensorflow/core/kernels/mlir_generated/gpu_binary_ops_test.cc
index a801d7a2b2c..6fc7428f16b 100644
--- a/tensorflow/core/kernels/mlir_generated/gpu_binary_ops_test.cc
+++ b/tensorflow/core/kernels/mlir_generated/gpu_binary_ops_test.cc
@@ -170,19 +170,13 @@ T baseline_div(T lhs, T rhs) {
   return lhs / rhs;
 }
 
-GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
-    Div,
-    /*test_name=*/Half, Eigen::half, Eigen::half,
-    test::DefaultInput<Eigen::half>(), test::DefaultInputNonZero<Eigen::half>(),
-    baseline_div);
-GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
-    Div,
-    /*test_name=*/Float, float, float, test::DefaultInput<float>(),
-    test::DefaultInputNonZero<float>(), baseline_div);
-GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
-    Div,
-    /*test_name=*/Double, double, double, test::DefaultInput<double>(),
-    test::DefaultInputNonZero<double>(), baseline_div);
+GENERATE_DEFAULT_TESTS(Div,
+                       /*test_name=*/Half, Eigen::half, Eigen::half,
+                       baseline_div);
+GENERATE_DEFAULT_TESTS(Div,
+                       /*test_name=*/Float, float, float, baseline_div);
+GENERATE_DEFAULT_TESTS(Div,
+                       /*test_name=*/Double, double, double, baseline_div);
 GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
     Div,
     /*test_name=*/Int16, int16, int16, test::DefaultInput<int16>(),
@@ -476,19 +470,13 @@ GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(Pow,
 
 /// Test `tf.RealDiv`.
 
-GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
-    RealDiv,
-    /*test_name=*/Half, Eigen::half, Eigen::half,
-    test::DefaultInput<Eigen::half>(), test::DefaultInputNonZero<Eigen::half>(),
-    baseline_div);
-GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
-    RealDiv,
-    /*test_name=*/Float, float, float, test::DefaultInput<float>(),
-    test::DefaultInputNonZero<float>(), baseline_div);
-GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
-    RealDiv,
-    /*test_name=*/Double, double, double, test::DefaultInput<double>(),
-    test::DefaultInputNonZero<double>(), baseline_div);
+GENERATE_DEFAULT_TESTS(RealDiv,
+                       /*test_name=*/Half, Eigen::half, Eigen::half,
+                       baseline_div);
+GENERATE_DEFAULT_TESTS(RealDiv,
+                       /*test_name=*/Float, float, float, baseline_div);
+GENERATE_DEFAULT_TESTS(RealDiv,
+                       /*test_name=*/Double, double, double, baseline_div);
 
 /// Test `tf.RightShift`.
 
diff --git a/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc b/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc
index 4d19ef01ed3..2496dcc89af 100644
--- a/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc
+++ b/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc
@@ -81,18 +81,11 @@ typename T::value_type baseline_angle(T x) {
   return std::arg(x);
 }
 
-GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES(
-    Angle, DT_COMPLEX64, DT_FLOAT,
-    test::ComplexInputFromValues<std::complex<float>>(
-        test::DefaultInputNonZero<float>(), test::DefaultInputNonZero<float>()),
-    baseline_angle, test::OpsTestConfig().AddTout().NoBufferReuse())
+GENERATE_DEFAULT_TEST(Angle, DT_COMPLEX64, DT_FLOAT, baseline_angle,
+                      test::OpsTestConfig().AddTout().NoBufferReuse())
 
-GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES(
-    Angle, DT_COMPLEX128, DT_DOUBLE,
-    test::ComplexInputFromValues<std::complex<double>>(
-        test::DefaultInputNonZero<double>(),
-        test::DefaultInputNonZero<double>()),
-    baseline_angle, test::OpsTestConfig().AddTout().NoBufferReuse())
+GENERATE_DEFAULT_TEST(Angle, DT_COMPLEX128, DT_DOUBLE, baseline_angle,
+                      test::OpsTestConfig().AddTout().NoBufferReuse())
 
 /// Test `tf.Asin`.
 
@@ -532,33 +525,24 @@ GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES_2(
 
 /// Test `tf.Log`.
 
-GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES(
-    Log, DT_FLOAT, DT_FLOAT, test::DefaultInputGreaterThanZero<float>(),
-    std::log, test::OpsTestConfig())
+GENERATE_DEFAULT_TEST(Log, DT_FLOAT, DT_FLOAT, std::log, test::OpsTestConfig())
 
-GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES(
-    Log, DT_DOUBLE, DT_DOUBLE, test::DefaultInputGreaterThanZero<double>(),
-    std::log, test::OpsTestConfig())
+GENERATE_DEFAULT_TEST(Log, DT_DOUBLE, DT_DOUBLE, std::log,
+                      test::OpsTestConfig())
 
-GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES_2(
-    Log, DT_HALF, DT_FLOAT, DT_HALF, DT_FLOAT,
-    test::DefaultInputGreaterThanZero<Eigen::half>(), std::log,
-    test::OpsTestConfig())
+GENERATE_DEFAULT_TEST_2(Log, DT_HALF, DT_FLOAT, DT_HALF, DT_FLOAT, std::log,
+                        test::OpsTestConfig())
 
 /// Test `tf.Log1p`.
 
-GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES(
-    Log1p, DT_FLOAT, DT_FLOAT, test::DefaultInputGreaterThanZero<float>(),
-    std::log1p, test::OpsTestConfig())
+GENERATE_DEFAULT_TEST(Log1p, DT_FLOAT, DT_FLOAT, std::log1p,
+                      test::OpsTestConfig())
 
-GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES(
-    Log1p, DT_DOUBLE, DT_DOUBLE, test::DefaultInputGreaterThanZero<double>(),
-    std::log1p, test::OpsTestConfig())
+GENERATE_DEFAULT_TEST(Log1p, DT_DOUBLE, DT_DOUBLE, std::log1p,
+                      test::OpsTestConfig())
 
-GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES_2(
-    Log1p, DT_HALF, DT_FLOAT, DT_HALF, DT_FLOAT,
-    test::DefaultInputGreaterThanZero<Eigen::half>(), std::log1p,
-    test::OpsTestConfig())
+GENERATE_DEFAULT_TEST_2(Log1p, DT_HALF, DT_FLOAT, DT_HALF, DT_FLOAT, std::log1p,
+                        test::OpsTestConfig())
 
 /// Test `tf.LogicalNot`
 
@@ -614,24 +598,21 @@ T baseline_rsqrt(T x) {
   return 1.0 / std::sqrt(x);
 }
 
-GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES(
-    Rsqrt, DT_FLOAT, DT_FLOAT, test::DefaultInputGreaterThanZero<float>(),
-    baseline_rsqrt, test::OpsTestConfig())
+GENERATE_DEFAULT_TEST(Rsqrt, DT_FLOAT, DT_FLOAT, baseline_rsqrt,
+                      test::OpsTestConfig())
 
-GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES(
-    Rsqrt, DT_DOUBLE, DT_DOUBLE, test::DefaultInputGreaterThanZero<double>(),
-    baseline_rsqrt, test::OpsTestConfig())
+GENERATE_DEFAULT_TEST(Rsqrt, DT_DOUBLE, DT_DOUBLE, baseline_rsqrt,
+                      test::OpsTestConfig())
 
-GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES_2(
-    Rsqrt, DT_HALF, DT_FLOAT, DT_HALF, DT_FLOAT,
-    test::DefaultInputGreaterThanZero<Eigen::half>(), baseline_rsqrt,
-    test::OpsTestConfig())
+GENERATE_DEFAULT_TEST_2(Rsqrt, DT_HALF, DT_FLOAT, DT_HALF, DT_FLOAT,
+                        baseline_rsqrt, test::OpsTestConfig())
 
 /// Test `tf.Sign`.
 
 // Reference implementation
 template <typename T>
 T baseline_sign(T x) {
+  if (isnan(x)) return x;
   if (x == 0) return 0;
   if (x < 0) return -1;
   return 1;
@@ -671,19 +652,14 @@ GENERATE_DEFAULT_TEST(Sinh, DT_DOUBLE, DT_DOUBLE, std::sinh,
 
 /// Test `tf.Sqrt`.
 
-GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES(
-    Sqrt, DT_FLOAT, DT_FLOAT, test::DefaultInputGreaterOrEqualToZero<float>(),
-    std::sqrt, test::OpsTestConfig())
+GENERATE_DEFAULT_TEST(Sqrt, DT_FLOAT, DT_FLOAT, std::sqrt,
+                      test::OpsTestConfig())
 
-GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES(
-    Sqrt, DT_DOUBLE, DT_DOUBLE,
-    test::DefaultInputGreaterOrEqualToZero<double>(), std::sqrt,
-    test::OpsTestConfig())
+GENERATE_DEFAULT_TEST(Sqrt, DT_DOUBLE, DT_DOUBLE, std::sqrt,
+                      test::OpsTestConfig())
 
-GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES_2(
-    Sqrt, DT_HALF, DT_FLOAT, DT_HALF, DT_FLOAT,
-    test::DefaultInputGreaterOrEqualToZero<Eigen::half>(), std::sqrt,
-    test::OpsTestConfig())
+GENERATE_DEFAULT_TEST_2(Sqrt, DT_HALF, DT_FLOAT, DT_HALF, DT_FLOAT, std::sqrt,
+                        test::OpsTestConfig())
 
 /// Test `tf.Tan`.