diff --git a/.ci/scripts/wheel/test_shared_libraries.py b/.ci/scripts/wheel/test_shared_libraries.py index a11726ca218..6ec28eeb70e 100644 --- a/.ci/scripts/wheel/test_shared_libraries.py +++ b/.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", diff --git a/backends/cuda/CMakeLists.txt b/backends/cuda/CMakeLists.txt index 06990692428..ba30393d1d7 100644 --- a/backends/cuda/CMakeLists.txt +++ b/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 diff --git a/setup.py b/setup.py index 2ed9cd61b90..beee589e809 100644 --- a/setup.py +++ b/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 diff --git a/tools/cmake/executorch-wheel-config.cmake b/tools/cmake/executorch-wheel-config.cmake index d6a6f5af5cc..553d3246049 100644 --- a/tools/cmake/executorch-wheel-config.cmake +++ b/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..so. This is the legacy contract used # to build custom-op extensions against the Python module, and is kept working