Ship the CUDA delegate in the wheel

.ci/scripts/wheel/test_shared_libraries.py

@@ -54,6 +54,16 @@
 # The quantized kernels, whose own library the wheel ships when they are built.
 # A separate group because they have a separate owner, and because a wheel built
 # without them ships neither the library nor these symbols.
+# The CUDA delegate and its stream helper, for a wheel built from a CUDA index. The
+# stream helper matters most: two copies means two notions of the caller's stream, so
+# work queued through one is invisible to the other.
+# A strongly defined symbol, chosen by reading the built library rather than guessed.
+# The CudaBackend methods are emitted weak, and the definer count only looks at strong
+# definitions, so naming one of those would report zero definers for a library that is
+# plainly present.
+_CUDA_BACKEND_SYMBOLS = ("executorch::backends::cuda::load_library",)
+_CUDA_STREAM_SYMBOLS = ("executorch::extension::cuda::getCallerStream",)
+
 _QUANTIZED_KERNEL_SYMBOLS = (
     "torch::executor::native::quantize_per_tensor_out",
     "torch::executor::native::dequantize_per_tensor_out",
@@ -361,6 +371,18 @@ def _assert_single_definer(symbols, what: str, owner: str | None = None) -> None
         "libexecutorch_kernels_quantized.so",
         False,
     ),
+    (
+        "CUDA delegate",
+        _CUDA_BACKEND_SYMBOLS,
+        "libexecutorch_backend_cuda.so",
+        False,
+    ),
+    (
+        "CUDA stream helper",
+        _CUDA_STREAM_SYMBOLS,
+        "libextension_cuda.so",
+        False,
+    ),
     # The third-party code these libraries bundle, checked separately from the
     # wrappers above. A wrapper can have a single owner while the implementation
     # underneath it is bundled into two of these, which is two real thread pools or
@@ -1240,6 +1262,8 @@ def test_shipped_library_names_are_expected() -> None:
         "libexecutorch",
         "libexecutorch_kernels_optimized",
         "libexecutorch_kernels_quantized",
+        "libexecutorch_backend_cuda",
+        "libextension_cuda",
         "libexecutorch_backend_xnnpack",
         "libexecutorch_threadpool",
         "libexecutorch_etdump",

backends/cuda/CMakeLists.txt

@@ -200,7 +200,17 @@ 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}
+)
 
 target_include_directories(
   aoti_cuda_backend
@@ -240,6 +250,22 @@ endif()
 
 executorch_target_link_options_shared_lib(aoti_cuda_backend)
 
+if(EXECUTORCH_BUILD_SHARED)
+  # Named after what the library provides rather than after the target that
+  # produces it, matching the other shipped delegates, so the file reads as
+  # libexecutorch_backend_cuda.so. The target name stays as it is because the
+  # rest of the build already refers to it.
+  set_target_properties(
+    aoti_cuda_backend PROPERTIES OUTPUT_NAME executorch_backend_cuda
+  )
+  executorch_target_soname_policy(aoti_cuda_backend)
+  # Resolve the runtime from the shared library rather than from the static
+  # core, so the delegate registers into the one registry the process has.
+  target_link_libraries(aoti_cuda_backend PUBLIC executorch_shared)
+  # Ships beside the runtime in the wheel's lib/ directory.
+  executorch_target_shipped_runtime_path(aoti_cuda_backend)
+endif()
+
 install(
   TARGETS aoti_cuda_backend
   EXPORT ExecuTorchTargets

setup.py

@@ -1138,6 +1138,10 @@ def run(self):  # noqa C901
             if cmake_cache.is_enabled("EXECUTORCH_BUILD_CUDA"):
                 cmake_build_args += ["--target", "aoti_cuda_backend"]
                 cmake_build_args += ["--target", "aoti_common_shims_slim"]
+                if cmake_cache.is_enabled("EXECUTORCH_BUILD_SHARED"):
+                    # The stream helper ships as its own library so a process has one
+                    # of it. Named because nothing else in a wheel build links it.
+                    cmake_build_args += ["--target", "extension_cuda"]
 
             if cmake_cache.is_enabled("EXECUTORCH_BUILD_EXTENSION_MODULE"):
                 cmake_build_args += ["--target", "extension_module"]
@@ -1278,6 +1282,28 @@ def run(self):  # noqa C901
                         "EXECUTORCH_BUILD_KERNELS_OPTIMIZED",
                     ],
                 ),
+                # The CUDA delegate and the process-wide CUDA stream helper, for a
+                # wheel built from a CUDA index. Only present when the build asks for
+                # CUDA, so packaging requires that rather than looking for files a
+                # CPU-only build never produced.
+                BuiltFile(
+                    src_dir="%CMAKE_CACHE_DIR%/backends/cuda/",
+                    src_name="libexecutorch_backend_cuda.so",
+                    dst="executorch/lib/libexecutorch_backend_cuda.so",
+                    dependent_cmake_flags=[
+                        "EXECUTORCH_BUILD_SHARED",
+                        "EXECUTORCH_BUILD_CUDA",
+                    ],
+                ),
+                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",
+                    ],
+                ),
                 # The quantized kernels, as their own library rather than code
                 # fused into the AOT-only extension beside the Python bindings.
                 # A C++ application running a quantized model could not link

tools/cmake/executorch-wheel-config.cmake

@@ -406,6 +406,11 @@ if(TARGET executorch::runtime)
 endif()
 
 executorch_define_component(backend_xnnpack executorch_backend_xnnpack)
+# The CUDA delegate and its stream helper, present only in a wheel built from a
+# CUDA index. A CPU wheel defines neither, so a consumer asking for one is told
+# while configuring.
+executorch_define_component(backend_cuda executorch_backend_cuda)
+executorch_define_component(cuda_stream extension_cuda)
 
 # Find prebuilt _portable_lib.<EXT_SUFFIX>.so. This is the legacy contract used
 # to build custom-op extensions against the Python module, and is kept working