Skip to content

Commit 488d3eb

Browse files
committed
[Executorch][threadpool] Add "private" method to reset num threads in threadpool
Pull Request resolved: #2351 This diff introduces new API to reset threadpool. This will be short-lived features. It is needed atm because we dont have a) a way to use only a subset of threads to distribute work in the threadpool and b) API to create threadpool of custom size and pass such a threadpool to backends like delegtes, cusotm ops, optimized cpu op lib etc. ghstack-source-id: 218210435 @exported-using-ghexport Differential Revision: [D54766070](https://our.internmc.facebook.com/intern/diff/D54766070/)
1 parent 289d709 commit 488d3eb

File tree

2 files changed

+23
-0
lines changed

2 files changed

+23
-0
lines changed

backends/xnnpack/threadpool/threadpool.cpp

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,18 @@ size_t ThreadPool::get_thread_count() const {
4444
return pthreadpool_get_threads_count(threadpool_.get());
4545
}
4646

47+
bool ThreadPool::_unsafe_reset_threadpool(uint32_t new_thread_count) {
48+
// No need to do anything if the count is same
49+
if (new_thread_count == get_thread_count()) {
50+
return true;
51+
}
52+
53+
std::lock_guard<std::mutex> lock{mutex_};
54+
55+
threadpool_.reset(pthreadpool_create(new_thread_count));
56+
return true;
57+
}
58+
4759
void ThreadPool::run(
4860
const std::function<void(size_t)>& fn,
4961
const size_t range) {

backends/xnnpack/threadpool/threadpool.h

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,17 @@ class ThreadPool final {
3434

3535
size_t get_thread_count() const;
3636

37+
/*
38+
* Resets the threadpool by creating a new threadpool with requested # of threads.
39+
* This is not a thread safe call. When calling this method, threads of the threadpool
40+
* might be doing some work.
41+
* Some other code may also be holding on to the threadpool pointer, that is no longer valid.
42+
* This is a private API, which will later be replaced by something that allows
43+
* creating of threadpool with requested size and use such a threadpool with backend
44+
* delegates, custom ops or optimized lib.
45+
*/
46+
bool _unsafe_reset_threadpool(uint32_t num_threads);
47+
3748
// Run, in parallel, function fn(task_id) over task_id in range [0, range).
3849
// This function is blocking. All input is processed by the time it returns.
3950
// NoThreadPoolGuard (see threadpool_guard.h) can used to disable

0 commit comments

Comments
 (0)