Limit the default thread pool to the allowed CPUs on Linux (not Android)

docs/source/using-executorch-faqs.md

@@ -41,7 +41,7 @@ If building the runtime from source, ensure that the build is done in release mo
 
 Ensure the model is delegated. If not targeting a specific accelerator, use the XNNPACK delegate for CPU performance. Undelegated operators will typically fall back to the ExecuTorch portable library, which is designed as a fallback, and is not intended for performance sensitive operators. To target XNNPACK, pass an `XnnpackPartitioner` to `to_edge_transform_and_lower`. See [Building and Running ExecuTorch with XNNPACK Backend](tutorial-xnnpack-delegate-lowering.md) for more information.
 
-Thread count can have a significant impact on CPU performance. The optimal thread count may depend on the model and application. By default, ExecuTorch will currently use as many threads as there are cores. Consider setting the thread count to cores / 2, or just set to 4 on mobile CPUs.
+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.
 ```

examples/models/gemma4/e2e_runner.cpp

@@ -11,7 +11,6 @@
 #include <gflags/gflags.h>
 
 #include <iostream>
-#include <thread>
 
 #include <executorch/examples/models/gemma4/image_utils.h>
 #include <executorch/examples/models/gemma4/runner/gemma4_runner.h>
@@ -57,7 +56,12 @@ int32_t main(int32_t argc, char** argv) {
       ? ::executorch::extension::cpuinfo::get_num_performant_cores()
       : static_cast<uint32_t>(FLAGS_cpu_threads);
   if (num_threads == 0) {
-    num_threads = std::min(std::thread::hardware_concurrency(), 8u);
+    // On Linux other than Android, the default pool size follows the CPUs
+    // this process may run on.
+    num_threads = static_cast<uint32_t>(std::min<size_t>(
+        ::executorch::extension::threadpool::get_threadpool()
+            ->get_thread_count(),
+        8));
   }
   ET_LOG(Info, "Setting threadpool to %d threads", (int)num_threads);
   ::executorch::extension::threadpool::get_threadpool()

extension/threadpool/cpuinfo_utils.cpp

@@ -20,6 +20,15 @@
 #include <sys/sysctl.h>
 #endif
 
+#if defined(__linux__)
+#include <sched.h>
+
+#include <algorithm>
+#include <cerrno>
+#include <cstdlib>
+#include <memory>
+#endif
+
 namespace executorch::extension::cpuinfo {
 
 // Ignore revisions (last digit (4 LSBs))
@@ -60,6 +69,7 @@ bool is_non_performant_core(const struct cpuinfo_uarch_info* uarch_info) {
   return false;
 }
 
+#if defined(__linux__) && (CPUINFO_ARCH_ARM || CPUINFO_ARCH_ARM64)
 std::vector<uint32_t>* get_static_cpu_midr_vector() {
   static std::vector<uint32_t> cpu_midrs;
   return &cpu_midrs;
@@ -103,7 +113,14 @@ bool populate_available_cpu_mids() {
     if (midr_file.is_open()) {
       std::string x;
       std::getline(midr_file, x);
-      tmp = std::stoi(x, nullptr, 16);
+      char* end = nullptr;
+      const unsigned long midr = std::strtoul(x.c_str(), &end, 16);
+      if (end == x.c_str()) {
+        ET_LOG(Info, "Failed to parse midr file %s", midr_file_path.c_str());
+        cpu_midrs->clear();
+        return false;
+      }
+      tmp = static_cast<uint32_t>(midr);
       (*cpu_midrs)[i] = tmp;
     } else {
       ET_LOG(Info, "Failed to open midr file %s", midr_file_path.c_str());
@@ -140,17 +157,18 @@ uint32_t _get_num_performant_cores() {
   }
   return num_possible_cores;
 }
+#endif
 
-} // namespace
-
-uint32_t get_num_performant_cores() {
+uint32_t count_performant_cores() {
   ET_CHECK_MSG(cpuinfo_initialize(), "cpuinfo cannot be initialized.");
+#if defined(__linux__) && (CPUINFO_ARCH_ARM || CPUINFO_ARCH_ARM64)
   // First try and see if we have number of cores profiled for this specific
   // device
   uint32_t model_specific_num_cores = _get_model_specific_num_cores();
   if (model_specific_num_cores > 0) {
     return model_specific_num_cores;
   }
+#endif
 
   // Else looks at either the # of litte cores if found
   // Or parse the midr in "Something seems wrong" section.
@@ -193,8 +211,48 @@ uint32_t get_num_performant_cores() {
       }
     }
 #endif
+#if defined(__linux__) && (CPUINFO_ARCH_ARM || CPUINFO_ARCH_ARM64)
     return _get_num_performant_cores();
+#else
+    // The MIDR files exist only on ARM Linux. Elsewhere the count is unknown.
+    return 0;
+#endif
+  }
+}
+
+} // namespace
+
+uint32_t get_num_allowed_cpus() {
+#if defined(__linux__)
+  // The kernel rejects a mask smaller than its possible CPU count, which can
+  // exceed the CPU_SETSIZE CPUs a cpu_set_t holds, so grow it until it fits.
+  for (int num_cpus = CPU_SETSIZE; num_cpus <= (1 << 20); num_cpus *= 2) {
+    std::unique_ptr<cpu_set_t, void (*)(cpu_set_t*)> cpus(
+        CPU_ALLOC(num_cpus), [](cpu_set_t* set) { CPU_FREE(set); });
+    if (!cpus) {
+      break;
+    }
+    const size_t size = CPU_ALLOC_SIZE(num_cpus);
+    if (sched_getaffinity(0, size, cpus.get()) == 0) {
+      return static_cast<uint32_t>(CPU_COUNT_S(size, cpus.get()));
+    }
+    if (errno != EINVAL) {
+      break;
+    }
   }
+#endif
+  return 0;
+}
+
+uint32_t get_num_performant_cores() {
+  const uint32_t count = count_performant_cores();
+#if defined(__linux__) && !defined(__ANDROID__)
+  const uint32_t allowed = get_num_allowed_cpus();
+  if (count > 0 && allowed > 0) {
+    return std::min(count, allowed);
+  }
+#endif
+  return count;
 }
 
 } // namespace executorch::extension::cpuinfo

extension/threadpool/cpuinfo_utils.h

@@ -12,8 +12,22 @@
 
 namespace executorch::extension::cpuinfo {
 
+/**
+ * Returns the number of performance cores, or 0 if it is unknown, for example
+ * on x86 or when every core is an efficiency core. On Linux other than
+ * Android, the count is capped at get_num_allowed_cpus(), without checking
+ * which cores the affinity mask allows.
+ */
 uint32_t get_num_performant_cores();
 
+/**
+ * Returns the number of CPUs in the calling thread's affinity mask on Linux,
+ * or 0 if it is not available. taskset and cgroup cpusets shrink the mask,
+ * CPU quotas do not. On Android the mask changes as an app moves between the
+ * foreground and the background, so it is a poor size for a long-lived pool.
+ */
+uint32_t get_num_allowed_cpus();
+
 } // namespace executorch::extension::cpuinfo
 
 namespace torch::executorch::cpuinfo { // DEPRECATED

extension/threadpool/test/targets.bzl

@@ -15,6 +15,7 @@ def define_common_targets():
         name = "threadpool_test",
         srcs = _THREADPOOL_TESTS,
         deps = [
+            "//executorch/extension/threadpool:cpuinfo_utils",
             "//executorch/extension/threadpool:threadpool",
         ],
     )

extension/threadpool/test/threadpool_test.cpp

@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <executorch/extension/threadpool/cpuinfo_utils.h>
 #include <executorch/extension/threadpool/threadpool.h>
 #include <executorch/runtime/platform/runtime.h>
 
@@ -17,6 +18,14 @@
 
 #include <gtest/gtest.h>
 
+#if defined(__linux__) && !defined(__ANDROID__)
+#include <sched.h>
+
+#include <cerrno>
+#include <cstdio>
+#include <cstdlib>
+#endif
+
 using namespace ::testing;
 
 namespace {
@@ -69,6 +78,59 @@ void run_lambda_with_size(
   auto threadpool = ::executorch::extension::threadpool::get_threadpool();
   threadpool->run(f, range);
 }
+
+// The pool has no more threads than the CPUs this process may run on.
+bool may_use_more_than_one_cpu() {
+#if defined(__linux__) && !defined(__ANDROID__)
+  // Grows the mask as get_num_allowed_cpus() does, for kernels with more
+  // possible CPUs than a cpu_set_t holds.
+  for (int num_cpus = CPU_SETSIZE; num_cpus <= (1 << 20); num_cpus *= 2) {
+    std::unique_ptr<cpu_set_t, void (*)(cpu_set_t*)> cpus(
+        CPU_ALLOC(num_cpus), [](cpu_set_t* set) { CPU_FREE(set); });
+    if (!cpus) {
+      break;
+    }
+    const size_t size = CPU_ALLOC_SIZE(num_cpus);
+    if (sched_getaffinity(0, size, cpus.get()) == 0) {
+      return CPU_COUNT_S(size, cpus.get()) > 1;
+    }
+    if (errno != EINVAL) {
+      break;
+    }
+  }
+#endif
+  return true;
+}
+
+#if defined(__linux__) && !defined(__ANDROID__)
+// Limits this thread to the CPU it runs on, then prints the default pool size
+// and the performance core count for the test to match. Failing to set the
+// limit aborts instead.
+[[noreturn]] void print_counts_on_one_cpu() {
+  const int cpu = sched_getcpu();
+  if (cpu < 0) {
+    std::abort();
+  }
+  std::unique_ptr<cpu_set_t, void (*)(cpu_set_t*)> cpus(
+      CPU_ALLOC(cpu + 1), [](cpu_set_t* set) { CPU_FREE(set); });
+  if (!cpus) {
+    std::abort();
+  }
+  const size_t size = CPU_ALLOC_SIZE(cpu + 1);
+  CPU_ZERO_S(size, cpus.get());
+  CPU_SET_S(cpu, size, cpus.get());
+  if (sched_setaffinity(0, size, cpus.get()) != 0) {
+    std::abort();
+  }
+  std::fprintf(
+      stderr,
+      "pool %zu perf %u end\n",
+      ::executorch::extension::threadpool::get_threadpool()->get_thread_count(),
+      static_cast<unsigned>(
+          ::executorch::extension::cpuinfo::get_num_performant_cores()));
+  std::exit(EXIT_SUCCESS);
+}
+#endif
 } // namespace
 
 TEST(ThreadPoolTest, ParallelAdd) {
@@ -88,7 +150,9 @@ TEST(ThreadPoolTest, ParallelAdd) {
   };
 
   auto threadpool = ::executorch::extension::threadpool::get_threadpool();
-  EXPECT_GT(threadpool->get_thread_count(), 1);
+  if (may_use_more_than_one_cpu()) {
+    EXPECT_GT(threadpool->get_thread_count(), 1);
+  }
 
   generate_add_test_inputs(a, b, c_ref, c, vector_size);
   run_lambda_with_size(add_lambda, vector_size, grain_size);
@@ -133,7 +197,9 @@ TEST(ThreadPoolTest, ParallelReduce) {
   };
 
   auto threadpool = ::executorch::extension::threadpool::get_threadpool();
-  EXPECT_GT(threadpool->get_thread_count(), 1);
+  if (may_use_more_than_one_cpu()) {
+    EXPECT_GT(threadpool->get_thread_count(), 1);
+  }
 
   generate_reduce_test_inputs(a, c_ref, vector_size);
   run_lambda_with_size(reduce_lambda, vector_size, grain_size);
@@ -201,3 +267,25 @@ TEST(TestNoThreadPoolGuard, TestRunWithGuard) {
   }
   ASSERT_EQ(inner, 6);
 }
+
+#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.
+TEST(ThreadPoolDeathTest, DefaultSizeHonorsCpuAffinity) {
+  // A later run, as with --gtest_repeat, forks after the other tests created
+  // the pool, and the child keeps that pool's size.
+  static int runs = 0;
+  if (runs++ > 0) {
+    GTEST_SKIP() << "The pool may already exist in this process";
+  }
+  executorch::runtime::runtime_init();
+
+  // A known performance core count drops to 1, and an unknown one stays 0.
+  const char* expected =
+      ::executorch::extension::cpuinfo::get_num_performant_cores() > 0
+      ? "pool 1 perf 1 end"
+      : "pool 1 perf 0 end";
+  EXPECT_EXIT(
+      print_counts_on_one_cpu(), ExitedWithCode(EXIT_SUCCESS), expected);
+}
+#endif

extension/threadpool/threadpool.cpp

@@ -140,6 +140,17 @@ ThreadPool* get_threadpool() {
     // Set threads equal to the number of performance cores.
     auto result = ::executorch::extension::cpuinfo::get_num_performant_cores();
 #endif
+#if defined(__linux__) && !defined(__ANDROID__)
+    // Never start more threads than the CPUs in this thread's affinity mask,
+    // which the pool's threads inherit, and use that number when the count is
+    // unknown (0), which pthreadpool would fill in with every online CPU, past
+    // the cap below.
+    const auto allowed =
+        ::executorch::extension::cpuinfo::get_num_allowed_cpus();
+    if (allowed > 0) {
+      result = result == 0 ? allowed : std::min(result, allowed);
+    }
+#endif
 
     /*
      * For llvm-tsan, holding limit for the number of locks for a single thread

extension/threadpool/threadpool.h

@@ -30,6 +30,15 @@
  * - EXECUTORCH_THREADPOOL_USE_ALL_LOGICAL_CORES (flag) - Sizes the threadpool
  * equal to the number of logical cores on system. This is the historical
  * behavior.
+ *
+ * On Linux other than Android, the threadpool starts with no more threads than
+ * the CPUs in the affinity mask of the thread that first calls
+ * get_threadpool(). taskset and cgroup cpusets shrink that mask, CPU quotas do
+ * not. The mask also sizes the threadpool when the performance core count is
+ * unknown, as on x86. It is read once: the threadpool does not follow later
+ * changes to the mask, and a child created by fork() keeps its parent's size.
+ * If the mask can grow later, set the thread count you need with
+ * _unsafe_reset_threadpool() before loading any model.
  */
 
 namespace executorch::extension::threadpool {