Keep a loaded XNNPACK method's thread pool alive, and ship threadpool.h in the wheel (#24770)

### What is wrong today

**From Python**, `_unsafe_reset_threadpool(n)` is the only way to set
the CPU thread count. Calling it after an XNNPACK method is loaded can
make later runs crash or hang:

```python
method = Runtime.get().load_program(pte.buffer).load_method("forward")
_unsafe_reset_threadpool(16)
for _ in range(20):
    method.execute([x])  # segfault, abort, or hang
```

The reset frees the pthreadpool, which XNNPACK got from
`get_pthreadpool()` at load time and keeps using. On Android, a second
`Module` with another thread count resets the pool and can crash the
first one.

**From C++**, an app built against the pip package cannot set the count:
`get_threadpool()->_unsafe_reset_threadpool(n)` is declared in
`threadpool.h`, which the wheel does not ship.

### What this change does

- `ThreadPool` keeps its pthreadpool in a `std::shared_ptr`, returned by
the new `get_shared_pthreadpool()`. The XNNPACK executor holds it, so
the XNNPACK part of a method loaded before a reset keeps the old pool
until the method is destroyed. Later methods and kernels use the new
pool.
- Both getters read the pointer under a new mutex that a reset also
takes, so the XNNPACK delegate can load a method during a reset on
another thread.
- A `fork()` child does not destroy inherited pools, whose threads
stayed in the parent.
- The wheel ships `threadpool.h`. It declares the pthreadpool type
itself and includes `pthreadpool.h` only where it exists, so it works
from the wheel, and existing code that calls pthreadpool keeps
compiling.
- The docs say to set the count before loading, how to reload a method,
and that resetting during a `CpuBackend` load or run is unsafe.

### Tests

- New `XNNCompilerTest.RuntimeKeepsItsThreadpoolAcrossReset` runs an
XNNPACK graph across a reset, checks the output, and checks that
destroying the delegate frees the old pool. New `threadpool_test.cpp`
tests cover a held pool across a reset, readers during resets, and a
release in a forked child. `.ci/scripts/wheel/test_cpp_sdk.py` now also
links a program calling `_unsafe_reset_threadpool` against the installed
wheel. Each fails without its part of the fix (the readers test under
ThreadSanitizer).
- Ran on Linux x86_64: thread pool and XNNPACK C++ tests (plain,
AddressSanitizer with UndefinedBehaviorSanitizer, ThreadSanitizer),
native CPU engine tests, Python tests for pybindings, XNNPACK and LLM
attention custom ops, and the C++ SDK checks on a wheel of this branch.
On Linux aarch64: thread pool C++ tests in the same modes.
- Every failing test also fails on main, or needs network access the
test machine lacked. No macOS or Windows wheel was built.
- On main, the Python code above crashed or hung in most runs. It passes
now.

Fixes the crash in #24765.

Author: shoumikhin

.ci/scripts/wheel/test_cpp_sdk.py

@@ -1233,6 +1233,14 @@ def test_every_shipped_header_compiles(work_dir: Path) -> None:
     assert (
         headers
     ), f"no headers found under {include_root}, so this check would prove nothing"
+    # Everything ships under include/executorch. A header at the top of the include directory is on
+    # every consumer's include path, where it can shadow the consumer's own copy of that library.
+    outside = [
+        header.relative_to(include_root).as_posix()
+        for header in headers
+        if header.relative_to(include_root).parts[0] != "executorch"
+    ]
+    assert not outside, f"the wheel ships headers outside include/executorch: {outside}"
 
     # The same include directories the CMake package exports, since that is what a consumer gets.
     includes = [
@@ -1499,6 +1507,12 @@ def test_shipped_headers_have_implementations(work_dir: Path) -> None:
             "int main() { return get_thread_count() > 0 && "
             "parallel_for(0, 1, 1, [](int64_t, int64_t) {}) ? 0 : 1; }\n"
         ),
