From 2cf1277b11784bd29ac3793713ac4fdb238f049b Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Thu, 30 Jul 2020 07:26:11 -0700
Subject: [PATCH] Only instantiate abs kernel templates if needed.
When MLIR generated GPU kernels are used, there is no need to instantiate the
GPU kernels based on Eigen. This would just lead to unused code being linked
in.
PiperOrigin-RevId: 323998013
Change-Id: Ia1a3057f39e5372ae2da15631f0c65a83a2ca0c8
---
tensorflow/core/kernels/cwise_op_gpu_abs.cu.cc | 4 ++++
1 file changed, 4 insertions(+)
diff --git a/tensorflow/core/kernels/cwise_op_gpu_abs.cu.cc b/tensorflow/core/kernels/cwise_op_gpu_abs.cu.cc
index f5b4f014c2c..ee14f2950d6 100644
--- a/tensorflow/core/kernels/cwise_op_gpu_abs.cu.cc
+++ b/tensorflow/core/kernels/cwise_op_gpu_abs.cu.cc
@@ -19,7 +19,11 @@ limitations under the License.
namespace tensorflow {
namespace functor {
+#ifdef MLIR_GENERATED_GPU_KERNELS_ENABLED
+DEFINE_UNARY2(abs, complex64, complex128);
+#else
DEFINE_UNARY6(abs, Eigen::half, float, double, int64, complex64, complex128);
+#endif
} // namespace functor
} // namespace tensorflow