[PIR] speed up set device id (#69676)

wanghuancoder · web-flow · commit 5c1124b8aa2a · 2024-11-28T10:36:51.000+08:00
* speed up set device id
diff --git a/paddle/phi/backends/gpu/cuda/cuda_info.cc b/paddle/phi/backends/gpu/cuda/cuda_info.cc
@@ -245,16 +245,36 @@ const gpuDeviceProp &GetDeviceProperties(int id) {
 }
 
 void SetDeviceId(int id) {
-  // TODO(qijun): find a better way to cache the cuda device count
-  PADDLE_ENFORCE_LT(id,
-                    GetGPUDeviceCount(),
-                    common::errors::InvalidArgument(
-                        "Device id must be less than GPU count, "
-                        "but received id is: %d. GPU count is: %d.",
-                        id,
-                        GetGPUDeviceCount()));
-  PADDLE_RETRY_CUDA_SUCCESS(cudaSetDevice(id));
-  VLOG(4) << "SetDeviceId " << id;
+  static thread_local bool first_call = true;
+  if (first_call) {
+    PADDLE_ENFORCE_LT(id,
+                      GetGPUDeviceCount(),
+                      common::errors::InvalidArgument(
+                          "Device id must be less than GPU count, "
+                          "but received id is: %d. GPU count is: %d.",
+                          id,
+                          GetGPUDeviceCount()));
+
+    PADDLE_RETRY_CUDA_SUCCESS(cudaSetDevice(id));
+    VLOG(4) << "SetDeviceId " << id;
+    first_call = false;
+    return;
+  }
+
+  int prev_id;
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaGetDevice(&prev_id));
+  if (prev_id != id) {
+    PADDLE_ENFORCE_LT(id,
+                      GetGPUDeviceCount(),
+                      common::errors::InvalidArgument(
+                          "Device id must be less than GPU count, "
+                          "but received id is: %d. GPU count is: %d.",
+                          id,
+                          GetGPUDeviceCount()));
+
+    PADDLE_RETRY_CUDA_SUCCESS(cudaSetDevice(id));
+    VLOG(4) << "SetDeviceId " << id;
+  }
 }
 
 void GpuMemcpyAsync(void *dst,
diff --git a/paddle/phi/backends/gpu/rocm/rocm_info.cc b/paddle/phi/backends/gpu/rocm/rocm_info.cc
@@ -241,15 +241,36 @@ const gpuDeviceProp &GetDeviceProperties(int id) {
 }
 
 void SetDeviceId(int id) {
-  // TODO(qijun): find a better way to cache the cuda device count
-  PADDLE_ENFORCE_LT(id,
-                    GetGPUDeviceCount(),
-                    common::errors::InvalidArgument(
-                        "Device id must be less than GPU count, "
-                        "but received id is: %d. GPU count is: %d.",
-                        id,
-                        GetGPUDeviceCount()));
-  PADDLE_RETRY_CUDA_SUCCESS(hipSetDevice(id));
+  static thread_local bool first_call = true;
+  if (first_call) {
+    PADDLE_ENFORCE_LT(id,
+                      GetGPUDeviceCount(),
+                      common::errors::InvalidArgument(
+                          "Device id must be less than GPU count, "
+                          "but received id is: %d. GPU count is: %d.",
+                          id,
+                          GetGPUDeviceCount()));
+
+    PADDLE_RETRY_CUDA_SUCCESS(hipSetDevice(id));
+    VLOG(4) << "SetDeviceId " << id;
+    first_call = false;
+    return;
+  }
+
+  int prev_id;
+  PADDLE_ENFORCE_GPU_SUCCESS(hipGetDevice(&prev_id));
+  if (prev_id != id) {
+    PADDLE_ENFORCE_LT(id,
+                      GetGPUDeviceCount(),
+                      common::errors::InvalidArgument(
+                          "Device id must be less than GPU count, "
+                          "but received id is: %d. GPU count is: %d.",
+                          id,
+                          GetGPUDeviceCount()));
+
+    PADDLE_RETRY_CUDA_SUCCESS(hipSetDevice(id));
+    VLOG(4) << "SetDeviceId " << id;
+  }
 }
 
 void GpuMemcpyAsync(void *dst,
