Enable the generated version of comparison operations.

PiperOrigin-RevId: 361104492 Change-Id: Icdd837233eb2772c03a94dce577424286caf1714
2021-03-05 02:35:31 -08:00 · 2021-03-05 02:35:31 -08:00 · 184cf2d642
commit 184cf2d642
parent a1b5beec58
15 changed files with 20 additions and 34 deletions
--- a/tensorflow/core/kernels/cwise_op_equal_to_1.cc
+++ b/tensorflow/core/kernels/cwise_op_equal_to_1.cc
@ -27,8 +27,7 @@ REGISTER_KERNEL_BUILDER(
    Name("ApproximateEqual").Device(DEVICE_CPU).TypeConstraint<double>("T"),
    ApproximateEqualOp<CPUDevice, double>);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
-    !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 REGISTER4(BinaryOp, GPU, "Equal", functor::equal_to, float, Eigen::half, double,
          uint8);
 #else
--- a/tensorflow/core/kernels/cwise_op_equal_to_2.cc
+++ b/tensorflow/core/kernels/cwise_op_equal_to_2.cc
@ -25,8 +25,7 @@ namespace tensorflow {
 REGISTER6(BinaryOp, CPU, "Equal", functor::equal_to, int32, int64, complex64,
          complex128, tstring, bool);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
-    !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 REGISTER6(BinaryOp, GPU, "Equal", functor::equal_to, int8, int16, int64,
          complex64, complex128, bool);
 #else
--- a/tensorflow/core/kernels/cwise_op_gpu_equal_to.cu.cc
+++ b/tensorflow/core/kernels/cwise_op_gpu_equal_to.cu.cc
@ -19,8 +19,7 @@ limitations under the License.

 namespace tensorflow {
 namespace functor {
-#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
-    !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 DEFINE_BINARY10(equal_to, float, Eigen::half, double, uint8, int8, int16, int64,
                complex64, complex128, bool);
 #else
--- a/tensorflow/core/kernels/cwise_op_gpu_greater.cu.cc
+++ b/tensorflow/core/kernels/cwise_op_gpu_greater.cu.cc
@ -19,8 +19,7 @@ limitations under the License.

 namespace tensorflow {
 namespace functor {
-#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
-    !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 DEFINE_BINARY10(greater, Eigen::half, float, double, int64, uint8, uint16,
                uint32, uint64, int8, int16);
 #else
--- a/tensorflow/core/kernels/cwise_op_gpu_greater_equal.cu.cc
+++ b/tensorflow/core/kernels/cwise_op_gpu_greater_equal.cu.cc
@ -19,8 +19,7 @@ limitations under the License.

 namespace tensorflow {
 namespace functor {
-#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
-    !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 DEFINE_BINARY10(greater_equal, Eigen::half, float, double, int64, uint8, uint16,
                uint32, uint64, int8, int16);
 #else
--- a/tensorflow/core/kernels/cwise_op_gpu_less.cu.cc
+++ b/tensorflow/core/kernels/cwise_op_gpu_less.cu.cc
@ -19,8 +19,7 @@ limitations under the License.

 namespace tensorflow {
 namespace functor {
-#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
-    !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 DEFINE_BINARY10(less, Eigen::half, float, double, int64, uint8, uint16, uint32,
                uint64, int8, int16);
 #else
--- a/tensorflow/core/kernels/cwise_op_gpu_less_equal.cu.cc
+++ b/tensorflow/core/kernels/cwise_op_gpu_less_equal.cu.cc
@ -19,8 +19,7 @@ limitations under the License.

 namespace tensorflow {
 namespace functor {
-#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
-    !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 DEFINE_BINARY10(less_equal, Eigen::half, float, double, int64, uint8, uint16,
                uint32, uint64, int8, int16);
 #else
--- a/tensorflow/core/kernels/cwise_op_gpu_not_equal_to.cu.cc
+++ b/tensorflow/core/kernels/cwise_op_gpu_not_equal_to.cu.cc
@ -19,8 +19,7 @@ limitations under the License.

 namespace tensorflow {
 namespace functor {
-#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
-    !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 DEFINE_BINARY7(not_equal_to, float, Eigen::half, double, int8, int16, int64,
               bool);
 #endif
--- a/tensorflow/core/kernels/cwise_op_greater.cc
+++ b/tensorflow/core/kernels/cwise_op_greater.cc
@ -20,8 +20,7 @@ REGISTER9(BinaryOp, CPU, "Greater", functor::greater, float, Eigen::half,
          double, int32, int64, uint8, uint16, uint32, uint64);
 REGISTER3(BinaryOp, CPU, "Greater", functor::greater, int8, int16, bfloat16);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
-    !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 REGISTER6(BinaryOp, GPU, "Greater", functor::greater, float, Eigen::half,
          double, int8, int16, int64);
 REGISTER4(BinaryOp, GPU, "Greater", functor::greater, uint8, uint16, uint32,
--- a/tensorflow/core/kernels/cwise_op_greater_equal.cc
+++ b/tensorflow/core/kernels/cwise_op_greater_equal.cc
@ -21,8 +21,7 @@ REGISTER9(BinaryOp, CPU, "GreaterEqual", functor::greater_equal, float,
 REGISTER3(BinaryOp, CPU, "GreaterEqual", functor::greater_equal, int8, int16,
          bfloat16);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
-    !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 REGISTER9(BinaryOp, GPU, "GreaterEqual", functor::greater_equal, float,
          Eigen::half, double, int64, uint8, uint16, uint32, uint64, int8);
 REGISTER(BinaryOp, GPU, "GreaterEqual", functor::greater_equal, int16);
--- a/tensorflow/core/kernels/cwise_op_less.cc
+++ b/tensorflow/core/kernels/cwise_op_less.cc
@ -22,8 +22,7 @@ REGISTER7(BinaryOp, CPU, "Less", functor::less, uint8, uint16, uint32, uint64,
          int8, int16, int64);

 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
-    !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 REGISTER9(BinaryOp, GPU, "Less", functor::less, float, Eigen::half, double,
          int64, uint8, uint16, uint32, uint64, int8);
 REGISTER(BinaryOp, GPU, "Less", functor::less, int16);
--- a/tensorflow/core/kernels/cwise_op_less_equal.cc
+++ b/tensorflow/core/kernels/cwise_op_less_equal.cc
@ -22,8 +22,7 @@ REGISTER7(BinaryOp, CPU, "LessEqual", functor::less_equal, int64, uint8, uint16,
          uint32, uint64, int8, int16);

 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
-    !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 REGISTER9(BinaryOp, GPU, "LessEqual", functor::less_equal, float, Eigen::half,
          double, int64, uint8, uint16, uint32, uint64, int8);
 REGISTER(BinaryOp, GPU, "LessEqual", functor::less_equal, int16);
--- a/tensorflow/core/kernels/cwise_op_not_equal_to_1.cc
+++ b/tensorflow/core/kernels/cwise_op_not_equal_to_1.cc
@ -21,8 +21,7 @@ REGISTER7(BinaryOp, CPU, "NotEqual", functor::not_equal_to, float, Eigen::half,
 REGISTER7(BinaryOp, CPU, "NotEqual", functor::not_equal_to, uint16, uint32,
          uint64, qint8, qint16, quint8, quint16);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
-    !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 REGISTER4(BinaryOp, GPU, "NotEqual", functor::not_equal_to, float, Eigen::half,
          double, uint8);
 #else
--- a/tensorflow/core/kernels/cwise_op_not_equal_to_2.cc
+++ b/tensorflow/core/kernels/cwise_op_not_equal_to_2.cc
@ -25,8 +25,7 @@ namespace tensorflow {
 REGISTER6(BinaryOp, CPU, "NotEqual", functor::not_equal_to, int32, int64,
          complex64, complex128, tstring, bool);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
-    !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 REGISTER6(BinaryOp, GPU, "NotEqual", functor::not_equal_to, int8, int16, int64,
          complex64, complex128, bool);
 #else
--- a/tensorflow/core/kernels/mlir_generated/BUILD
+++ b/tensorflow/core/kernels/mlir_generated/BUILD
@ -107,10 +107,16 @@ filegroup(
        "gpu_op_bitwise_or.cc",
        "gpu_op_bitwise_xor.cc",
        "gpu_op_complex.cc",
+        "gpu_op_equal.cc",
+        "gpu_op_greater.cc",
+        "gpu_op_greater_equal.cc",
+        "gpu_op_less.cc",
+        "gpu_op_less_equal.cc",
        "gpu_op_logical_and.cc",
        "gpu_op_logical_or.cc",
        "gpu_op_maximum.cc",
        "gpu_op_minimum.cc",
+        "gpu_op_not_equal.cc",
    ],
    compatible_with = get_compatible_with_cloud(),
 )
@ -120,15 +126,9 @@ filegroup(
    srcs = [
        "gpu_op_atan2.cc",
        "gpu_op_div.cc",
-        "gpu_op_equal.cc",
        "gpu_op_floor_div.cc",
-        "gpu_op_greater.cc",
-        "gpu_op_greater_equal.cc",
        "gpu_op_left_shift.cc",
-        "gpu_op_less.cc",
-        "gpu_op_less_equal.cc",
        "gpu_op_mul.cc",
-        "gpu_op_not_equal.cc",
        "gpu_op_pow.cc",
        "gpu_op_right_shift.cc",
        "gpu_op_squared_difference.cc",