From 55a4089e40d0469e0baf468ff0b4debb6a7e0b3c Mon Sep 17 00:00:00 2001
From: Dynamitos <dynamitos15@gmail.com>
Date: Sat, 25 Jan 2025 13:28:52 +0100
Subject: [PATCH] More efficient threadpool

---
 src/Minimal.h       | 21 ++++++++--
 src/ThreadPool.cpp  | 94 +++++++++++++++++++++++----------------------
 src/ThreadPool.h    |  4 +-
 src/scene/Scene.cpp | 10 +++--
 4 files changed, 74 insertions(+), 55 deletions(-)

diff --git a/src/Minimal.h b/src/Minimal.h
index 3e0ac92..65b3c69 100644
--- a/src/Minimal.h
+++ b/src/Minimal.h
@@ -1,9 +1,22 @@
 #pragma once
+#include <coroutine>
 #include <memory>
 
-#define DEFINE_REF(x)                                                                                                                      \
-    typedef ::std::unique_ptr<x> P##x;                                                                                                   \
+#define DEFINE_REF(x) typedef ::std::unique_ptr<x> P##x;
 
 #define DECLARE_REF(x)                                                                                                                     \
-    class x;                                                                                                                               \
-    typedef ::std::unique_ptr<x> P##x;                                                                                                   \
+  class 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;
+};
diff --git a/src/ThreadPool.cpp b/src/ThreadPool.cpp
index b8e9266..b61170b 100644
--- a/src/ThreadPool.cpp
+++ b/src/ThreadPool.cpp
@@ -3,64 +3,66 @@
 
 ThreadPool::ThreadPool(uint32_t numThreads)
 {
-    for (uint32_t i = 0; i < numThreads; ++i)
-    {
-        workers.push_back(std::thread(&ThreadPool::work, this));
-    }
+  for (uint32_t i = 0; i < numThreads; ++i)
+  {
+    workers.push_back(std::thread(&ThreadPool::work, this));
+  }
 }
 
 ThreadPool::~ThreadPool()
 {
-    running.store(false);
-    {
-        std::unique_lock l(queueLock);
-        queueCV.notify_all();
-    }
-    for (auto& worker : workers)
-    {
-        worker.join();
-    }
+  running.store(false);
+  {
+    std::unique_lock l(queueLock);
+    queueCV.notify_all();
+  }
+  for (auto& worker : workers)
+  {
+    worker.join();
+  }
 }
 
 void ThreadPool::runBatch(Batch&& batch)
 {
-    {
-        std::unique_lock l(queueLock);
-        taskQueue.push_back(batch);
-        queueCV.notify_one();
-    }
-    while (true)
-    {
-        std::unique_lock l(queueLock);
-        if (taskQueue.empty())
-            return;
-        completedCV.wait(l);
-    }
+  {
+    std::unique_lock l(queueLock);
+    numRemaining = batch.jobs.size();
+    taskQueue.push_back(batch);
+    queueCV.notify_all();
+  }
+  while (true)
+  {
+    std::unique_lock l(queueLock);
+    if (taskQueue.empty())
+      return;
+    completedCV.wait(l);
+  }
 }
 
 void ThreadPool::work()
 {
-    while (running)
+  while (running)
+  {
+    Task job;
     {
-        std::function<void()> job;
-        {
-            std::unique_lock l(queueLock);
-            if (taskQueue.empty() || taskQueue.front().jobs.empty())
-            {
-                queueCV.wait(l);
-                continue;
-            }
-            job = taskQueue.front().jobs.front();
-            taskQueue.front().jobs.pop_front();
-        }
-        job();
-        {
-            std::unique_lock l(queueLock);
-            if (taskQueue.front().jobs.empty())
-            {
-                taskQueue.pop_front();
-                completedCV.notify_one();
-            }
-        }
+      std::unique_lock l(queueLock);
+      if (taskQueue.empty() || taskQueue.front().jobs.empty())
+      {
+        queueCV.wait(l);
+        continue;
+      }
+      job = taskQueue.front().jobs.front();
+      taskQueue.front().jobs.pop_front();
     }
+    job.handle();
+    {
+      std::unique_lock l(queueLock);
+      numRemaining--;
+      if (numRemaining == 0)
+      {
+        taskQueue.pop_front();
+        completedCV.notify_one();
+      }
+    }
+  }
 }
diff --git a/src/ThreadPool.h b/src/ThreadPool.h
index dba5b67..b5832b3 100644
--- a/src/ThreadPool.h
+++ b/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;
 };
\ No newline at end of file
diff --git a/src/scene/Scene.cpp b/src/scene/Scene.cpp
index 87e7ed2..cd8e0e6 100644
--- a/src/scene/Scene.cpp
+++ b/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),
-                                                                 glm::vec4(0.0f, 0.0f, 1.0f, 0.0f), glm::vec4(0.0f, 0.0f, 0.0f, 1.0f)));
+  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), 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));