MLX
Loading...
Searching...
No Matches
threadpool.h
Go to the documentation of this file.
1// This code was modified from https://github.com/progschj/ThreadPool
2// The original License is copied below:
3//
4// Copyright (c) 2012 Jakob Progsch, Václav Zeman
5// This software is provided 'as-is', without any express or implied
6// warranty. In no event will the authors be held liable for any damages
7// arising from the use of this software.
8//
9// Permission is granted to anyone to use this software for any purpose,
10// including commercial applications, and to alter it and redistribute it
11// freely, subject to the following restrictions:
12//
13// 1. The origin of this software must not be misrepresented; you must not
14// claim that you wrote the original software. If you use this software
15// in a product, an acknowledgment in the product documentation would be
16// appreciated but is not required.
17//
18// 2. Altered source versions must be plainly marked as such, and must not be
19// misrepresented as being the original software.
20//
21// 3. This notice may not be removed or altered from any source
22// distribution.
23#pragma once
24
25#include <condition_variable>
26#include <functional>
27#include <future>
28#include <memory>
29#include <mutex>
30#include <queue>
31#include <stdexcept>
32#include <thread>
33#include <vector>
34
36 public:
37 ThreadPool(size_t);
38 template <class F, class... Args>
39 auto enqueue(F&& f, Args&&... args)
40 -> std::future<typename std::invoke_result_t<F, Args...>>;
42
43 private:
44 std::vector<std::thread> workers;
45 std::queue<std::function<void()>> tasks;
46 std::mutex queue_mutex;
47 std::condition_variable condition;
48 bool stop;
49};
50
51inline ThreadPool::ThreadPool(size_t threads) : stop(false) {
52 for (size_t i = 0; i < threads; ++i)
53 workers.emplace_back([this] {
54 for (;;) {
55 std::function<void()> task;
56
57 {
58 std::unique_lock<std::mutex> lock(this->queue_mutex);
59 this->condition.wait(
60 lock, [this] { return this->stop || !this->tasks.empty(); });
61 if (this->stop && this->tasks.empty())
62 return;
63 task = std::move(this->tasks.front());
64 this->tasks.pop();
65 }
66
67 task();
68 }
69 });
70}
71
72template <class F, class... Args>
73auto ThreadPool::enqueue(F&& f, Args&&... args)
74 -> std::future<typename std::invoke_result_t<F, Args...>> {
75 using return_type = typename std::invoke_result_t<F, Args...>;
76
77 auto task = std::make_shared<std::packaged_task<return_type()>>(
78 std::bind(std::forward<F>(f), std::forward<Args>(args)...));
79
80 std::future<return_type> res = task->get_future();
81 {
82 std::unique_lock<std::mutex> lock(queue_mutex);
83
84 if (stop) {
85 throw std::runtime_error(
86 "[ThreadPool::enqueue] Not allowed on stopped ThreadPool");
87 }
88
89 tasks.emplace([task]() { (*task)(); });
90 }
91 condition.notify_one();
92 return res;
93}
94
96 {
97 std::unique_lock<std::mutex> lock(queue_mutex);
98 stop = true;
99 }
100 condition.notify_all();
101 for (std::thread& worker : workers)
102 worker.join();
103}
Definition threadpool.h:35
auto enqueue(F &&f, Args &&... args) -> std::future< typename std::invoke_result_t< F, Args... > >
Definition threadpool.h:73
~ThreadPool()
Definition threadpool.h:95
ThreadPool(size_t)
Definition threadpool.h:51
float f
Definition bf16.h:16