Fix: Fix crash in Debug build with multi-GPU due to forced cudaSetDevice(0) (#6498)

Signed-off-by：Tianxiang Wang<[email protected]>, Contributed under MetaX Integrated Circuits (Shanghai) Co., Ltd.

Author: wangtianxiang

source/source_base/module_device/device.h

@@ -65,8 +65,8 @@ std::string get_device_flag(const std::string& device,
 /**
  * @brief Get the rank of current node
  *        Note that GPU can only be binded with CPU in the same node
- * 
- * @return int 
+ *
+ * @return int
  */
 int get_node_rank();
 int get_node_rank_with_mpi_shared(const MPI_Comm mpi_comm = MPI_COMM_WORLD);
@@ -91,6 +91,14 @@ void record_device_memory(const Device* dev, std::ofstream& ofs_device, std::str
     return;
 }
 
+#if defined(__CUDA) || defined(__ROCM)
+template <>
+void print_device_info<base_device::DEVICE_GPU>(const base_device::DEVICE_GPU *ctx, std::ofstream &ofs_device);
+
+template <>
+void record_device_memory<base_device::DEVICE_GPU>(const base_device::DEVICE_GPU* dev, std::ofstream& ofs_device, std::string str, size_t size);
+#endif
+
 } // end of namespace information
 } // end of namespace base_device
 

source/source_base/module_device/output_device.cpp

@@ -190,7 +190,7 @@ void print_device_info<base_device::DEVICE_GPU>(
     ofs_device << "Detected " << deviceCount << " CUDA Capable device(s)\n";
   }
   int dev = 0, driverVersion = 0, runtimeVersion = 0;
-  cudaErrcheck(cudaSetDevice(dev));
+  cudaErrcheck(cudaGetDevice(&dev));
   cudaDeviceProp deviceProp;
   cudaErrcheck(cudaGetDeviceProperties(&deviceProp, dev));
   ofs_device << "\nDevice " << dev << ":\t " << deviceProp.name << std::endl;
@@ -429,7 +429,7 @@ void print_device_info<base_device::DEVICE_GPU>(
     ofs_device << "Detected " << deviceCount << " CUDA Capable device(s)\n";
   }
   int dev = 0, driverVersion = 0, runtimeVersion = 0;
-  hipErrcheck(hipSetDevice(dev));
+  hipErrcheck(hipGetDevice(&dev));
   hipDeviceProp_t deviceProp;
   hipErrcheck(hipGetDeviceProperties(&deviceProp, dev));
   ofs_device << "\nDevice " << dev << ":\t " << deviceProp.name << std::endl;