From cd0ef5f0a24d8482f66c47200541fdfec07200a8 Mon Sep 17 00:00:00 2001
From: Tres Popp <tpopp@google.com>
Date: Wed, 4 Nov 2020 12:13:34 -0800
Subject: [PATCH] Lower all StandardOps to LLVMIR in kernel_gen.

PiperOrigin-RevId: 340702811
Change-Id: I525313971f23516454273f944c82f2a74ff01289
---
 .../tools/kernel_gen/transforms/tf_kernel_to_llvm_pass.cc     | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/tf_kernel_to_llvm_pass.cc b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/tf_kernel_to_llvm_pass.cc
index dff2d470b93..ccb040599c9 100644
--- a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/tf_kernel_to_llvm_pass.cc
+++ b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/tf_kernel_to_llvm_pass.cc
@@ -56,8 +56,8 @@ class TFKernelToLLVMPass : public TFKernelToLLVMPassBase<TFKernelToLLVMPass> {
     // Set target.
     ConversionTarget target(getContext());
     target.addLegalDialect<LLVM::LLVMDialect>();
-    target
-        .addIllegalDialect<gpu::GPUDialect, tf_framework::TFFrameworkDialect>();
+    target.addIllegalDialect<gpu::GPUDialect, StandardOpsDialect,
+                             tf_framework::TFFrameworkDialect>();
     target.addIllegalOp<LLVM::DialectCastOp>();
 
     if (failed(applyPartialConversion(m, target, std::move(patterns)))) {