Add EXECUTORCH_THREADPOOL_SIZE option and heuristic

Author: GregoryComer

extension/threadpool/targets.bzl

@@ -34,7 +34,10 @@ def define_common_targets():
         ],
         exported_preprocessor_flags = [
             "-DET_USE_THREADPOOL",
-        ],
+        ] + (
+            # Default to perf heuristic (0) in OSS, all cores (-1) otherwise.
+            ["-DEXECUTORCH_THREADPOOL_SIZE=0"] if runtime.is_oss else ["-DEXECUTORCH_THREADPOOL_SIZE=-1"]
+        ),
         visibility = [
             "//executorch/...",
             "//executorch/backends/...",

extension/threadpool/threadpool.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 <algorithm>
@@ -101,7 +102,21 @@ ThreadPool* get_threadpool() {
     return nullptr; // NOLINT(facebook-hte-NullableReturn)
   }
 
+  // Choose the number of threads according to the EXECUTORCH_THREADPOOL_SIZE
+  // value. See the description in threadpool.h.
+
+#if defined(EXECUTORCH_THREADPOOL_SIZE) && ((EXECUTORCH_THREADPOOL_SIZE) > 0)
+  // Use an explicit threadpool size.
+  int num_threads = EXECUTORCH_THREADPOOL_SIZE;
+#elif defined(EXECUTORCH_THREADPOOL_SIZE) && \
+    ((EXECUTORCH_THREADPOOL_SIZE) == -1)
+  // Use threads=cores.
   int num_threads = cpuinfo_get_processors_count();
+#else
+  // Use a performance heuristic.
+  int num_threads = ::executorch::extension::cpuinfo::get_num_performant_cores();
+#endif
+
   /*
    * For llvm-tsan, holding limit for the number of locks for a single thread
    * is 63 (because of comparison < 64 instead of <=). pthreadpool's worst

extension/threadpool/threadpool.h

@@ -14,6 +14,32 @@
 
 #include <pthreadpool.h>
 
+/*
+ * Threadpool Options:
+ *
+ * Threadpool size has a sizble affect on performance. The following
+ * options are exposed to control the threadpool size.
+ *
+ * EXECUTORCH_THREADPOOL_SIZE: int - Set the size of the threadpool,
+ * in number of threads.
+ *
+ *  Special Values:
+ *   -  0: Use a perforance heuristic to determine the default size,
+ *          based on the active hardware. This is the default mode
+ *          for CMake.
+ *   - -1: Set the thread count equal to the number of cores on the
+ *          active hardware.
+ *
+ *   Any other positive value will be interpreted as a thread count.
+ *   For example, setting EXECUTORCH_THREADPOOL_SIZE=4 will default
+ *   the threadpool to use 4 threads.
+ */
+
+#ifndef EXECUTORCH_THREADPOOL_SIZE
+// Default to using a runtime heuristic.
+#define EXECUTORCH_THREADPOOL_SIZE 0
+#endif
+
 namespace executorch::extension::threadpool {
 
 class ThreadPool final {