Ship the merged CPU kernels as their own library

.ci/scripts/wheel/test_cpp_sdk.py

@@ -40,6 +40,19 @@
 # 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",)
+
+# The registry entry points, kept separate from the kernel implementations above.
+# A library that carries its own copy of these has its own registration code, which
+# is what this split is meant to prevent: one owner of the operator table. Checking
+# only a kernel implementation would miss that entirely.
+_KERNEL_REGISTRY_SYMBOLS = (
+    "executorch::runtime::register_kernels",
+    "executorch::runtime::get_registered_kernels",
+)
+
 # `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>.+)$")
@@ -110,6 +123,32 @@ def _defines_symbol(library: Path, symbol: str) -> bool:
     return False
 
 
+def _report_definers(symbols, what: str, limit: int) -> None:
+    """Print which shipped libraries define each symbol and hold a ceiling.
+
+    Used where the desired invariant is one owner but the wheel does not reach it
+    yet for reasons outside this change. Printing alone would let a regression
+    from a few definers to many stay green, so the count is also capped at the
+    number that exists today. The cap is meant to be lowered, never raised.
+    """
+    package_dir = _installed_package_dir()
+    libraries = _shipped_shared_objects(package_dir)
+    for symbol in symbols:
+        definers = [
+            str(library.relative_to(package_dir))
+            for library in libraries
+            if _defines_symbol(library, symbol)
+        ]
+        print(f"- {what}: {symbol.split('::')[-1]} defined by {len(definers)}")
+        for definer in definers:
+            print(f"    {definer}")
+        assert len(definers) <= limit, (
+            f"{what}: {symbol} is defined by {len(definers)} shipped libraries, "
+            f"more than the {limit} that existed when this check was added, so a "
+            f"new duplicate crept in: {definers}"
+        )
+
+
 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"
@@ -139,6 +178,17 @@ 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")
+    # Ownership of the operator table, not just of a kernel implementation. This is
+    # reported rather than asserted because two extension modules still link the
+    # static core and carry their own copy, which predates this split: a released
+    # wheel has five definers where this one has three. Asserting here would fail
+    # on those pre-existing copies rather than on anything this change introduced.
+    _report_definers(_KERNEL_REGISTRY_SYMBOLS, "operator registry", limit=2)
+
+
 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 +289,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

@@ -1263,6 +1263,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,7 +1120,7 @@ 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."
@@ -1136,6 +1136,27 @@ def run(self):  # noqa C901
                         "EXECUTORCH_BUILD_CPUINFO",
                     ],
                 ),
+                # 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()}"
+                    ),
+                    # The target is only created when the optimized kernels are
+                    # enabled, so packaging has to require that too rather than
+                    # looking for a file a shared build may never have produced.
+                    dependent_cmake_flags=[
+                        "EXECUTORCH_BUILD_SHARED",
+                        "EXECUTORCH_BUILD_KERNELS_OPTIMIZED",
+                    ],
+                ),
                 # 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
@@ -293,6 +314,13 @@ function(gen_operators_lib)
   endif()
   target_sources(${GEN_LIB_NAME} PRIVATE ${_srcs_list})
   target_link_libraries(${GEN_LIB_NAME} PRIVATE ${GEN_DEPS})
+  if(GEN_SHARED)
+    # Resolve the runtime from the shared library rather than from the static
+    # core in GEN_DEPS. Linking the static core gives this library its own copy
+    # of the operator table, so its static initializer registers into a table
+    # nothing else reads and the operators appear missing at run time.
+    executorch_target_link_shared_runtime(${GEN_LIB_NAME})
+  endif()
   set(portable_kernels_check "portable_kernels")
   if(GEN_KERNEL_LIBS)