Ship the merged CPU kernels as their own library

.ci/scripts/wheel/test_cpp_sdk.py

@@ -43,6 +43,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>.+)$")
@@ -171,6 +184,16 @@ 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. A
+    # second copy means a second table, and a static initializer registering into
+    # a table nothing else reads shows up as an operator that is missing at run
+    # time rather than as a link error.
+    _assert_single_definer(_KERNEL_REGISTRY_SYMBOLS, "operator registry")
+
+
 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"
@@ -761,7 +784,7 @@ def test_no_absolute_runtime_paths() -> None:
 
 # Link every component this wheel offers, and report which ones those are so the test
 # can check the result. Guarded individually because the set depends on the wheel.
-foreach(_component threadpool)
+foreach(_component threadpool kernels_optimized)
   if(TARGET executorch::${_component})
     target_link_libraries(component_consumer PRIVATE executorch::${_component})
     # Report the library file, not just the target name: the two differ, and the test
@@ -929,6 +952,7 @@ def run_tests(work_dir: Path) -> None:
     test_custom_op_compiles(work_dir)
     test_no_absolute_runtime_paths()
     test_single_threadpool()
+    test_single_kernel_registration()
     test_cpp_consumer(work_dir)
     test_documented_example_compiles(work_dir)
     test_component_targets_link(work_dir)

CMakeLists.txt

@@ -1299,6 +1299,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
@@ -65,4 +73,17 @@ if(EXECUTORCH_BUILD_KERNELS_OPTIMIZED)
     EXPORT ExecuTorchTargets
     DESTINATION ${CMAKE_INSTALL_LIBDIR}
   )
+  if(EXECUTORCH_BUILD_SHARED)
+    # The shipped file is named after the component a consumer asks for, rather
+    # than after the code generation target, so the pair reads together:
+    # executorch::kernels_optimized resolves to
+    # libexecutorch_kernels_optimized.so. The target name stays as it is because
+    # a source build already refers to it.
+    set_target_properties(
+      optimized_native_cpu_ops_lib
+      PROPERTIES OUTPUT_NAME executorch_kernels_optimized
+                 VERSION "${PROJECT_VERSION}"
+                 SOVERSION "${PROJECT_VERSION_MAJOR}"
+    )
+  endif()
 endif()

docs/source/using-executorch-cpp.md

@@ -108,6 +108,7 @@ each one is defined only when the installed wheel actually ships it:
 | --- | --- |
 | `executorch::runtime` | The core runtime. Always present. |
 | `executorch::threadpool` | The shared thread pool the kernels and backends use. |
+| `executorch::kernels_optimized` | CPU operator kernels, for any operator not taken by a backend. |
 
 Each target already carries what it needs: the runtime dependency, the include
 directories, the runtime search paths, and the linker options that keep a
@@ -131,7 +132,7 @@ target_link_libraries(
     PRIVATE executorch
     executorch::backends
     executorch::extensions
-    executorch::kernels)
+    executorch::kernels_optimized)
 ```
 
 See [Building from Source](using-executorch-building-from-source.md) for more information on the CMake build process.

setup.py

@@ -1288,7 +1288,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."
@@ -1304,6 +1304,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_kernels_optimized.so."
+                        f"{get_runtime_soname_major()}.*"
+                    ),
+                    dst=(
+                        "executorch/lib/"
+                        "libexecutorch_kernels_optimized.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

@@ -274,9 +274,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}")
@@ -289,7 +293,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
@@ -299,6 +320,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)
 

tools/cmake/executorch-wheel-config.cmake

@@ -44,8 +44,8 @@
 # dependency and, for a registration-only library, the link options that keep it
 # from being dropped. The names, when present, are:
 #
-# executorch::threadpool executorch::kernels executorch::xnnpack_backend
-# executorch::cuda_backend
+# executorch::threadpool executorch::kernels_optimized
+# executorch::xnnpack_backend executorch::cuda_backend
 #
 # Check with if(TARGET executorch::<name>) rather than assuming one exists. A
 # namespaced name that was never defined is a configure-time error that names
@@ -375,6 +375,11 @@ endfunction()
 
 executorch_define_component(threadpool executorch_threadpool)
 
+# The merged CPU kernels. Documented as a component and asserted by the release
+# checks, so it has to be defined here or a consumer following the documentation
+# gets a bare name that CMake hands to the linker as a literal flag.
+executorch_define_component(kernels_optimized executorch_kernels_optimized)
+
 # A consumer that links the thread pool has to see the same switch a source
 # build sets, or the parallel helpers in the runtime headers compile their
 # serial fallback instead and the library they linked is never used.