From 3306d2127513facc617a14da10e9669392f7d217 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 6 Apr 2018 08:58:14 -0700
Subject: [PATCH] Fix up the support for the case where a given array name
 occurs multiple times in the inputs/outputs list of an op. The
 (non-essential) computation of the optimal workspace size had not been
 updated for that case, causing it to fail on a simple test case. Moreover,
 the initial implementation had some redundant usage of std::find that this CL
 moves to a shared helper function.

PiperOrigin-RevId: 191894081
---
 .../contrib/lite/toco/allocate_transient_arrays.cc | 36 ++++++++++------------
 1 file changed, 17 insertions(+), 19 deletions(-)

diff --git a/tensorflow/contrib/lite/toco/allocate_transient_arrays.cc b/tensorflow/contrib/lite/toco/allocate_transient_arrays.cc
index 621fbcb..1f3ea2e 100644
--- a/tensorflow/contrib/lite/toco/allocate_transient_arrays.cc
+++ b/tensorflow/contrib/lite/toco/allocate_transient_arrays.cc
@@ -200,6 +200,12 @@ void DeallocateTransientArray(const Model& model, const string& array_name,
   allocator->Deallocate(*array->alloc);
 }
 
+void PushBackIfNotFound(const string& s, std::vector<string>* v) {
+  if (std::find(v->begin(), v->end(), s) == v->end()) {
+    v->push_back(s);
+  }
+}
+
 }  // namespace
 
 void AllocateTransientArrays(Model* model,
@@ -251,18 +257,12 @@ void AllocateTransientArrays(Model* model,
     std::vector<string> arrays_to_allocate;
     for (const auto& input : op->inputs) {
       if (StartsAt(array_lifespans[input], op_index)) {
-        if (std::find(arrays_to_allocate.begin(), arrays_to_allocate.end(),
-                      input) == arrays_to_allocate.end()) {
-          arrays_to_allocate.push_back(input);
-        }
+        PushBackIfNotFound(input, &arrays_to_allocate);
       }
     }
     for (const auto& output : op->outputs) {
       if (StartsAt(array_lifespans[output], op_index)) {
-        if (std::find(arrays_to_allocate.begin(), arrays_to_allocate.end(),
-                      output) == arrays_to_allocate.end()) {
-          arrays_to_allocate.push_back(output);
-        }
+        PushBackIfNotFound(output, &arrays_to_allocate);
       }
     }
     for (const string& array : arrays_to_allocate) {
@@ -274,18 +274,12 @@ void AllocateTransientArrays(Model* model,
     std::vector<string> arrays_to_deallocate;
     for (const auto& input : op->inputs) {
       if (EndsAt(array_lifespans[input], op_index)) {
-        if (std::find(arrays_to_deallocate.begin(), arrays_to_deallocate.end(),
-                      input) == arrays_to_deallocate.end()) {
-          arrays_to_deallocate.push_back(input);
-        }
+        PushBackIfNotFound(input, &arrays_to_deallocate);
       }
     }
     for (const auto& output : op->outputs) {
       if (EndsAt(array_lifespans[output], op_index)) {
-        if (std::find(arrays_to_deallocate.begin(), arrays_to_deallocate.end(),
-                      output) == arrays_to_deallocate.end()) {
-          arrays_to_deallocate.push_back(output);
-        }
+        PushBackIfNotFound(output, &arrays_to_deallocate);
       }
     }
     for (const string& array : arrays_to_deallocate) {
@@ -310,17 +304,21 @@ void AllocateTransientArrays(Model* model,
     // for each operator, compute the sum of the sizes of the array that must
     // be live during the execution of this operator, plus the size of
     // persistent arrays that must be live at all times.
-    std::size_t size = persistent_alloc_size;
+    std::vector<string> non_persistent_edges;
     for (const auto& input : op->inputs) {
       if (!array_lifespans[input].persistent) {
-        size += TransientArraySize(*model, input, transient_data_alignment);
+        PushBackIfNotFound(input, &non_persistent_edges);
       }
     }
     for (const auto& output : op->outputs) {
       if (!array_lifespans[output].persistent) {
-        size += TransientArraySize(*model, output, transient_data_alignment);
+        PushBackIfNotFound(output, &non_persistent_edges);
       }
     }
+    std::size_t size = persistent_alloc_size;
+    for (const string& edge : non_persistent_edges) {
+      size += TransientArraySize(*model, edge, transient_data_alignment);
+    }
     // The optimal total size is the maximum of all operator-specific sizes.
     optimal_transient_alloc_size = std::max(optimal_transient_alloc_size, size);
   }
-- 
2.7.4