[MLIR][KernelGen] Add unranked ceil kernel and microbenchmarks
PiperOrigin-RevId: 341360359 Change-Id: I5717b9731b4fd5171d708bfee7c94e56079d61f3
This commit is contained in:
parent
ae5f78b6c4
commit
7978215e25
@ -20,7 +20,10 @@ REGISTER4(UnaryOp, CPU, "Ceil", functor::ceil, float, Eigen::half, bfloat16,
|
|||||||
double);
|
double);
|
||||||
|
|
||||||
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
|
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
|
||||||
|
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
|
||||||
|
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
|
||||||
REGISTER3(UnaryOp, GPU, "Ceil", functor::ceil, float, Eigen::half, double);
|
REGISTER3(UnaryOp, GPU, "Ceil", functor::ceil, float, Eigen::half, double);
|
||||||
#endif
|
#endif
|
||||||
|
#endif
|
||||||
|
|
||||||
} // namespace tensorflow
|
} // namespace tensorflow
|
||||||
|
@ -41,6 +41,7 @@ filegroup(
|
|||||||
srcs = if_mlir_unranked_kernels_enabled(
|
srcs = if_mlir_unranked_kernels_enabled(
|
||||||
[
|
[
|
||||||
"unranked_op_gpu_abs.cc",
|
"unranked_op_gpu_abs.cc",
|
||||||
|
"unranked_op_gpu_ceil.cc",
|
||||||
"unranked_op_gpu_tanh.cc",
|
"unranked_op_gpu_tanh.cc",
|
||||||
"unranked_op_gpu_base.h",
|
"unranked_op_gpu_base.h",
|
||||||
"unranked_op_gpu_base.cc",
|
"unranked_op_gpu_base.cc",
|
||||||
@ -60,6 +61,7 @@ cc_library(
|
|||||||
[
|
[
|
||||||
":abs_unranked_kernels",
|
":abs_unranked_kernels",
|
||||||
":addv2_unranked_kernels",
|
":addv2_unranked_kernels",
|
||||||
|
":ceil_unranked_kernels",
|
||||||
":tanh_unranked_kernels",
|
":tanh_unranked_kernels",
|
||||||
"//tensorflow/compiler/mlir/tools/kernel_gen:tf_cuda_runtime_wrappers",
|
"//tensorflow/compiler/mlir/tools/kernel_gen:tf_cuda_runtime_wrappers",
|
||||||
"//tensorflow/compiler/mlir/tools/kernel_gen:tf_framework_c_interface",
|
"//tensorflow/compiler/mlir/tools/kernel_gen:tf_framework_c_interface",
|
||||||
|
@ -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(Ceil, f16, DT_HALF, Eigen::half);
|
||||||
|
REGISTER_AND_GENERATE_KERNEL(Ceil, f32, DT_FLOAT, float);
|
||||||
|
REGISTER_AND_GENERATE_KERNEL(Ceil, f64, DT_DOUBLE, double);
|
||||||
|
|
||||||
|
} // namespace tensorflow
|
Loading…
x
Reference in New Issue
Block a user