Add non-cooperative jobs to new ThreadPool (#6245)

* This PR adds more lightweight non-cooperative jobs to new thread pool.
* The default Job and IncrementalJob become non-cooperative
* Cooperative[Incremental]Job can be used for reentrant behavior.

---------
Signed-off-by: Michal Zientkiewicz <michalz@nvidia.com>

Author: mzient

dali/core/exec/thread_pool_base.cc

@@ -18,34 +18,43 @@
 
 namespace dali {
 
-JobBase::~JobBase() noexcept(false) {
+template <bool cooperative>
+JobBase<cooperative>::~JobBase() noexcept(false) {
   if (total_tasks_ > 0 && !wait_completed_) {
     throw std::logic_error("The job is not empty, but hasn't been discarded or waited for.");
   }
   while (running_)
     std::this_thread::yield();
 }
 
-void JobBase::DoWait() {
+template <bool cooperative>
+void JobBase<cooperative>::DoWait() {
   if (wait_started_)
     throw std::logic_error("This job has already been waited for.");
-  wait_started_ = true;
 
   if (total_tasks_ == 0) {
+    // If there are no tasks, it's legal to skip a call to Run, therefore executor_ can be null.
+    wait_started_ = true;
     wait_completed_ = true;
     return;
   }
 
-  if (executor_ == nullptr)
+  if (this->executor_ == nullptr)
     throw std::logic_error("This job hasn't been run - cannot wait for it.");
 
-  auto ready = [&]() { return num_pending_tasks_ == 0; };
-  if (ThreadPoolBase::this_thread_pool() != nullptr) {
-    bool result = ThreadPoolBase::this_thread_pool()->WaitOrRunTasks(cv_, ready);
-    wait_completed_ = true;
-    if (!result)
-      throw std::logic_error("The thread pool was stopped");
+  if (ThreadPoolBase::this_thread_pool() == this->executor_) {
+    if constexpr (cooperative) {
+      auto ready = [&]() { return num_pending_tasks_ == 0; };
+      wait_started_ = true;
+      bool result = ThreadPoolBase::this_thread_pool()->WaitOrRunTasks(this->cv_, ready);
+      wait_completed_ = true;
+      if (!result)
+        throw std::logic_error("The thread pool was stopped");
+    } else {
+      throw std::logic_error("Cannot wait for this job from inside the thread pool.");
+    }
   } else {
+    wait_started_ = true;
     int old = num_pending_tasks_.load();
     while (old != 0) {
       num_pending_tasks_.wait(old);
@@ -56,10 +65,13 @@ void JobBase::DoWait() {
   }
 }
 
-void JobBase::DoNotify() {
+template <bool cooperative>
+void JobBase<cooperative>::DoNotify() {
   num_pending_tasks_.notify_all();
-  (void)std::lock_guard(mtx_);
-  cv_.notify_all();
+  if constexpr (cooperative) {
+    (void)std::lock_guard(this->mtx_);
+    this->cv_.notify_all();
+  }
   // We need this second flag to avoid a race condition where the destructor is called between
   // decrementing num_pending_tasks_ and notification_ without excessive use of mutexes.
   // This must be the very last operation in the task function that touches `this`.
@@ -68,29 +80,35 @@ void JobBase::DoNotify() {
 
 // Job ////////////////////////////////////////////////////////////////////
 
-void Job::Run(ThreadPoolBase &tp, bool wait) {
-  if (executor_ != nullptr)
+template <bool cooperative>
+void JobImpl<cooperative>::Run(ThreadPoolBase &tp, bool wait) {
+  if (this->executor_ != nullptr)
     throw std::logic_error("This job has already been started.");
-  executor_ = &tp;
-  running_ = !tasks_.empty();
+
+  if (!cooperative && &tp == ThreadPoolBase::this_thread_pool())
+    throw std::logic_error("Cannot run this job from inside the thread pool.");
+
+  this->executor_ = &tp;
+  this->running_ = !tasks_.empty();
   {
     auto batch = tp.BeginBulkAdd();
     for (auto &x : tasks_) {
       batch.Add(std::move(x.second.func));
     }
     int added = batch.Size();
     if (added) {
-      num_pending_tasks_ += added;
-      running_ = true;
+      this->num_pending_tasks_ += added;
+      this->running_ = true;
     }
     batch.Submit();
   }
   if (wait && !tasks_.empty())
     Wait();
 }
 
-void Job::Wait() {
-  DoWait();
+template <bool cooperative>
+void JobImpl<cooperative>::Wait() {
+  this->DoWait();
 
   // note - this vector is not allocated unless there were exceptions thrown
   std::vector<std::exception_ptr> errors;
@@ -104,19 +122,25 @@ void Job::Wait() {
     throw MultipleErrors(std::move(errors));
 }
 
-void Job::Discard() {
-  if (executor_ != nullptr)
+template <bool cooperative>
+void JobImpl<cooperative>::Discard() {
+  if (this->executor_ != nullptr)
     throw std::logic_error("Cannot discard a job that has already been started");
   tasks_.clear();
-  total_tasks_ = 0;
+  this->total_tasks_ = 0;
 }
 
 // IncrementalJob /////////////////////////////////////////////////////////
 
-void IncrementalJob::Run(ThreadPoolBase &tp, bool wait) {
-  if (executor_ && executor_ != &tp)
+template <bool cooperative>
+void IncrementalJobImpl<cooperative>::Run(ThreadPoolBase &tp, bool wait) {
+  if (this->executor_ && this->executor_ != &tp)
     throw std::logic_error("This job is already running in a different executor.");
-  executor_ = &tp;
+
+  if (!cooperative && &tp == ThreadPoolBase::this_thread_pool())
+    throw std::logic_error("Cannot run this job from inside the thread pool.");
+
+  this->executor_ = &tp;
   {
     auto it = last_task_run_.has_value() ? std::next(*last_task_run_) : tasks_.begin();
     auto batch = tp.BeginBulkAdd();
@@ -126,24 +150,26 @@ void IncrementalJob::Run(ThreadPoolBase &tp, bool wait) {
     }
     int added = batch.Size();
     if (added) {
-      num_pending_tasks_ += added;
-      running_ = true;
+      this->num_pending_tasks_ += added;
+      this->running_ = true;
     }
     batch.Submit();
   }
   if (wait && !tasks_.empty())
     Wait();
 }
 
-void IncrementalJob::Discard() {
-  if (executor_)
+template <bool cooperative>
+void IncrementalJobImpl<cooperative>::Discard() {
+  if (this->executor_)
     throw std::logic_error("Cannot discard a job that has already been started");
   tasks_.clear();
-  total_tasks_ = 0;
+  this->total_tasks_ = 0;
 }
 
-void IncrementalJob::Wait() {
-  DoWait();
+template <bool cooperative>
+void IncrementalJobImpl<cooperative>::Wait() {
+  this->DoWait();
   // note - this vector is not allocated unless there were exceptions thrown
   std::vector<std::exception_ptr> errors;
   for (auto &x : tasks_) {
@@ -158,8 +184,15 @@ void IncrementalJob::Wait() {
 
 ///////////////////////////////////////////////////////////////////////////
 
+template class JobBase<true>;
+template class JobImpl<true>;
+template class JobBase<false>;
+template class JobImpl<false>;
+template class IncrementalJobImpl<true>;
+template class IncrementalJobImpl<false>;
+
 thread_local ThreadPoolBase *ThreadPoolBase::this_thread_pool_ = nullptr;
-thread_local int ThreadPoolBase::this_thread_idx_ = -1;
+thread_local int ThisThreadIdx::this_thread_idx_ = -1;
 
 void ThreadPoolBase::Init(int num_threads, const std::function<OnThreadStartFn> &on_thread_start) {
   if (shutdown_pending_)

dali/core/exec/thread_pool_base_test.cc

@@ -153,9 +153,21 @@ TEST(NewThreadPool, RunLargeIncrementalJobInThreadPool) {
 template <typename JobType>
 class NewThreadPoolJobTest : public ::testing::Test {};
 
-using JobTypes = ::testing::Types<Job, IncrementalJob>;
+template <typename JobType>
+class NewThreadPoolCooperativeJobTest : public ::testing::Test {};
+
+template <typename JobType>
+class NewThreadPoolNonCooperativeJobTest : public ::testing::Test {};
+
+using JobTypes = ::testing::Types<Job, IncrementalJob, CooperativeJob, CooperativeIncrementalJob>;
 TYPED_TEST_SUITE(NewThreadPoolJobTest, JobTypes);
 
+using CooperativeJobTypes = ::testing::Types<CooperativeJob, CooperativeIncrementalJob>;
+TYPED_TEST_SUITE(NewThreadPoolCooperativeJobTest, CooperativeJobTypes);
+
+using NonCooperativeJobTypes = ::testing::Types<Job, IncrementalJob>;
+TYPED_TEST_SUITE(NewThreadPoolNonCooperativeJobTest, NonCooperativeJobTypes);
+
 
 TYPED_TEST(NewThreadPoolJobTest, RunJobInSeries) {
   TypeParam job;
@@ -184,7 +196,7 @@ TYPED_TEST(NewThreadPoolJobTest, Discard) {
   });
 }
 
-TYPED_TEST(NewThreadPoolJobTest, ErrorIncrementalJobNotStarted) {
+TYPED_TEST(NewThreadPoolJobTest, ErrorJobNotStarted) {
   try {
     TypeParam job;
     job.AddTask([]() {});
@@ -195,6 +207,19 @@ TYPED_TEST(NewThreadPoolJobTest, ErrorIncrementalJobNotStarted) {
   GTEST_FAIL() << "Expected a logic error.";
 }
 
+TYPED_TEST(NewThreadPoolJobTest, ErrorWaitBeforeRun) {
+  TypeParam job;
+  try {
+    job.AddTask([]() {});
+    job.Wait();
+  } catch (std::logic_error &e) {
+    EXPECT_NE(nullptr, strstr(e.what(), "hasn't been run"));
+    job.Discard();
+    return;
+  }
+  GTEST_FAIL() << "Expected a logic error.";
+}
+
 TYPED_TEST(NewThreadPoolJobTest, RethrowMultipleErrors) {
   TypeParam job;
   ThreadPoolBase tp(4);
@@ -210,7 +235,7 @@ TYPED_TEST(NewThreadPoolJobTest, RethrowMultipleErrors) {
   EXPECT_THROW(job.Run(tp, true), MultipleErrors);
 }
 
-TYPED_TEST(NewThreadPoolJobTest, Reentrant) {
+TYPED_TEST(NewThreadPoolCooperativeJobTest, Reentrant) {
   TypeParam job;
   ThreadPoolBase tp(1);  // must not hang with just one thread
   std::atomic_int outer{0}, inner{0};
@@ -221,7 +246,7 @@ TYPED_TEST(NewThreadPoolJobTest, Reentrant) {
   }
 
   job.AddTask([&]() {
-    Job innerJob;
+    TypeParam innerJob;
 
     for (int i = 0; i < 10; i++)
       innerJob.AddTask([&, i]() {
@@ -241,6 +266,25 @@ TYPED_TEST(NewThreadPoolJobTest, Reentrant) {
   job.Run(tp, true);
 }
 
+TYPED_TEST(NewThreadPoolNonCooperativeJobTest, Reentrant) {
+  TypeParam job;
+  ThreadPoolBase tp(1);  // must not hang with just one thread
+  job.AddTask([&]() {
+    TypeParam innerJob;
+    innerJob.AddTask([]() {});
+
+    try {
+      innerJob.Run(tp, true);
+    } catch (...) {
+      innerJob.Discard();
+      throw;
+    }
+  });
+
+  job.Run(tp, false);
+  EXPECT_THROW(job.Wait(), std::logic_error);
+}
+
 TYPED_TEST(NewThreadPoolJobTest, JobPerf) {
   using JobType = TypeParam;
   ThreadPoolBase tp(4);

include/dali/core/exec/thread_idx.h

@@ -0,0 +1,39 @@
+// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef DALI_CORE_EXEC_THREAD_IDX_H_
+#define DALI_CORE_EXEC_THREAD_IDX_H_
+
+#include "dali/core/api_helper.h"
+
+namespace dali {
+
+class DLL_PUBLIC ThisThreadIdx {
+ public:
+  /**
+   * @brief Returns the index of the current thread within the current thread pool
+   *
+   * @return the thread index or -1 if the calling thread does not belong to a thread pool
+   */
+  static int this_thread_idx() {
+    return this_thread_idx_;
+  }
+
+ protected:
+  static thread_local int this_thread_idx_;
+};
+
+}  // namespace dali
+
+#endif  // DALI_CORE_EXEC_THREAD_IDX_H_