+        # The call the documentation gives for setting the thread count.
+        "extension/threadpool/threadpool.h": (
+            "#include <executorch/extension/threadpool/threadpool.h>\n"
+            "using namespace executorch::extension::threadpool;\n"
+            "int main() { return get_threadpool()->_unsafe_reset_threadpool(2) ? 0 : 1; }\n"
+        ),
     }
 
     includes = [

backends/xnnpack/runtime/XNNCompiler.cpp

@@ -2490,13 +2490,15 @@ ET_NODISCARD Error XNNCompiler::compileModel(
         : nullptr;
   }
 
+  executor->threadpool_ =
+      ::executorch::extension::threadpool::get_shared_pthreadpool();
   // NOLINTBEGIN(facebook-hte-NullableDereference) - weights cache is allowed to
   // be null
   status = xnn_create_runtime_v4(
       subgraph.get(),
       weights_cache_ptr,
       workspace,
-      ::executorch::extension::threadpool::get_pthreadpool(),
+      executor->threadpool_.get(),
       runtime_flags,
       &runtime_ptr);
   // NOLINTEND(facebook-hte-NullableDereference)

backends/xnnpack/runtime/XNNExecutor.h

@@ -35,6 +35,10 @@ class XNNExecutor {
   // we need to hold onto the buffers to keep them alive.
   std::vector<executorch::runtime::FreeableBuffer> unpacked_buffers_;
 
+  // Keeps runtime_'s threadpool alive across thread count resets. Declared
+  // before runtime_, so it outlives it.
+  std::shared_ptr<pthreadpool> threadpool_;
+
   std::unique_ptr<xnn_runtime, decltype(&xnn_delete_runtime)> runtime_{
       nullptr,
       &xnn_delete_runtime};

backends/xnnpack/test/runtime/test_xnnexecutor.cpp

@@ -9,6 +9,7 @@
 #include <executorch/backends/xnnpack/runtime/XNNExecutor.h>
 #include <executorch/backends/xnnpack/runtime/XNNPACKBackend.h>
 #include <executorch/backends/xnnpack/serialization/schema_generated.h>
+#include <executorch/extension/threadpool/threadpool.h>
 #include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
 #include <executorch/runtime/platform/runtime.h>
 #include <flatbuffers/flatbuffers.h>
@@ -205,6 +206,81 @@ TEST(XNNCompilerTest, Bf16BatchMatrixMultiplyOutputWithoutDims) {
   backend->destroy(handle.get());
 }
 
+TEST(XNNCompilerTest, RuntimeKeepsItsThreadpoolAcrossReset) {
+  runtime::runtime_init();
+  auto* backend = runtime::get_backend_class(xnnpack_backend_key);
+  ASSERT_NE(backend, nullptr);
+  auto* threadpool = extension::threadpool::get_threadpool();
+  struct RestoreThreadCount {
+    extension::threadpool::ThreadPool* threadpool;
+    uint32_t threads;
+    ~RestoreThreadCount() {
+      // @lint-ignore CLANGTIDY facebook-hte-Deprecated
+      threadpool->_unsafe_reset_threadpool(threads);
+    }
+  } restore{threadpool, static_cast<uint32_t>(threadpool->get_thread_count())};
+  // @lint-ignore CLANGTIDY facebook-hte-Deprecated
+  ASSERT_TRUE(threadpool->_unsafe_reset_threadpool(2));
+
+  flatbuffers::FlatBufferBuilder builder;
+  std::vector<flatbuffers::Offset<fb_xnnpack::XValue>> values;
+  for (uint32_t id = 0; id < 3; ++id) {
+    const auto tensor = fb_xnnpack::CreateXNNTensorValue(
+        builder,
+        fb_xnnpack::XNNDatatype::xnn_datatype_fp32,
+        1,
+        builder.CreateVector<uint32_t>({4}),
+        0,
+        id,
+        id == 2 ? XNN_VALUE_FLAG_EXTERNAL_OUTPUT
+                : XNN_VALUE_FLAG_EXTERNAL_INPUT,
+        id);
+    values.push_back(fb_xnnpack::CreateXValue(
+        builder, fb_xnnpack::XValueUnion::XNNTensorValue, tensor.Union()));
+  }
+  const auto add = fb_xnnpack::Create_XNNNode2x1(builder, 0, 1, 2);
+  const auto node = fb_xnnpack::CreateXNode(
+      builder, fb_xnnpack::XNodeUnion::XNNAdd, add.Union());
+  const auto graph = fb_xnnpack::CreateXNNGraph(
+      builder,
+      builder.CreateString("0"),
+      builder.CreateVector(&node, 1),
+      builder.CreateVector(values),
+      3,
+      builder.CreateVector<uint32_t>({0, 1}),
+      builder.CreateVector<uint32_t>({2}));
+  fb_xnnpack::FinishXNNGraphBuffer(builder, graph);
+
+  std::array<uint8_t, 4096> memory{};
+  runtime::MemoryAllocator allocator(memory.size(), memory.data());
+  runtime::BackendInitContext context(&allocator);
+  runtime::FreeableBuffer processed(
+      builder.GetBufferPointer(), builder.GetSize(), nullptr);
+  auto handle = backend->init(context, &processed, {});
+  ASSERT_TRUE(handle.ok());
+  const std::weak_ptr<pthreadpool> loaded_with =
+      extension::threadpool::get_shared_pthreadpool();
+
+  // @lint-ignore CLANGTIDY facebook-hte-Deprecated
+  ASSERT_TRUE(threadpool->_unsafe_reset_threadpool(3));
+  ASSERT_FALSE(loaded_with.expired());
+  TensorFactory<executorch::aten::ScalarType::Float> tf;
+  EValue a(tf.make({4}, {1, 2, 3, 4}));
+  EValue b(tf.make({4}, {10, 20, 30, 40}));
+  EValue out(tf.zeros({4}));
+  std::array<EValue*, 3> args{&a, &b, &out};
+  runtime::BackendExecutionContext execution_context;
+  EXPECT_EQ(
+      backend->execute(execution_context, handle.get(), {args.data(), 3}),
+      Error::Ok);
+  const float* sum = out.toTensor().const_data_ptr<float>();
+  EXPECT_EQ(
+      std::vector<float>(sum, sum + 4), (std::vector<float>{11, 22, 33, 44}));
+
+  backend->destroy(handle.get());
+  EXPECT_TRUE(loaded_with.expired());
+}
+
 } // namespace executorch::backends::xnnpack::delegate
 
 TEST(XNNExecutorTest, ArgumentWithTooManyDimensions) {

backends/xnnpack/test/targets.bzl

@@ -32,6 +32,7 @@ def define_common_targets():
             "//executorch/runtime/core/exec_aten/util:scalar_type_util",
             "//executorch/backends/xnnpack:xnnpack_backend",
             "//executorch/backends/xnnpack/serialization:xnnpack_flatbuffer_header",
+            "//executorch/extension/threadpool:threadpool",
         ],
     )
 

