thread_pool.hpp

Go to the documentation of this file.

#pragma once
 
#define THREAD_POOL_VERSION "v2.0.0 (2021-08-14)"
 
#include <atomic>      // std::atomic
#include <chrono>      // std::chrono
#include <cstdint>     // std::int_fast64_t, std::uint_fast32_t
#include <functional>  // std::function
#include <future>      // std::future, std::promise
#include <iostream>    // std::cout, std::ostream
#include <memory>      // std::shared_ptr, std::unique_ptr
#include <mutex>       // std::mutex, std::scoped_lock
#include <queue>       // std::queue
#include <thread>      // std::this_thread, std::thread
#include <type_traits>  // std::common_type_t, std::decay_t, std::enable_if_t, std::is_void_v, std::invoke_result_t
#include <utility>  // std::move
 
// ============================================================================================= //
//                                    Begin class thread_pool                                    //
 
namespace kvikio::third_party {
 
class thread_pool {
  typedef std::uint_fast32_t ui32;
  typedef std::uint_fast64_t ui64;
 
 public:
  // ============================
  // Constructors and destructors
  // ============================
 
  thread_pool(const ui32& _thread_count = std::thread::hardware_concurrency())
    : thread_count(_thread_count ? _thread_count : std::thread::hardware_concurrency()),
      threads(new std::thread[_thread_count ? _thread_count : std::thread::hardware_concurrency()])
  {
    create_threads();
  }
 
  ~thread_pool()
  {
    wait_for_tasks();
    running = false;
    destroy_threads();
  }
 
  // =======================
  // Public member functions
  // =======================
 
  ui64 get_tasks_queued() const
  {
    const std::scoped_lock lock(queue_mutex);
    return tasks.size();
  }
 
  ui32 get_tasks_running() const { return tasks_total - (ui32)get_tasks_queued(); }
 
  ui32 get_tasks_total() const { return tasks_total; }
 
  ui32 get_thread_count() const { return thread_count; }
 
  template <typename T1, typename T2, typename F>
  void parallelize_loop(const T1& first_index,
                        const T2& index_after_last,
                        const F& loop,
                        ui32 num_blocks = 0)
  {
    typedef std::common_type_t<T1, T2> T;
    T the_first_index = (T)first_index;
    T last_index      = (T)index_after_last;
    if (the_first_index == last_index) return;
    if (last_index < the_first_index) {
      T temp          = last_index;
      last_index      = the_first_index;
      the_first_index = temp;
    }
    last_index--;
    if (num_blocks == 0) num_blocks = thread_count;
    ui64 total_size = (ui64)(last_index - the_first_index + 1);
    ui64 block_size = (ui64)(total_size / num_blocks);
    if (block_size == 0) {
      block_size = 1;
      num_blocks = (ui32)total_size > 1 ? (ui32)total_size : 1;
    }
    std::atomic<ui32> blocks_running = 0;
    for (ui32 t = 0; t < num_blocks; t++) {
      T start = ((T)(t * block_size) + the_first_index);
      T end =
        (t == num_blocks - 1) ? last_index + 1 : ((T)((t + 1) * block_size) + the_first_index);
      blocks_running++;
      push_task([start, end, &loop, &blocks_running] {
        loop(start, end);
        blocks_running--;
      });
    }
    while (blocks_running != 0) {
      sleep_or_yield();
    }
  }
 
  template <typename F>
  void push_task(const F& task)
  {
    tasks_total++;
    {
      const std::scoped_lock lock(queue_mutex);
      tasks.push(std::function<void()>(task));
    }
  }
 
  template <typename F, typename... A>
  void push_task(const F& task, const A&... args)
  {
    push_task([task, args...] { task(args...); });
  }
 
  void reset(const ui32& _thread_count = std::thread::hardware_concurrency())
  {
    bool was_paused = paused;
    paused          = true;
    wait_for_tasks();
    running = false;
    destroy_threads();
    thread_count = _thread_count ? _thread_count : std::thread::hardware_concurrency();
    threads.reset(new std::thread[thread_count]);
    paused  = was_paused;
    running = true;
    create_threads();
  }
 
  template <typename F,
            typename... A,
            typename = std::enable_if_t<
              std::is_void_v<std::invoke_result_t<std::decay_t<F>, std::decay_t<A>...>>>>
  std::future<bool> submit(const F& task, const A&... args)
  {
    std::shared_ptr<std::promise<bool>> task_promise(new std::promise<bool>);
    std::future<bool> future = task_promise->get_future();
    push_task([task, args..., task_promise] {
      try {
        task(args...);
        task_promise->set_value(true);
      } catch (...) {
        try {
          task_promise->set_exception(std::current_exception());
        } catch (...) {
        }
      }
    });
    return future;
  }
 
  template <typename F,
            typename... A,
            typename R = std::invoke_result_t<std::decay_t<F>, std::decay_t<A>...>,
            typename   = std::enable_if_t<!std::is_void_v<R>>>
  std::future<R> submit(const F& task, const A&... args)
  {
    std::shared_ptr<std::promise<R>> task_promise(new std::promise<R>);
    std::future<R> future = task_promise->get_future();
    push_task([task, args..., task_promise] {
      try {
        task_promise->set_value(task(args...));
      } catch (...) {
        try {
          task_promise->set_exception(std::current_exception());
        } catch (...) {
        }
      }
    });
    return future;
  }
 
  void wait_for_tasks()
  {
    while (true) {
      if (!paused) {
        if (tasks_total == 0) break;
      } else {
        if (get_tasks_running() == 0) break;
      }
      sleep_or_yield();
    }
  }
 
  // ===========
  // Public data
  // ===========
 
  std::atomic<bool> paused = false;
 
  ui32 sleep_duration = 1000;
 
 private:
  // ========================
  // Private member functions
  // ========================
 
  void create_threads()
  {
    for (ui32 i = 0; i < thread_count; i++) {
      threads[i] = std::thread(&thread_pool::worker, this);
    }
  }
 
  void destroy_threads()
  {
    for (ui32 i = 0; i < thread_count; i++) {
      threads[i].join();
    }
  }
 
  bool pop_task(std::function<void()>& task)
  {
    const std::scoped_lock lock(queue_mutex);
    if (tasks.empty())
      return false;
    else {
      task = std::move(tasks.front());
      tasks.pop();
      return true;
    }
  }
 
  void sleep_or_yield()
  {
    if (sleep_duration)
      std::this_thread::sleep_for(std::chrono::microseconds(sleep_duration));
    else
      std::this_thread::yield();
  }
 
  void worker()
  {
    while (running) {
      std::function<void()> task;
      if (!paused && pop_task(task)) {
        task();
        tasks_total--;
      } else {
        sleep_or_yield();
      }
    }
  }
 
  // ============
  // Private data
  // ============
 
  mutable std::mutex queue_mutex = {};
 
  std::atomic<bool> running = true;
 
  std::queue<std::function<void()>> tasks = {};
 
  ui32 thread_count;
 
  std::unique_ptr<std::thread[]> threads;
 
  std::atomic<ui32> tasks_total = 0;
};
}  // namespace kvikio::third_party
 
//                                     End class thread_pool                                     //
// ============================================================================================= //