[Runtime] Use cudaGetDeviceCount to check if device exists (#16377)

Using `cudaDeviceGetAttribute` will set the global error code when the device doesn't exist and will impact subsequent CUDA API calls.
apache · Jan 10, 2024 · 524ec5f · 524ec5f
1 parent 3166366
commit 524ec5f
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/runtime/cuda/cuda_device_api.cc b/src/runtime/cuda/cuda_device_api.cc
@@ -42,8 +42,9 @@ class CUDADeviceAPI final : public DeviceAPI {
     int value = 0;
     switch (kind) {
       case kExist:
-        value = (cudaDeviceGetAttribute(&value, cudaDevAttrMaxThreadsPerBlock, dev.device_id) ==
-                 cudaSuccess);
+        int count;
+        CUDA_CALL(cudaGetDeviceCount(&count));
+        value = static_cast<int>(dev.device_id < count);
         break;
       case kMaxThreadsPerBlock: {
         CUDA_CALL(cudaDeviceGetAttribute(&value, cudaDevAttrMaxThreadsPerBlock, dev.device_id));