docs/source/using-executorch-faqs.md

@@ -43,13 +43,12 @@ Ensure the model is delegated. If not targeting a specific accelerator, use the
 
 Thread count can have a significant impact on CPU performance. The optimal thread count may depend on the model and application. By default, ExecuTorch uses one thread per performance core where it can tell core types apart, and one per logical core otherwise. On Linux other than Android, it also uses no more threads than the CPUs the process may run on, for example under `taskset` or in a container with a CPU set, and currently at most 63. It reads that limit once, from the thread that first uses the thread pool. CPU quotas, such as `docker --cpus`, are not detected. Consider setting the thread count to cores / 2, or just set to 4 on mobile CPUs.
 
-Thread count can be set with the following function. Ensure this is done prior to loading or running a model.
+Thread count can be set with the following function. Call it before loading a model. Methods loaded earlier keep their XNNPACK work on the old thread count, and their `CpuBackend` work fails to run. To load such a method again, call `Module::unload_method()` before `load_method()`, or create a new `Module`. Changing the thread count while a method that uses `CpuBackend` is loading or running is not safe.
 ```
 ::executorch::extension::threadpool::get_threadpool()->_unsafe_reset_threadpool(num_threads);
 ```
-This needs `extension/threadpool/threadpool.h`, which a source build provides. The pip
-package does not publish it, because the header holds a `pthreadpool` member by value and
-the wheel vendors no pthreadpool headers to compile it against.
+It is declared in `executorch/extension/threadpool/threadpool.h`. The pip package ships
+this header, so a C++ application built against the package can call it.
 
 For a deeper investigation into model performance, ExecuTorch supports operator-level performance profiling. See [Using the ExecuTorch Developer Tools to Profile a Model](tutorials/devtools-integration-tutorial) <!-- @lint-ignore --> for more information.
 

