[MLIR][KernelGen] Add unranked MLIR-generated kernels and microbenchmarks
Add unranked MLIR-generated versions for the following simple unary kernels:
- cos
- exp
- floor
- log
- rsqrt
- sqrt
PiperOrigin-RevId: 341382592
Change-Id: I7b6c79204f793262fb1ec1e6cebcd4a944857e86
diff --git a/tensorflow/core/kernels/cwise_op_cos.cc b/tensorflow/core/kernels/cwise_op_cos.cc
index 64e9fab..5bf127f 100644
--- a/tensorflow/core/kernels/cwise_op_cos.cc
+++ b/tensorflow/core/kernels/cwise_op_cos.cc
@@ -20,7 +20,10 @@
double, complex64, complex128);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
+ !defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
REGISTER3(UnaryOp, GPU, "Cos", functor::cos, float, Eigen::half, double);
#endif
+#endif
} // namespace tensorflow
diff --git a/tensorflow/core/kernels/cwise_op_exp.cc b/tensorflow/core/kernels/cwise_op_exp.cc
index 28ace80..24d098a 100644
--- a/tensorflow/core/kernels/cwise_op_exp.cc
+++ b/tensorflow/core/kernels/cwise_op_exp.cc
@@ -20,8 +20,11 @@
double, complex64, complex128);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
+ !defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
REGISTER5(UnaryOp, GPU, "Exp", functor::exp, float, Eigen::half, double,
complex64, complex128);
#endif
+#endif
} // namespace tensorflow
diff --git a/tensorflow/core/kernels/cwise_op_floor.cc b/tensorflow/core/kernels/cwise_op_floor.cc
index da5619b..57296f9 100644
--- a/tensorflow/core/kernels/cwise_op_floor.cc
+++ b/tensorflow/core/kernels/cwise_op_floor.cc
@@ -20,6 +20,9 @@
double);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
+ !defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
REGISTER3(UnaryOp, GPU, "Floor", functor::floor, float, Eigen::half, double);
#endif
+#endif
} // namespace tensorflow
diff --git a/tensorflow/core/kernels/cwise_op_log.cc b/tensorflow/core/kernels/cwise_op_log.cc
index 236f95d..f0ece6c 100644
--- a/tensorflow/core/kernels/cwise_op_log.cc
+++ b/tensorflow/core/kernels/cwise_op_log.cc
@@ -20,7 +20,10 @@
bfloat16, complex64, complex128);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
+ !defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
REGISTER3(UnaryOp, GPU, "Log", functor::log, float, Eigen::half, double);
#endif
+#endif
} // namespace tensorflow
diff --git a/tensorflow/core/kernels/cwise_op_rsqrt.cc b/tensorflow/core/kernels/cwise_op_rsqrt.cc
index 21e3bf4..cb6c1ef 100644
--- a/tensorflow/core/kernels/cwise_op_rsqrt.cc
+++ b/tensorflow/core/kernels/cwise_op_rsqrt.cc
@@ -20,8 +20,11 @@
complex64, complex128);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
+ !defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
REGISTER3(UnaryOp, GPU, "Rsqrt", functor::rsqrt, float, Eigen::half, double);
#endif
+#endif
REGISTER5(SimpleBinaryOp, CPU, "RsqrtGrad", functor::rsqrt_grad, float,
Eigen::half, double, complex64, complex128);
diff --git a/tensorflow/core/kernels/cwise_op_sqrt.cc b/tensorflow/core/kernels/cwise_op_sqrt.cc
index 2e33297..32c78e4 100644
--- a/tensorflow/core/kernels/cwise_op_sqrt.cc
+++ b/tensorflow/core/kernels/cwise_op_sqrt.cc
@@ -20,9 +20,11 @@
bfloat16, complex64, complex128);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
+ !defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
REGISTER3(UnaryOp, GPU, "Sqrt", functor::sqrt, float, Eigen::half, double);
#endif
-
+#endif
REGISTER6(SimpleBinaryOp, CPU, "SqrtGrad", functor::sqrt_grad, float,
Eigen::half, bfloat16, double, complex64, complex128);
diff --git a/tensorflow/core/kernels/mlir_generated/BUILD b/tensorflow/core/kernels/mlir_generated/BUILD
index 46d3e30..c3e700d 100644
--- a/tensorflow/core/kernels/mlir_generated/BUILD
+++ b/tensorflow/core/kernels/mlir_generated/BUILD
@@ -42,6 +42,12 @@
[
"unranked_op_gpu_abs.cc",
"unranked_op_gpu_ceil.cc",
+ "unranked_op_gpu_cos.cc",
+ "unranked_op_gpu_exp.cc",
+ "unranked_op_gpu_floor.cc",
+ "unranked_op_gpu_log.cc",
+ "unranked_op_gpu_rsqrt.cc",
+ "unranked_op_gpu_sqrt.cc",
"unranked_op_gpu_tanh.cc",
"unranked_op_gpu_base.h",
"unranked_op_gpu_base.cc",
@@ -62,6 +68,12 @@
":abs_unranked_kernels",
":addv2_unranked_kernels",
":ceil_unranked_kernels",
+ ":cos_unranked_kernels",
+ ":exp_unranked_kernels",
+ ":floor_unranked_kernels",
+ ":log_unranked_kernels",
+ ":rsqrt_unranked_kernels",
+ ":sqrt_unranked_kernels",
":tanh_unranked_kernels",
"//tensorflow/compiler/mlir/tools/kernel_gen:tf_cuda_runtime_wrappers",
"//tensorflow/compiler/mlir/tools/kernel_gen:tf_framework_c_interface",
diff --git a/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_cos.cc b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_cos.cc
new file mode 100644
index 0000000..b21f835
--- /dev/null
+++ b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_cos.cc
@@ -0,0 +1,25 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
+#include "tensorflow/core/kernels/mlir_generated/unranked_op_gpu_base.h"
+
+namespace tensorflow {
+
+REGISTER_AND_GENERATE_KERNEL(Cos, f16, DT_HALF, Eigen::half);
+REGISTER_AND_GENERATE_KERNEL(Cos, f32, DT_FLOAT, float);
+REGISTER_AND_GENERATE_KERNEL(Cos, f64, DT_DOUBLE, double);
+
+} // namespace tensorflow
diff --git a/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_exp.cc b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_exp.cc
new file mode 100644
index 0000000..70f3306
--- /dev/null
+++ b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_exp.cc
@@ -0,0 +1,25 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
+#include "tensorflow/core/kernels/mlir_generated/unranked_op_gpu_base.h"
+
+namespace tensorflow {
+
+REGISTER_AND_GENERATE_KERNEL(Exp, f16, DT_HALF, Eigen::half);
+REGISTER_AND_GENERATE_KERNEL(Exp, f32, DT_FLOAT, float);
+REGISTER_AND_GENERATE_KERNEL(Exp, f64, DT_DOUBLE, double);
+
+} // namespace tensorflow
diff --git a/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_floor.cc b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_floor.cc
new file mode 100644
index 0000000..80b05d0
--- /dev/null
+++ b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_floor.cc
@@ -0,0 +1,25 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
+#include "tensorflow/core/kernels/mlir_generated/unranked_op_gpu_base.h"
+
+namespace tensorflow {
+
+REGISTER_AND_GENERATE_KERNEL(Floor, f16, DT_HALF, Eigen::half);
+REGISTER_AND_GENERATE_KERNEL(Floor, f32, DT_FLOAT, float);
+REGISTER_AND_GENERATE_KERNEL(Floor, f64, DT_DOUBLE, double);
+
+} // namespace tensorflow
diff --git a/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_log.cc b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_log.cc
new file mode 100644
index 0000000..5314e4a
--- /dev/null
+++ b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_log.cc
@@ -0,0 +1,25 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
+#include "tensorflow/core/kernels/mlir_generated/unranked_op_gpu_base.h"
+
+namespace tensorflow {
+
+REGISTER_AND_GENERATE_KERNEL(Log, f16, DT_HALF, Eigen::half);
+REGISTER_AND_GENERATE_KERNEL(Log, f32, DT_FLOAT, float);
+REGISTER_AND_GENERATE_KERNEL(Log, f64, DT_DOUBLE, double);
+
+} // namespace tensorflow
diff --git a/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_rsqrt.cc b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_rsqrt.cc
new file mode 100644
index 0000000..448b242
--- /dev/null
+++ b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_rsqrt.cc
@@ -0,0 +1,25 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
+#include "tensorflow/core/kernels/mlir_generated/unranked_op_gpu_base.h"
+
+namespace tensorflow {
+
+REGISTER_AND_GENERATE_KERNEL(Rsqrt, f16, DT_HALF, Eigen::half);
+REGISTER_AND_GENERATE_KERNEL(Rsqrt, f32, DT_FLOAT, float);
+REGISTER_AND_GENERATE_KERNEL(Rsqrt, f64, DT_DOUBLE, double);
+
+} // namespace tensorflow
diff --git a/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_sqrt.cc b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_sqrt.cc
new file mode 100644
index 0000000..98f1342
--- /dev/null
+++ b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_sqrt.cc
@@ -0,0 +1,25 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
+#include "tensorflow/core/kernels/mlir_generated/unranked_op_gpu_base.h"
+
+namespace tensorflow {
+
+REGISTER_AND_GENERATE_KERNEL(Sqrt, f16, DT_HALF, Eigen::half);
+REGISTER_AND_GENERATE_KERNEL(Sqrt, f32, DT_FLOAT, float);
+REGISTER_AND_GENERATE_KERNEL(Sqrt, f64, DT_DOUBLE, double);
+
+} // namespace tensorflow
