Include outputs in liveness analysis.

Signed-off-by: Michal Zientkiewicz <michalz@nvidia.com>
NVIDIA · Oct 3, 2023 · d9197ef · d9197ef
1 parent 2c3c120
commit d9197ef
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 3 deletions.
diff --git a/dali/pipeline/executor/executor.h b/dali/pipeline/executor/executor.h
@@ -580,7 +580,7 @@ void Executor<WorkspacePolicy, QueuePolicy>::Build(OpGraph *graph, vector<string
 
   // Create corresponding storage type for TensorNodes in graph
   tensor_to_store_queue_ =
-      CreateBackingStorageForTensorNodes(*graph_, max_batch_size_, queue_sizes);
+      CreateBackingStorageForTensorNodes(*graph_, max_batch_size_, queue_sizes, output_names_);
 
   InitializeLivenessInfo(liveness_info_, tensor_to_store_queue_);
 

diff --git a/dali/pipeline/graph/op_graph_storage.cc b/dali/pipeline/graph/op_graph_storage.cc
@@ -12,28 +12,38 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <set>
+#include <string>
 #include <vector>
 
 #include "dali/pipeline/graph/op_graph_storage.h"
 
 namespace dali {
 
 std::vector<tensor_data_store_queue_t> CreateBackingStorageForTensorNodes(
-    const OpGraph &op_graph, int batch_size, const std::vector<int> &queue_sizes) {
+    const OpGraph &op_graph, int batch_size, const std::vector<int> &queue_sizes,
+    const std::vector<std::string> &output_names) {
   DALI_ENFORCE(static_cast<int>(queue_sizes.size()) == op_graph.NumTensor(),
                "Data queue sizes undefined for some Tensor nodes.");
   std::vector<tensor_data_store_queue_t> result;
   result.resize(op_graph.NumTensor());
 
+  std::set<int64_t> outputs;
+  auto output_ids = op_graph.GetOutputs(output_names);;
+  outputs.insert(output_ids.begin(), output_ids.end());
+
   // Assign data to each Tensor node in graph
   for (int i = 0; i < op_graph.NumTensor(); i++) {
     const auto &tensor = op_graph.Tensor(i);
     auto producer_op_type = op_graph.Node(tensor.producer.node).op_type;
     result[i] =
         BatchFactory(producer_op_type, tensor.producer.storage_device, batch_size, queue_sizes[i]);
 
+    bool is_output = outputs.count(tensor.id) > 0;
     tuple_for_each(result[i], [&](auto &x) {
       x.num_consumers = tensor.consumers.size();
+      if (is_output)
+        x.num_consumers++;
     });
   }
   return result;

diff --git a/dali/pipeline/graph/op_graph_storage.h b/dali/pipeline/graph/op_graph_storage.h
@@ -15,6 +15,7 @@
 #ifndef DALI_PIPELINE_GRAPH_OP_GRAPH_STORAGE_H_
 #define DALI_PIPELINE_GRAPH_OP_GRAPH_STORAGE_H_
 
+#include <string>
 #include <vector>
 
 #include "dali/pipeline/graph/op_graph.h"
@@ -28,7 +29,8 @@ namespace dali {
 using MixedOpEventMap = std::vector<std::vector<cudaEvent_t>>;
 
 DLL_PUBLIC std::vector<tensor_data_store_queue_t> CreateBackingStorageForTensorNodes(
-    const OpGraph& op_graph, int batch_size, const std::vector<int>& queue_sizes);
+    const OpGraph& op_graph, int batch_size, const std::vector<int>& queue_sizes,
+    const std::vector<std::string> &output_names);
 
 // Mapping from MixedOp partition id to queue of corresponding events
 DLL_PUBLIC MixedOpEventMap CreateEventsForMixedOps(EventPool& event_pool, const OpGraph& op_graph,