Stop ignoring a failed cudaSetDevice in CudaAllocator

backends/cuda/runtime/cuda_allocator.cpp

@@ -61,12 +61,36 @@ Error copy_impl(
   }
 
   int prev_device = 0;
-  cudaError_t prev_device_err = cudaSuccess;
+  bool switched_device = false;
 
   if (index >= 0) {
-    prev_device_err = cudaGetDevice(&prev_device);
-    if (prev_device_err == cudaSuccess) {
-      (void)cudaSetDevice(index);
+    // Without the current device there is no way to switch to `index` and no
+    // way to restore afterwards, so copying would run against whatever device
+    // happens to be current and report success. Fail instead, as
+    // cuda_mutable_state.cpp does for the same call.
+    cudaError_t prev_device_err = cudaGetDevice(&prev_device);
+    if (prev_device_err != cudaSuccess) {
+      ET_LOG(
+          Error,
+          "%s: cudaGetDevice failed: %s",
+          method,
+          cudaGetErrorString(prev_device_err));
+      return Error::Internal;
+    }
+    if (static_cast<int>(index) != prev_device) {
+      cudaError_t set_err = cudaSetDevice(index);
+      if (set_err != cudaSuccess) {
+        // Nothing was switched, so there is nothing to restore. Copying now
+        // would silently run against whatever device is still current.
+        ET_LOG(
+            Error,
+            "%s: cudaSetDevice(%d) failed: %s",
+            method,
+            static_cast<int>(index),
+            cudaGetErrorString(set_err));
+        return Error::Internal;
+      }
+      switched_device = true;
     }
   }
   cudaError_t err = cudaSuccess;
@@ -79,7 +103,7 @@ Error copy_impl(
     err = cudaMemcpy(dst, src, nbytes, kind);
   }
 
-  if (index >= 0 && prev_device_err == cudaSuccess) {
+  if (switched_device) {
     (void)cudaSetDevice(prev_device);
   }
 
@@ -130,14 +154,41 @@ CudaAllocator::allocate(size_t nbytes, DeviceIndex index, size_t alignment) {
 
   void* ptr = nullptr;
   int prev_device = 0;
-  cudaError_t prev_device_err = cudaGetDevice(&prev_device);
+  bool switch_device = false;
+
+  // If index == -1, fall back to the current device and skip the set/restore
+  // round-trip.
+  if (index >= 0) {
+    // Without the current device there is no way to switch to `index` and no
+    // way to restore afterwards, so the allocation would land on whatever
+    // device happens to be current while the caller records it as living on
+    // `index`. Fail instead.
+    cudaError_t prev_device_err = cudaGetDevice(&prev_device);
+    if (prev_device_err != cudaSuccess) {
+      ET_LOG(
+          Error,
+          "CudaAllocator::allocate: cudaGetDevice failed: %s",
+          cudaGetErrorString(prev_device_err));
+      return Error::Internal;
+    }
+    switch_device = static_cast<int>(index) != prev_device;
+  }
 
-  // If index == -1, fall back to the current device returned by cudaGetDevice
-  // and skip the set/restore round-trip.
-  const bool switch_device = index >= 0 && prev_device_err == cudaSuccess &&
-      static_cast<int>(index) != prev_device;
   if (switch_device) {
-    (void)cudaSetDevice(index);
+    cudaError_t set_err = cudaSetDevice(index);
+    if (set_err != cudaSuccess) {
+      // Allocating now would return a pointer on the current device while the
+      // caller records it as living on the requested one. cudaSetDevice reports
+      // more than a bad ordinal here (a valid device can be unavailable or in
+      // prohibited mode), so report it the way the rest of the CUDA runtime
+      // code does rather than blaming the caller's argument.
+      ET_LOG(
+          Error,
+          "CudaAllocator::allocate: cudaSetDevice(%d) failed: %s",
+          static_cast<int>(index),
+          cudaGetErrorString(set_err));
+      return Error::Internal;
+    }
   }
 
   cudaError_t err = cudaMalloc(&ptr, nbytes);
@@ -178,18 +229,37 @@ void CudaAllocator::deallocate(void* ptr, DeviceIndex index) {
   }
 
   int prev_device = 0;
-  cudaError_t prev_device_err = cudaSuccess;
+  bool switched_device = false;
 
   if (index >= 0) {
-    prev_device_err = cudaGetDevice(&prev_device);
-    if (prev_device_err == cudaSuccess) {
-      (void)cudaSetDevice(index);
+    cudaError_t prev_device_err = cudaGetDevice(&prev_device);
+    if (prev_device_err != cudaSuccess) {
+      // cudaFree accepts a pointer from any device under unified addressing, so
+      // free it anyway rather than leak, but do not try to restore a device we
+      // never read.
+      ET_LOG(
+          Error,
+          "CudaAllocator::deallocate: cudaGetDevice failed: %s",
+          cudaGetErrorString(prev_device_err));
+    } else if (static_cast<int>(index) != prev_device) {
+      cudaError_t set_err = cudaSetDevice(index);
+      if (set_err != cudaSuccess) {
+        // Same reasoning: keep going rather than leak it, but do not stay
+        // silent about it.
+        ET_LOG(
+            Error,
+            "CudaAllocator::deallocate: cudaSetDevice(%d) failed: %s",
+            static_cast<int>(index),
+            cudaGetErrorString(set_err));
+      } else {
+        switched_device = true;
+      }
     }
   }
 
   cudaError_t err = cudaFree(ptr);
 
-  if (index >= 0 && prev_device_err == cudaSuccess) {
+  if (switched_device) {
     (void)cudaSetDevice(prev_device);
   }
 

backends/cuda/runtime/test/test_cuda_allocator.cpp

@@ -11,6 +11,7 @@
 #include <executorch/extension/cuda/runtime_api.h>
 
 #include <cstdint>
+#include <limits>
 #include <vector>
 
 #include <executorch/backends/cuda/runtime/cuda_allocator.h>
@@ -20,18 +21,35 @@
 
 using executorch::backends::cuda::CudaAllocator;
 using executorch::runtime::Error;
+using executorch::runtime::etensor::DeviceIndex;
 
 class CudaAllocatorTest : public testing::Test {
  protected:
   void SetUp() override {
     et_pal_init();
 
-    int device_count = 0;
-    cudaError_t err = cudaGetDeviceCount(&device_count);
-    if (err != cudaSuccess || device_count == 0) {
+    cudaError_t err = cudaGetDeviceCount(&device_count_);
+    if (err != cudaSuccess || device_count_ == 0) {
       GTEST_SKIP() << "CUDA not available";
     }
   }
+
+  // One past the last valid device ordinal, so switching to it always fails.
+  // Only the tests that need such an ordinal call this, so the fit check lives
+  // here rather than in SetUp, where it would also skip the device-0 tests.
+  DeviceIndex missing_device() const {
+    return static_cast<DeviceIndex>(device_count_);
+  }
+
+  // missing_device() has to stay a valid-but-absent ordinal. DeviceIndex is
+  // int8_t, so on a host with more than 127 visible GPUs the count wraps to a
+  // negative index (which the >= -1 argument check rejects for a different
+  // reason) or, at 256, back onto real device 0.
+  bool missing_device_fits() const {
+    return device_count_ <= std::numeric_limits<DeviceIndex>::max();
+  }
+
+  int device_count_ = 0;
 };
 
 TEST_F(CudaAllocatorTest, CopyRoundtrip) {
@@ -112,3 +130,54 @@ TEST_F(CudaAllocatorTest, CopyDeviceToHostNullSrcReturnsInvalidArgument) {
       << "expected InvalidArgument for null src, got "
       << static_cast<uint32_t>(e);
 }
+
+TEST_F(CudaAllocatorTest, AllocateOnMissingDeviceFails) {
+  if (!missing_device_fits()) {
+    GTEST_SKIP() << "device count " << device_count_
+                 << " leaves no absent ordinal in DeviceIndex";
+  }
+  CudaAllocator& a = CudaAllocator::instance();
+  auto res = a.allocate(1024, missing_device());
+  ASSERT_FALSE(res.ok()) << "allocate must not report success for device "
+                         << static_cast<int>(missing_device())
+                         << ", which does not exist";
+  EXPECT_EQ(res.error(), Error::Internal);
+}
+
+TEST_F(CudaAllocatorTest, CopyHostToDeviceOnMissingDeviceFails) {
+  if (!missing_device_fits()) {
+    GTEST_SKIP() << "device count " << device_count_
+                 << " leaves no absent ordinal in DeviceIndex";
+  }
+  CudaAllocator& a = CudaAllocator::instance();
+  constexpr size_t N = 64;
+  auto res = a.allocate(N, 0);
+  ASSERT_TRUE(res.ok());
+  void* dptr = res.get();
+
+  std::vector<uint8_t> h(N, 7);
+  EXPECT_EQ(
+      a.copy_host_to_device(dptr, h.data(), N, missing_device()),
+      Error::Internal);
+
+  a.deallocate(dptr, 0);
+}
+
+TEST_F(CudaAllocatorTest, CopyDeviceToHostOnMissingDeviceFails) {
+  if (!missing_device_fits()) {
+    GTEST_SKIP() << "device count " << device_count_
+                 << " leaves no absent ordinal in DeviceIndex";
+  }
+  CudaAllocator& a = CudaAllocator::instance();
+  constexpr size_t N = 64;
+  auto res = a.allocate(N, 0);
+  ASSERT_TRUE(res.ok());
+  void* dptr = res.get();
+
+  std::vector<uint8_t> h(N, 0);
+  EXPECT_EQ(
+      a.copy_device_to_host(h.data(), dptr, N, missing_device()),
+      Error::Internal);
+
+  a.deallocate(dptr, 0);
+}