From 944ea677f833c8e3f897898f00cd3e02c40bbb2f Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Fri, 31 Jul 2026 12:20:27 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .ci/scripts/wheel/test_cpp_sdk.py | 25 +++++++++++++---- extension/threadpool/CMakeLists.txt | 43 +++++++++++++++++++++++++---- setup.py | 14 ++++++++++ 3 files changed, 72 insertions(+), 10 deletions(-) diff --git a/.ci/scripts/wheel/test_cpp_sdk.py b/.ci/scripts/wheel/test_cpp_sdk.py index c2e729eb1f2..6de24259f9f 100644 --- a/.ci/scripts/wheel/test_cpp_sdk.py +++ b/.ci/scripts/wheel/test_cpp_sdk.py @@ -36,6 +36,10 @@ "executorch::runtime::get_backend_class", ) +# The thread pool accessor. A second definer means a second pool, which +# oversubscribes the CPU because each pool sizes itself to all cores. +_THREADPOOL_SYMBOLS = ("executorch::extension::threadpool::get_threadpool",) + # `nm -DC` prints " " for a definition and # " U " for an undefined reference. _DEFINED = re.compile(r"^[0-9a-fA-F]+\s+(?P[A-Za-z])\s+(?P.+)$") @@ -99,23 +103,33 @@ def _defines_symbol(library: Path, symbol: str) -> bool: return False -def test_single_backend_registry() -> None: - """Exactly one shipped library may define the backend registry.""" +def _assert_single_definer(symbols, what: str) -> None: + """Exactly one shipped library may define each of `symbols`.""" assert shutil.which("nm") is not None, "nm is required to inspect the wheel" package_dir = _installed_package_dir() libraries = _shipped_shared_objects(package_dir) assert libraries, f"no shared libraries found under {package_dir}" - for symbol in _REGISTRY_SYMBOLS: + for symbol in symbols: definers = [lib for lib in libraries if _defines_symbol(lib, symbol)] pretty = [str(lib.relative_to(package_dir)) for lib in definers] assert len(definers) == 1, ( f"expected exactly one library to define {symbol}, found " f"{len(definers)}: {pretty}. More than one definition means the " - f"process has more than one backend registry." + f"process has more than one {what}." ) - print(f"✓ single backend registry across {len(libraries)} shipped libraries") + print(f"✓ single {what} across {len(libraries)} shipped libraries") + + +def test_single_backend_registry() -> None: + """Exactly one shipped library may define the backend registry.""" + _assert_single_definer(_REGISTRY_SYMBOLS, "backend registry") + + +def test_single_threadpool() -> None: + """Exactly one shipped library may define the thread pool accessor.""" + _assert_single_definer(_THREADPOOL_SYMBOLS, "thread pool") def test_cpp_consumer(work_dir: Path) -> None: @@ -172,4 +186,5 @@ def test_cpp_consumer(work_dir: Path) -> None: def run_tests(work_dir: Path) -> None: test_single_backend_registry() + test_single_threadpool() test_cpp_consumer(work_dir) diff --git a/extension/threadpool/CMakeLists.txt b/extension/threadpool/CMakeLists.txt index 3b9c7c66ddb..33a13286bf4 100644 --- a/extension/threadpool/CMakeLists.txt +++ b/extension/threadpool/CMakeLists.txt @@ -30,13 +30,46 @@ else() set(_threadpool_size_flag "EXECUTORCH_THREADPOOL_USE_PERFORMANCE_CORES") endif() +# The thread pool is a process-wide singleton held in a function-local static, +# so every library that links it statically gets its own copy. Build it shared +# for the wheel, where several extensions are loaded into one interpreter, and +# keep it static everywhere else so no other build changes. +if(EXECUTORCH_BUILD_SHARED) + set(_threadpool_library_type SHARED) +else() + set(_threadpool_library_type STATIC) +endif() + add_library( - extension_threadpool threadpool.cpp threadpool_guard.cpp thread_parallel.cpp - cpuinfo_utils.cpp -) -target_link_libraries( - extension_threadpool PUBLIC executorch_core cpuinfo pthreadpool + extension_threadpool + ${_threadpool_library_type} threadpool.cpp threadpool_guard.cpp + thread_parallel.cpp cpuinfo_utils.cpp ) +if(EXECUTORCH_BUILD_SHARED) + set_target_properties( + extension_threadpool + PROPERTIES OUTPUT_NAME executorch_threadpool + VERSION "${PROJECT_VERSION}" + SOVERSION "${PROJECT_VERSION_MAJOR}" + ) + if(NOT APPLE) + # Ships beside libexecutorch.so in the wheel's lib/ directory. + set_target_properties( + extension_threadpool PROPERTIES BUILD_RPATH "$ORIGIN" INSTALL_RPATH + "$ORIGIN" + ) + endif() + # cpuinfo and pthreadpool are forced static, so bundle them inside this + # library instead of making every consumer supply them. + executorch_target_whole_archive(extension_threadpool cpuinfo) + executorch_target_whole_archive(extension_threadpool pthreadpool) + target_link_libraries(extension_threadpool PRIVATE cpuinfo pthreadpool) + target_link_libraries(extension_threadpool PUBLIC executorch_shared) +else() + target_link_libraries( + extension_threadpool PUBLIC executorch_core cpuinfo pthreadpool + ) +endif() target_include_directories( extension_threadpool PUBLIC ${_common_include_directories} ) diff --git a/setup.py b/setup.py index f33e5478476..d0b3d0448ff 100644 --- a/setup.py +++ b/setup.py @@ -1114,6 +1114,20 @@ def run(self): # noqa C901 dst=f"executorch/lib/libexecutorch.so.{get_runtime_soname_major()}", dependent_cmake_flags=["EXECUTORCH_BUILD_SHARED"], ), + # Install the shared thread pool next to it. It is a separate + # library so that a process has one pool rather than one per + # component that uses it. + BuiltFile( + src_dir="%CMAKE_CACHE_DIR%/extension/threadpool/", + src_name=( + "libexecutorch_threadpool.so." f"{get_runtime_soname_major()}.*" + ), + dst=( + "executorch/lib/libexecutorch_threadpool.so." + f"{get_runtime_soname_major()}" + ), + dependent_cmake_flags=["EXECUTORCH_BUILD_SHARED"], + ), # Install the prebuilt pybindings extension wrapper for the runtime, # portable kernels, and a selection of backends. This lets users # load and execute .pte files from python.