[MLIR][KernelGen] Add lgamma kernels

PiperOrigin-RevId: 354519407
Change-Id: I22d6981835acc26f4e9f7c287bec1378c47bfb28

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 <typename OpTy>
 inline void AddLegalOpOnRankedTensor(ConversionTarget *target) {

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],

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

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",

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

tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc

@@ -17,6 +17,7 @@ limitations under the License.
 #include <complex>
 #include <functional>
 #include <initializer_list>
+#include <limits>
 #include <memory>
 #include <numeric>
 #include <vector>
@@ -575,6 +576,52 @@ GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES_2(
     test::NearZeroInfAndNanInput<Eigen::half>(), std::isnan,
     test::GpuOpsTestConfig().ExpectStrictlyEqual().NoBufferReuse());
 
+/// Test `tf.Lgamma`.
+
+static constexpr std::initializer_list<double> kLgammaValues = {
+    -std::numeric_limits<double>::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<double>::infinity(),
+};
+
+GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES(
+    Lgamma, DT_FLOAT, DT_FLOAT, test::InputAsVector<float>(kLgammaValues),
+    std::lgamma, test::GpuOpsTestConfig())
+
+GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES(
+    Lgamma, DT_DOUBLE, DT_DOUBLE, test::InputAsVector<double>(kLgammaValues),
+    std::lgamma, test::GpuOpsTestConfig())
+
+GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES_2(
+    Lgamma, DT_HALF, DT_FLOAT, DT_HALF, DT_FLOAT,
+    test::InputAsVector<Eigen::half>(kLgammaValues), std::lgamma,
+    test::GpuOpsTestConfig())
+
 /// Test `tf.Log`.
 
 GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES(