diff --git a/tensorflow/core/kernels/cwise_op_gpu_imag.cu.cc b/tensorflow/core/kernels/cwise_op_gpu_imag.cu.cc index ca7fea8a79e..d6653b55042 100644 --- a/tensorflow/core/kernels/cwise_op_gpu_imag.cu.cc +++ b/tensorflow/core/kernels/cwise_op_gpu_imag.cu.cc @@ -19,7 +19,9 @@ limitations under the License. namespace tensorflow { namespace functor { +#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) DEFINE_UNARY2(get_imag, complex64, complex128); +#endif } // namespace functor } // namespace tensorflow diff --git a/tensorflow/core/kernels/cwise_op_gpu_real.cu.cc b/tensorflow/core/kernels/cwise_op_gpu_real.cu.cc index 9b7bc624ed7..564fbe9aea3 100644 --- a/tensorflow/core/kernels/cwise_op_gpu_real.cu.cc +++ b/tensorflow/core/kernels/cwise_op_gpu_real.cu.cc @@ -19,7 +19,9 @@ limitations under the License. namespace tensorflow { namespace functor { +#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) DEFINE_UNARY2(get_real, complex64, complex128); +#endif } // namespace functor } // namespace tensorflow diff --git a/tensorflow/core/kernels/cwise_op_imag.cc b/tensorflow/core/kernels/cwise_op_imag.cc index 9e34c09b936..14e163e5123 100644 --- a/tensorflow/core/kernels/cwise_op_imag.cc +++ b/tensorflow/core/kernels/cwise_op_imag.cc @@ -27,8 +27,7 @@ REGISTER_COMPLEX(CPU, float, complex64); REGISTER_COMPLEX(CPU, double, complex128); #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM -#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ - !defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED) +#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) REGISTER_COMPLEX(GPU, float, complex64); REGISTER_COMPLEX(GPU, double, complex128); #endif diff --git a/tensorflow/core/kernels/cwise_op_real.cc b/tensorflow/core/kernels/cwise_op_real.cc index 351376875b6..4b1e4de667b 100644 --- a/tensorflow/core/kernels/cwise_op_real.cc +++ b/tensorflow/core/kernels/cwise_op_real.cc @@ -28,8 +28,7 @@ REGISTER_COMPLEX(CPU, float, complex64); REGISTER_COMPLEX(CPU, double, complex128); #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM -#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ - !defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED) +#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) REGISTER_COMPLEX(GPU, float, complex64); REGISTER_COMPLEX(GPU, double, complex128); #endif diff --git a/tensorflow/core/kernels/mlir_generated/BUILD b/tensorflow/core/kernels/mlir_generated/BUILD index 613628b82a0..6b10e4e6054 100644 --- a/tensorflow/core/kernels/mlir_generated/BUILD +++ b/tensorflow/core/kernels/mlir_generated/BUILD @@ -50,6 +50,8 @@ filegroup( "gpu_op_erf.cc", "gpu_op_erfc.cc", "gpu_op_floor.cc", + "gpu_op_imag.cc", + "gpu_op_real.cc", "gpu_op_sqrt.cc", "gpu_op_square.cc", "gpu_op_tan.cc", @@ -67,7 +69,6 @@ filegroup( "gpu_op_asinh.cc", "gpu_op_exp.cc", "gpu_op_expm1.cc", - "gpu_op_imag.cc", "gpu_op_invert.cc", "gpu_op_is_finite.cc", "gpu_op_is_inf.cc", @@ -77,7 +78,6 @@ filegroup( "gpu_op_log1p.cc", "gpu_op_logical_not.cc", "gpu_op_neg.cc", - "gpu_op_real.cc", "gpu_op_rsqrt.cc", "gpu_op_sign.cc", "gpu_op_sin.cc",