diff --git a/backends/cuda/runtime/cuda_allocator.cpp b/backends/cuda/runtime/cuda_allocator.cpp index 3d81d058684..4c7d6aec288 100644 --- a/backends/cuda/runtime/cuda_allocator.cpp +++ b/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(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(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(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) { - (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(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(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(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 baf9889df77..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 @@ -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(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; }; TEST_F(CudaAllocatorTest, CopyRoundtrip) { @@ -112,3 +130,54 @@ TEST_F(CudaAllocatorTest, CopyDeviceToHostNullSrcReturnsInvalidArgument) { << "expected InvalidArgument for null src, got " << static_cast(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(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 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 h(N, 0); + EXPECT_EQ( + a.copy_device_to_host(h.data(), dptr, N, missing_device()), + Error::Internal); + + a.deallocate(dptr, 0); +}