From de5620b74cb1688dcc81f3e50ad8fc24b18beb9b Mon Sep 17 00:00:00 2001
From: Stephan Herhut <herhut@google.com>
Date: Thu, 18 Jun 2020 04:33:29 -0700
Subject: [PATCH] Add option in kernel_lowering to use the tanh approximation.

PiperOrigin-RevId: 317077338
Change-Id: I5cc7bfb84d6defba05439186377f90e422247d68
---
 tensorflow/compiler/xla/service/mlir_gpu/BUILD              | 1 +
 tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc | 5 +++++
 tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h  | 1 +
 3 files changed, 7 insertions(+)

diff --git a/tensorflow/compiler/xla/service/mlir_gpu/BUILD b/tensorflow/compiler/xla/service/mlir_gpu/BUILD
index ce45d937424..efe69450846 100644
--- a/tensorflow/compiler/xla/service/mlir_gpu/BUILD
+++ b/tensorflow/compiler/xla/service/mlir_gpu/BUILD
@@ -167,6 +167,7 @@ cc_library(
         "//tensorflow/compiler/mlir/xla:lhlo_legalize_to_affine",
         "//tensorflow/compiler/mlir/xla:lhlo_legalize_to_gpu",
         "//tensorflow/compiler/mlir/xla:xla_dialect_registration",
+        "//tensorflow/compiler/mlir/xla:xla_legalize_tanh_to_approximation",
         "//tensorflow/compiler/mlir/xla:xla_legalize_to_linalg",
         "//tensorflow/compiler/xla:status",
         "//tensorflow/compiler/xla:statusor",
diff --git a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
index b0cbddcdb92..9d5b52df010 100644
--- a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
+++ b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
@@ -505,6 +505,11 @@ Status LowerLHLOToGPU(mlir::ModuleOp module, LowerLHLOToGPUOptions options) {
   // Some basic cleanup.
   pm.addNestedPass<::mlir::FuncOp>(::mlir::createCanonicalizerPass());
   pm.addNestedPass<::mlir::FuncOp>(::mlir::createCSEPass());
+  // Approximate of requested.
+  if (options.use_approximations) {
+    pm.addNestedPass<::mlir::FuncOp>(
+        ::mlir::xla::createLegalizeTanhToApproximationPass());
+  }
   // Move scalar operations into the launch to ensure smaller signatures.
   pm.addPass(absl::make_unique<MoveScalarComputationsIntoGpuLaunch>());
   // Take launches to launches with kernels.
diff --git a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h
index 77cf75b9e47..bd633bb06cb 100644
--- a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h
+++ b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h
@@ -28,6 +28,7 @@ struct LowerLHLOToGPUOptions {
   llvm::ArrayRef<unsigned> unroll_factors = {};
   bool collapse_parallel_loops = true;
   bool rewrite_signature = true;
+  bool use_approximations = false;
 };
 
 Status LowerLHLOToGPU(mlir::ModuleOp module,