[MLIR][KernelGen] Enable JIT-compiled GPU kernels for min and max

Enable the JIT-compiled GPU kernels for the following ops and dtypes: - Maximum i8, ui16, ui32, ui64 - Minimum i8, ui16, ui32, ui64 PiperOrigin-RevId: 407580629 Change-Id: I4ec68881a2b29779697e98a200e216b23a5fcbf0
Marbck · Nov 4, 2021 · 61e4fd4 · 61e4fd4
1 parent 89fc3d2
commit 61e4fd4
Show file tree

Hide file tree

Showing 3 changed files with 0 additions and 10 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
@@ -810,8 +810,6 @@ GENERATE_DEFAULT_TESTS(Maximum, /*test_name=*/UInt8, uint8_t, uint8_t,
                        test::OpsTestConfig().ExpectStrictlyEqual())
 
 /// Test the experimental JIT-compiled kernels.
-#if defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) && \
-    defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
 GENERATE_DEFAULT_TESTS(Maximum, /*test_name=*/Int8, int8_t, int8_t,
                        baseline_maximum,
                        test::OpsTestConfig().ExpectStrictlyEqual())
@@ -825,7 +823,6 @@ GENERATE_DEFAULT_TESTS(Maximum, /*test_name=*/UInt32, uint32_t, uint32_t,
 GENERATE_DEFAULT_TESTS(Maximum, /*test_name=*/UInt64, uint64_t, uint64_t,
                        baseline_maximum,
                        test::OpsTestConfig().ExpectStrictlyEqual())
-#endif
 
 /// Test `tf.Minmum`.
 
@@ -854,8 +851,6 @@ GENERATE_DEFAULT_TESTS(Minimum, /*test_name=*/UInt8, uint8_t, uint8_t,
                        test::OpsTestConfig().ExpectStrictlyEqual())
 
 /// Test the experimental JIT-compiled kernels.
-#if defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) && \
-    defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
 GENERATE_DEFAULT_TESTS(Minimum, /*test_name=*/Int8, int8_t, int8_t,
                        baseline_minimum,
                        test::OpsTestConfig().ExpectStrictlyEqual())
@@ -869,7 +864,6 @@ GENERATE_DEFAULT_TESTS(Minimum, /*test_name=*/UInt32, uint32_t, uint32_t,
 GENERATE_DEFAULT_TESTS(Minimum, /*test_name=*/UInt64, uint64_t, uint64_t,
                        baseline_minimum,
                        test::OpsTestConfig().ExpectStrictlyEqual())
-#endif
 
 /// Test `tf.Mul`.
 

diff --git a/tensorflow/core/kernels/mlir_generated/gpu_op_maximum.cc b/tensorflow/core/kernels/mlir_generated/gpu_op_maximum.cc
@@ -27,11 +27,9 @@ GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Maximum, DT_INT64);
 GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Maximum, DT_UINT8);
 
 // These kernels are JIT-compiled.
-#if defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
 GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Maximum, DT_INT8);
 GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Maximum, DT_UINT16);
 GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Maximum, DT_UINT32);
 GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Maximum, DT_UINT64);
-#endif
 
 }  // namespace tensorflow
diff --git a/tensorflow/core/kernels/mlir_generated/gpu_op_minimum.cc b/tensorflow/core/kernels/mlir_generated/gpu_op_minimum.cc
@@ -27,11 +27,9 @@ GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Minimum, DT_INT64);
 GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Minimum, DT_UINT8);
 
 // These kernels are JIT-compiled.
-#if defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
 GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Minimum, DT_INT8);
 GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Minimum, DT_UINT16);
 GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Minimum, DT_UINT32);
 GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Minimum, DT_UINT64);
-#endif
 
 }  // namespace tensorflow