extension/pybindings/pybindings.pyi

@@ -462,7 +462,16 @@ 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.
+
+    Call it before loading methods. A method loaded earlier keeps its XNNPACK
+    parts on the old thread pool until the method is destroyed, while its other
+    kernels use the new count and its ``CpuBackend`` parts fail to run. A
+    ``Program`` from ``executorch.runtime`` and an ``ExecuTorchModule`` keep the
+    methods they loaded, so load the method from a new one to use the new count.
+    Calling it while a ``CpuBackend`` part is loading or running is not safe.
+    A count of 0, or the current count, changes nothing.
+
     .. warning::
 
         This API is experimental and subject to change without notice.

extension/threadpool/test/threadpool_test.cpp

@@ -10,16 +10,20 @@
 #include <executorch/extension/threadpool/threadpool.h>
 #include <executorch/runtime/platform/runtime.h>
 
+#include <atomic>
 #include <mutex>
 #include <numeric>
 #include <random>
+#include <thread>
 
 #include <executorch/extension/threadpool/threadpool_guard.h>
 
 #include <gtest/gtest.h>
 
 #if defined(__linux__) && !defined(__ANDROID__)
 #include <sched.h>
+#include <sys/wait.h>
+#include <unistd.h>
 
 #include <cerrno>
 #include <cstdio>
@@ -131,6 +135,16 @@ bool may_use_more_than_one_cpu() {
   std::exit(EXIT_SUCCESS);
 }
 #endif
