From d7f96c409b4a67e89fed9289284c56d2d3ce9d14 Mon Sep 17 00:00:00 2001
From: Sean Silva <silvasean@google.com>
Date: Mon, 11 Jan 2021 16:22:25 -0800
Subject: [PATCH] Fix handling of tf_saved_model.asset in
 tf-saved-model-optimize-global-tensors

PiperOrigin-RevId: 351259338
Change-Id: Idfa548a3a20b1b5f9e77e5c59010060d85d4578d
---
 .../tf_saved_model_optimize_global_tensors.mlir  | 16 ++++++++++++++++
 .../transforms/optimize_global_tensors.cc        |  3 +++
 2 files changed, 19 insertions(+)

diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_optimize_global_tensors.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_optimize_global_tensors.mlir
index 0c68cf0cf64..4d3938a9570 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_optimize_global_tensors.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_optimize_global_tensors.mlir
@@ -228,3 +228,19 @@ module attributes {tf_saved_model.semantics} {
     return
   }
 }
+
+// -----
+
+// Test presence of tf_saved_model.asset's as bound inputs.
+// It should not crash.
+module attributes {tf_saved_model.semantics}  {
+
+  "tf_saved_model.session_initializer"() {initializers = [@legacy_init_op]} : () -> ()
+  "tf_saved_model.asset"() {filename = "assets/foo.txt", sym_name = "asset"} : () -> ()
+  // CHECK: @legacy_init_op
+  func @legacy_init_op(%arg0: tensor<!tf.string> {tf_saved_model.bound_input = @asset}) attributes {tf_saved_model.exported_names = ["f"]} {
+    %0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf.string, shared_name = "hash_table_shared_name", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf.resource>
+    "tf.InitializeTableFromTextFileV2"(%0, %arg0) {delimiter = "\09", device = "", key_index = -2 : i64, value_index = -1 : i64, vocab_size = 205 : i64} : (tensor<!tf.resource>, tensor<!tf.string>) -> ()
+    return
+  }
+}
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/optimize_global_tensors.cc b/tensorflow/compiler/mlir/tensorflow/transforms/optimize_global_tensors.cc
index e0e0c45d63a..4c824e8807d 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/optimize_global_tensors.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/optimize_global_tensors.cc
@@ -212,6 +212,9 @@ GlobalTensorUsesMap CreateGlobalTensorUsesMap(ModuleOp module) {
       }
       auto global_tensor = symbol_table.lookup<GlobalTensorOp>(
           sym.cast<FlatSymbolRefAttr>().getValue());
+      if (!global_tensor) {
+        continue;
+      }
       global_tensor_uses[global_tensor].push_back({func, i});
     }
   }