From 087bf0db2c4a5db2544dffa58d012405f387fb16 Mon Sep 17 00:00:00 2001 From: PyTorch Bot Date: Mon, 21 Sep 2026 09:23:05 -0700 Subject: [PATCH] Share the CUDA device guard through the extension so other backends can use it Selecting a CUDA device and putting the caller's back is written out by hand in two places in the CUDA backend: a struct private to one source file, and the same save and restore spelled inline in a destructor. A comment on a third copy says it follows the first. The TensorRT delegate, which lives in another repository and loads into the same process, had written a fourth. The tree already had the right guard, in a backend directory. It restores on destruction, is built through a factory returning a Result so a failure cannot be ignored by accident, logs a failed restore, neuters a moved-from copy so only one of the two restores, and carries tests. What it lacked was a home any other backend could reach. So it moves to the CUDA extension, beside the caller-stream guard, for the same reason that one is there: one program can run several backends whose engines sit on different devices, so each has to leave the device as it found it or the next inherits a selection it never made. Kernel launches follow the stream they are given, but allocation follows the current device, so a backend that allocates has to select one even when the caller chose the stream. Only the guard moves. The per-device stream helpers stay where their callers are, because the extension already answers the stream question differently and two mechanisms side by side would leave a reader guessing. Their header keeps the name in scope, which it needs anyway: its own stream guard holds one. Two details carried across rather than rewritten. The destructor decides whether to restore from what the caller asked for, not from whether the switch reported success, because a failed selection can still move the calling thread once an error is pending on the device. And a negative index is refused rather than ignored, so the one call site that can reach one asks first, the way it did before. The class is marked for export the same way the stream guard beside it is, so its symbols leave the shared library on a platform that requires saying so. The build that runs here does not hide symbols, so the omission would only have shown up for a consumer on another platform. The header reaches for no CUDA header at all, which is what lets a wheel ship it. Nothing in the interface names a CUDA type, so the bodies live in a source file beside it and only a consumer that actually uses the guard links that. A wheel check compiles every shipped header against the wheel and nothing else, deliberately without a toolkit include path, and a header that only compiles with one fails in a consumer project rather than in that check. The device index is a plain int here. The old signature used an alias from the backend's compatibility layer, and that alias was the only thing tying this guard to that backend. Test plan: Measured on a two card machine: a selection that needs a switch restores the caller's device, one that does not changes nothing, a negative index and an absent device are both refused, and moving a guard restores once rather than twice. That last one matters because a defaulted move would have restored in both scopes. The guard's tests, and the stream guard's tests beside them, had a target in one build system and none in the other, so the build a pull request runs never compiled either file. Both have a target now, and the job that runs the other CUDA runtime tests builds and runs them. Twenty nine cases that compiled nowhere now compile and run. The guard's existing tests still exercise it through the header that re-exports it, including the cases for a negative index and for move assignment being unavailable. The moved header and the file it came from both compile against the real runtime headers on Arm. The definitions are compiled into the library this directory already ships, rather than into one of their own. A separate archive would have had to be shipped and declared as its own component for a consumer to link what the header declares, and it would have needed position independent code of its own, because shared objects link it and the global setting for that covers only one build configuration. Folding it in avoids all of that. The library is only built for an accelerator, so it can call the runtime, which the stream helper beside it deliberately does not. It also takes the runtime now, because the guard logs and the logging lives there, and before this source existed that library called nothing at all. It takes it through the helper the tree provides rather than as an ordinary dependency, because the ordinary form puts the static core ahead of the shared runtime on the link line and leaves the library with its own copy of the registries. It also declares the runtime path every shipped library beside it declares. The CUDA delegate takes the runtime through the same helper. It asked for the shared runtime as an ordinary dependency, which its comment says was the intent, but that form loses to the ordering the helper exists to correct, and the delegate ended up carrying its own copy of the backend registry. It went unnoticed because a library that no longer exists used to sit on that link line and happened to order it correctly. It takes those from the shared runtime where one is built, not from the static core, because this library ships in the wheel and pulling the static core into it would give a process two platform layers. The CUDA platform library beside it chooses the same way, and privately, so that a library linking this one does not inherit the runtime and end up carrying a second copy of the registries. Every place that reaches the header is wired: a target of its own in the extension, a dependency from the guard's old target, and one from each backend target whose sources include it. The Buck graph is not built by any public job, so this was checked by reading rather than by a green tick. --- .github/workflows/cuda.yml | 4 +- backends/aoti/CMakeLists.txt | 13 ++++ backends/aoti/slim/cuda/guard.cpp | 44 ----------- backends/aoti/slim/cuda/guard.h | 72 +---------------- backends/aoti/slim/cuda/targets.bzl | 1 + .../aoti/slim/cuda/test/test_cuda_guard.cpp | 36 +++++++++ backends/cuda/CMakeLists.txt | 13 +++- backends/cuda/runtime/cuda_delegate_handle.h | 13 ++-- backends/cuda/runtime/cuda_mutable_state.cpp | 69 ++++++----------- extension/cuda/CMakeLists.txt | 19 +++-- extension/cuda/device_guard.cpp | 76 ++++++++++++++++++ extension/cuda/device_guard.h | 77 +++++++++++++++++++ extension/cuda/targets.bzl | 3 + setup.py | 14 +++- tools/cmake/executorch-wheel-config.cmake | 3 +- 15 files changed, 274 insertions(+), 183 deletions(-) create mode 100644 extension/cuda/device_guard.cpp create mode 100644 extension/cuda/device_guard.h diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index d2e5926d386..92c282d21cd 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -508,10 +508,12 @@ jobs: -v -o "addopts=" cmake --preset llm-release-cuda -DEXECUTORCH_BUILD_TESTS=ON - cmake --build cmake-out --target test_cuda_allocator test_cuda_mutable_state test_cuda_weight_cache -j$(nproc) + cmake --build cmake-out --target test_cuda_allocator test_cuda_mutable_state test_cuda_weight_cache test_cuda_guard test_cuda_stream_guard -j$(nproc) ctest --test-dir cmake-out -R test_cuda_allocator --output-on-failure -V ctest --test-dir cmake-out -R test_cuda_mutable_state --output-on-failure -V ctest --test-dir cmake-out -R test_cuda_weight_cache --output-on-failure -V + ctest --test-dir cmake-out -R test_cuda_guard --output-on-failure -V + ctest --test-dir cmake-out -R test_cuda_stream_guard --output-on-failure -V test-model-cuda-e2e: name: test-model-cuda-e2e-${{ matrix.model.name }}-${{ matrix.quant }} diff --git a/backends/aoti/CMakeLists.txt b/backends/aoti/CMakeLists.txt index 5ba98c67f20..550021e33a0 100644 --- a/backends/aoti/CMakeLists.txt +++ b/backends/aoti/CMakeLists.txt @@ -148,6 +148,19 @@ if(MSVC) ) endif() +# These two files had a Buck target and no CMake one, so the CMake build never +# compiled them. Declared beside the library that carries their subject. +if(BUILD_TESTING AND (EXECUTORCH_BUILD_CUDA OR EXECUTORCH_BUILD_ROCM)) + include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) + et_cxx_test( + test_cuda_guard SOURCES slim/cuda/test/test_cuda_guard.cpp EXTRA_LIBS + aoti_common_shims_slim extension_cuda + ) + et_cxx_test( + test_cuda_stream_guard SOURCES slim/cuda/test/test_cuda_stream_guard.cpp + EXTRA_LIBS aoti_common_shims_slim extension_cuda + ) +endif() install( TARGETS aoti_common_shims_slim EXPORT ExecuTorchTargets diff --git a/backends/aoti/slim/cuda/guard.cpp b/backends/aoti/slim/cuda/guard.cpp index 0d73b414c2d..fb78c40bc49 100644 --- a/backends/aoti/slim/cuda/guard.cpp +++ b/backends/aoti/slim/cuda/guard.cpp @@ -80,50 +80,6 @@ void clearCurrentCUDAStream(DeviceIndex device_index) { current_streams_.erase(device_index); } -CUDAGuard::CUDAGuard(CUDAGuard&& other) noexcept - : original_device_index_(other.original_device_index_), - current_device_index_(other.current_device_index_) { - // Mark the moved-from object as "already restored" so its destructor doesn't - // try to restore the device - other.original_device_index_ = other.current_device_index_; -} - -CUDAGuard::~CUDAGuard() { - if (original_device_index_ != current_device_index_) { - // DeviceIndex (int8_t) implicitly widens to int for cudaSetDevice - cudaError_t err = cudaSetDevice(original_device_index_); - if (err != cudaSuccess) { - ET_LOG( - Error, - "~CUDAGuard: Failed to restore device to %d: %s", - static_cast(original_device_index_), - cudaGetErrorString(err)); - } - } -} - -Error CUDAGuard::set_index(DeviceIndex device_index) { - // CUDA API returns int, explicit cast to DeviceIndex (int8_t) following ATen - int tmp_device = -1; - ET_CUDA_CHECK_OR_RETURN_ERROR(cudaGetDevice(&tmp_device)); - - original_device_index_ = static_cast(tmp_device); - current_device_index_ = device_index; - - if (current_device_index_ != original_device_index_) { - // DeviceIndex (int8_t) implicitly widens to int for cudaSetDevice - ET_CUDA_CHECK_OR_RETURN_ERROR(cudaSetDevice(current_device_index_)); - } - - return Error::Ok; -} - -Result CUDAGuard::create(DeviceIndex device_index) { - CUDAGuard guard; // Fixed: Removed () to create a variable, not a function - ET_CHECK_OK_OR_RETURN_ERROR(guard.set_index(device_index)); - return guard; -} - CUDAStreamGuard::CUDAStreamGuard(CUDAStreamGuard&& other) noexcept : device_guard_(std::move(other.device_guard_)), original_stream_(other.original_stream_), diff --git a/backends/aoti/slim/cuda/guard.h b/backends/aoti/slim/cuda/guard.h index f8a71873ff2..5a75fc01260 100644 --- a/backends/aoti/slim/cuda/guard.h +++ b/backends/aoti/slim/cuda/guard.h @@ -12,12 +12,15 @@ #include #include +#include #include #include #include namespace executorch::backends::cuda { +using ::executorch::extension::cuda::CUDAGuard; + using executorch::runtime::Error; using executorch::runtime::Result; @@ -66,75 +69,6 @@ std::optional peekCurrentCUDAStream( */ void clearCurrentCUDAStream(DeviceIndex device_index = -1); -/** - * RAII guard that sets the current CUDA device and restores it on destruction. - * This ensures that the device is properly restored even if an exception - * occurs. - * - */ -class CUDAGuard { - private: - /** - * Private constructor - use create() factory method instead. - */ - explicit CUDAGuard() - : original_device_index_(-1), current_device_index_(-1) {} - - public: - /** - * Factory method to create a CUDAGuard. - * - * @param device_index The device index to set as current - * @return Result containing the guard on success, or an error code on failure - */ - static Result create(DeviceIndex device_index); - - // Copy is not allowed - CUDAGuard(const CUDAGuard&) = delete; - CUDAGuard& operator=(const CUDAGuard&) = delete; - - // Move constructor and assignment - CUDAGuard(CUDAGuard&& other) noexcept; - CUDAGuard& operator=(CUDAGuard&& other) = delete; - - /** - * Destructor that restores the original device if necessary. - */ - ~CUDAGuard(); - - /** - * Sets the CUDA device to the given device index. - * - * @param device_index The device index to set as current - * @return Error code indicating success or failure - */ - Error set_index(DeviceIndex device_index); - - /** - * Get the original device index before the guard was created. - * - * @return The original device index - */ - DeviceIndex original_device() const { - return original_device_index_; - } - - /** - * Get the current device index. - * - * @return The current device index - */ - DeviceIndex current_device() const { - return current_device_index_; - } - - private: - /// The original device before this guard was created - DeviceIndex original_device_index_; - /// The current device managed by this guard - DeviceIndex current_device_index_; -}; - /** * RAII guard that sets the current CUDA device and stream, restoring both on * destruction. This is useful for temporarily switching to a different device diff --git a/backends/aoti/slim/cuda/targets.bzl b/backends/aoti/slim/cuda/targets.bzl index 9585cd71635..80ff0013b33 100644 --- a/backends/aoti/slim/cuda/targets.bzl +++ b/backends/aoti/slim/cuda/targets.bzl @@ -18,6 +18,7 @@ def define_common_targets(): exported_deps = [ "//executorch/backends/aoti/slim/c10/core:device", "//executorch/backends/aoti/slim/c10/cuda:exception", + "//executorch/extension/cuda:caller_stream", "//executorch/extension/cuda:runtime_api", "//executorch/runtime/core:core", "//executorch/runtime/core/exec_aten:lib", diff --git a/backends/aoti/slim/cuda/test/test_cuda_guard.cpp b/backends/aoti/slim/cuda/test/test_cuda_guard.cpp index 70da3108aba..1f5c370ee24 100644 --- a/backends/aoti/slim/cuda/test/test_cuda_guard.cpp +++ b/backends/aoti/slim/cuda/test/test_cuda_guard.cpp @@ -97,6 +97,42 @@ TEST_F(CUDAGuardTest, NegativeDeviceIndex) { // Compile-time type-trait checks. These do not need a CUDA device, so they // live outside the CUDAGuardTest fixture (whose SetUp() calls GTEST_SKIP // when no CUDA device is available). +TEST_F(CUDAGuardTest, MovingLeavesOneRestorer) { + int entry_device = -1; + ASSERT_EQ(cudaGetDevice(&entry_device), cudaSuccess); + int other_device = entry_device == 0 ? 1 : 0; + int device_count = 0; + ASSERT_EQ(cudaGetDeviceCount(&device_count), cudaSuccess); + if (other_device >= device_count) { + GTEST_SKIP() << "needs a second device to tell one restore from two"; + } + + { + auto created = CUDAGuard::create(other_device); + ASSERT_TRUE(created.ok()); + { + CUDAGuard moved(std::move(created.get())); + int inside = -1; + ASSERT_EQ(cudaGetDevice(&inside), cudaSuccess); + EXPECT_EQ(inside, other_device); + } + // The guard that took over the borrow has gone, so the entry device is + // back. + int after_move_target = -1; + ASSERT_EQ(cudaGetDevice(&after_move_target), cudaSuccess); + EXPECT_EQ(after_move_target, entry_device); + // Move somewhere else, so a second restore would be visible rather than + // landing on the value that is already current. + ASSERT_EQ(cudaSetDevice(other_device), cudaSuccess); + } + // The moved-from guard has now gone too. It must not have restored again. + int after_moved_from = -1; + ASSERT_EQ(cudaGetDevice(&after_moved_from), cudaSuccess); + EXPECT_EQ(after_moved_from, other_device) + << "the moved-from guard restored as well, so the device was put back twice"; + ASSERT_EQ(cudaSetDevice(entry_device), cudaSuccess); +} + TEST(CUDAGuardCompileTimeTest, CopyConstructorDeleted) { static_assert( !std::is_copy_constructible_v, diff --git a/backends/cuda/CMakeLists.txt b/backends/cuda/CMakeLists.txt index 6ed8f83678f..3ce9c527a53 100644 --- a/backends/cuda/CMakeLists.txt +++ b/backends/cuda/CMakeLists.txt @@ -409,10 +409,17 @@ endif() executorch_target_link_options_shared_lib(aoti_cuda_backend) -if(EXECUTORCH_BUILD_SHARED) +if(_executorch_cuda_ships_libraries) # Resolve the runtime from the shared library rather than from the static - # core, so the delegate registers into the one registry the process has. True - # of a static delegate in a shared build too, which is what ROCm builds. + # core, so the delegate registers into the one registry the process has. + # Through the helper, not as an ordinary dependency: CMake orders an archive + # ahead of what it depends on, so the static core would satisfy the runtime + # symbols first and this library would carry its own registry. + executorch_target_link_shared_runtime(aoti_cuda_backend) +elseif(EXECUTORCH_BUILD_SHARED) + # Static delegate in a shared build, which is what ROCm produces. It carries + # no link step of its own, so the helper cannot apply and its consumer + # resolves the runtime instead. target_link_libraries(aoti_cuda_backend PUBLIC executorch_shared) endif() diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index d4af1ed0741..0a9fb42a4e8 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -10,6 +10,7 @@ #include #include +#include #include #include #include @@ -49,15 +50,11 @@ struct CudaWeightStorage { std::free(data); return; } - int previous_device = 0; - const cudaError_t get_device_error = cudaGetDevice(&previous_device); - if (get_device_error == cudaSuccess && previous_device != device_index) { - (void)cudaSetDevice(device_index); - } + // A destructor cannot report, and the guard logs a failed restore itself. + const auto guard = + ::executorch::extension::cuda::CUDAGuard::create(device_index); + (void)guard; (void)cudaFree(data); - if (get_device_error == cudaSuccess && previous_device != device_index) { - (void)cudaSetDevice(previous_device); - } } CudaWeightStorage(const CudaWeightStorage&) = delete; diff --git a/backends/cuda/runtime/cuda_mutable_state.cpp b/backends/cuda/runtime/cuda_mutable_state.cpp index 0cfbd0253cb..49b694b6a4f 100644 --- a/backends/cuda/runtime/cuda_mutable_state.cpp +++ b/backends/cuda/runtime/cuda_mutable_state.cpp @@ -13,12 +13,14 @@ #include #include #include +#include #include #include #include #include #include +#include #include #include @@ -30,6 +32,7 @@ namespace aoti = ::executorch::backends::aoti; namespace slimc10 = ::executorch::backends::aoti::slim::c10; using ::executorch::backends::aoti::slim::from_blob; using ::executorch::backends::aoti::slim::SlimTensor; +using ::executorch::extension::cuda::CUDAGuard; using ::executorch::runtime::Error; using ::executorch::runtime::Result; @@ -103,42 +106,6 @@ bool handle_has_symbols(CudaDelegateHandle* h) { h->update_user_managed_constant_buffer_pairs; } -struct CudaDeviceGuard { - int prev_device{0}; - bool restore{false}; - - Error set(int device) { - if (device < 0) { - return Error::Ok; - } - cudaError_t err = cudaGetDevice(&prev_device); - if (err != cudaSuccess) { - ET_LOG(Error, "mutable_state: cudaGetDevice failed"); - return Error::Internal; - } - if (prev_device == device) { - return Error::Ok; - } - err = cudaSetDevice(device); - if (err != cudaSuccess) { - ET_LOG( - Error, - "mutable_state: cudaSetDevice(%d) failed: %s", - device, - cudaGetErrorString(err)); - return Error::Internal; - } - restore = true; - return Error::Ok; - } - - ~CudaDeviceGuard() { - if (restore) { - (void)cudaSetDevice(prev_device); - } - } -}; - Result tensor_cuda_device_index(const SlimTensor& t) { const slimc10::Device device = t.device(); ET_CHECK_OR_RETURN_ERROR( @@ -174,13 +141,21 @@ void cuda_free_on_pointer_device(void* ptr, bool synchronize) { (void)cudaGetLastError(); } - CudaDeviceGuard guard; - if (device >= 0 && guard.set(device) != Error::Ok) { - ET_LOG( - Error, - "mutable_state: freeing pointer %p without switching to device %d", - ptr, - device); + // The guard refuses a negative index rather than ignoring it, and this path + // reaches one whenever the pointer's attributes could not be read. Freeing is + // best effort, so stay quiet. + std::optional guard; + if (device >= 0) { + auto created = CUDAGuard::create(device); + if (created.ok()) { + guard.emplace(std::move(created.get())); + } else { + ET_LOG( + Error, + "mutable_state: freeing pointer %p without switching to device %d", + ptr, + device); + } } if (synchronize) { const cudaError_t sync_err = cudaDeviceSynchronize(); @@ -333,8 +308,8 @@ Error build_descriptors(Context& c, CudaDelegateHandle* h) { c.discovered_fqns.insert(fqn); if (c.template_ptr.find(fqn) == c.template_ptr.end()) { - CudaDeviceGuard guard; - ET_CHECK_OK_OR_RETURN_ERROR(guard.set(device)); + auto guard = CUDAGuard::create(device); + ET_CHECK_OK_OR_RETURN_ERROR(guard.error()); void* tpl = nullptr; if (cudaMalloc(&tpl, t->nbytes()) != cudaSuccess) { @@ -372,8 +347,8 @@ Error ensure_session_buffers(Context& c, int token) { ET_LOG(Error, "mutable_state: no template device for '%s'", fqn.c_str()); return Error::Internal; } - CudaDeviceGuard guard; - ET_CHECK_OK_OR_RETURN_ERROR(guard.set(device_it->second)); + auto guard = CUDAGuard::create(device_it->second); + ET_CHECK_OK_OR_RETURN_ERROR(guard.error()); void* p = nullptr; if (cudaMalloc(&p, nbytes) != cudaSuccess) { diff --git a/extension/cuda/CMakeLists.txt b/extension/cuda/CMakeLists.txt index 599186f5383..c059dac6384 100644 --- a/extension/cuda/CMakeLists.txt +++ b/extension/cuda/CMakeLists.txt @@ -24,7 +24,7 @@ endif() # definition across every shared object in the process (see export.h). A static # copy linked into multiple shared libraries would create multiple thread-locals # and silently break the caller-stream handshake. -add_library(extension_cuda SHARED caller_stream.cpp) +add_library(extension_cuda SHARED caller_stream.cpp device_guard.cpp) if(EXECUTORCH_BUILD_SHARED) # Named after what it provides rather than after the target, matching the # other libraries the wheel ships. The target name stays as it is because the @@ -33,15 +33,22 @@ if(EXECUTORCH_BUILD_SHARED) extension_cuda PROPERTIES OUTPUT_NAME executorch_extension_cuda ) executorch_target_soname_policy(extension_cuda) + # Ships beside libexecutorch.so in the wheel's lib/ directory. + executorch_target_shipped_runtime_path(extension_cuda) endif() -# No CUDA headers or libraries: caller_stream.cpp uses cudaStream_t as an opaque -# handle and calls no CUDA function, so it compiles without cuda_runtime.h and -# needs no libcudart. Anything else that links this library and does call CUDA -# links the runtime itself. ROCm is the exception: the handle is a hip type -# there, so the header needs hip's own. +# device_guard.cpp calls into the runtime, so this needs it. Only built under an +# accelerator build, where the toolkit is present by definition. +# caller_stream.cpp still calls nothing and treats its handle as opaque. if(EXECUTORCH_BUILD_ROCM) target_link_libraries(extension_cuda PUBLIC hip::host) +else() + target_link_libraries(extension_cuda PUBLIC CUDA::cudart) endif() +# device_guard.cpp logs, so this needs the ExecuTorch runtime. Taken through the +# helper rather than as an ordinary dependency, because CMake would order the +# static core ahead of the shared runtime and leave this library with its own +# copy of the registries. +executorch_target_link_shared_runtime(extension_cuda) target_include_directories(extension_cuda PUBLIC ${_common_include_directories}) target_compile_options( extension_cuda PUBLIC "$<$:${_common_compile_options}>" diff --git a/extension/cuda/device_guard.cpp b/extension/cuda/device_guard.cpp new file mode 100644 index 00000000000..4a1b1571ed1 --- /dev/null +++ b/extension/cuda/device_guard.cpp @@ -0,0 +1,76 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +#if defined(EXECUTORCH_USE_HIP) +#include +#else +#include +#endif + +namespace executorch::extension::cuda { + +using ::executorch::runtime::Error; +using ::executorch::runtime::Result; + +Result CUDAGuard::create(int device_index) { + CUDAGuard guard; + ET_CHECK_OK_OR_RETURN_ERROR(guard.set_index(device_index)); + return guard; +} + +CUDAGuard::CUDAGuard(CUDAGuard&& other) noexcept + : original_device_index_(other.original_device_index_), + current_device_index_(other.current_device_index_) { + other.original_device_index_ = other.current_device_index_; +} + +CUDAGuard::~CUDAGuard() { + // Keyed on what was asked for, not on whether the selection reported success. + // A failed cudaSetDevice can still move the calling thread, which is what + // happens once an error is already pending on the device. + if (original_device_index_ != current_device_index_) { + cudaError_t err = cudaSetDevice(original_device_index_); + if (err != cudaSuccess) { + ET_LOG( + Error, + "~CUDAGuard: failed to restore device to %d: %s", + original_device_index_, + cudaGetErrorString(err)); + } + } +} + +Error CUDAGuard::set_index(int device_index) { + int current = -1; + cudaError_t err = cudaGetDevice(¤t); + if (err != cudaSuccess) { + ET_LOG( + Error, "CUDAGuard: cudaGetDevice failed: %s", cudaGetErrorString(err)); + return Error::Internal; + } + original_device_index_ = current; + current_device_index_ = device_index; + if (current_device_index_ != original_device_index_) { + err = cudaSetDevice(current_device_index_); + if (err != cudaSuccess) { + ET_LOG( + Error, + "CUDAGuard: cudaSetDevice(%d) failed: %s", + device_index, + cudaGetErrorString(err)); + return Error::Internal; + } + } + return Error::Ok; +} + +} // namespace executorch::extension::cuda diff --git a/extension/cuda/device_guard.h b/extension/cuda/device_guard.h new file mode 100644 index 00000000000..a67cf1d5923 --- /dev/null +++ b/extension/cuda/device_guard.h @@ -0,0 +1,77 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include + +namespace executorch::extension::cuda { + +/** + * Selects a CUDA device for its own lifetime and puts the caller's device back + * on destruction. + * + * The current device belongs to whoever called in, so a backend that needs its + * own device borrows it rather than keeping it. Kernel launches follow the + * stream they are given, but allocation follows the current device, so a + * backend that allocates has to select one even when the caller chose the + * stream. + * + * Backend-neutral for the reason CallerStreamGuard is: one program may run + * several backends whose engines sit on different devices, so each has to leave + * the device as it found it, or the next one inherits a selection it never + * made. + * + * No CUDA header is reached for here. Nothing in this interface names a CUDA + * type, so a consumer can include this against the wheel alone, the same as + * every other header it ships. + */ +class EXECUTORCH_EXTENSION_CUDA_API CUDAGuard { + public: + /** + * Borrows @p device_index for the returned guard's lifetime. The result + * carries the failure instead of the guard, so a caller cannot use a guard + * that never selected anything. + */ + static ::executorch::runtime::Result create(int device_index); + + /** Restores the device that was current when this guard selected its own. */ + ~CUDAGuard(); + + /** + * Takes over the borrow. The moved-from guard is left looking already + * restored, so only one of the two puts the device back. + */ + CUDAGuard(CUDAGuard&& other) noexcept; + CUDAGuard& operator=(CUDAGuard&&) = delete; + CUDAGuard(const CUDAGuard&) = delete; + CUDAGuard& operator=(const CUDAGuard&) = delete; + + /** Selects @p device_index, recording the one it replaced. */ + ::executorch::runtime::Error set_index(int device_index); + + /** The device that was current before this guard selected its own. */ + int original_device() const { + return original_device_index_; + } + + /** The device this guard selected. */ + int current_device() const { + return current_device_index_; + } + + private: + CUDAGuard() = default; + + int original_device_index_{-1}; + int current_device_index_{-1}; +}; + +} // namespace executorch::extension::cuda diff --git a/extension/cuda/targets.bzl b/extension/cuda/targets.bzl index a3b06c3d47d..a63fc7616c9 100644 --- a/extension/cuda/targets.bzl +++ b/extension/cuda/targets.bzl @@ -29,13 +29,16 @@ def define_common_targets(): name = "caller_stream", srcs = [ "caller_stream.cpp", + "device_guard.cpp", ], exported_headers = [ "caller_stream.h", + "device_guard.h", "export.h", ], exported_deps = [ ":runtime_api", + "//executorch/runtime/core:core", ], # Opt out of the OSS force_static default so consumers *can* link one # shared instance and keep the thread-local unique (see above); the diff --git a/setup.py b/setup.py index 8135d5ba190..2f8e750f349 100644 --- a/setup.py +++ b/setup.py @@ -2300,15 +2300,21 @@ def run(self): "devtools/etdump/utils.h", "devtools/etdump/data_sinks/", ] + ( - # The CUDA stream helper's public header, and the export macros it includes. Its library is - # shared so the process has one copy of the caller-stream state, and that is a handshake the - # caller takes part in, so a consumer needs the declarations to take part at all. + # The CUDA stream helper's public header, the device guard beside it, and the export macros + # they include. The stream helper's library is shared so the process has one copy of the + # caller-stream state, and that is a handshake the caller takes part in, so a consumer needs + # the declarations to take part at all. The device guard's own definitions are compiled into that + # same library, so a consumer needs this header to reach them. # # Only when this wheel carries the CUDA delegate, and decided from the same CMake cache the # libraries ship on. Keying it off the release row's CUDA version instead meant a build on # an unrecognised toolkit shipped both CUDA libraries and both CMake components with no # header, so a consumer got a component it could link and not include. - ["extension/cuda/caller_stream.h", "extension/cuda/export.h"] + [ + "extension/cuda/caller_stream.h", + "extension/cuda/device_guard.h", + "extension/cuda/export.h", + ] if _cuda_libraries_built(cmake_cache_dir) else [] ): diff --git a/tools/cmake/executorch-wheel-config.cmake b/tools/cmake/executorch-wheel-config.cmake index c668f97becb..0cfdf80b370 100644 --- a/tools/cmake/executorch-wheel-config.cmake +++ b/tools/cmake/executorch-wheel-config.cmake @@ -78,7 +78,8 @@ # executorch::kernels_torchao The TorchAO kernels. Linux and macOS on # aarch64 only. # executorch::backend_cuda The CUDA delegate. Linux only. -# executorch::extension_cuda The CUDA stream extension. Linux only. +# executorch::extension_cuda The CUDA stream and device helpers. Linux +# only. # executorch::backend_openvino The OpenVINO delegate. Linux only. Opens the # OpenVINO runtime by name, which a C++ program # installs and points OPENVINO_LIB_PATH at.