Dump operator stats

dali/c_api/c_api.cc

@@ -109,7 +109,8 @@ void daliCreatePipeline(daliPipelineHandle *pipe_handle,
                         int separated_execution,
                         int prefetch_queue_depth,
                         int cpu_prefetch_queue_depth,
-                        int gpu_prefetch_queue_depth) {
+                        int gpu_prefetch_queue_depth,
+                        int get_memory_stats) {
   bool se = separated_execution != 0;
   auto pipeline = std::make_unique<dali::Pipeline>(std::string(serialized_pipeline, length),
                                                    batch_size, num_threads, device_id, true,
@@ -118,6 +119,7 @@ void daliCreatePipeline(daliPipelineHandle *pipe_handle,
   if (se) {
     pipeline->SetQueueSizes(cpu_prefetch_queue_depth, gpu_prefetch_queue_depth);
   }
+  pipeline->EnableOperatorOutputMemoryStatistics(get_memory_stats);
   pipeline->Build();
   auto ws = std::make_unique<dali::DeviceWorkspace>();
   auto stream = dali::CUDAStream::Create(true);
@@ -441,3 +443,34 @@ void daliGetReaderMetadata(daliPipelineHandle* pipe_handle, const char *reader_n
   meta->pad_last_batch = returned_meta.pad_last_batch;
   meta->stick_to_shard = returned_meta.stick_to_shard;
 }
+
+void daliGetExecutorMetadata(daliPipelineHandle* pipe_handle, daliExecutorMetadata **operator_meta,
+                             size_t *operator_meta_num) {
+  dali::Pipeline* pipeline = reinterpret_cast<dali::Pipeline*>(pipe_handle->pipe);
+  auto returned_meta = pipeline->GetExecutorMeta();
+  *operator_meta_num = returned_meta.size();
+  *operator_meta = static_cast<daliExecutorMetadata*>(malloc(sizeof(daliExecutorMetadata) *
+                                                     returned_meta.size()));
+
+  int i = 0;
+  for (const auto &stat : returned_meta) {
+    auto op_name_size = stat.first.size();
+    (*operator_meta)[i].operator_name = static_cast<char*>(malloc(sizeof(char) *
+                                                                 (op_name_size + 1)));
+    stat.first.copy((*operator_meta)[i].operator_name, op_name_size);
+    (*operator_meta)[i].operator_name[op_name_size] = '\0';
+
+    auto num_outputs = stat.second.size();
+    (*operator_meta)[i].out_num = num_outputs;
+    (*operator_meta)[i].real_size = static_cast<size_t*>(malloc(sizeof(size_t) * num_outputs));
+    (*operator_meta)[i].reserved = static_cast<size_t*>(malloc(sizeof(size_t) * num_outputs));
+
+    for (size_t j = 0; j < num_outputs; ++j) {
+      const auto &entry = stat.second[j];
+      (*operator_meta)[i].real_size[j] = entry.real_size;
+      (*operator_meta)[i].reserved[j] = entry.reserved;
+    }
+    ++i;
+  }
+}
+

dali/c_api/c_api_test.cc

@@ -166,7 +166,7 @@ TYPED_TEST(CApiTest, FileReaderPipe) {
   daliPipelineHandle handle;
   daliCreatePipeline(&handle, serialized.c_str(), serialized.size(), batch_size, num_thread,
                      device_id, false, prefetch_queue_depth, prefetch_queue_depth,
-                     prefetch_queue_depth);
+                     prefetch_queue_depth, false);
   daliPrefetchUniform(&handle, prefetch_queue_depth);
 
   dali::DeviceWorkspace ws;
@@ -223,7 +223,7 @@ TYPED_TEST(CApiTest, ExternalSourceSingleAllocPipe) {
   daliPipelineHandle handle;
   daliCreatePipeline(&handle, serialized.c_str(), serialized.size(), batch_size, num_thread,
                      device_id, false, prefetch_queue_depth, prefetch_queue_depth,
-                     prefetch_queue_depth);
+                     prefetch_queue_depth, false);
 
   for (int i = 0; i < prefetch_queue_depth; i++) {
     SequentialFill(view<uint8_t>(input_cpu), 42 * i);
@@ -280,7 +280,7 @@ TYPED_TEST(CApiTest, ExternalSourceMultipleAllocPipe) {
   daliPipelineHandle handle;
   daliCreatePipeline(&handle, serialized.c_str(), serialized.size(), batch_size, num_thread,
                      device_id, false, prefetch_queue_depth, prefetch_queue_depth,
-                     prefetch_queue_depth);
+                     prefetch_queue_depth, false);
 
   for (int i = 0; i < prefetch_queue_depth; i++) {
     SequentialFill(view<uint8_t>(input_cpu), 42 * i);
@@ -340,7 +340,7 @@ TYPED_TEST(CApiTest, ExternalSourceSingleAllocDifferentBackendsTest) {
   daliPipelineHandle handle;
   daliCreatePipeline(&handle, serialized.c_str(), serialized.size(), batch_size, num_thread,
                      device_id, false, prefetch_queue_depth, prefetch_queue_depth,
-                     prefetch_queue_depth);
+                     prefetch_queue_depth, false);
 
   for (int i = 0; i < prefetch_queue_depth; i++) {
     SequentialFill(view<uint8_t>(input_cpu), 42 * i);
@@ -404,7 +404,7 @@ TYPED_TEST(CApiTest, ExternalSourceMultipleAllocDifferentBackendsTest) {
   daliPipelineHandle handle;
   daliCreatePipeline(&handle, serialized.c_str(), serialized.size(), batch_size, num_thread,
                      device_id, false, prefetch_queue_depth, prefetch_queue_depth,
-                     prefetch_queue_depth);
+                     prefetch_queue_depth, false);
 
   for (int i = 0; i < prefetch_queue_depth; i++) {
     SequentialFill(view<uint8_t>(input_cpu), 42 * i);
@@ -445,4 +445,30 @@ TYPED_TEST(CApiTest, ExternalSourceMultipleAllocDifferentBackendsTest) {
   ComparePipelinesOutputs<OpBackend>(handle, *pipe_ptr);
 }
 
+TYPED_TEST(CApiTest, TestExecutorMeta) {
+  auto pipe_ptr = GetTestPipeline<TypeParam>(true, this->output_device_);
+  auto serialized = pipe_ptr->SerializeToProtobuf();
+
+  pipe_ptr.reset();
+  daliPipelineHandle handle;
+  daliCreatePipeline(&handle, serialized.c_str(), serialized.size(), batch_size, num_thread,
+                     device_id, false, prefetch_queue_depth, prefetch_queue_depth,
+                     prefetch_queue_depth, true);
+
+  daliRun(&handle);
+  daliOutput(&handle);
+
+  size_t N;
+  daliExecutorMetadata *meta;
+  daliGetExecutorMetadata(&handle, &meta, &N);
+  EXPECT_EQ(N, 4);
+  for (size_t i = 0; i < N; ++i) {
+    free(meta[i].operator_name);
+    free(meta[i].real_size);
+    free(meta[i].reserved);
+  }
+  free(meta);
+  CUDA_CALL(cudaDeviceSynchronize());
+}
+
 }  // namespace dali

dali/pipeline/data/tensor_vector.h

@@ -113,6 +113,22 @@ class TensorVector {
     return tensors_.size();
   }
 
+  size_t nbytes() const noexcept {
+    size_t total_nbytes = 0;
+    for (const auto &t : tensors_) {
+      total_nbytes += t->nbytes();
+    }
+    return total_nbytes;
+  }
+
+  size_t capacity() const noexcept {
+    size_t total_capacity = 0;
+    for (const auto &t : tensors_) {
+      total_capacity += t->capacity();
+    }
+    return total_capacity;
+  }
+
   TensorListShape<> shape() const {
     if (state_ == State::contiguous) {
       return tl_->shape();

dali/pipeline/executor/executor.h

@@ -21,6 +21,8 @@
 #include <string>
 #include <utility>
 #include <vector>
+#include <unordered_map>
+#include <mutex>
 
 #include "dali/core/common.h"
 #include "dali/core/error_handling.h"
@@ -38,9 +40,15 @@
 #include "dali/pipeline/workspace/host_workspace.h"
 #include "dali/pipeline/workspace/mixed_workspace.h"
 #include "dali/pipeline/workspace/workspace_data_factory.h"
+#include "dali/pipeline/data/backend.h"
 
 namespace dali {
 
+struct DLL_PUBLIC ExecutorMeta {
+  size_t real_size;
+  size_t reserved;
+};
+using ExecutorMetaMap = std::unordered_map<std::string, std::vector<ExecutorMeta>>;
 
 namespace detail {
 // This is stream callback used on GPU stream to indicate that GPU work for this
@@ -62,6 +70,8 @@ class DLL_PUBLIC ExecutorBase {
   DLL_PUBLIC virtual void ShareOutputs(DeviceWorkspace *ws) = 0;
   DLL_PUBLIC virtual void ReleaseOutputs() = 0;
   DLL_PUBLIC virtual void SetCompletionCallback(ExecutorCallback cb) = 0;
+  DLL_PUBLIC virtual void EnableMemoryStats(bool get_memory_stats = false) = 0;
+  DLL_PUBLIC virtual ExecutorMetaMap GetExecutorMeta() = 0;
 
  protected:
   // virtual to allow the TestPruneWholeGraph test in gcc
@@ -101,6 +111,9 @@ class DLL_PUBLIC Executor : public ExecutorBase, public WorkspacePolicy, public
     stage_queue_depths_ = QueuePolicy::GetQueueSizes(prefetch_queue_depth);
   }
 
+  DLL_PUBLIC void EnableMemoryStats(bool get_memory_stats = false) override {
+    get_memory_stats_ = get_memory_stats;
+  }
   DLL_PUBLIC void Build(OpGraph *graph, vector<string> output_names) override;
   DLL_PUBLIC void Init() override {}
   DLL_PUBLIC void RunCPU() override;
@@ -110,6 +123,7 @@ class DLL_PUBLIC Executor : public ExecutorBase, public WorkspacePolicy, public
   DLL_PUBLIC void ShareOutputs(DeviceWorkspace *ws) override;
   DLL_PUBLIC void ReleaseOutputs() override;
   DLL_PUBLIC void SetCompletionCallback(ExecutorCallback cb) override;
+  DLL_PUBLIC ExecutorMetaMap GetExecutorMeta() override;
 
   DLL_PUBLIC void ShutdownQueue() {
     QueuePolicy::SignalStop();
@@ -118,6 +132,31 @@ class DLL_PUBLIC Executor : public ExecutorBase, public WorkspacePolicy, public
   DISABLE_COPY_MOVE_ASSIGN(Executor);
 
  protected:
+  template <typename W>
+  inline void FillStats(ExecutorMetaMap &memory_stats, W ws, std::string op_name,
+                       std::mutex &write_mutex) {
+    if (get_memory_stats_) {
+        size_t out_size = 0;
+        size_t reserved_size = 0;
+        std::lock_guard<std::mutex> lck(write_mutex);
+        auto &stats = memory_stats[op_name];
+        stats.resize(ws.NumOutput(), {0, 0});
+        for (int i = 0; i < ws.NumOutput(); ++i) {
+          out_size = 0;
+          reserved_size = 0;
+          if (ws.template OutputIsType<CPUBackend>(i)) {
+            out_size = ws.template OutputRef<CPUBackend>(i).nbytes();
+            reserved_size = ws.template OutputRef<CPUBackend>(i).capacity();
+          } else {
+            out_size = ws.template OutputRef<GPUBackend>(i).nbytes();
+            reserved_size = ws.template OutputRef<GPUBackend>(i).capacity();
+          }
+          stats[i].real_size = std::max(out_size, stats[i].real_size);
+          stats[i].reserved = std::max(reserved_size, stats[i].reserved);
+        }
+      }
+  }
+
   void HandleError(const char *message = "Unknown exception") {
     exec_error_ = true;
     ShutdownQueue();
@@ -220,6 +259,12 @@ class DLL_PUBLIC Executor : public ExecutorBase, public WorkspacePolicy, public
   // in some edge cases where there are no operators
   std::vector<cudaEvent_t> mixed_callback_events_;
 
+  std::atomic<bool> get_memory_stats_ = ATOMIC_VAR_INIT(false);;
+  ExecutorMetaMap cpu_memory_stats_, mixed_memory_stats_, gpu_memory_stats_;
+  std::mutex cpu_memory_stats_mutex_;
+  std::mutex mixed_memory_stats_mutex_;
+  std::mutex gpu_memory_stats_mutex_;
+
  private:
   template <typename Workspace>
   void RunHelper(OpNode &op_node, Workspace &ws) {
@@ -266,6 +311,23 @@ void Executor<WorkspacePolicy, QueuePolicy>::SetCompletionCallback(ExecutorCallb
   }
 }
 
+template<typename map>
+void AppendToMap(map &ret, map &in_stats, std::mutex &mutex) {
+  const std::lock_guard<std::mutex> lock(mutex);
+  for (auto const& stats : in_stats) {
+    ret.emplace(stats);
+  }
+}
+
+template <typename WorkspacePolicy, typename QueuePolicy>
+ExecutorMetaMap Executor<WorkspacePolicy, QueuePolicy>::GetExecutorMeta() {
+  ExecutorMetaMap ret;
+  AppendToMap(ret, cpu_memory_stats_, cpu_memory_stats_mutex_);
+  AppendToMap(ret, mixed_memory_stats_, mixed_memory_stats_mutex_);
+  AppendToMap(ret, gpu_memory_stats_, gpu_memory_stats_mutex_);
+  return ret;
+}
+
 template <typename WorkspacePolicy, typename QueuePolicy>
 void Executor<WorkspacePolicy, QueuePolicy>::Build(OpGraph *graph, vector<string> output_names) {
   DALI_ENFORCE(graph != nullptr, "Input graph is nullptr.");
@@ -347,6 +409,7 @@ void Executor<WorkspacePolicy, QueuePolicy>::RunCPU() {
 
     try {
       RunHelper(op_node, ws);
+      FillStats(cpu_memory_stats_, ws, "CPU_" + op_node.instance_name, cpu_memory_stats_mutex_);
     } catch (std::exception &e) {
       HandleError(e.what());
     } catch (...) {
@@ -381,6 +444,8 @@ void Executor<WorkspacePolicy, QueuePolicy>::RunMixed() {
       TimeRange tr("[Executor] Run Mixed op " + op_node.instance_name,
           TimeRange::kOrange);
       RunHelper(op_node, ws);
+      FillStats(mixed_memory_stats_, ws,  "MIXED_" + op_node.instance_name,
+                mixed_memory_stats_mutex_);
       if (ws.has_stream() && ws.has_event()) {
         CUDA_CALL(cudaEventRecord(ws.event(), ws.stream()));
       }
@@ -438,6 +503,7 @@ void Executor<WorkspacePolicy, QueuePolicy>::RunGPU() {
       TimeRange tr("[Executor] Run GPU op " + op_node.instance_name,
           TimeRange::knvGreen);
       RunHelper(op_node, ws);
+      FillStats(gpu_memory_stats_, ws, "GPU_" + op_node.instance_name, gpu_memory_stats_mutex_);
       if (ws.has_event()) {
         CUDA_CALL(cudaEventRecord(ws.event(), ws.stream()));
       }

dali/pipeline/pipeline.cc

@@ -167,11 +167,11 @@ int Pipeline::AddOperator(const OpSpec &spec, const std::string& inst_name) {
 }
 
 int Pipeline::AddOperator(const OpSpec &spec, int logical_id) {
-  return AddOperator(spec, "<no name>", logical_id);
+  return AddOperator(spec, make_string("<no name>_", logical_id), logical_id);
 }
 
 int Pipeline::AddOperator(const OpSpec &spec) {
-  return AddOperator(spec, "<no name>", GetNextLogicalId());
+  return AddOperator(spec, GetNextLogicalId());
 }
 
 
@@ -431,6 +431,7 @@ void Pipeline::Build(vector<std::pair<string, string>> output_names) {
   executor_ = GetExecutor(pipelined_execution_, separated_execution_, async_execution_, batch_size_,
                           num_threads_, device_id_, bytes_per_sample_hint_, set_affinity_,
                           max_num_stream_, default_cuda_stream_priority_, prefetch_queue_depth_);
+  executor_->EnableMemoryStats(get_memory_stats_);
   executor_->Init();
 
   // Creating the graph

dali/pipeline/pipeline.h

@@ -95,7 +95,6 @@ class DLL_PUBLIC Pipeline {
          default_cuda_stream_priority, QueueSizes{prefetch_queue_depth});
   }
 
-
   DLL_PUBLIC Pipeline(const string &serialized_pipe, int batch_size = -1, int num_threads = -1,
                       int device_id = -1, bool pipelined_execution = true,
                       int prefetch_queue_depth = 2, bool async_execution = true,
@@ -310,6 +309,29 @@ class DLL_PUBLIC Pipeline {
     async_execution_ = async_execution;
   }
 
+  /**
+   * @brief Set if the DALI pipeline should operator output buffer statistics
+   *
+   * @param get_memory_stats If DALI should print operator output buffer statistics.
+   * Usefull for `bytes_per_sample_hint` operator parameter.
+   */
+  DLL_PUBLIC void EnableOperatorOutputMemoryStatistics(bool get_memory_stats = true) {
+    get_memory_stats_ = get_memory_stats;
+    if (executor_) {
+      executor_->EnableMemoryStats(get_memory_stats_);
+    }
+  }
+
+  /**
+   * @brief Obtains the executor statistics
+   */
+  DLL_PUBLIC ExecutorMetaMap GetExecutorMeta() {
+    if (executor_) {
+      return  executor_->GetExecutorMeta();
+    } else {
+      return {};
+    }
+  }
 
   /**
    * @brief Set queue sizes for Pipeline using Separated Queues
@@ -492,6 +514,7 @@ class DLL_PUBLIC Pipeline {
   int next_logical_id_ = 0;
   int next_internal_logical_id_ = -1;
   QueueSizes prefetch_queue_depth_;
+  bool get_memory_stats_ = false;
 
   std::vector<int64_t> seed_;
   int original_seed_;