[Runtime] Fix duplicate references to const tensor (DeepRec-AI#119)

maciekpac · Dec 7, 2021 · a8adcc9 · a8adcc9
1 parent 4e7dbf3
commit a8adcc9
Show file tree

Hide file tree

Showing 4 changed files with 2 additions and 16 deletions.
diff --git a/tensorflow/core/common_runtime/executor.cc b/tensorflow/core/common_runtime/executor.cc
@@ -689,7 +689,7 @@ void ExecutorState<PropagatorStateType>::ProcessConstTensor(
   nodestats::SetOpEnd(stats);
   Entry& output = (*outputs)[0];
   output.state = Entry::State::HAS_CONST_TENSOR;
-  output.const_tensor = item.const_tensor;
+  output.const_tensor = item.kernel->const_tensor();
   output.alloc_attr = item.output_attrs()[0];
 }
 
@@ -800,7 +800,7 @@ void ExecutorState<PropagatorStateType>::Process(TaggedNode tagged_node,
       if (outputs.size() < item.num_outputs) outputs.resize(item.num_outputs);
     } else if (TF_PREDICT_FALSE(item.is_noop)) {
       ProcessNoop(stats);
-    } else if (item.const_tensor != nullptr && !params.track_allocations) {
+    } else if (item.kernel->const_tensor() != nullptr && !params.track_allocations) {
       ProcessConstTensor(item, &outputs, stats);
     } else {
       // Prepares inputs.

diff --git a/tensorflow/core/common_runtime/graph_view.h b/tensorflow/core/common_runtime/graph_view.h
@@ -89,9 +89,6 @@ struct NodeItem {
   // The kernel for this node.
   OpKernel* kernel = nullptr;
 
-  // If the kernel is a Const op, this containts points to the constant tensor.
-  const Tensor* const_tensor = nullptr;
-
   // Cached values of node->num_inputs() and node->num_outputs(), to
   // avoid levels of indirection.
   int num_inputs;

diff --git a/tensorflow/core/common_runtime/immutable_executor_state.cc b/tensorflow/core/common_runtime/immutable_executor_state.cc
@@ -144,14 +144,6 @@ Status ImmutableExecutorState::Initialize() {
         break;
       }
     }
-    const Tensor* const_tensor = item->kernel->const_tensor();
-    if (const_tensor) {
-      // Hold onto a shallow copy of the constant tensor in `*this` so that the
-      // reference count does not drop to 1. This prevents the constant tensor
-      // from being forwarded, and its buffer reused.
-      const_tensors_.emplace_back(*const_tensor);
-    }
-    item->const_tensor = const_tensor;
     item->is_noop = (item->kernel->type_string() == "NoOp");
     item->is_enter = IsEnter(n);
     if (item->is_enter) {

diff --git a/tensorflow/core/common_runtime/immutable_executor_state.h b/tensorflow/core/common_runtime/immutable_executor_state.h
@@ -154,9 +154,6 @@ class ImmutableExecutorState {
   // pending counts for the nodes in the graph, indexed by node ID.
   std::unique_ptr<std::atomic<int32>[]> atomic_pending_counts_;
 
-  // Shallow copies of the constant tensors used in the graph.
-  std::vector<Tensor> const_tensors_;
-
   TF_DISALLOW_COPY_AND_ASSIGN(ImmutableExecutorState);
 };