From 0e2545d934abf542dfda1b06cdc61153a2d4cab1 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 18 Jan 2021 05:13:02 -0800 Subject: [PATCH] [MLIR][KernelGen] Add `tf.Atanh` kernels PiperOrigin-RevId: 352393602 Change-Id: I2431e39759a12735241e9efb9ff778bdb287e6d3 --- .../mlir-hlo/Dialect/mhlo/IR/chlo_ops.td | 14 +++++++++++ .../chlo_legalize_to_hlo_patterns.td | 23 ++++++++++++++++++ .../mhlo/transforms/transform_unranked_hlo.cc | 7 +++--- .../xla/transforms/legalize_tf_patterns.td | 1 + tensorflow/core/kernels/cwise_op_atanh.cc | 5 +++- tensorflow/core/kernels/mlir_generated/BUILD | 12 ++++++++++ .../kernels/mlir_generated/gpu_op_atanh.cc | 24 +++++++++++++++++++ .../mlir_generated/gpu_unary_ops_test.cc | 10 ++++++++ .../op_definitions/atanh.mlir.tmpl | 5 ++++ 9 files changed, 97 insertions(+), 4 deletions(-) create mode 100644 tensorflow/core/kernels/mlir_generated/gpu_op_atanh.cc create mode 100644 tensorflow/core/kernels/mlir_generated/op_definitions/atanh.mlir.tmpl diff --git a/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.td b/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.td index 558da58fa08..c633bd2cd81 100644 --- a/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.td +++ b/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.td @@ -397,6 +397,20 @@ def HLOClient_AtanOp : HLOClient_UnaryElementwiseOp<"atan", [], }]; } +def HLOClient_AtanhOp : HLOClient_UnaryElementwiseOp<"atanh", [], + HLO_FpOrComplexTensor> { + let summary = "Atanh operator"; + + let description = [{ + Returns `Atanh(operand)` element-wise. + + $$ + \atanh(x) = 0.5 * log((1 + x) / (1 - x)) if abs(x) <= 1 + = nan otherwise + $$ + }]; +} + def HLOClient_ConjOp : HLOClient_UnaryElementwiseOp<"conj", [], HLO_FpOrComplexTensor> { let summary = "Conj operator"; diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/chlo_legalize_to_hlo_patterns.td b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/chlo_legalize_to_hlo_patterns.td index b8b6abb72d3..a2b97a8c432 100644 --- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/chlo_legalize_to_hlo_patterns.td +++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/chlo_legalize_to_hlo_patterns.td @@ -175,6 +175,29 @@ def : Pat<(HLOClient_AtanOp $input), (HLO_ConstantLike<"1"> $input) )>; +// Express `atanh` as follows: +// atanh(x) = 0.5 * log((1 + x) / (1 - x)) if abs(x) <= 1 +// atanh(x) = nan otherwise +def : Pat<(HLOClient_AtanhOp NonComplexElementType:$input), + (HLO_SelectOp + (HLO_CompareOp + (HLO_AbsOp $input), + (HLO_ConstantLike<"1"> $input), + HLO_COMPARISON_DIRECTION_GT, + (HLO_DEFAULT_COMPARISON_TYPE) + ), + (HLO_ConstantLike<"NAN"> $input), + (HLO_MulOp + (HLO_SubOp + (HLO_Log1pOp $input), + (HLO_Log1pOp + (HLO_NegOp $input) + ) + ), + (HLO_ConstantLike<"0.5"> $input) + ) + )>; + // Express `conj` as // conj(x) = (re(x), -im(x)). def : Pat<(HLOClient_ConjOp $v), 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 bd6d8918d4f..70d5d380597 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 @@ -50,9 +50,10 @@ namespace { sep fn(ShiftRightLogicalOp) sep fn(SubOp) sep fn(XorOp) // TODO(herhut): Generate these out of op definitions. -#define MAP_CHLO_OPERATION_CWISE_UNARY(fn, sep) \ - fn(AcosOp) sep fn(AsinOp) sep fn(AsinhOp) sep fn(AtanOp) sep fn(ConjOp) \ - sep fn(CoshOp) sep fn(ErfOp) sep fn(ErfcOp) sep fn(SinhOp) sep fn(TanOp) +#define MAP_CHLO_OPERATION_CWISE_UNARY(fn, sep) \ + fn(AcosOp) 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) 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 7707b4f8081..9603bf0a754 100644 --- a/tensorflow/compiler/mlir/xla/transforms/legalize_tf_patterns.td +++ b/tensorflow/compiler/mlir/xla/transforms/legalize_tf_patterns.td @@ -588,6 +588,7 @@ foreach Mapping = [ [TF_AcosOp, HLOClient_AcosOp], [TF_AsinOp, HLOClient_AsinOp], [TF_AtanOp, HLOClient_AtanOp], + [TF_AtanhOp, HLOClient_AtanhOp], [TF_CeilOp, HLO_CeilOp], [TF_CoshOp, HLOClient_CoshOp], [TF_ComplexAbsOp, HLO_AbsOp], diff --git a/tensorflow/core/kernels/cwise_op_atanh.cc b/tensorflow/core/kernels/cwise_op_atanh.cc index 2404cd19646..def20136f8b 100644 --- a/tensorflow/core/kernels/cwise_op_atanh.cc +++ b/tensorflow/core/kernels/cwise_op_atanh.cc @@ -20,8 +20,11 @@ namespace tensorflow { REGISTER4(UnaryOp, CPU, "Atanh", functor::atanh, float, double, complex64, complex128); - #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM +#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ + !defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED) REGISTER2(UnaryOp, GPU, "Atanh", functor::atanh, float, double); #endif +#endif + } // namespace tensorflow diff --git a/tensorflow/core/kernels/mlir_generated/BUILD b/tensorflow/core/kernels/mlir_generated/BUILD index e45360adb0d..e78e00cc305 100644 --- a/tensorflow/core/kernels/mlir_generated/BUILD +++ b/tensorflow/core/kernels/mlir_generated/BUILD @@ -50,6 +50,7 @@ filegroup( "gpu_op_asin.cc", "gpu_op_asinh.cc", "gpu_op_atan.cc", + "gpu_op_atanh.cc", "gpu_op_ceil.cc", "gpu_op_complex.cc", "gpu_op_complex_abs.cc", @@ -118,6 +119,7 @@ tf_kernel_library( ":asin_kernels", ":asinh_kernels", ":atan_kernels", + ":atanh_kernels", ":ceil_kernels", ":complex_abs_kernels", ":complex_kernels", @@ -349,6 +351,16 @@ gen_kernel_library( unroll_factors = "4", ) +gen_kernel_library( + name = "atanh", + tile_size = "256", + types = [ + "f32", + "f64", + ], + unroll_factors = "4", +) + gen_kernel_library( name = "conj", tile_size = "256", diff --git a/tensorflow/core/kernels/mlir_generated/gpu_op_atanh.cc b/tensorflow/core/kernels/mlir_generated/gpu_op_atanh.cc new file mode 100644 index 00000000000..5f162188784 --- /dev/null +++ b/tensorflow/core/kernels/mlir_generated/gpu_op_atanh.cc @@ -0,0 +1,24 @@ +/* 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(Atanh, f32, DT_FLOAT, float); +GENERATE_AND_REGISTER_UNARY_KERNEL(Atanh, 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 fc6f76746e9..78543917bf1 100644 --- a/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc +++ b/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc @@ -207,6 +207,16 @@ GENERATE_DEFAULT_TEST(Atan, DT_FLOAT, DT_FLOAT, std::atan, GENERATE_DEFAULT_TEST(Atan, DT_DOUBLE, DT_DOUBLE, std::atan, test::GpuOpsTestConfig()) +/// Test `tf.Atanh`. + +GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES( + Atanh, DT_FLOAT, DT_FLOAT, test::DefaultInputBetweenZeroAndOne(), + std::atanh, test::GpuOpsTestConfig()) + +GENERATE_DEFAULT_TEST_WITH_SPECIFIC_INPUT_VALUES( + Atanh, DT_DOUBLE, DT_DOUBLE, test::DefaultInputBetweenZeroAndOne(), + std::atanh, test::GpuOpsTestConfig()) + /// Test `tf.Ceil`. GENERATE_DEFAULT_TEST(Ceil, DT_FLOAT, DT_FLOAT, std::ceil, diff --git a/tensorflow/core/kernels/mlir_generated/op_definitions/atanh.mlir.tmpl b/tensorflow/core/kernels/mlir_generated/op_definitions/atanh.mlir.tmpl new file mode 100644 index 00000000000..560483385e9 --- /dev/null +++ b/tensorflow/core/kernels/mlir_generated/op_definitions/atanh.mlir.tmpl @@ -0,0 +1,5 @@ +func @Atanh_elem_type(%arg0: tensor<*xelem_type>) + -> tensor<*xelem_type> attributes {tf_entry, llvm.emit_c_interface} { + %0 = "tf.Atanh"(%arg0) : (tensor<*xelem_type>) -> tensor<*xelem_type> + return %0 : tensor<*xelem_type> +}