Ship the CUDA delegate as its own library

.ci/scripts/wheel/test_cpp_sdk.py

@@ -50,6 +50,10 @@
     "executorch::backends::xnnpack::XnnpackBackendOptions::workspace_manager",
 )
 
+# A representative symbol from the CUDA delegate's shim layer. The delegate's own
+# methods are weak symbols, so this checks a strong one instead.
+_CUDA_SYMBOLS = ("executorch::backends::cuda::clearCurrentCUDAStream",)
+
 # `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>.+)$")
@@ -120,16 +124,117 @@ def _defines_symbol(library: Path, symbol: str) -> bool:
     return False
 
 
-def _assert_single_definer(symbols, what: str) -> None:
-    """Exactly one shipped library may define each of `symbols`."""
+def report_wheel_composition() -> None:
+    """Print what the wheel ships and what each library needs.
+
+    Not an assertion. A size jump or an unexpected external dependency is the
+    first visible sign that a component got statically duplicated again, so the
+    numbers are worth having in the log of every run.
+    """
+    package_dir = _installed_package_dir()
+    libraries = _shipped_shared_objects(package_dir)
+
+    print("shipped libraries:")
+    total = 0
+    for library in sorted(libraries, key=lambda path: path.name):
+        size = library.stat().st_size
+        total += size
+        print(f"  {size / 1024:9.1f} KiB  {library.relative_to(package_dir)}")
+    print(f"  {total / 1024:9.1f} KiB  total")
+
+    if shutil.which("readelf") is None:
+        return
+    # Anything the libraries need that the wheel does not itself ship has to be
+    # present on the user's machine, so it belongs in the report. Compare against
+    # the shipped file names rather than guessing from name prefixes.
+    shipped = {library.name for library in libraries}
+    external = set()
+    for library in libraries:
+        dynamic = subprocess.run(
+            ["readelf", "-d", str(library)],
+            capture_output=True,
+            text=True,
+            check=False,
+        ).stdout
+        for line in dynamic.splitlines():
+            if "(NEEDED)" not in line or "[" not in line:
+                continue
+            name = line.split("[", 1)[1].rstrip("]").strip()
+            if name not in shipped:
+                external.add(name)
+    if external:
+        print("external dependencies expected from the environment:")
+        for name in sorted(external):
+            print(f"  {name}")
+
+
+def test_shipped_libraries_load() -> None:
+    """Every shipped library must depend only on things that exist.
+
+    The symbol checks prove each component is defined exactly once, but a library
+    can still be unloadable if it needs something nothing provides, which is a
+    packaging bug rather than a duplication bug.
+
+    A dependency the wheel ships elsewhere is fine even when `ldd` cannot resolve
+    it: some extensions are loaded after `import torch` has already brought their
+    dependencies into the process, so they intentionally carry no path to them.
+    Only a name nothing in the wheel provides is a real problem.
+    """
+    if shutil.which("ldd") is None:
+        print("- ldd not available, skipping the load check")
+        return
+
+    package_dir = _installed_package_dir()
+    libraries = _shipped_shared_objects(package_dir)
+    shipped = {library.name for library in libraries}
+
+    broken = {}
+    for library in libraries:
+        resolved = subprocess.run(
+            ["ldd", str(library)], capture_output=True, text=True, check=False
+        ).stdout
+        missing = [
+            name
+            for name in (
+                line.split("=>")[0].strip()
+                for line in resolved.splitlines()
+                if "not found" in line
+            )
+            if name not in shipped
+        ]
+        if missing:
+            broken[str(library.relative_to(package_dir))] = missing
+
+    assert not broken, (
+        "shipped libraries need dependencies that nothing provides, so they will "
+        f"fail to load: {broken}"
+    )
+    print("✓ every shipped library depends only on things that exist")
+
+
+def _assert_single_definer(symbols, what: str, optional: bool = False) -> None:
+    """Exactly one shipped library may define each of `symbols`.
+
+    `optional` allows a component that is only present in some wheel flavors,
+    such as an accelerator delegate, to be absent without failing.
+    """
     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 symbols:
-        definers = [lib for lib in libraries if _defines_symbol(lib, symbol)]
+    # Resolve every symbol first, so a component that is only half present is
+    # reported rather than being mistaken for one that is absent entirely.
+    found = {
+        symbol: [lib for lib in libraries if _defines_symbol(lib, symbol)]
+        for symbol in symbols
+    }
+    if optional and not any(found.values()):
+        print(f"- no {what} in this wheel, skipping")
+        return
+
+    for symbol, definers in found.items():
         pretty = [str(lib.relative_to(package_dir)) for lib in definers]
         assert len(definers) == 1, (
             f"expected exactly one library to define {symbol}, found "
@@ -159,6 +264,11 @@ def test_single_xnnpack_delegate() -> None:
     _assert_single_definer(_XNNPACK_SYMBOLS, "XNNPACK delegate")
 
 
+def test_single_cuda_delegate() -> None:
+    """Exactly one shipped library may define the CUDA delegate, if present."""
+    _assert_single_definer(_CUDA_SYMBOLS, "CUDA delegate", optional=True)
+
+
 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"
@@ -257,8 +367,11 @@ def _assert_runs_relocated(consumer, package_dir, work_dir, environment) -> None
 
 
 def run_tests(work_dir: Path) -> None:
+    report_wheel_composition()
+    test_shipped_libraries_load()
     test_single_backend_registry()
     test_single_threadpool()
     test_single_kernel_registration()
     test_single_xnnpack_delegate()
+    test_single_cuda_delegate()
     test_cpp_consumer(work_dir)

CMakeLists.txt

@@ -1269,7 +1269,9 @@ if(EXECUTORCH_BUILD_PYBIND)
   # initializer, so nothing here references a symbol from them and some linkers
   # drop them from DT_NEEDED. That surfaces at runtime as a missing kernel or an
   # unregistered backend rather than as a link error.
-  foreach(_retained_component optimized_native_cpu_ops_lib xnnpack_backend)
+  foreach(_retained_component optimized_native_cpu_ops_lib xnnpack_backend
+                              aoti_cuda_backend
+  )
     if(TARGET ${_retained_component})
       executorch_target_retain_shared_library(
         portable_lib ${_retained_component}

backends/cuda/CMakeLists.txt

@@ -93,8 +93,17 @@ target_compile_options(
   PUBLIC "$<$<COMPILE_LANGUAGE:CXX>:${_cuda_cxx_compile_options}>"
 )
 
-# Link against ExecuTorch core libraries
-target_link_libraries(cuda_platform PRIVATE executorch_core ${CMAKE_DL_LIBS})
+# Link against ExecuTorch core libraries. Resolve them from the shared runtime
+# when there is one, so this does not carry a second copy of the backend
+# registry.
+if(EXECUTORCH_BUILD_SHARED)
+  target_link_libraries(
+    cuda_platform PRIVATE executorch_shared ${CMAKE_DL_LIBS}
+  )
+  executorch_target_link_shared_runtime(cuda_platform)
+else()
+  target_link_libraries(cuda_platform PRIVATE executorch_core ${CMAKE_DL_LIBS})
+endif()
 
 install(
   TARGETS cuda_platform
@@ -169,14 +178,9 @@ if(_cuda_is_msvc_toolchain)
 else()
   target_link_libraries(
     aoti_cuda_shims
-    PRIVATE cuda_platform
-    PUBLIC -Wl,--whole-archive
-           aoti_common_shims_slim
-           -Wl,--no-whole-archive
-           CUDA::cudart
-           CUDA::curand
-           extension_cuda
-           ${CMAKE_DL_LIBS}
+    PRIVATE cuda_platform -Wl,--whole-archive aoti_common_shims_slim
+            -Wl,--no-whole-archive
+    PUBLIC CUDA::cudart CUDA::curand extension_cuda ${CMAKE_DL_LIBS}
   )
 endif()
 
@@ -200,7 +204,35 @@ if(_cuda_is_msvc_toolchain)
   list(APPEND _aoti_cuda_backend_sources runtime/cuda_allocator.cpp)
 endif()
 
-add_library(aoti_cuda_backend STATIC ${_aoti_cuda_backend_sources})
+# Build the delegate as a shared library for the wheel so a process has one copy
+# of it, and keep it static everywhere else so no other build changes.
+if(EXECUTORCH_BUILD_SHARED)
+  set(_aoti_cuda_backend_library_type SHARED)
+else()
+  set(_aoti_cuda_backend_library_type STATIC)
+endif()
+add_library(
+  aoti_cuda_backend ${_aoti_cuda_backend_library_type}
+                    ${_aoti_cuda_backend_sources}
+)
+if(EXECUTORCH_BUILD_SHARED)
+  set_target_properties(
+    aoti_cuda_backend
+    PROPERTIES OUTPUT_NAME executorch_cuda_backend
+               VERSION "${PROJECT_VERSION}"
+               SOVERSION "${PROJECT_VERSION_MAJOR}"
+  )
+  if(NOT APPLE)
+    # Ships in the wheel's lib/ directory, but the CUDA shim library it links
+    # lives under backends/cuda, so both locations have to be searchable. The
+    # CUDA runtime itself comes from the environment and is not bundled.
+    set(_cuda_backend_rpath "$ORIGIN:$ORIGIN/../backends/cuda")
+    set_target_properties(
+      aoti_cuda_backend PROPERTIES BUILD_RPATH "${_cuda_backend_rpath}"
+                                   INSTALL_RPATH "${_cuda_backend_rpath}"
+    )
+  endif()
+endif()
 
 target_include_directories(
   aoti_cuda_backend

setup.py

@@ -1160,6 +1160,36 @@ def run(self):  # noqa C901
                         "EXECUTORCH_BUILD_XNNPACK",
                     ],
                 ),
+                # Install the CUDA delegate beside them when it is built. The CUDA
+                # runtime itself is not bundled; it comes from the environment.
+                BuiltFile(
+                    src_dir="%CMAKE_CACHE_DIR%/backends/cuda/",
+                    src_name=(
+                        "libexecutorch_cuda_backend.so."
+                        f"{get_runtime_soname_major()}.*"
+                    ),
+                    dst=(
+                        "executorch/lib/libexecutorch_cuda_backend.so."
+                        f"{get_runtime_soname_major()}"
+                    ),
+                    dependent_cmake_flags=[
+                        "EXECUTORCH_BUILD_SHARED",
+                        "EXECUTORCH_BUILD_CUDA",
+                    ],
+                ),
+                # The CUDA delegate calls into this for stream handling, so an
+                # application that links the delegate from the wheel cannot
+                # resolve it unless this ships too. It carries no SONAME version,
+                # so the name is used as built.
+                BuiltFile(
+                    src_dir="%CMAKE_CACHE_DIR%/extension/cuda/",
+                    src_name="libextension_cuda.so",
+                    dst="executorch/lib/libextension_cuda.so",
+                    dependent_cmake_flags=[
+                        "EXECUTORCH_BUILD_SHARED",
+                        "EXECUTORCH_BUILD_CUDA",
+                    ],
+                ),
                 # 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.