More efficient threadpool

src/Minimal.h

@@ -1,9 +1,22 @@
 #pragma once
+#include <coroutine>
 #include <memory>
 
-#define DEFINE_REF(x)                                                                                                                      \
+#define DEFINE_REF(x) typedef ::std::unique_ptr<x> P##x;
-    typedef ::std::unique_ptr<x> P##x;                                                                                                   \
 
 #define DECLARE_REF(x)                                                                                                                     \
   class x;                                                                                                                                 \
-    typedef ::std::unique_ptr<x> P##x;                                                                                                   \
+  typedef ::std::unique_ptr<x> P##x;
+
+struct Task
+{
+  struct promise_type
+  {
+    Task get_return_object() { return {std::coroutine_handle<promise_type>::from_promise(*this)}; }
+    std::suspend_always initial_suspend() noexcept { return {}; }
+    std::suspend_never final_suspend() noexcept { return {}; }
+    void return_void() {}
+    void unhandled_exception() {}
+  };
+  std::coroutine_handle<promise_type> handle;
+};

src/ThreadPool.cpp

@@ -26,8 +26,9 @@ void ThreadPool::runBatch(Batch&& batch)
 {
   {
     std::unique_lock l(queueLock);
+    numRemaining = batch.jobs.size();
     taskQueue.push_back(batch);
-        queueCV.notify_one();
+    queueCV.notify_all();
   }
   while (true)
   {
@@ -42,7 +43,7 @@ void ThreadPool::work()
 {
   while (running)
   {
-        std::function<void()> job;
+    Task job;
     {
       std::unique_lock l(queueLock);
       if (taskQueue.empty() || taskQueue.front().jobs.empty())
@@ -53,10 +54,11 @@ void ThreadPool::work()
       job = taskQueue.front().jobs.front();
       taskQueue.front().jobs.pop_front();
     }
-        job();
+    job.handle();
     {
       std::unique_lock l(queueLock);
-            if (taskQueue.front().jobs.empty())
+      numRemaining--;
+      if (numRemaining == 0)
       {
         taskQueue.pop_front();
         completedCV.notify_one();

src/ThreadPool.h

@@ -3,10 +3,11 @@
 #include <functional>
 #include <list>
 #include <thread>
+#include "Minimal.h"
 
 struct Batch
 {
-    std::list<std::function<void()>> jobs;
+    std::list<Task> jobs;
 };
 
 class ThreadPool
@@ -21,6 +22,7 @@ private:
     std::mutex queueLock;
     std::condition_variable queueCV;
     std::condition_variable completedCV;
+    uint32_t numRemaining;
     std::list<Batch> taskQueue;
     std::vector<std::thread> workers;
 };

src/scene/Scene.cpp

@@ -5,8 +5,9 @@
 
 Scene::Scene()
 {
-  bvh.addModels(ModelLoader::loadModel("../res/models/cube.fbx"), glm::mat4(glm::vec4(1.0f, 0.0f, 0.0f, 0.0f), glm::vec4(0.0f, 1.0f, 0.0f, 0.0f),
+  bvh.addModels(ModelLoader::loadModel("../res/models/cube.fbx"),
-                                                                 glm::vec4(0.0f, 0.0f, 1.0f, 0.0f), glm::vec4(0.0f, 0.0f, 0.0f, 1.0f)));
+                glm::mat4(glm::vec4(1.0f, 0.0f, 0.0f, 0.0f), glm::vec4(0.0f, 1.0f, 0.0f, 0.0f), glm::vec4(0.0f, 0.0f, 1.0f, 0.0f),
+                          glm::vec4(0.0f, 0.0f, 0.0f, 1.0f)));
   bvh.generate();
 }
 
@@ -39,7 +40,7 @@ void Scene::render(Camera cam, RenderParameter params)
     for (int w = 0; w < params.width; ++w)
     {
       batch.jobs.push_back(
-          [&, w]()
+          [&](int w) -> Task
           {
             for (int h = 0; h < params.height; ++h)
             {
@@ -51,7 +52,8 @@ void Scene::render(Camera cam, RenderParameter params)
               accumulator[w + h * params.width] +=
                   glm::vec3(w / float(params.width * params.numSamples), h / float(params.height * params.numSamples), 0);
             }
-          });
+            co_return;
+          }(w));
     }
     auto start = std::chrono::high_resolution_clock::now();
     threadPool.runBatch(std::move(batch));