[MLIR][KernelGen] Add tan kernels
PiperOrigin-RevId: 352746800 Change-Id: I1961a4f25b7941aaa7f4e494fb07db10905d8dbc
This commit is contained in:
parent
531c3b720d
commit
c3e7d0dac0
@ -20,7 +20,10 @@ REGISTER6(UnaryOp, CPU, "Tan", functor::tan, Eigen::half, bfloat16, float,
|
|||||||
double, complex64, complex128);
|
double, complex64, complex128);
|
||||||
|
|
||||||
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
|
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
|
||||||
|
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
|
||||||
|
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
|
||||||
REGISTER3(UnaryOp, GPU, "Tan", functor::tan, Eigen::half, float, double);
|
REGISTER3(UnaryOp, GPU, "Tan", functor::tan, Eigen::half, float, double);
|
||||||
#endif
|
#endif
|
||||||
|
#endif
|
||||||
|
|
||||||
} // namespace tensorflow
|
} // namespace tensorflow
|
||||||
|
@ -74,6 +74,7 @@ filegroup(
|
|||||||
"gpu_op_sign.cc",
|
"gpu_op_sign.cc",
|
||||||
"gpu_op_sin.cc",
|
"gpu_op_sin.cc",
|
||||||
"gpu_op_sqrt.cc",
|
"gpu_op_sqrt.cc",
|
||||||
|
"gpu_op_tan.cc",
|
||||||
],
|
],
|
||||||
compatible_with = get_compatible_with_cloud(),
|
compatible_with = get_compatible_with_cloud(),
|
||||||
)
|
)
|
||||||
@ -145,6 +146,7 @@ tf_kernel_library(
|
|||||||
":sign_kernels",
|
":sign_kernels",
|
||||||
":sin_kernels",
|
":sin_kernels",
|
||||||
":sqrt_kernels",
|
":sqrt_kernels",
|
||||||
|
":tan_kernels",
|
||||||
":tanh_kernels",
|
":tanh_kernels",
|
||||||
":gpu_ops_base",
|
":gpu_ops_base",
|
||||||
"//third_party/eigen3",
|
"//third_party/eigen3",
|
||||||
@ -703,5 +705,6 @@ gen_kernel_library(
|
|||||||
for name in [
|
for name in [
|
||||||
"cos",
|
"cos",
|
||||||
"sin",
|
"sin",
|
||||||
|
"tan",
|
||||||
]
|
]
|
||||||
]
|
]
|
||||||
|
25
tensorflow/core/kernels/mlir_generated/gpu_op_tan.cc
Normal file
25
tensorflow/core/kernels/mlir_generated/gpu_op_tan.cc
Normal file
@ -0,0 +1,25 @@
|
|||||||
|
/* Copyright 2021 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/gpu_ops_base.h"
|
||||||
|
|
||||||
|
namespace tensorflow {
|
||||||
|
|
||||||
|
GENERATE_AND_REGISTER_UNARY_KERNEL(Tan, f16, DT_HALF, Eigen::half);
|
||||||
|
GENERATE_AND_REGISTER_UNARY_KERNEL(Tan, f32, DT_FLOAT, float);
|
||||||
|
GENERATE_AND_REGISTER_UNARY_KERNEL(Tan, f64, DT_DOUBLE, double);
|
||||||
|
|
||||||
|
} // namespace tensorflow
|
@ -539,6 +539,17 @@ GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES_2(
|
|||||||
test::DefaultInputGreaterOrEqualToZero<Eigen::half>(), std::sqrt,
|
test::DefaultInputGreaterOrEqualToZero<Eigen::half>(), std::sqrt,
|
||||||
test::GpuOpsTestConfig())
|
test::GpuOpsTestConfig())
|
||||||
|
|
||||||
|
/// Test `tf.Tan`.
|
||||||
|
|
||||||
|
GENERATE_DEFAULT_TEST(Tan, DT_FLOAT, DT_FLOAT, std::tan,
|
||||||
|
test::GpuOpsTestConfig())
|
||||||
|
|
||||||
|
GENERATE_DEFAULT_TEST(Tan, DT_DOUBLE, DT_DOUBLE, std::tan,
|
||||||
|
test::GpuOpsTestConfig())
|
||||||
|
|
||||||
|
GENERATE_DEFAULT_TEST_2(Tan, DT_HALF, DT_FLOAT, DT_HALF, DT_FLOAT, std::tan,
|
||||||
|
test::GpuOpsTestConfig())
|
||||||
|
|
||||||
/// Test `tf.Tanh`.
|
/// Test `tf.Tanh`.
|
||||||
|
|
||||||
GENERATE_DEFAULT_TEST(Tanh, DT_FLOAT, DT_FLOAT, std::tanh,
|
GENERATE_DEFAULT_TEST(Tanh, DT_FLOAT, DT_FLOAT, std::tanh,
|
||||||
|
Loading…
x
Reference in New Issue
Block a user