[WIP]

Signed-off-by: Michał Zientkiewicz <[email protected]>

Author: mzient

dali/pipeline/util/new_thread_pool.cc

@@ -0,0 +1,22 @@
+// Copyright (c) 2022, 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.
+
+#include "dali/pipeline/util/new_thread_pool.h"
+
+namespace dali {
+namespace experimental {
+
+}  // namespace experimental
+}  // namespace dali
+

dali/pipeline/util/new_thread_pool.h

@@ -23,6 +23,8 @@
 #include <string>
 #include <utility>
 #include <vector>
+#include <mutex>
+#include <condition_variable>
 #include "dali/core/call_at_exit.h"
 #include "dali/core/error_handling.h"
 #include "dali/core/multi_error.h"
@@ -53,7 +55,7 @@ class Job {
   std::enable_if_t<std::is_convertible_v<Runnable, std::function<void(int)>>>
   AddTask(Runnable &&runnable, priority_t priority = {}) {
     if (started_)
-      throw std::logic_error("This has already been started - cannot add more tasks to it");
+      throw std::logic_error("This job has already been started - cannot add more tasks to it");
     auto it = tasks_.emplace(priority, Task());
     try {
       it->second.func = [this, task = &it->second, f = std::move(runnable)](int tid) {
@@ -150,19 +152,21 @@ class ThreadPoolBase {
     Stop();
   }
 
-  void AddTask(TaskFunc f) {
-    {
-      std::lock_guard<std::mutex> g(mtx_);
-      if (stop_requested_)
-        throw std::logic_error("The thread pool is stopped and no longer accepts new tasks.");
-      tasks_.push(std::move(f));
-    }
-    cv_.notify_one();
+  void AddTask(TaskFunc f);
+
+  static ThreadPoolBase *this_thread_pool() {
+    return this_thread_pool_;
+  }
+
+  static int this_thread_idx() {
+    return thread_idx_;
   }
 
  protected:
-  virtual void OnThreadStart(int thread_idx) {}
-  virtual void OnThreadStop(int thread_idx) {}
+  virtual void OnThreadStart(int thread_idx) noexcept {}
+  virtual void OnThreadStop(int thread_idx) noexcept {}
+
+  friend class Job;
 
   void Stop() {
     {
@@ -176,33 +180,66 @@ class ThreadPoolBase {
 
     {
       std::lock_guard<std::mutex> g(mtx_);
+      for (auto &task : tasks_) {
+      }
       threads_.clear();
     }
   }
 
-  void Run(int index) noexcept {
-    OnThreadStart(index);
-    detail::CallAtExit([&]() { OnThreadStop(index); });
+  template <typename Condition>
+  bool WaitOrRunTasks(std::condition_variable cv, Condition &&condition) {
+    assert(this_thread_pool() == this);
     std::unique_lock lock(mtx_);
-    while (!stop_requested_) {
+    do {
       cv_.wait(lock, [&]() { return stop_requested_ || !tasks_.empty(); });
-      if (stop_requested_)
-        break;
-      TaskFunc t = std::move(tasks_.front());
-      tasks_.pop();
-      lock.unlock();
-      t(index);
-      lock.lock();
     }
+
   }
 
+  static thread_local ThreadPoolBase *this_thread_pool_;
+  static int thread_idx_;
+
+  void Run(int index) noexcept;
+
   std::mutex mtx_;
   std::condition_variable cv_;
   bool stop_requested_ = false;
   std::queue<TaskFunc> tasks_;
   std::vector<std::thread> threads_;
 };
 
+
+thread_local ThreadPoolBase *ThreadPoolBase::this_thread_pool_ = nullptr;
+thread_local int ThreadPoolBase::this_thread_index_ = -1;;
+
+inline void ThreadPoolBase::AddTask(TaskFunc f) {
+  {
+    std::lock_guard<std::mutex> g(mtx_);
+    if (stop_requested_)
+    throw std::logic_error("The thread pool is stopped and no longer accepts new tasks.");
+    tasks_.push(std::move(f));
+  }
+  cv_.notify_one();
+}
+
+inline void ThreadPoolBase::Run(int index) noexcept {
+  ThreadPoolBase *this_thread_pool_ = this;
+  this_thread_idx_ = index;
+  OnThreadStart(index);
+  detail::CallAtExit([&]() { OnThreadStop(index); });
+  std::unique_lock lock(mtx_);
+  while (!stop_requested_) {
+    cv_.wait(lock, [&]() { return stop_requested_ || !tasks_.empty(); });
+    if (stop_requested_)
+      break;
+    TaskFunc t = std::move(tasks_.front());
+    tasks_.pop();
+    lock.unlock();
+    t(index);
+    lock.lock();
+  }
+}
+
 }  // namespace experimental
 }  // namespace dali