Use a separate threadpool library

CMakeLists.txt

@@ -637,6 +637,13 @@ if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/runner_util)
 endif()
 
+if(EXECUTORCH_BUILD_PTHREADPOOL
+   AND EXECUTORCH_BUILD_CPUINFO
+   AND CMAKE_CXX_STANDARD GREATER_EQUAL 14
+)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/threadpool)
+endif()
+
 if(EXECUTORCH_BUILD_PYBIND)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third-party/pybind11)
 

backends/xnnpack/test/CMakeLists.txt

@@ -34,6 +34,7 @@ et_cxx_test(
   SOURCES
   ${_test_srcs}
   EXTRA_LIBS
+  extension_threadpool
   xnnpack_backend
   XNNPACK
   pthreadpool

build/Test.cmake

@@ -25,9 +25,11 @@ find_package(executorch CONFIG REQUIRED)
 enable_testing()
 find_package(GTest CONFIG REQUIRED)
 
+target_link_options_shared_lib(cpuinfo)
 target_link_options_shared_lib(extension_data_loader)
 target_link_options_shared_lib(portable_kernels)
 target_link_options_shared_lib(portable_ops_lib)
+target_link_options_shared_lib(pthreadpool)
 target_link_options_shared_lib(quantized_ops_lib)
 
 # Add code coverage flags to supported compilers

build/executorch-config.cmake

@@ -46,6 +46,7 @@ set(lib_list
     extension_module
     extension_module_static
     extension_runner_util
+    extension_threadpool
     xnnpack_backend
     XNNPACK
     cpuinfo

examples/models/llama2/CMakeLists.txt

@@ -144,22 +144,15 @@ set(XNNPACK_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../backends/xnnpack)
 # Extra compile option and include dir for pthreadpool
 if(EXECUTORCH_BUILD_PTHREADPOOL)
   list(APPEND _common_compile_options -DET_USE_THREADPOOL)
-  list(APPEND link_libraries pthreadpool)
-  # These 2 source files are included in xnnpack_backend
-  if(NOT TARGET xnnpack_backend)
-    list(APPEND _srcs ${EXECUTORCH_ROOT}/extension/threadpool/threadpool.cpp
-         ${EXECUTORCH_ROOT}/extension/threadpool/threadpool_guard.cpp
-    )
-  endif()
+  list(APPEND link_libraries extension_threadpool pthreadpool)
   list(APPEND _common_include_directories
        ${XNNPACK_ROOT}/third-party/pthreadpool/include
   )
 endif()
 
 # Extra sources for cpuinfo
 if(EXECUTORCH_BUILD_CPUINFO)
-  list(APPEND link_libraries cpuinfo)
-  list(APPEND _srcs ${EXECUTORCH_ROOT}/extension/threadpool/cpuinfo_utils.cpp)
+  list(APPEND link_libraries extension_threadpool cpuinfo)
   list(APPEND _common_include_directories
        ${XNNPACK_ROOT}/third-party/cpuinfo/include
   )

examples/models/llava/CMakeLists.txt

@@ -127,22 +127,15 @@ set(XNNPACK_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../backends/xnnpack)
 # Extra compile option and include dir for pthreadpool
 if(EXECUTORCH_BUILD_PTHREADPOOL)
   list(APPEND _common_compile_options -DET_USE_THREADPOOL)
-  list(APPEND link_libraries pthreadpool)
-  # These 2 source files are included in xnnpack_backend
-  if(NOT TARGET xnnpack_backend)
-    list(APPEND _srcs ${EXECUTORCH_ROOT}/extension/threadpool/threadpool.cpp
-         ${EXECUTORCH_ROOT}/extension/threadpool/threadpool_guard.cpp
-    )
-  endif()
+  list(APPEND link_libraries extension_threadpool pthreadpool)
   list(APPEND _common_include_directories
        ${XNNPACK_ROOT}/third-party/pthreadpool/include
   )
 endif()
 
 # Extra sources for cpuinfo
 if(EXECUTORCH_BUILD_CPUINFO)
-  list(APPEND link_libraries cpuinfo)
-  list(APPEND _srcs ${EXECUTORCH_ROOT}/extension/threadpool/cpuinfo_utils.cpp)
+  list(APPEND link_libraries extension_threadpool cpuinfo)
   list(APPEND _common_include_directories
        ${XNNPACK_ROOT}/third-party/cpuinfo/include
   )

extension/android/CMakeLists.txt

@@ -32,7 +32,7 @@ find_package(executorch CONFIG REQUIRED)
 target_link_options_shared_lib(executorch)
 
 set(link_libraries)
-list(APPEND link_libraries extension_data_loader extension_module executorch
+list(APPEND link_libraries extension_data_loader extension_module extension_threadpool executorch
      fbjni
 )
 
@@ -93,13 +93,7 @@ if(EXECUTORCH_BUILD_LLAMA_JNI)
 
   target_link_options_shared_lib(quantized_ops_lib)
 
-  if(TARGET pthreadpool)
-    set(LLAMA_JNI_SRCS jni/jni_layer_llama.cpp
-                       ../../extension/threadpool/cpuinfo_utils.cpp
-    )
-  else()
-    set(LLAMA_JNI_SRCS jni/jni_layer_llama.cpp)
-  endif()
+  set(LLAMA_JNI_SRCS jni/jni_layer_llama.cpp)
   add_library(executorch_llama_jni SHARED ${LLAMA_JNI_SRCS})
   if(TARGET pthreadpool)
     target_compile_definitions(executorch_llama_jni PRIVATE ET_USE_THREADPOOL=1)

extension/llm/custom_ops/CMakeLists.txt

@@ -47,17 +47,10 @@ list(APPEND custom_ops_libs eigen_blas)
 
 list(TRANSFORM _custom_ops__srcs PREPEND "${EXECUTORCH_ROOT}/")
 
-# TODO: Consider moving xnnpack/threadpool in a separate lib since it's now used
-# by custom ops too.
 if(NOT EXECUTORCH_BUILD_XNNPACK)
-  list(
-    APPEND
-    _custom_ops__srcs
-    "${CMAKE_CURRENT_SOURCE_DIR}/../../../extension/threadpool/threadpool.cpp"
-    "${CMAKE_CURRENT_SOURCE_DIR}/../../../extension/threadpool/threadpool_guard.cpp"
-  )
+  list(APPEND custom_ops_libs extension_threadpool)
 else()
-  list(APPEND custom_ops_libs xnnpack_backend)
+  list(APPEND custom_ops_libs extension_threadpool xnnpack_backend)
 endif()
 
 add_library(custom_ops ${_custom_ops__srcs})

extension/parallel/test/CMakeLists.txt

@@ -23,12 +23,10 @@ include(${EXECUTORCH_ROOT}/build/Test.cmake)
 
 set(_test_srcs
     thread_parallel_test.cpp ../thread_parallel.cpp
-    ${EXECUTORCH_ROOT}/extension/threadpool/threadpool.cpp
-    ${EXECUTORCH_ROOT}/extension/threadpool/threadpool_guard.cpp
 )
 
 et_cxx_test(
-  extension_parallel_test SOURCES ${_test_srcs} EXTRA_LIBS pthreadpool cpuinfo
+  extension_parallel_test SOURCES ${_test_srcs} EXTRA_LIBS pthreadpool cpuinfo extension_threadpool
 )
 target_include_directories(
   extension_parallel_test

extension/threadpool/CMakeLists.txt

@@ -0,0 +1,40 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Please this file formatted by running:
+# ~~~
+# cmake-format -i CMakeLists.txt
+# ~~~
+
+cmake_minimum_required(VERSION 3.19)
+
+# Source root directory for executorch.
+if(NOT EXECUTORCH_ROOT)
+  set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../..)
+endif()
+
+if(NOT CMAKE_CXX_STANDARD)
+  set(CMAKE_CXX_STANDARD 17)
+endif()
+
+add_library(extension_threadpool threadpool.cpp threadpool_guard.cpp cpuinfo_utils.cpp)
+target_link_libraries(extension_threadpool PUBLIC executorch cpuinfo pthreadpool)
+target_include_directories(extension_threadpool PUBLIC ${EXECUTORCH_ROOT}/..)
+target_include_directories(
+  extension_threadpool
+  PUBLIC
+    ${EXECUTORCH_ROOT}/backends/xnnpack/third-party/cpuinfo/include
+    ${EXECUTORCH_ROOT}/backends/xnnpack/third-party/pthreadpool/include
+)
+target_compile_options(extension_threadpool PUBLIC ${_common_compile_options})
+
+# Install libraries
+install(
+  TARGETS extension_threadpool
+  DESTINATION lib
+  INCLUDES
+  DESTINATION ${_common_include_directories}
+)

extension/threadpool/threadpool.cpp

@@ -14,6 +14,7 @@
 #include <cpuinfo.h>
 
 #include <atomic>
+#include <memory>
 
 namespace torch {
 namespace executorch {