From ce80ef8846f1181b040ba407c791ffaba916da70 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 6 Nov 2026 18:48:59 -0800 Subject: [PATCH] Keep a shared thread pool alive when the thread count is reset `_unsafe_reset_threadpool()` destroyed the current pthreadpool and created a new one. Delegates such as XNNPACK get that pthreadpool from `get_pthreadpool()` when a method is loaded and keep using it, so after a reset the loaded method ran on a freed pool and crashed or hung. From Python the reset is the only way to set the thread count, so calling it after `load_method` was enough to hit this. Now a pthreadpool that `get_pthreadpool()` has handed out stays alive with its size when a reset replaces it. XNNPACK methods loaded before the reset keep running on it, and methods loaded after it use the new pool. A reset back to a kept size reuses that pool, so at most one pool per size is kept, and `_unsafe_destroy_threadpool()` frees the kept pools too. A pool that was never handed out is destroyed as before, so a reset before loading costs nothing extra. Resizing a pool in place is not an option, because XNNPACK sizes per-thread buffers by the thread count when it plans a method. Fixes #24765 --- extension/pybindings/pybindings.pyi | 7 ++- extension/threadpool/test/threadpool_test.cpp | 57 +++++++++++++++++++ extension/threadpool/threadpool.cpp | 28 ++++++++- extension/threadpool/threadpool.h | 30 ++++++---- 4 files changed, 108 insertions(+), 14 deletions(-) diff --git a/extension/pybindings/pybindings.pyi b/extension/pybindings/pybindings.pyi index b78a657434e..0849adcfdff 100644 --- a/extension/pybindings/pybindings.pyi +++ b/extension/pybindings/pybindings.pyi @@ -462,7 +462,12 @@ def _reset_profile_results() -> None: @experimental("This API is experimental and subject to change without notice.") def _unsafe_reset_threadpool(num_threads: int) -> None: - """ + """Sets the number of threads in the CPU thread pool. + + Delegates such as XNNPACK keep using the thread pool that a method was + loaded with, so load the method after this call for the new count to apply + to it. + .. warning:: This API is experimental and subject to change without notice. diff --git a/extension/threadpool/test/threadpool_test.cpp b/extension/threadpool/test/threadpool_test.cpp index 00a380ba0dd..9ab96c55c67 100644 --- a/extension/threadpool/test/threadpool_test.cpp +++ b/extension/threadpool/test/threadpool_test.cpp @@ -10,6 +10,8 @@ #include #include +#include +#include #include #include #include @@ -268,6 +270,61 @@ TEST(TestNoThreadPoolGuard, TestRunWithGuard) { ASSERT_EQ(inner, 6); } +TEST(ThreadPoolTest, ResetKeepsHandedOutPthreadpoolUsable) { + executorch::runtime::runtime_init(); + + auto* threadpool = ::executorch::extension::threadpool::get_threadpool(); + const auto thread_count = threadpool->get_thread_count(); + // @lint-ignore CLANGTIDY facebook-hte-Deprecated + ASSERT_TRUE(threadpool->_unsafe_reset_threadpool(2)); + // Delegates such as XNNPACK hold this pointer while a loaded method lives. + const auto held = ::executorch::extension::threadpool::get_pthreadpool(); + // @lint-ignore CLANGTIDY facebook-hte-Deprecated + ASSERT_TRUE(threadpool->_unsafe_reset_threadpool(3)); + EXPECT_EQ(threadpool->get_thread_count(), 3); + // Going back to 2 threads reuses the held threadpool, which must then be + // kept again by the next reset. + // @lint-ignore CLANGTIDY facebook-hte-Deprecated + ASSERT_TRUE(threadpool->_unsafe_reset_threadpool(2)); + // @lint-ignore CLANGTIDY facebook-hte-Deprecated + ASSERT_TRUE(threadpool->_unsafe_reset_threadpool(3)); + EXPECT_EQ(pthreadpool_get_threads_count(held), 2); + + // Each item waits for the other, so both finish their waits only if the + // held threadpool still has its worker thread. + struct Rendezvous { + std::mutex mutex; + std::condition_variable arrived; + size_t count = 0; + bool met = true; + } rendezvous; + pthreadpool_parallelize_1d( + held, + [](void* context, size_t) { + auto& rendezvous = *static_cast(context); + std::unique_lock lock(rendezvous.mutex); + ++rendezvous.count; + rendezvous.arrived.notify_all(); + if (!rendezvous.arrived.wait_for( + lock, std::chrono::seconds(10), [&rendezvous] { + return rendezvous.count == 2; + })) { + rendezvous.met = false; + } + }, + &rendezvous, + 2, + 0); + EXPECT_TRUE(rendezvous.met); + + // @lint-ignore CLANGTIDY facebook-hte-Deprecated + ASSERT_TRUE(threadpool->_unsafe_reset_threadpool(2)); + EXPECT_EQ(::executorch::extension::threadpool::get_pthreadpool(), held); + + // @lint-ignore CLANGTIDY facebook-hte-Deprecated + threadpool->_unsafe_reset_threadpool(thread_count); +} + #if defined(__linux__) && !defined(__ANDROID__) // gtest runs death tests before the other tests, so the forked child creates // the pool itself, and the CPU limit it sets does not reach the other tests. diff --git a/extension/threadpool/threadpool.cpp b/extension/threadpool/threadpool.cpp index 036a6c10e6f..8ce73cb4354 100644 --- a/extension/threadpool/threadpool.cpp +++ b/extension/threadpool/threadpool.cpp @@ -72,9 +72,28 @@ bool ThreadPool::_unsafe_reset_threadpool(uint32_t new_thread_count) { std::lock_guard lock{mutex_}; - threadpool_.reset(pthreadpool_create(new_thread_count)); - thread_count_ = - threadpool_ ? pthreadpool_get_threads_count(threadpool_.get()) : 0; + // Holders of a shared threadpool, such as XNNPACK, size their buffers by its + // thread count and are not told about a new one, so it is kept as is. + if (threadpool_shared_ && threadpool_) { + retired_threadpools_.emplace_back(thread_count_, std::move(threadpool_)); + } + // Reusing a kept threadpool of the requested size bounds how many are kept. + const auto retired = std::find_if( + retired_threadpools_.begin(), + retired_threadpools_.end(), + [new_thread_count](const auto& entry) { + return entry.first == new_thread_count; + }); + threadpool_shared_ = retired != retired_threadpools_.end(); + if (threadpool_shared_) { + threadpool_ = std::move(retired->second); + thread_count_ = retired->first; + retired_threadpools_.erase(retired); + } else { + threadpool_.reset(pthreadpool_create(new_thread_count)); + thread_count_ = + threadpool_ ? pthreadpool_get_threads_count(threadpool_.get()) : 0; + } return true; } @@ -83,6 +102,7 @@ void ThreadPool::_unsafe_destroy_threadpool() { ET_LOG(Info, "Destroying threadpool."); threadpool_.reset(); thread_count_ = 0; + retired_threadpools_.clear(); } void ThreadPool::run( @@ -199,6 +219,8 @@ pthreadpool_t get_pthreadpool() { } ThreadPool* const threadpool = get_threadpool(); ET_CHECK_MSG(threadpool, "Failed to acquire an instance of ThreadPool!"); + std::lock_guard lock{threadpool->mutex_}; + threadpool->threadpool_shared_ = true; return threadpool->threadpool_.get(); } diff --git a/extension/threadpool/threadpool.h b/extension/threadpool/threadpool.h index 6a22cd286fb..ccf66fcbf59 100644 --- a/extension/threadpool/threadpool.h +++ b/extension/threadpool/threadpool.h @@ -11,6 +11,8 @@ #include #include #include +#include +#include #include @@ -66,10 +68,13 @@ class ThreadPool final { size_t get_thread_count() const; /** - * INTERNAL: Resets the threadpool by creating a new threadpool with requested - * # of threads. This is not a thread safe call. When calling this method, - * threads of the threadpool might be doing some work. Some other code may - * also be holding on to the threadpool pointer, that is no longer valid. This + * INTERNAL: Resets the threadpool to the requested # of threads. This is not + * a thread safe call. A threadpool that get_pthreadpool() returned is not + * destroyed but kept with its size, because delegates such as XNNPACK keep + * using it for the methods they loaded. So the new size applies to such + * methods only when they are loaded after this call. A later reset to a kept + * size reuses that threadpool. Kept threadpools live until + * _unsafe_destroy_threadpool() or the end of the process. This * is a private API, which will later be replaced by something that allows * creating of threadpool with requested size and use such a threadpool with * backend delegates, custom ops or optimized lib. @@ -104,14 +109,18 @@ class ThreadPool final { friend pthreadpool_t get_pthreadpool(); private: - // This mutex is used inside get_thread_count API but it is not really needed - // since data members of ThreadPool objects are not really mutable. - // TODO(kimishpatel): Figure out if we will allow set_num_threads API, in - // which case this mutex will be useful. Otherwise remove it. + using PthreadpoolPtr = + std::unique_ptr; + + // Guards the data members below. mutable std::mutex mutex_; - std::unique_ptr threadpool_; + PthreadpoolPtr threadpool_; // The number of threads in the threadpool. size_t thread_count_; + // Whether get_pthreadpool() has returned threadpool_. + bool threadpool_shared_ = false; + // Shared threadpools that a reset replaced, with their thread counts. + std::vector> retired_threadpools_; }; /** @@ -123,7 +132,8 @@ ThreadPool* get_threadpool(); * Returns the underlying pthreadpool instance used by the implementation of * ThreadPool returned by `get_threadpool()`. Only for use in external libraries * so as to unify threading across internal (i.e. ATen, etc.) and external (e.g. - * NNPACK, QNNPACK, XNNPACK) use cases. + * NNPACK, QNNPACK, XNNPACK) use cases. The returned instance stays valid after + * ThreadPool::_unsafe_reset_threadpool(). */ pthreadpool_t get_pthreadpool();