MCPcopy Create free account
hub / github.com/ml-explore/mlx-data / ThreadPool

Method ThreadPool

mlx/data/core/ThreadPool.cpp:11–51  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

9std::shared_ptr<ThreadController> ThreadPool::thread_controller = nullptr;
10
11ThreadPool::ThreadPool(size_t thread_count) {
12 if (!thread_controller) {
13 thread_controller = std::make_shared<ThreadController>();
14 }
15 for (size_t i = 0; i < thread_count; ++i) {
16 // start waiting threads. Workers listen for changes through
17 // the ThreadPool member condition_variable
18 threads_.emplace_back(std::thread([&]() {
19 std::unique_lock<std::mutex> queue_lock(task_mutex_, std::defer_lock);
20
21 while (true) {
22 queue_lock.lock();
23 task_cv_.wait(queue_lock, [&]() -> bool {
24 return !tasks_.empty() || stop_threads_;
25 });
26
27 // used by dtor to stop all threads without having to
28 // unceremoniously stop tasks. The tasks must all be
29 // finished, lest we break a promise and risk a `future`
30 // object throwing an exception.
31 if (stop_threads_ && tasks_.empty())
32 return;
33
34 // to initialize temp_task, we must move the unique_ptr
35 // from the queue to the local stack. Since a unique_ptr
36 // cannot be copied (obviously), it must be explicitly
37 // moved. This transfers ownership of the pointed-to
38 // object to *this, as specified in 20.11.1.2.1
39 // [unique.ptr.single.ctor].
40 auto temp_task = std::move(tasks_.front());
41
42 tasks_.pop();
43 queue_lock.unlock();
44
45 auto thread_state = thread_controller->limit();
46 (*temp_task)();
47 thread_controller->restore(thread_state);
48 }
49 }));
50 }
51}
52
53ThreadPool::~ThreadPool() {
54 stop_threads_ = true;

Callers

nothing calls this directly

Calls 2

limitMethod · 0.80
restoreMethod · 0.80

Tested by

no test coverage detected