Ship the merged CPU kernels as their own library

.ci/scripts/wheel/test_cpp_sdk.py

@@ -40,6 +40,10 @@
 # oversubscribes the CPU because each pool sizes itself to all cores.
 _THREADPOOL_SYMBOLS = ("executorch::extension::threadpool::get_threadpool",)
 
+# A representative operator from the merged CPU kernels. A second definer means
+# the operators are registered twice, which aborts at startup.
+_KERNEL_SYMBOLS = ("torch::executor::native::abs_out",)
+
 # `nm -DC` prints "<hexaddr> <kind> <name>" for a definition and
 # "                 U <name>" for an undefined reference.
 _DEFINED = re.compile(r"^[0-9a-fA-F]+\s+(?P<kind>[A-Za-z])\s+(?P<name>.+)$")
@@ -139,6 +143,11 @@ def test_single_threadpool() -> None:
     _assert_single_definer(_THREADPOOL_SYMBOLS, "thread pool")
 
 
+def test_single_kernel_registration() -> None:
+    """Exactly one shipped library may define the merged CPU kernels."""
+    _assert_single_definer(_KERNEL_SYMBOLS, "set of CPU kernels")
+
+
 def test_cpp_consumer(work_dir: Path) -> None:
     """A standalone C++ app builds and runs against the installed wheel."""
     assert shutil.which("cmake") is not None, "cmake is required to build a consumer"
@@ -239,4 +248,5 @@ def _assert_runs_relocated(consumer, package_dir, work_dir, environment) -> None
 def run_tests(work_dir: Path) -> None:
     test_single_backend_registry()
     test_single_threadpool()
+    test_single_kernel_registration()
     test_cpp_consumer(work_dir)

CMakeLists.txt

@@ -1257,6 +1257,14 @@ if(EXECUTORCH_BUILD_PYBIND)
   target_compile_options(portable_lib PUBLIC ${_pybind_compile_options})
   target_link_libraries(portable_lib PRIVATE ${_dep_libs})
   executorch_target_link_shared_runtime(portable_lib)
+  # The operators register themselves from a static initializer, so nothing here
+  # references a symbol from the kernels library and some linkers drop it, which
+  # surfaces at runtime as a missing kernel rather than a link error.
+  if(TARGET optimized_native_cpu_ops_lib)
+    executorch_target_retain_shared_library(
+      portable_lib optimized_native_cpu_ops_lib
+    )
+  endif()
 
   # Set RPATH to find PyTorch and backend libraries relative to the installation
   # location. This goes from executorch/extension/pybindings up to

configurations/CMakeLists.txt

@@ -50,7 +50,15 @@ if(EXECUTORCH_BUILD_KERNELS_OPTIMIZED)
   else()
     set(_optimized_native_cpu_ops_lib_portable_kernels_lib portable_kernels)
   endif()
+  # Ship this as a shared library in the wheel so the kernels are registered
+  # once per process instead of once per component that links them.
+  if(EXECUTORCH_BUILD_SHARED)
+    set(_merged_cpu_ops_shared SHARED)
+  else()
+    set(_merged_cpu_ops_shared "")
+  endif()
   gen_operators_lib(
+    ${_merged_cpu_ops_shared}
     LIB_NAME
     "optimized_native_cpu_ops_lib"
     KERNEL_LIBS

setup.py

@@ -1120,14 +1120,29 @@ def run(self):  # noqa C901
                 BuiltFile(
                     src_dir="%CMAKE_CACHE_DIR%/extension/threadpool/",
                     src_name=(
-                        "libexecutorch_threadpool.so." f"{get_runtime_soname_major()}.*"
+                        f"libexecutorch_threadpool.so.{get_runtime_soname_major()}.*"
                     ),
                     dst=(
                         "executorch/lib/libexecutorch_threadpool.so."
                         f"{get_runtime_soname_major()}"
                     ),
                     dependent_cmake_flags=["EXECUTORCH_BUILD_SHARED"],
                 ),
+                # Install the merged CPU kernels beside them, so the operators are
+                # registered once per process rather than once per component.
+                BuiltFile(
+                    src_dir="%CMAKE_CACHE_DIR%/configurations/",
+                    src_name=(
+                        "libexecutorch_optimized_native_cpu_ops_lib.so."
+                        f"{get_runtime_soname_major()}.*"
+                    ),
+                    dst=(
+                        "executorch/lib/"
+                        "libexecutorch_optimized_native_cpu_ops_lib.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.

tools/cmake/Codegen.cmake

@@ -268,9 +268,13 @@ function(gen_custom_ops_aot_lib)
 endfunction()
 
 # Generate a runtime lib for registering operators in Executorch
+#
+# SHARED opts this library into being a shared object. It is opt-in because most
+# callers want the default static library, and only the one shipped in the wheel
+# needs to be shared so a process has a single copy of the kernels.
 function(gen_operators_lib)
   set(multi_arg_names LIB_NAME KERNEL_LIBS DEPS DTYPE_SELECTIVE_BUILD)
-  cmake_parse_arguments(GEN "" "" "${multi_arg_names}" ${ARGN})
+  cmake_parse_arguments(GEN "SHARED" "" "${multi_arg_names}" ${ARGN})
 
   message(STATUS "Generating operator lib:")
   message(STATUS "  LIB_NAME: ${GEN_LIB_NAME}")
@@ -283,7 +287,24 @@ function(gen_operators_lib)
     set(_opvariant_h ${_out_dir}/selected_op_variants.h)
   endif()
 
-  add_library(${GEN_LIB_NAME})
+  if(GEN_SHARED)
+    add_library(${GEN_LIB_NAME} SHARED)
+    set_target_properties(
+      ${GEN_LIB_NAME}
+      PROPERTIES OUTPUT_NAME executorch_${GEN_LIB_NAME}
+                 VERSION "${PROJECT_VERSION}"
+                 SOVERSION "${PROJECT_VERSION_MAJOR}"
+    )
+    if(NOT APPLE)
+      # Ships beside the runtime in the wheel's lib/ directory.
+      set_target_properties(
+        ${GEN_LIB_NAME} PROPERTIES BUILD_RPATH "$ORIGIN" INSTALL_RPATH
+                                                         "$ORIGIN"
+      )
+    endif()
+  else()
+    add_library(${GEN_LIB_NAME})
+  endif()
 
   set(_srcs_list ${_out_dir}/RegisterCodegenUnboxedKernelsEverything.cpp
                  ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h