Keep a shared thread pool alive when the thread count is reset

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.

extension/threadpool/test/threadpool_test.cpp

@@ -10,6 +10,8 @@
 #include <executorch/extension/threadpool/threadpool.h>
 #include <executorch/runtime/platform/runtime.h>
 
+#include <chrono>
+#include <condition_variable>
 #include <mutex>
 #include <numeric>
 #include <random>
@@ -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<Rendezvous*>(context);
+        std::unique_lock<std::mutex> 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.

extension/threadpool/threadpool.cpp

@@ -72,9 +72,28 @@ bool ThreadPool::_unsafe_reset_threadpool(uint32_t new_thread_count) {
 
   std::lock_guard<std::mutex> 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<std::mutex> lock{threadpool->mutex_};
+  threadpool->threadpool_shared_ = true;
   return threadpool->threadpool_.get();
 }
 

extension/threadpool/threadpool.h

@@ -11,6 +11,8 @@
 #include <functional>
 #include <memory>
 #include <mutex>
+#include <utility>
+#include <vector>
 
 #include <pthreadpool.h>
 
@@ -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<pthreadpool, decltype(&pthreadpool_destroy)>;
+
+  // Guards the data members below.
   mutable std::mutex mutex_;
-  std::unique_ptr<pthreadpool, decltype(&pthreadpool_destroy)> 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<std::pair<size_t, PthreadpoolPtr>> 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();