Support chunked output buffers (#74)

Support chunked output buffers.

Signed-off-by: Rafal <[email protected]>

Author: banasraf

src/dali_executor/dali_executor.cc

@@ -44,40 +44,82 @@ void DaliExecutor::SetupInputs(const std::vector<IDescr>& inputs) {
       assert(inp_size <= c_inputs.back().buffers[0].size);
     }
   }
-  RunInputCopy();
+  WaitForCopies();
   for (auto& inp : c_inputs) {
     pipeline_.SetInput(inp);
   }
 }
 
 
-IDescr DaliExecutor::ScheduleInputCopy(const IDescr& input) {
-  assert(input.buffers.size() > 0);
+IOBufferI* DaliExecutor::GetInputBuffer(const std::string& name, device_type_t device) {
   IOBufferI* buffer;
-  if (input.buffers[0].device == device_type_t::CPU) {
-    buffer = &cpu_buffers_[input.meta.name];
+  if (device == device_type_t::CPU) {
+    buffer = &cpu_buffers_[name + "_inp"];
   } else {
-    buffer = &gpu_buffers_[input.meta.name];
+    buffer = &gpu_buffers_[name + "_inp"];
   }
+  return buffer;
+}
+
+
+IOBufferI* DaliExecutor::GetOutputBuffer(const std::string& name, device_type_t device) {
+  IOBufferI* buffer;
+  if (device == device_type_t::CPU) {
+    buffer = &cpu_buffers_[name + "_out"];
+  } else {
+    buffer = &gpu_buffers_[name + "_out"];
+  }
+  return buffer;
+}
+
+
+IDescr DaliExecutor::ScheduleInputCopy(const IDescr& input) {
+  assert(input.buffers.size() > 0);
+  IOBufferI* buffer = GetInputBuffer(input.meta.name, input.buffers[0].device);
   size_t size = 0;
   for (auto& buf : input.buffers)
     size += buf.size;
   buffer->resize(size);
   auto descriptor = buffer->get_descr();
   char* dst = reinterpret_cast<char*>(descriptor.data);
+  auto stream = pipeline_.CopyStream();
   for (auto& buf : input.buffers) {
     thread_pool_.AddWork(
-        [descriptor, dst, buf](int) {
-          MemCopy(descriptor.device, dst, buf.device, buf.data, buf.size);
+        [stream, descriptor, dst, buf](int) {
+          MemCopy(descriptor.device, dst, buf.device, buf.data, buf.size, stream);
         },
         buf.size, true);
     dst += buf.size;
   }
   return IDescr{input.meta, {descriptor}};
 }
 
-void DaliExecutor::RunInputCopy() {
+void DaliExecutor::ScheduleOutputCopy(const ODescr& output, int output_idx) {
+  const auto& name = output.meta.name;
+  const auto& out_buffers = output.buffers;
+  size_t size = 0;
+  for (auto& out_buff : out_buffers) {
+    size += out_buff.size;
+  }
+  IOBufferI* interm_buffer = GetOutputBuffer(name, pipeline_.GetOutputDevice(output_idx));
+  interm_buffer->resize(size);
+  auto interm_descr = interm_buffer->get_descr();
+  pipeline_.PutOutput(interm_descr.data, output_idx, interm_descr.device);
+  char* src = reinterpret_cast<char*>(interm_descr.data);
+  auto stream = pipeline_.CopyStream();
+  for (auto& buf : out_buffers) {
+    thread_pool_.AddWork(
+        [stream, src, buf, interm_descr](int) {
+          MemCopy(buf.device, buf.data, interm_descr.device, src, buf.size, stream);
+        },
+        buf.size);
+    src += buf.size;
+  }
+}
+
+void DaliExecutor::WaitForCopies() {
   thread_pool_.RunAll();
+  pipeline_.SyncStream();
 }
 
 
@@ -106,14 +148,14 @@ std::vector<OutputInfo> DaliExecutor::Run(const std::vector<IDescr>& inputs) {
 
 void DaliExecutor::PutOutputs(const std::vector<ODescr>& outputs) {
   for (uint32_t output_idx = 0; output_idx < outputs.size(); ++output_idx) {
-    ENFORCE(outputs[output_idx].buffers.size() == 1,
-            "Ouptut can be copied only to a single buffer");
-    auto buffer = outputs[output_idx].buffers[0];
-    auto data = buffer.data;
-    auto device_type = buffer.device;
-    pipeline_.PutOutput(data, output_idx, device_type);
+    if (outputs[output_idx].buffers.size() == 1) {
+      auto buffer = outputs[output_idx].buffers[0];
+      pipeline_.PutOutput(buffer.data, output_idx, buffer.device);
+    } else {
+      ScheduleOutputCopy(outputs[output_idx], output_idx);
+    }
   }
-  pipeline_.SyncOutputStream();
+  WaitForCopies();
 }
 
 }}}  // namespace triton::backend::dali

src/dali_executor/dali_executor.h

@@ -61,16 +61,23 @@ class DaliExecutor {
   void SetupInputs(const std::vector<IDescr>& inputs);
 
   /**
-   * @brief Schedule a copy off all buffers within input IDescr to a continuous buffer.
-   *        The copy will be performed after calling RunInputCopy().
+   * @brief Schedule a copy of all buffers within input IDescr to a continuous buffer.
+   *        Call WaitForCopies() to wait for the copy to finish.
    * @return IDecr to the new, continuous, buffer.
    */
   IDescr ScheduleInputCopy(const IDescr& buffers);
 
   /**
-   * @brief Run copies scheduled by ScheduleInputCopy and wait for them to finish.
+   * @brief Schedule a copy to a chunked output through an intermediate buffer.
+   *        Call WaitForCopies() to wait for the copy to finish.
    */
-  void RunInputCopy();
+  void ScheduleOutputCopy(const ODescr& output, int output_idx);
+
+  /**
+   * @brief Wait for the copies scheduled by ScheduleInputCopy or ScheduleOutputCopy
+   *        and wait for them to finish.
+   */
+  void WaitForCopies();
 
   /**
    * @brief Check if an input can be used without a copy.
@@ -82,6 +89,16 @@ class DaliExecutor {
     return (n_threads < 1) ? 1 : n_threads;
   }
 
+  /**
+   * @brief Get an intermediate buffer located on the \p device for an input with a given \p name
+   */
+  IOBufferI* GetInputBuffer(const std::string& name, device_type_t device);
+
+  /**
+   * @brief Get an intermediate buffer located on the \p device for an output with a given \p name
+   */
+  IOBufferI* GetOutputBuffer(const std::string& name, device_type_t device);
+
   DaliPipeline pipeline_;
   ThreadPool thread_pool_;
   std::map<std::string, IOBuffer<CPU>> cpu_buffers_;

src/dali_executor/dali_pipeline.cc

@@ -77,7 +77,7 @@ void DaliPipeline::SetInput(const IDescr& io_descr) {
   SetInput(buffer.data, meta.name.c_str(), buffer.device, meta.type, meta.shape);
 }
 
-void DaliPipeline::SyncOutputStream() {
+void DaliPipeline::SyncStream() {
   if (NoGpu())
     return;
   DeviceGuard dg(device_id_);

src/dali_executor/dali_pipeline.h

@@ -125,11 +125,15 @@ class DaliPipeline {
   void PutOutput(void* destination, int output_idx, device_type_t destination_device);
 
   /**
-   * @brief Wait for all output copies.
+   * @brief Wait for the work scheduled on the copy stream.
    *
    * This should be always called after copying all of the pipeline outputs.
    */
-  void SyncOutputStream();
+  void SyncStream();
+
+  cudaStream_t CopyStream() {
+    return output_stream_;
+  }
 
   void Reset() {
     ReleasePipeline();

src/dali_executor/executor.test.cc

@@ -29,15 +29,49 @@
 
 namespace triton { namespace backend { namespace dali { namespace test {
 
+template<typename T, typename Op>
+void coalesced_compare(const std::vector<OBufferDescr> &obuffers,
+                       const std::vector<std::vector<T>> &ibuffers, size_t inp_size, const Op &op) {
+  size_t inp_buff_i = 0;
+  size_t inp_i = 0;
+  size_t out_buff_i = 0;
+  size_t out_i = 0;
+  std::vector<T> obuffer;
+  for (size_t i = 0; i < inp_size; ++i) {
+    if (inp_i == ibuffers[inp_buff_i].size()) {
+      inp_i = 0;
+      inp_buff_i++;
+    }
+    if (out_i == obuffers[out_buff_i].size / sizeof(T)) {
+      out_i = 0;
+      out_buff_i++;
+    }
+    if (out_i == 0) {
+      auto descr = obuffers[out_buff_i];
+      REQUIRE(descr.size % sizeof(T) == 0);
+      obuffer.resize(descr.size / sizeof(T));
+      MemCopy(CPU, obuffer.data(), descr.device, descr.data, descr.size);
+    }
+    REQUIRE(obuffer[out_i] == op(ibuffers[inp_buff_i][inp_i]));
+    out_i++;
+    inp_i++;
+  }
+}
+
 TEST_CASE("Scaling Pipeline") {
   std::string pipeline_s((const char *)pipelines::scale_pipeline_str,
                          pipelines::scale_pipeline_len);
-  DaliPipeline pipeline(pipeline_s, 8, 4, 0);
+  DaliPipeline pipeline(pipeline_s, 256, 4, 0);
   DaliExecutor executor(std::move(pipeline));
   std::mt19937 rand(1217);
   std::uniform_real_distribution<float> dist(-1.f, 1.f);
   const std::string inp_name = "INPUT0";
-  auto scaling_test = [&](const std::vector<int> &batch_sizes) {
+  auto scaling_test = [&](const std::vector<int> &batch_sizes,
+                          const std::vector<int> &out_batch_sizes,
+                          const std::vector<device_type_t> &out_devs) {
+    REQUIRE(std::accumulate(batch_sizes.begin(), batch_sizes.end(), 0) ==
+            std::accumulate(out_batch_sizes.begin(), out_batch_sizes.end(), 0));
+    REQUIRE(out_devs.size() == out_batch_sizes.size());
     std::vector<TensorListShape<>> shapes;
     for (auto batch_size : batch_sizes) {
       TensorListShape<> shape(batch_size, 2);
@@ -53,33 +87,39 @@ TEST_CASE("Scaling Pipeline") {
     size_t inp_size = 0;
     for (auto &inp_buffer : input_buffers)
       inp_size += inp_buffer.size();
-    std::vector<float> output_buffer(inp_size);
+    std::vector<std::unique_ptr<IOBufferI>> output_buffers;
+    int ti = 0;
+    for (size_t out_i = 0; out_i < out_batch_sizes.size(); ++out_i) {
+      int64_t buffer_vol = 0;
+      for (int i = 0; i < out_batch_sizes[out_i]; ++i) {
+        buffer_vol += volume(output[0].shape[ti]) * sizeof(float);
+        ti++;
+      }
+      if (out_devs[out_i] == device_type_t::CPU) {
+        output_buffers.emplace_back(std::make_unique<IOBuffer<CPU>>(buffer_vol));
+      } else {
+        output_buffers.emplace_back(std::make_unique<IOBuffer<GPU>>(buffer_vol));
+      }
+    }
     std::vector<ODescr> output_vec(1);
     auto &outdesc = output_vec[0];
-    OBufferDescr buf_descr;
-    buf_descr.device = device_type_t::CPU;
-    buf_descr.data = output_buffer.data();
-    buf_descr.size = output_buffer.size() * sizeof(decltype(output_buffer)::size_type);
-    outdesc.buffers = {buf_descr};
-    executor.PutOutputs(output_vec);
-    size_t out_i = 0;
-    int i = 0;
-    for (auto &inp_buffer : input_buffers) {
-      for (size_t i = 0; i < inp_buffer.size(); ++i) {
-        REQUIRE(output_buffer[out_i] == inp_buffer[i] * 2);
-        ++out_i;
-      }
+    for (auto &out_buffer : output_buffers) {
+      outdesc.buffers.push_back(out_buffer->get_descr());
     }
+    executor.PutOutputs(output_vec);
+    coalesced_compare(outdesc.buffers, input_buffers, inp_size, [](float a) { return a * 2; });
   };
 
   SECTION("Simple execute") {
-    scaling_test({3, 2, 1});
-    scaling_test({5});
+    scaling_test({3, 2, 1}, {6}, {CPU});
+    scaling_test({5}, {5}, {GPU});
   }
 
-  SECTION("Repeat batch size") {
-    scaling_test({3, 3});
-    scaling_test({6});
+  SECTION("Chunked output") {
+    scaling_test({3, 3}, {3, 3}, {CPU, CPU});
+    scaling_test({6}, {2, 4}, {GPU, GPU});
+    scaling_test({8}, {6, 2}, {CPU, GPU});
+    scaling_test({64}, {32, 16, 16}, {CPU, GPU, GPU});
   }
 }
 
@@ -110,7 +150,7 @@ TEST_CASE("RN50 pipeline") {
     obuffer.device = device_type_t::CPU;
     obuffer.device_id = 0;
     obuffer.data = output_buffer.data();
-    obuffer.size = output_buffer.size() * sizeof(decltype(output_buffer)::size_type);
+    obuffer.size = output_buffer.size() * sizeof(decltype(output_buffer)::value_type);
     outdesc.buffers = {obuffer};
     executor.PutOutputs(output_vec);
     for (int c = 0; c < output_c; ++c) {