From 21928d2d023488b0d0af54bd89095149de1d455a Mon Sep 17 00:00:00 2001 From: Stephan Herhut Date: Mon, 1 Feb 2021 12:08:15 -0800 Subject: [PATCH] Enable the generated versions of invert, logical_not and neg. PiperOrigin-RevId: 354995894 Change-Id: I563774705838bd04e03f17f878fcef8dfcc981c6 --- tensorflow/core/kernels/cwise_op_gpu_invert.cu.cc | 4 ++++ tensorflow/core/kernels/cwise_op_gpu_logical_not.cu.cc | 2 ++ tensorflow/core/kernels/cwise_op_gpu_neg.cu.cc | 5 ++++- tensorflow/core/kernels/cwise_op_invert.cc | 3 +-- tensorflow/core/kernels/cwise_op_logical_not.cc | 3 +-- tensorflow/core/kernels/cwise_op_neg_1.cc | 3 +-- tensorflow/core/kernels/cwise_op_neg_2.cc | 9 +++------ tensorflow/core/kernels/mlir_generated/BUILD | 6 +++--- 8 files changed, 19 insertions(+), 16 deletions(-) diff --git a/tensorflow/core/kernels/cwise_op_gpu_invert.cu.cc b/tensorflow/core/kernels/cwise_op_gpu_invert.cu.cc index db720372667..3526bcf65cb 100644 --- a/tensorflow/core/kernels/cwise_op_gpu_invert.cu.cc +++ b/tensorflow/core/kernels/cwise_op_gpu_invert.cu.cc @@ -19,7 +19,11 @@ limitations under the License. namespace tensorflow { namespace functor { +#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) DEFINE_UNARY8(invert, int8, int16, int32, int64, uint8, uint16, uint32, uint64); +#else +DEFINE_UNARY4(invert, uint8, uint16, uint32, uint64); +#endif } // namespace functor } // namespace tensorflow diff --git a/tensorflow/core/kernels/cwise_op_gpu_logical_not.cu.cc b/tensorflow/core/kernels/cwise_op_gpu_logical_not.cu.cc index f619b9b5d77..f7517987aed 100644 --- a/tensorflow/core/kernels/cwise_op_gpu_logical_not.cu.cc +++ b/tensorflow/core/kernels/cwise_op_gpu_logical_not.cu.cc @@ -19,7 +19,9 @@ limitations under the License. namespace tensorflow { namespace functor { +#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) template struct UnaryFunctor; +#endif } // namespace functor } // namespace tensorflow diff --git a/tensorflow/core/kernels/cwise_op_gpu_neg.cu.cc b/tensorflow/core/kernels/cwise_op_gpu_neg.cu.cc index 4f7bb9b2075..23ca6d61a93 100644 --- a/tensorflow/core/kernels/cwise_op_gpu_neg.cu.cc +++ b/tensorflow/core/kernels/cwise_op_gpu_neg.cu.cc @@ -19,8 +19,11 @@ limitations under the License. namespace tensorflow { namespace functor { +#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) DEFINE_UNARY4(neg, int8, int16, int32, int64); -DEFINE_UNARY6(neg, Eigen::half, float, double, bfloat16, complex64, complex128); +DEFINE_UNARY3(neg, Eigen::half, float, double); +#endif +DEFINE_UNARY3(neg, bfloat16, complex64, complex128); } // namespace functor } // namespace tensorflow diff --git a/tensorflow/core/kernels/cwise_op_invert.cc b/tensorflow/core/kernels/cwise_op_invert.cc index bdf325f5661..9f7bcbfadbe 100644 --- a/tensorflow/core/kernels/cwise_op_invert.cc +++ b/tensorflow/core/kernels/cwise_op_invert.cc @@ -21,8 +21,7 @@ REGISTER8(UnaryOp, CPU, "Invert", functor::invert, int8, int16, int32, int64, #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) REGISTER8(UnaryOp, GPU, "Invert", functor::invert, int8, int16, int32, int64, uint8, uint16, uint32, uint64); #else diff --git a/tensorflow/core/kernels/cwise_op_logical_not.cc b/tensorflow/core/kernels/cwise_op_logical_not.cc index d38876590ce..ac0863c0707 100644 --- a/tensorflow/core/kernels/cwise_op_logical_not.cc +++ b/tensorflow/core/kernels/cwise_op_logical_not.cc @@ -19,8 +19,7 @@ namespace tensorflow { REGISTER_KERNEL_BUILDER(Name("LogicalNot").Device(DEVICE_CPU), UnaryOp); #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_KERNEL_BUILDER(Name("LogicalNot").Device(DEVICE_GPU), UnaryOp); #endif diff --git a/tensorflow/core/kernels/cwise_op_neg_1.cc b/tensorflow/core/kernels/cwise_op_neg_1.cc index 9551d845e35..19000abdfc2 100644 --- a/tensorflow/core/kernels/cwise_op_neg_1.cc +++ b/tensorflow/core/kernels/cwise_op_neg_1.cc @@ -20,8 +20,7 @@ REGISTER4(UnaryOp, CPU, "Neg", functor::neg, int8, int16, int32, int64); #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) REGISTER3(UnaryOp, GPU, "Neg", functor::neg, int8, int16, int64); #endif diff --git a/tensorflow/core/kernels/cwise_op_neg_2.cc b/tensorflow/core/kernels/cwise_op_neg_2.cc index a2857fa1337..02bf64c1462 100644 --- a/tensorflow/core/kernels/cwise_op_neg_2.cc +++ b/tensorflow/core/kernels/cwise_op_neg_2.cc @@ -20,12 +20,9 @@ REGISTER6(UnaryOp, CPU, "Neg", functor::neg, Eigen::half, float, double, bfloat16, complex64, complex128); #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM -#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ - !defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED) -REGISTER6(UnaryOp, GPU, "Neg", functor::neg, Eigen::half, float, double, - bfloat16, complex64, complex128); -#else +#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) +REGISTER3(UnaryOp, GPU, "Neg", functor::neg, Eigen::half, float, double); +#endif REGISTER3(UnaryOp, GPU, "Neg", functor::neg, bfloat16, complex64, complex128); #endif -#endif } // namespace tensorflow diff --git a/tensorflow/core/kernels/mlir_generated/BUILD b/tensorflow/core/kernels/mlir_generated/BUILD index 6b10e4e6054..b6ea132315d 100644 --- a/tensorflow/core/kernels/mlir_generated/BUILD +++ b/tensorflow/core/kernels/mlir_generated/BUILD @@ -51,6 +51,9 @@ filegroup( "gpu_op_erfc.cc", "gpu_op_floor.cc", "gpu_op_imag.cc", + "gpu_op_invert.cc", + "gpu_op_logical_not.cc", + "gpu_op_neg.cc", "gpu_op_real.cc", "gpu_op_sqrt.cc", "gpu_op_square.cc", @@ -69,15 +72,12 @@ filegroup( "gpu_op_asinh.cc", "gpu_op_exp.cc", "gpu_op_expm1.cc", - "gpu_op_invert.cc", "gpu_op_is_finite.cc", "gpu_op_is_inf.cc", "gpu_op_is_nan.cc", "gpu_op_lgamma.cc", "gpu_op_log.cc", "gpu_op_log1p.cc", - "gpu_op_logical_not.cc", - "gpu_op_neg.cc", "gpu_op_rsqrt.cc", "gpu_op_sign.cc", "gpu_op_sin.cc",