From bb5339d628107dde7fa4789a0adfd41f8e14ee72 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 24 Aug 2026 10:38:26 -0700 Subject: [PATCH 1/2] Stop ignoring a failed cudaSetDevice in CudaAllocator 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. 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 only shows up much later as a wrong result or an unrelated CUDA error. Now allocate() and the copy helpers log the CUDA error and return Error::MemoryAllocationFailed and Error::Internal instead of going ahead. deallocate() returns void and cudaFree works on a pointer from any device under unified addressing, so it logs the error and still frees rather than leaking. Example, on a machine with one GPU: ``` auto r = CudaAllocator::instance().allocate(1024, /*index=*/1); // before: r.ok() == true, pointer actually on device 0 // after: r.ok() == false, r.error() == Error::MemoryAllocationFailed ``` Test plan: Added three tests to test_cuda_allocator.cpp that 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 Built and ran the suite on one NVIDIA H100. With this change: [ PASSED ] 9 tests. 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 showed the old behavior directly: the call reported success for a device that does not exist. The logged message on the fixed build: CudaAllocator::allocate: cudaSetDevice(1) failed: invalid device ordinal clang-format reports no changes needed on either touched file. Not covered: the deallocate path only logs, so no test asserts on it. --- backends/cuda/runtime/cuda_allocator.cpp | 36 +++++++++++-- .../cuda/runtime/test/test_cuda_allocator.cpp | 52 +++++++++++++++++-- 2 files changed, 82 insertions(+), 6 deletions(-) diff --git a/backends/cuda/runtime/cuda_allocator.cpp b/backends/cuda/runtime/cuda_allocator.cpp index 3d81d058684..b0a930acd23 100644 --- a/backends/cuda/runtime/cuda_allocator.cpp +++ b/backends/cuda/runtime/cuda_allocator.cpp @@ -66,7 +66,18 @@ Error copy_impl( if (index >= 0) { prev_device_err = cudaGetDevice(&prev_device); if (prev_device_err == cudaSuccess) { - (void)cudaSetDevice(index); + 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(index), + cudaGetErrorString(set_err)); + return Error::Internal; + } } } cudaError_t err = cudaSuccess; @@ -137,7 +148,17 @@ CudaAllocator::allocate(size_t nbytes, DeviceIndex index, size_t alignment) { const bool switch_device = index >= 0 && prev_device_err == cudaSuccess && static_cast(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. + ET_LOG( + Error, + "CudaAllocator::allocate: cudaSetDevice(%d) failed: %s", + static_cast(index), + cudaGetErrorString(set_err)); + return Error::MemoryAllocationFailed; + } } cudaError_t err = cudaMalloc(&ptr, nbytes); @@ -183,7 +204,16 @@ void CudaAllocator::deallocate(void* ptr, DeviceIndex index) { if (index >= 0) { prev_device_err = cudaGetDevice(&prev_device); if (prev_device_err == cudaSuccess) { - (void)cudaSetDevice(index); + cudaError_t set_err = cudaSetDevice(index); + if (set_err != cudaSuccess) { + // cudaFree accepts a pointer from any device under unified addressing, + // so keep going rather than leak it, but do not stay silent about it. + ET_LOG( + Error, + "CudaAllocator::deallocate: cudaSetDevice(%d) failed: %s", + static_cast(index), + cudaGetErrorString(set_err)); + } } } diff --git a/backends/cuda/runtime/test/test_cuda_allocator.cpp b/backends/cuda/runtime/test/test_cuda_allocator.cpp index baf9889df77..18353846d0e 100644 --- a/backends/cuda/runtime/test/test_cuda_allocator.cpp +++ b/backends/cuda/runtime/test/test_cuda_allocator.cpp @@ -20,18 +20,25 @@ 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. + DeviceIndex missing_device() const { + return static_cast(device_count_); + } + + int device_count_ = 0; }; TEST_F(CudaAllocatorTest, CopyRoundtrip) { @@ -112,3 +119,42 @@ TEST_F(CudaAllocatorTest, CopyDeviceToHostNullSrcReturnsInvalidArgument) { << "expected InvalidArgument for null src, got " << static_cast(e); } + +TEST_F(CudaAllocatorTest, AllocateOnMissingDeviceFails) { + CudaAllocator& a = CudaAllocator::instance(); + auto res = a.allocate(1024, missing_device()); + ASSERT_FALSE(res.ok()) << "allocate must not report success for device " + << static_cast(missing_device()) + << ", which does not exist"; + EXPECT_EQ(res.error(), Error::MemoryAllocationFailed); +} + +TEST_F(CudaAllocatorTest, CopyHostToDeviceOnMissingDeviceFails) { + 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 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) { + 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 h(N, 0); + EXPECT_EQ( + a.copy_device_to_host(h.data(), dptr, N, missing_device()), + Error::Internal); + + a.deallocate(dptr, 0); +} From afa82ae94f4e622d8a3eb5b74716059389d65576 Mon Sep 17 00:00:00 2001 From: r Date: Mon, 24 Aug 2026 13:05:15 -0700 Subject: [PATCH 2/2] Report a failed device switch as Internal, and gate the missing-device tests cudaSetDevice and cudaGetDevice failures were reported as MemoryAllocationFailed and InvalidArgument. Neither fits: a device that cannot be selected is not out of memory, and a well-formed ordinal that does not exist is an operational failure rather than a bad argument. Use Internal, which is what Exception.h maps every CUDA error to and what cuda_mutable_state.cpp already returns for this same call. Also fail instead of continuing when cudaGetDevice fails, since without the current device there is no way to switch or restore, and the allocation would land on whatever device happened to be current. The three tests that need a valid-but-absent ordinal now check that one exists in DeviceIndex (int8_t) rather than gating in SetUp, which skipped the six device-0 tests too. The bound is <= max, so 127 runs and 128 skips. --- backends/cuda/runtime/cuda_allocator.cpp | 74 ++++++++++++++----- .../cuda/runtime/test/test_cuda_allocator.cpp | 25 ++++++- 2 files changed, 81 insertions(+), 18 deletions(-) diff --git a/backends/cuda/runtime/cuda_allocator.cpp b/backends/cuda/runtime/cuda_allocator.cpp index b0a930acd23..4c7d6aec288 100644 --- a/backends/cuda/runtime/cuda_allocator.cpp +++ b/backends/cuda/runtime/cuda_allocator.cpp @@ -61,11 +61,23 @@ 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) { + // 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(index) != prev_device) { cudaError_t set_err = cudaSetDevice(index); if (set_err != cudaSuccess) { // Nothing was switched, so there is nothing to restore. Copying now @@ -78,6 +90,7 @@ Error copy_impl( cudaGetErrorString(set_err)); return Error::Internal; } + switched_device = true; } } cudaError_t err = cudaSuccess; @@ -90,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); } @@ -141,23 +154,40 @@ 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(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(index) != prev_device; if (switch_device) { 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. + // 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(index), cudaGetErrorString(set_err)); - return Error::MemoryAllocationFailed; + return Error::Internal; } } @@ -199,27 +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) { + 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(index) != prev_device) { cudaError_t set_err = cudaSetDevice(index); if (set_err != cudaSuccess) { - // cudaFree accepts a pointer from any device under unified addressing, - // so keep going rather than leak it, but do not stay silent about it. + // 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(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); } diff --git a/backends/cuda/runtime/test/test_cuda_allocator.cpp b/backends/cuda/runtime/test/test_cuda_allocator.cpp index 18353846d0e..a872099c4f7 100644 --- a/backends/cuda/runtime/test/test_cuda_allocator.cpp +++ b/backends/cuda/runtime/test/test_cuda_allocator.cpp @@ -11,6 +11,7 @@ #include #include +#include #include #include @@ -34,10 +35,20 @@ class CudaAllocatorTest : public testing::Test { } // 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(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::max(); + } + int device_count_ = 0; }; @@ -121,15 +132,23 @@ TEST_F(CudaAllocatorTest, CopyDeviceToHostNullSrcReturnsInvalidArgument) { } 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(missing_device()) << ", which does not exist"; - EXPECT_EQ(res.error(), Error::MemoryAllocationFailed); + 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); @@ -145,6 +164,10 @@ TEST_F(CudaAllocatorTest, CopyHostToDeviceOnMissingDeviceFails) { } 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);