diff --git a/extension/threadpool/targets.bzl b/extension/threadpool/targets.bzl index 4a7185ce97..149f43247d 100644 --- a/extension/threadpool/targets.bzl +++ b/extension/threadpool/targets.bzl @@ -1,6 +1,16 @@ load("@fbsource//xplat/executorch/backends/xnnpack/third-party:third_party_libs.bzl", "third_party_dep") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +def get_threadpool_size(): + return native.read_config("executorch", "threadpool_size") + +def get_threadpool_flags(): + flags = [] + threadpool_size = get_threadpool_size() + if threadpool_size != None: + flags += ["-DET_THREADPOOL_SIZE=" + threadpool_size] + return flags + def define_common_targets(): """Defines targets that should be shared between fbcode and xplat. @@ -32,6 +42,7 @@ def define_common_targets(): exported_preprocessor_flags = [ "-DET_USE_THREADPOOL", ], + preprocessor_flags = get_threadpool_flags(), visibility = [ "//executorch/...", "//executorch/backends/...", diff --git a/extension/threadpool/threadpool.cpp b/extension/threadpool/threadpool.cpp index 4134bb8669..d7d68324a9 100644 --- a/extension/threadpool/threadpool.cpp +++ b/extension/threadpool/threadpool.cpp @@ -97,7 +97,13 @@ void ThreadPool::run( // Make this part threadsafe: TODO(kimishpatel) ThreadPool* get_threadpool() { ET_CHECK_MSG(cpuinfo_initialize(), "cpuinfo initialization failed"); + +#ifdef ET_THREADPOOL_SIZE + int num_threads = ET_THREADPOOL_SIZE; +#else int num_threads = cpuinfo_get_processors_count(); +#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