diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/transform_unranked_hlo.cc b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/transform_unranked_hlo.cc index 0f654217efa..f1bb5b2dddc 100644 --- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/transform_unranked_hlo.cc +++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/transform_unranked_hlo.cc @@ -54,7 +54,7 @@ namespace { #define MAP_CHLO_OPERATION_CWISE_UNARY(fn, sep) \ fn(AcosOp) sep fn(AcoshOp) sep fn(AsinOp) sep fn(AsinhOp) sep fn(AtanOp) \ sep fn(AtanhOp) sep fn(ConjOp) sep fn(CoshOp) sep fn(ErfOp) \ - sep fn(ErfcOp) sep fn(SinhOp) sep fn(TanOp) + sep fn(ErfcOp) sep fn(LgammaOp) sep fn(SinhOp) sep fn(TanOp) template inline void AddLegalOpOnRankedTensor(ConversionTarget *target) { diff --git a/tensorflow/compiler/mlir/xla/transforms/legalize_tf_patterns.td b/tensorflow/compiler/mlir/xla/transforms/legalize_tf_patterns.td index aed9549edb2..4f835aebea5 100644 --- a/tensorflow/compiler/mlir/xla/transforms/legalize_tf_patterns.td +++ b/tensorflow/compiler/mlir/xla/transforms/legalize_tf_patterns.td @@ -602,6 +602,7 @@ foreach Mapping = [ [TF_ImagOp, HLO_ImagOp], [TF_InvertOp, HLO_NotOp], [TF_IsFiniteOp, HLO_IsFiniteOp], + [TF_LgammaOp, HLOClient_LgammaOp], [TF_LogOp, HLO_LogOp], [TF_Log1pOp, HLO_Log1pOp], [TF_LogicalNotOp, HLO_NotOp], diff --git a/tensorflow/core/kernels/cwise_op_lgamma.cc b/tensorflow/core/kernels/cwise_op_lgamma.cc index 1446393921b..86cc2479aa9 100644 --- a/tensorflow/core/kernels/cwise_op_lgamma.cc +++ b/tensorflow/core/kernels/cwise_op_lgamma.cc @@ -16,9 +16,14 @@ limitations under the License. #include "tensorflow/core/kernels/cwise_ops_common.h" 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_GPU_KERNELS_ENABLED) REGISTER3(UnaryOp, GPU, "Lgamma", functor::lgamma, float, Eigen::half, double); #endif +#endif } // namespace tensorflow diff --git a/tensorflow/core/kernels/mlir_generated/BUILD b/tensorflow/core/kernels/mlir_generated/BUILD index ec1b688feff..809d51244c9 100644 --- a/tensorflow/core/kernels/mlir_generated/BUILD +++ b/tensorflow/core/kernels/mlir_generated/BUILD @@ -72,6 +72,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", @@ -148,6 +149,7 @@ tf_kernel_library( ":is_finite_kernels", ":is_inf_kernels", ":is_nan_kernels", + ":lgamma_kernels", ":log_kernels", ":log1p_kernels", ":logical_not_kernels", @@ -745,6 +747,7 @@ gen_kernel_library( "is_finite", "is_inf", "is_nan", + "lgamma", "log", "log1p", "rsqrt", diff --git a/tensorflow/core/kernels/mlir_generated/gpu_op_lgamma.cc b/tensorflow/core/kernels/mlir_generated/gpu_op_lgamma.cc new file mode 100644 index 00000000000..8641a6ba35e --- /dev/null +++ b/tensorflow/core/kernels/mlir_generated/gpu_op_lgamma.cc @@ -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(Lgamma, f16, DT_HALF, Eigen::half); +GENERATE_AND_REGISTER_UNARY_KERNEL(Lgamma, f32, DT_FLOAT, float); +GENERATE_AND_REGISTER_UNARY_KERNEL(Lgamma, f64, DT_DOUBLE, double); + +} // namespace tensorflow diff --git a/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc b/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc index b0ddc6ac413..2b4645e131e 100644 --- a/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc +++ b/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include #include +#include #include #include #include @@ -575,6 +576,52 @@ GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES_2( test::NearZeroInfAndNanInput(), std::isnan, test::GpuOpsTestConfig().ExpectStrictlyEqual().NoBufferReuse()); +/// Test `tf.Lgamma`. + +static constexpr std::initializer_list kLgammaValues = { + -std::numeric_limits::infinity(), + -9.0, + -8.5, + -8.3, + -2.0, + -1.5, + -1.3, + -1.0, + -0.5, + -0.3, + 0.0, + 0.1, + 0.2, + 0.3, + 0.4, + 0.5, + 0.6, + 0.7, + 0.8, + 0.9, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 5.0, + 100.0, + std::numeric_limits::infinity(), +}; + +GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES( + Lgamma, DT_FLOAT, DT_FLOAT, test::InputAsVector(kLgammaValues), + std::lgamma, test::GpuOpsTestConfig()) + +GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES( + Lgamma, DT_DOUBLE, DT_DOUBLE, test::InputAsVector(kLgammaValues), + std::lgamma, test::GpuOpsTestConfig()) + +GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES_2( + Lgamma, DT_HALF, DT_FLOAT, DT_HALF, DT_FLOAT, + test::InputAsVector(kLgammaValues), std::lgamma, + test::GpuOpsTestConfig()) + /// Test `tf.Log`. GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES(