Combine setup into one task

Signed-off-by: Joaquin Anton Guirao <[email protected]>

Author: jantonguirao

dali/operators/imgcodec/image_decoder.h

@@ -19,6 +19,7 @@
 #include <vector>
 #include "dali/core/call_at_exit.h"
 #include "dali/core/mm/memory.h"
+#include "dali/core/semaphore.h"
 #include "dali/operators.h"
 #include "dali/operators/decoder/cache/cached_decoder_impl.h"
 #include "dali/operators/generic/slice/slice_attr.h"
@@ -677,7 +678,7 @@ class ImageDecoder : public StatelessOperator<Backend> {
     auto setup_block = [&](int block_idx, int nblocks, int tid) {
       int i_start = nsamples * block_idx / nblocks;
       int i_end = nsamples * (block_idx + 1) / nblocks;
-      DomainTimeRange tr("Setup #" + std::to_string(block_idx) + "/" + std::to_string(nblocks),
+      DomainTimeRange tr("Parse #" + std::to_string(block_idx) + "/" + std::to_string(nblocks),
                           DomainTimeRange::kOrange);
       for (int i = i_start; i < i_end; i++) {
         auto *st = state_[i].get();
@@ -729,40 +730,6 @@ class ImageDecoder : public StatelessOperator<Backend> {
       }
     };
 
-    int nsamples_per_block = 16;
-    int nblocks = std::max(1, nsamples / nsamples_per_block);
-    int ntasks = std::min<int>(nblocks, std::min<int>(8, tp_->NumThreads() + 1));
-
-    if (ntasks < 2) {
-      DomainTimeRange tr("Setup", DomainTimeRange::kOrange);
-      setup_block(0, 1, -1);  // run all in current thread
-    } else {
-      int block_idx = 0;
-      atomic_idx_.store(0);
-      auto setup_task = [&, nblocks](int tid) {
-        DomainTimeRange tr("Setup", DomainTimeRange::kOrange);
-        int block_idx;
-        while ((block_idx = atomic_idx_.fetch_add(1)) < nblocks) {
-          setup_block(block_idx, nblocks, tid);
-        }
-      };
-
-      for (int task_idx = 0; task_idx < ntasks - 1; task_idx++) {
-        tp_->AddWork(setup_task, -task_idx);
-      }
-      assert(ntasks >= 2);
-      tp_->RunAll(false);  // start work but not wait
-      setup_task(-1);      // last task in current thread
-      tp_->WaitForWork();  // wait for the other threads
-    }
-
-    // Allocate the memory for the outputs...
-    {
-      DomainTimeRange tr("Alloc output", DomainTimeRange::kOrange);
-      output.Resize(out_shape);
-    }
-    // ... and create image descriptors.
-
     // The image descriptors are created in parallel, in block-wise fashion.
     auto init_desc_task = [&](int start_sample, int end_sample) {
       DomainTimeRange tr(
@@ -795,30 +762,55 @@ class ImageDecoder : public StatelessOperator<Backend> {
       }
     };
 
-    // Just one task? Run it in this thread!
+    int nsamples_per_block = 16;
+    int nblocks = std::max(1, nsamples / nsamples_per_block);
+    int ntasks = std::min<int>(nblocks, std::min<int>(8, tp_->NumThreads() + 1));
+
     if (ntasks < 2) {
-      DomainTimeRange tr("Create images", DomainTimeRange::kOrange);
+      // run all in current thread
+      DomainTimeRange tr("Setup", DomainTimeRange::kOrange);
+      setup_block(0, 1, -1);
+      {
+        DomainTimeRange tr("Alloc output", DomainTimeRange::kOrange);
+        output.Resize(out_shape);
+      }
       init_desc_task(0, nsamples);
     } else {
-      DomainTimeRange tr("Create images", DomainTimeRange::kOrange);
-      // Many tasks? Run in thread pool.
+      // run in parallel
       int block_idx = 0;
-      atomic_idx_.store(0);
-      auto create_images_task = [&, nblocks](int tid) {
+      // relaxed, only need atomicity, not ordering
+      atomic_idx_.store(0, std::memory_order_relaxed);
+      parse_barrier_.Reset(ntasks);
+      alloc_output_barrier_.Reset(ntasks);
+      auto setup_task = [&, nblocks](int tid) {
         int block_idx;
-        while ((block_idx = atomic_idx_.fetch_add(1)) < nblocks) {
+        DomainTimeRange tr("Setup", DomainTimeRange::kOrange);
+        while ((block_idx = atomic_idx_.fetch_add(1, std::memory_order_relaxed)) < nblocks) {
+          setup_block(block_idx, nblocks, tid);
+        }
+        parse_barrier_.Wait();  // wait until parsing is done
+        
+        if (tid == -1) {
+          DomainTimeRange tr("Alloc output", DomainTimeRange::kOrange);
+          output.Resize(out_shape);
+          atomic_idx_.store(0, std::memory_order_relaxed);
+        }
+
+        alloc_output_barrier_.Wait();  // wait until allocation is done
+        // Create image descriptors
+        while ((block_idx = atomic_idx_.fetch_add(1, std::memory_order_relaxed)) < nblocks) {
           int64_t start = nsamples * block_idx / nblocks;
           int64_t end = nsamples * (block_idx + 1) / nblocks;
           init_desc_task(start, end);
         }
       };
 
       for (int task_idx = 0; task_idx < ntasks - 1; task_idx++) {
-        tp_->AddWork(create_images_task, -task_idx);
+        tp_->AddWork(setup_task, -task_idx);
       }
       assert(ntasks >= 2);
       tp_->RunAll(false);  // start work but not wait
-      create_images_task(-1);
+      setup_task(-1);      // last task in current thread
       tp_->WaitForWork();  // wait for the other threads
     }
 
@@ -985,6 +977,40 @@ class ImageDecoder : public StatelessOperator<Backend> {
   std::vector<nvimgcodecExtension_t> extensions_;
 
   std::vector<std::function<void(int)>> nvimgcodec_scheduled_tasks_;
+
+  class ThreadBarrier {
+   public:
+    explicit ThreadBarrier(std::size_t count) : count_(count), current_(count) {}
+    void Wait(bool reset = false) {
+      std::unique_lock<std::mutex> lock(mutex_);
+      if (current_ == 0) {
+        throw std::logic_error("barrier is already completed");
+      }
+      current_--;
+      if (current_ == 0 || count_ == 0) {
+        if (reset)
+          current_ = count_;
+        cv_.notify_all();
+      } else {
+        cv_.wait(lock, [this] { 
+          return current_ == 0; });
+      }
+    }
+    void Reset(std::size_t count) {
+      std::lock_guard<std::mutex> lock(mutex_);
+      count_ = count;
+      current_ = count;
+    }
+
+   private:
+    std::mutex mutex_;
+    std::condition_variable cv_;
+    size_t count_;
+    size_t current_;
+  };
+
+  ThreadBarrier parse_barrier_{1};
+  ThreadBarrier alloc_output_barrier_{1};
 };
 
 }  // namespace imgcodec