commit a163ecfbe2cded83fcf960516754bf78b061e9c0
author Adrian Kuegel <akuegel@google.com> Mon Sep 27 05:28:03 2021 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> Mon Sep 27 05:33:28 2021 -0700
tree f2034844f70fbfa6583dae2f89b845728f530ad7
parent 4d5a1c23b8b5a972172e6cf470d84e8dd85572d7

Enable MLIR generated Sigmoid GPU kernel by default.

PiperOrigin-RevId: 399166253
Change-Id: If71a566f63108ad5067417e77b2e394287a1f5b8

tensorflow/core/kernels/cwise_op_gpu_sigmoid.cu.cc

diff --git a/tensorflow/core/kernels/cwise_op_gpu_sigmoid.cu.cc b/tensorflow/core/kernels/cwise_op_gpu_sigmoid.cu.cc
index 1cd1129..eacb304 100644
--- a/tensorflow/core/kernels/cwise_op_gpu_sigmoid.cu.cc
+++ b/tensorflow/core/kernels/cwise_op_gpu_sigmoid.cu.cc
@@ -20,8 +20,7 @@
 
 namespace tensorflow {
 namespace functor {
-#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
-    !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 DEFINE_UNARY3(sigmoid, Eigen::half, float, double);
 #endif
 DEFINE_SIMPLE_BINARY3(sigmoid_grad, Eigen::half, float, double);

tensorflow/core/kernels/cwise_op_sigmoid.cc

diff --git a/tensorflow/core/kernels/cwise_op_sigmoid.cc b/tensorflow/core/kernels/cwise_op_sigmoid.cc
index c24a4bc..fabab75 100644
--- a/tensorflow/core/kernels/cwise_op_sigmoid.cc
+++ b/tensorflow/core/kernels/cwise_op_sigmoid.cc
@@ -20,8 +20,7 @@
 REGISTER6(UnaryOp, CPU, "Sigmoid", functor::sigmoid, bfloat16, float,
           Eigen::half, double, complex64, complex128);
 #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, "Sigmoid", functor::sigmoid, float, Eigen::half,
           double);
 #endif

tensorflow/core/kernels/mlir_generated/BUILD

diff --git a/tensorflow/core/kernels/mlir_generated/BUILD b/tensorflow/core/kernels/mlir_generated/BUILD
index b9c37ec..e04adef 100644
--- a/tensorflow/core/kernels/mlir_generated/BUILD
+++ b/tensorflow/core/kernels/mlir_generated/BUILD
@@ -132,6 +132,7 @@
         "gpu_op_rint.cc",
         "gpu_op_round.cc",
         "gpu_op_rsqrt.cc",
+        "gpu_op_sigmoid.cc",
         "gpu_op_sign.cc",
         "gpu_op_sin.cc",
         "gpu_op_sinh.cc",
@@ -139,8 +140,6 @@
         "gpu_op_square.cc",
         "gpu_op_tan.cc",
         "gpu_op_tanh.cc",
-    ]) + if_mlir_generated_experimental_kernels_enabled([
-        "gpu_op_sigmoid.cc",
     ]),
     copts = if_mlir_generated_experimental_kernels_enabled([
         "-DMLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED",
@@ -182,6 +181,7 @@
         ":gpu_rint_kernels",
         ":gpu_round_kernels",
         ":gpu_rsqrt_kernels",
+        ":gpu_sigmoid_kernels",
         ":gpu_sign_kernels",
         ":gpu_sin_kernels",
         ":gpu_sinh_kernels",
@@ -190,8 +190,6 @@
         ":gpu_tan_kernels",
         ":gpu_tanh_kernels",
         "//third_party/eigen3",
-    ]) + if_mlir_generated_experimental_kernels_enabled([
-        ":gpu_sigmoid_kernels",
     ]),
 )