Share the CUDA device guard through the extension so other backends can use it

.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 }}

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

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<int>(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<DeviceIndex>(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> 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_),

backends/aoti/slim/cuda/guard.h

@@ -12,12 +12,15 @@
 
 #include <executorch/backends/aoti/slim/c10/core/Device.h>
 #include <executorch/backends/aoti/slim/c10/cuda/Exception.h>
+#include <executorch/extension/cuda/device_guard.h>
 #include <executorch/extension/cuda/runtime_api.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/result.h>
 
 namespace executorch::backends::cuda {
 
+using ::executorch::extension::cuda::CUDAGuard;
+
 using executorch::runtime::Error;
 using executorch::runtime::Result;
 
@@ -66,75 +69,6 @@ std::optional<cudaStream_t> 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<CUDAGuard> 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

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",

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<CUDAGuard>,

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()
 

backends/cuda/runtime/cuda_delegate_handle.h

@@ -10,6 +10,7 @@
 
 #include <executorch/backends/aoti/aoti_delegate_handle.h>
 #include <executorch/backends/aoti/slim/core/slim_tensor.h>
+#include <executorch/extension/cuda/device_guard.h>
 #include <executorch/extension/cuda/runtime_api.h>
 #include <cstdint>
 #include <cstdlib>
@@ -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;

backends/cuda/runtime/cuda_mutable_state.cpp

@@ -13,12 +13,14 @@
 #include <executorch/backends/aoti/slim/core/slim_tensor.h>
 #include <executorch/backends/aoti/slim/factory/from_blob.h>
 #include <executorch/backends/cuda/runtime/cuda_delegate_handle.h>
+#include <executorch/extension/cuda/device_guard.h>
 #include <executorch/extension/cuda/runtime_api.h>
 #include <executorch/runtime/platform/log.h>
 
 #include <iterator>
 #include <memory>
 #include <mutex>
+#include <optional>
 #include <unordered_map>
 #include <unordered_set>
 
@@ -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<int> 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<CUDAGuard> 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) {