From 8e7be6f71aaa97f8c2ef806da8945a9d11f99830 Mon Sep 17 00:00:00 2001
From: Kuangyuan Chen <chky@google.com>
Date: Wed, 17 Jun 2020 13:56:34 -0700
Subject: [PATCH] Populate side-effect decorators for HashTable ops.

PiperOrigin-RevId: 316958908
Change-Id: I8e75fe538bb3dfc4b6d0742f63072fe87f133a4d
---
 tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td | 6 +++---
 tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td       | 3 +++
 tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h   | 4 ++++
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
index d403462e6a6..7ed5a215ab8 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
@@ -3862,7 +3862,7 @@ table will be immutable.
   );
 
   let results = (outs
-    TF_ResourceTensor:$table_handle
+    Res<TF_ResourceTensor, "", [TF_TableRead]>:$table_handle
   );
 }
 
@@ -4814,7 +4814,7 @@ table. It must also be of the same type as the table values.
   }];
 
   let arguments = (ins
-    TF_ResourceTensor:$table_handle,
+    Arg<TF_ResourceTensor, "", [TF_TableRead]>:$table_handle,
     TF_Tensor:$keys,
     TF_Tensor:$default_value
   );
@@ -4838,7 +4838,7 @@ The tensor `values` must be of the type of the table values.
   }];
 
   let arguments = (ins
-    TF_ResourceTensor:$table_handle,
+    Arg<TF_ResourceTensor, "", [TF_TableWrite]>:$table_handle,
     TF_Tensor:$keys,
     TF_Tensor:$values
   );
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td
index 17424b54fc2..aac03061718 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td
@@ -91,14 +91,17 @@ class TF_ResourceBase<string resourceKind> :
 def TF_VariableResource : TF_ResourceBase<"Variable">;
 def TF_StackResource : TF_ResourceBase<"Stack">;
 def TF_TensorArrayResource : TF_ResourceBase<"TensorArray">;
+def TF_TableResource : TF_ResourceBase<"Table">;
 
 def TF_VariableRead : MemRead<TF_VariableResource>;
 def TF_StackRead : MemRead<TF_StackResource>;
 def TF_TensorArrayRead : MemRead<TF_TensorArrayResource>;
+def TF_TableRead : MemRead<TF_TableResource>;
 
 def TF_VariableWrite : MemWrite<TF_VariableResource>;
 def TF_StackWrite : MemWrite<TF_StackResource>;
 def TF_TensorArrayWrite : MemWrite<TF_TensorArrayResource>;
+def TF_TableWrite: MemWrite<TF_TableResource>;
 
 //===----------------------------------------------------------------------===//
 // TensorFlow op definitions
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h b/tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h
index 9be61b1db39..ab1d7935bad 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h
@@ -35,6 +35,10 @@ struct TensorArray : ::mlir::SideEffects::Resource::Base<TensorArray> {
   StringRef getName() final { return "TensorArray"; }
 };
 
+struct Table : ::mlir::SideEffects::Resource::Base<Table> {
+  StringRef getName() final { return "Table"; }
+};
+
 }  // namespace ResourceEffects
 }  // namespace TF
 }  // namespace mlir