Stop ignoring a failed cudaSetDevice in CudaAllocator (pytorch#22093)

### Summary

`CudaAllocator` switches the current CUDA device before it allocates,
frees or
copies, and switches back afterwards. The result of that switch was
thrown away
in all three places:

```cpp
(void)cudaSetDevice(index);
```

When the switch fails, for example because the requested device index
does not
exist on this machine, the work still went ahead on whatever device
happened to
be current. `allocate` then returned success with a pointer that lives
on a
different device than the caller asked for. The caller stores the
requested
index next to that pointer, so the pointer and its recorded device
disagree
from then on, and the mistake surfaces much later as a wrong result or
an
unrelated CUDA error somewhere else.

On a machine with one GPU:

```cpp
auto r = CudaAllocator::instance().allocate(1024, /*index=*/1);
// before: r.ok() == true, and the pointer is really on device 0
// after:  r.ok() == false, r.error() == Error::Internal
```

What changes:

* `allocate` logs the CUDA error and returns `Error::Internal` instead
of
allocating on the wrong device. A device that cannot be selected is not
an
out-of-memory condition, and `Error::Internal` is what `Exception.h`
maps
every CUDA error to and what `cuda_mutable_state.cpp` already returns
for
  this same call.
* The host to device and device to host copy helpers log the CUDA error
and
  return `Error::Internal` instead of copying against the wrong device.
* All three paths also fail when `cudaGetDevice` fails, instead of
carrying on
without switching. Without the current device there is no way to switch
to the
requested one or restore afterwards, so the work would run against
whatever
device happened to be current. This is the branch the tests cannot
reach: on a
  working CUDA host `cudaGetDevice` does not fail, so it is verified by
  inspection rather than by the new tests.
* `deallocate` returns `void`, and `cudaFree` accepts a pointer from any
device
under unified addressing, so it logs the error and still frees the
pointer
  rather than leaking it. The failure is no longer silent.

Only the switch to the requested device is checked. The switch back to
the
previous device stays best effort, because there is nothing useful to do
if
restoring fails.

### Test plan

Three new tests in `backends/cuda/runtime/test/test_cuda_allocator.cpp`
ask for
device index `device_count`, which is one past the last valid ordinal on
any
machine, so the switch always fails:

* `AllocateOnMissingDeviceFails`
* `CopyHostToDeviceOnMissingDeviceFails`
* `CopyDeviceToHostOnMissingDeviceFails`

`DeviceIndex` is `int8_t`, so each of the three first checks that
one-past-the-last
ordinal still fits in it: at 128 or more visible GPUs the count wraps to
a negative
index, or at 256 back onto real device 0, and neither is a valid absent
ordinal. The
check is per-test rather than in `SetUp` so it does not also skip the
six device-0
tests. 127 devices still runs; 128 skips.

Built and ran the suite on one NVIDIA H100.

With this change:

```
[  PASSED  ] 9 tests.
```

Confirmed the tests really catch the bug. With the allocator reverted to
the
old code and the same three tests:

```
[  FAILED  ] CudaAllocatorTest.AllocateOnMissingDeviceFails
[  FAILED  ] CudaAllocatorTest.CopyHostToDeviceOnMissingDeviceFails
[  FAILED  ] CudaAllocatorTest.CopyDeviceToHostOnMissingDeviceFails
```

The allocate failure shows the old behavior directly: the call reported
success
for a device that does not exist.

The message the fixed build logs:

```
CudaAllocator::allocate: cudaSetDevice(1) failed: invalid device ordinal
```

The six tests already in that file still pass, so the change does not
disturb
the normal single device path.

`clang-format` reports no changes needed on either touched file.

### Not covered

The `cudaGetDevice` failure branch is not covered by these tests. On a
working CUDA
host that call does not fail, and there is no injection point in the
suite, so
reverting just that branch leaves all three tests passing. Reaching it
would need a
function-pointer seam or an injected fake CUDA runtime, both larger than
the guard
they would protect.

`deallocate` only logs, so no test asserts on it. Everything here ran on
a
single GPU machine, so the switch always failed with an invalid ordinal.
A
device that exists but is unavailable, for example one in exclusive
compute
mode, takes the same code path but was not exercised.

---------

Co-authored-by: r <r@e>

Author: shoumikhin

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);
+}