From 4d27598a22656a2430b3feb63ad00b994c5035e8 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Fri, 13 Nov 2020 05:24:45 -0800
Subject: [PATCH] Fix sin kernel for ROCM.

It should be lowered to an intrinsic call.
Re-enable generating a GPU kernel for sin, now that the
problem is fixed.

PiperOrigin-RevId: 342243974
Change-Id: If151fedbdde3146c01ee5f80f002e224d9e8e3b6
---
 .../xla/service/mlir_gpu/kernel_lowering.cc      |  8 ++++----
 tensorflow/core/kernels/mlir_generated/BUILD     | 16 +---------------
 2 files changed, 5 insertions(+), 19 deletions(-)

diff --git a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
index c698ec72bae..b738c40c374 100644
--- a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
+++ b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
@@ -236,10 +236,10 @@ class LowerToROCDLPass
 
     ::mlir::ConversionTarget target(getContext());
     target.addIllegalDialect<::mlir::gpu::GPUDialect>();
-    target
-        .addIllegalOp<mlir::LLVM::CosOp, mlir::LLVM::ExpOp, mlir::LLVM::FAbsOp,
-                      mlir::LLVM::FCeilOp, mlir::LLVM::LogOp,
-                      mlir::LLVM::Log10Op, mlir::LLVM::Log2Op>();
+    target.addIllegalOp<mlir::LLVM::CosOp, mlir::LLVM::ExpOp,
+                        mlir::LLVM::FAbsOp, mlir::LLVM::FCeilOp,
+                        mlir::LLVM::LogOp, mlir::LLVM::Log10Op,
+                        mlir::LLVM::Log2Op, mlir::LLVM::SinOp>();
     target.addIllegalOp<mlir::FuncOp>();
     target.addLegalDialect<::mlir::LLVM::LLVMDialect>();
     target.addLegalDialect<::mlir::ROCDL::ROCDLDialect>();
diff --git a/tensorflow/core/kernels/mlir_generated/BUILD b/tensorflow/core/kernels/mlir_generated/BUILD
index 38cfc5db641..b47fdbee1df 100644
--- a/tensorflow/core/kernels/mlir_generated/BUILD
+++ b/tensorflow/core/kernels/mlir_generated/BUILD
@@ -302,21 +302,6 @@ gen_kernel_library(
     unroll_factors = "4",
 )
 
-# Temporarily disabled due to failure on ROCM, fix will come soon.
-# gen_kernel_library(
-#    name = "sin",
-#    generate_ranked = False,
-#    generate_unranked = True,
-#    tags = ["no_rocm"],
-#    tile_size = "256",
-#    types = [
-#        "f16",
-#        "f32",
-#        "f64",
-#    ],
-#    unroll_factors = "4",
-# )
-
 gen_kernel_library(
     name = "addv2",
     generate_ranked = False,
@@ -352,6 +337,7 @@ gen_kernel_library(
         "log",
         "neg",
         "rsqrt",
+        "sin",
         "sqrt",
         "tanh",
     ]