Create images per blocks

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

Author: jantonguirao

dali/operators/imgcodec/image_decoder.h

@@ -765,6 +765,9 @@ class ImageDecoder : public StatelessOperator<Backend> {
 
     // The image descriptors are created in parallel, in block-wise fashion.
     auto init_desc_task = [&](int start_sample, int end_sample) {
+      DomainTimeRange tr(
+          "Create images " + std::to_string(start_sample) + ".." + std::to_string(end_sample),
+          DomainTimeRange::kOrange);
       for (int orig_idx = start_sample; orig_idx < end_sample; orig_idx++) {
         auto &st = *state_[orig_idx];
         if (use_cache && st.load_from_cache) {
@@ -799,19 +802,24 @@ class ImageDecoder : public StatelessOperator<Backend> {
     } else {
       DomainTimeRange tr("Create images", DomainTimeRange::kOrange);
       // Many tasks? Run in thread pool.
-      // The first span of tasks is processed in the main operator thread.
-      for (int i = 1; i < ntasks; i++) {
-        int start = i * nsamples / ntasks;
-        int end = (i + 1) * nsamples / ntasks;
-        tp_->AddWork([&, start, end](int) {
+      int block_idx = 0;
+      atomic_idx_.store(0);
+      auto create_images_task = [&, nblocks](int tid) {
+        int block_idx;
+        while ((block_idx = atomic_idx_.fetch_add(1)) < 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);
       }
-      // Start processing of subsequent segments...
-      tp_->RunAll(false);
-      // ...and process the 1st segment in this thread.
-      init_desc_task(0, nsamples / ntasks);
-      tp_->WaitForWork();
+      assert(ntasks >= 2);
+      tp_->RunAll(false);  // start work but not wait
+      create_images_task(-1);
+      tp_->WaitForWork();  // wait for the other threads
     }
 
     bool any_need_processing = false;

dali/pipeline/util/thread_pool.cc

@@ -60,7 +60,7 @@ ThreadPool::~ThreadPool() {
 
 void ThreadPool::AddWork(Work work, int64_t priority, bool start_immediately) {
   bool started_before = started_;
-  outstanding_work_.fetch_add(1);
+  outstanding_work_.fetch_add(1, std::memory_order_relaxed);
   if (started_before) {
     std::lock_guard lock(queue_lock_);
     work_queue_.push({priority, std::move(work)});
@@ -81,10 +81,12 @@ void ThreadPool::AddWork(Work work, int64_t priority, bool start_immediately) {
 
 // Blocks until all work issued to the thread pool is complete
 void ThreadPool::WaitForWork(bool checkForErrors) {
-  if (outstanding_work_.load()) {
+  // Wait for outstanding work to complete
+  if (outstanding_work_.load(std::memory_order_relaxed) > 0) {
+    // If still not complete, use condition variable
     std::unique_lock lock(completed_mutex_);
     completed_.wait(lock, [&, this] {
-      return this->outstanding_work_ == 0;
+      return this->outstanding_work_.load(std::memory_order_relaxed) == 0;
     });
   }
   started_ = false;
@@ -184,7 +186,7 @@ void ThreadPool::ThreadMain(int thread_id, int device_id, bool set_affinity,
     // If it reaches zero, we must safely notify the potential threads waiting for the work
     // to complete.
     // NOTE: We don't have to acquire the mutex until the number of waiting threads reaches 0.
-    if (--outstanding_work_ == 0) {
+    if (outstanding_work_.fetch_sub(1, std::memory_order_relaxed) == 1) {
       // We don't need to guard the modification of the atomic value with a mutex -
       // however, we need to lock it briefly to make sure we don't have this scenario:
       //