From ed365f7817eeb4b6320232a1591ea77439826390 Mon Sep 17 00:00:00 2001
From: Stephan Herhut <herhut@google.com>
Date: Mon, 8 Feb 2021 11:24:51 -0800
Subject: [PATCH] Enable generated versions of lgamma and digamma.

PiperOrigin-RevId: 356310920
Change-Id: Icb3af5e20b260e5d86f62f3e5b0e3f55f94b3451
---
 tensorflow/core/kernels/cwise_op_digamma.cc        | 3 +--
 tensorflow/core/kernels/cwise_op_gpu_digamma.cu.cc | 2 ++
 tensorflow/core/kernels/cwise_op_gpu_lgamma.cu.cc  | 2 ++
 tensorflow/core/kernels/cwise_op_lgamma.cc         | 3 +--
 tensorflow/core/kernels/mlir_generated/BUILD       | 4 ++--
 5 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/tensorflow/core/kernels/cwise_op_digamma.cc b/tensorflow/core/kernels/cwise_op_digamma.cc
index 879c3f2938b..0c3588b5e63 100644
--- a/tensorflow/core/kernels/cwise_op_digamma.cc
+++ b/tensorflow/core/kernels/cwise_op_digamma.cc
@@ -21,8 +21,7 @@ REGISTER3(UnaryOp, CPU, "Digamma", functor::digamma, float, Eigen::half,
           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)
 REGISTER3(UnaryOp, GPU, "Digamma", functor::digamma, float, Eigen::half,
           double);
 #endif
diff --git a/tensorflow/core/kernels/cwise_op_gpu_digamma.cu.cc b/tensorflow/core/kernels/cwise_op_gpu_digamma.cu.cc
index beb9475e003..010b1358336 100644
--- a/tensorflow/core/kernels/cwise_op_gpu_digamma.cu.cc
+++ b/tensorflow/core/kernels/cwise_op_gpu_digamma.cu.cc
@@ -19,7 +19,9 @@ limitations under the License.
 
 namespace tensorflow {
 namespace functor {
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 DEFINE_UNARY3(digamma, Eigen::half, float, double);
+#endif
 }  // namespace functor
 }  // namespace tensorflow
 
diff --git a/tensorflow/core/kernels/cwise_op_gpu_lgamma.cu.cc b/tensorflow/core/kernels/cwise_op_gpu_lgamma.cu.cc
index ddcef1497e7..d2e29352dfd 100644
--- a/tensorflow/core/kernels/cwise_op_gpu_lgamma.cu.cc
+++ b/tensorflow/core/kernels/cwise_op_gpu_lgamma.cu.cc
@@ -19,7 +19,9 @@ limitations under the License.
 
 namespace tensorflow {
 namespace functor {
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 DEFINE_UNARY3(lgamma, Eigen::half, float, double);
+#endif
 }  // namespace functor
 }  // namespace tensorflow
 
diff --git a/tensorflow/core/kernels/cwise_op_lgamma.cc b/tensorflow/core/kernels/cwise_op_lgamma.cc
index 15734f2bb87..b84997dade5 100644
--- a/tensorflow/core/kernels/cwise_op_lgamma.cc
+++ b/tensorflow/core/kernels/cwise_op_lgamma.cc
@@ -20,8 +20,7 @@ namespace tensorflow {
 REGISTER3(UnaryOp, CPU, "Lgamma", functor::lgamma, float, Eigen::half, 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)
 REGISTER3(UnaryOp, GPU, "Lgamma", functor::lgamma, float, Eigen::half, double);
 #endif
 #endif
diff --git a/tensorflow/core/kernels/mlir_generated/BUILD b/tensorflow/core/kernels/mlir_generated/BUILD
index 924142dfe1a..04d9241d88b 100644
--- a/tensorflow/core/kernels/mlir_generated/BUILD
+++ b/tensorflow/core/kernels/mlir_generated/BUILD
@@ -54,6 +54,7 @@ filegroup(
         "gpu_op_conj.cc",
         "gpu_op_cos.cc",
         "gpu_op_cosh.cc",
+        "gpu_op_digamma.cc",
         "gpu_op_erf.cc",
         "gpu_op_erfc.cc",
         "gpu_op_floor.cc",
@@ -62,6 +63,7 @@ filegroup(
         "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",
@@ -81,10 +83,8 @@ filegroup(
 filegroup(
     name = "experimental_unary_gpu_kernel_srcs",
     srcs = [
-        "gpu_op_digamma.cc",
         "gpu_op_exp.cc",
         "gpu_op_expm1.cc",
-        "gpu_op_lgamma.cc",
         "gpu_op_sign.cc",
     ],
     compatible_with = get_compatible_with_cloud(),