+
+// Puts back the thread count a test changed, also when the test fails early.
+struct RestoreThreadCount {
+  ::executorch::extension::threadpool::ThreadPool* threadpool;
+  uint32_t threads;
+  ~RestoreThreadCount() {
+    // @lint-ignore CLANGTIDY facebook-hte-Deprecated
+    threadpool->_unsafe_reset_threadpool(threads);
+  }
+};
 } // namespace
 
 TEST(ThreadPoolTest, ParallelAdd) {
@@ -268,6 +282,64 @@ TEST(TestNoThreadPoolGuard, TestRunWithGuard) {
   ASSERT_EQ(inner, 6);
 }
 
+TEST(ThreadPoolTest, ResetKeepsSharedPthreadpoolUntilReleased) {
+  executorch::runtime::runtime_init();
+  auto* threadpool = ::executorch::extension::threadpool::get_threadpool();
+  const RestoreThreadCount restore{
+      threadpool, static_cast<uint32_t>(threadpool->get_thread_count())};
+  // @lint-ignore CLANGTIDY facebook-hte-Deprecated
+  ASSERT_TRUE(threadpool->_unsafe_reset_threadpool(2));
+  auto held = ::executorch::extension::threadpool::get_shared_pthreadpool();
+  ASSERT_EQ(held.get(), ::executorch::extension::threadpool::get_pthreadpool());
+  EXPECT_EQ(held.use_count(), 2);
+  const std::weak_ptr<pthreadpool> released = held;
+
+  // @lint-ignore CLANGTIDY facebook-hte-Deprecated
+  ASSERT_TRUE(threadpool->_unsafe_reset_threadpool(3));
+  EXPECT_NE(
+      ::executorch::extension::threadpool::get_shared_pthreadpool().get(),
+      held.get());
+  std::vector<uint8_t> ran(64);
+  pthreadpool_parallelize_1d(
+      held.get(),
+      [](void* context, size_t item) {
+        static_cast<uint8_t*>(context)[item] = 1;
+      },
+      ran.data(),
+      ran.size(),
+      0);
+  EXPECT_EQ(ran, std::vector<uint8_t>(64, 1));
+
+  held.reset();
+  EXPECT_TRUE(released.expired());
+}
+
+TEST(ThreadPoolTest, PthreadpoolGettersDuringReset) {
+  executorch::runtime::runtime_init();
+  auto* threadpool = ::executorch::extension::threadpool::get_threadpool();
+  const RestoreThreadCount restore{
+      threadpool, static_cast<uint32_t>(threadpool->get_thread_count())};
+  // Reads the pool on other threads, as delegate loads do, during resets.
+  std::atomic<bool> done{false};
+  std::thread shared_loader([&done] {
+    while (!done) {
+      ::executorch::extension::threadpool::get_shared_pthreadpool();
+    }
+  });
+  std::thread raw_loader([&done] {
+    while (!done) {
+      ::executorch::extension::threadpool::get_pthreadpool();
+    }
+  });
+  for (uint32_t i = 0; i < 100; ++i) {
+    // @lint-ignore CLANGTIDY facebook-hte-Deprecated
+    EXPECT_TRUE(threadpool->_unsafe_reset_threadpool(2 + i % 2));
+  }
+  done = true;
+  shared_loader.join();
+  raw_loader.join();
+}
+
 #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.
@@ -288,4 +360,31 @@ TEST(ThreadPoolDeathTest, DefaultSizeHonorsCpuAffinity) {
   EXPECT_EXIT(
       print_counts_on_one_cpu(), ExitedWithCode(EXIT_SUCCESS), expected);
 }
+
+TEST(ThreadPoolTest, ForkedChildReleasesInheritedSharedPthreadpool) {
+  executorch::runtime::runtime_init();
+  auto* threadpool = ::executorch::extension::threadpool::get_threadpool();
+  const RestoreThreadCount restore{
+      threadpool, static_cast<uint32_t>(threadpool->get_thread_count())};
+  // Destroying an inherited pool crashes once its thread stacks overflow
+  // glibc's 40 MiB stack cache; 32 threads do with stacks as small as 1 MiB.
+  // @lint-ignore CLANGTIDY facebook-hte-Deprecated
+  ASSERT_TRUE(threadpool->_unsafe_reset_threadpool(32));
+  auto held = ::executorch::extension::threadpool::get_shared_pthreadpool();
+  // @lint-ignore CLANGTIDY facebook-hte-Deprecated
+  ASSERT_TRUE(threadpool->_unsafe_reset_threadpool(2));
+  // The child holds the last reference to a threadpool whose threads stayed
+  // in the parent. Not a death test, whose "threadsafe" style would start the
+  // child over with a threadpool of its own.
+  const pid_t child = fork();
+  ASSERT_NE(child, -1);
+  if (child == 0) {
+    held.reset();
+    std::_Exit(EXIT_SUCCESS);
+  }
+  int status = 0;
+  ASSERT_EQ(waitpid(child, &status, 0), child);
+  EXPECT_TRUE(WIFEXITED(status) && WEXITSTATUS(status) == EXIT_SUCCESS)
+      << "wait status " << status;
+}
 #endif