Implement device=auto with safe GPU probe and MPI broadcast

- Add probe_gpu_availability() function that safely checks GPU without exiting
- Refactor get_device_flag() to support device=auto with clear user feedback
- Implement MPI broadcast for consistent device selection across all ranks
- Add safety check in get_device_kpar() to prevent GPU calls when unavailable
- Add informative messages for device selection (INFO/WARNING)

Co-authored-by: dzzz2001 <153698752+dzzz2001@users.noreply.github.com>

Author: Copilot

source/source_base/module_device/device.cpp

@@ -147,58 +147,110 @@ int set_device_by_rank(const MPI_Comm mpi_comm) {
 
 #endif
 
-std::string get_device_flag(const std::string &device,
-                            const std::string &basis_type) {
-if (device == "cpu") {
-  return "cpu"; // no extra checks required
-}
-std::string error_message;
-if (device != "auto" and device != "gpu")
-{
-  error_message += "Parameter \"device\" can only be set to \"cpu\" or \"gpu\"!";
-  ModuleBase::WARNING_QUIT("device", error_message);
-}
-
-// Get available GPU count
-int device_count = -1;
-#if ((defined __CUDA) || (defined __ROCM))
+bool probe_gpu_availability() {
 #if defined(__CUDA)
-cudaGetDeviceCount(&device_count);
+    int device_count = 0;
+    // Directly call cudaGetDeviceCount without cudaErrcheck to prevent program exit
+    cudaError_t error_id = cudaGetDeviceCount(&device_count);
+    if (error_id == cudaSuccess && device_count > 0) {
+        return true;
+    }
+    return false;
 #elif defined(__ROCM)
-hipGetDeviceCount(&device_count);
-/***auto start_time = std::chrono::high_resolution_clock::now();
-std::cout << "Starting hipGetDeviceCount.." << std::endl;
-auto end_time = std::chrono::high_resolution_clock::now();
-auto duration = std::chrono::duration_cast<std::chrono::duration<double>>(end_time - start_time);
-std::cout << "hipGetDeviceCount took " << duration.count() << "seconds" << std::endl;***/
+    int device_count = 0;
+    hipError_t error_id = hipGetDeviceCount(&device_count);
+    if (error_id == hipSuccess && device_count > 0) {
+        return true;
+    }
+    return false;
+#else
+    // If not compiled with GPU support, GPU is not available
+    return false;
 #endif
-if (device_count <= 0)
-{
-  error_message += "Cannot find GPU on this computer!\n";
 }
-#else // CPU only
-error_message += "ABACUS is built with CPU support only. Please rebuild with GPU support.\n";
+
+std::string get_device_flag(const std::string &device,
+                            const std::string &basis_type) {
+    // 1. Validate input string
+    if (device != "cpu" && device != "gpu" && device != "auto") {
+        ModuleBase::WARNING_QUIT("device", "Parameter \"device\" can only be set to \"cpu\", \"gpu\", or \"auto\"!");
+    }
+    
+    int decision = 0; // 0 for CPU, 1 for GPU
+
+#ifdef __MPI
+    int world_rank = 0;
+    MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
+    
+    if (world_rank == 0) {
+        // Rank 0 makes the decision
+        if (device == "gpu") {
+            if (probe_gpu_availability()) {
+                decision = 1;
+                std::cout << " INFO: 'device=gpu' specified. GPU will be used." << std::endl;
+            } else {
+                ModuleBase::WARNING_QUIT("device", "Device is set to 'gpu', but no available GPU was found. Please check your hardware/drivers or set 'device=cpu'.");
+            }
+        } else if (device == "auto") {
+            if (probe_gpu_availability()) {
+                decision = 1;
+                std::cout << " INFO: 'device=auto' specified. GPU detected and will be used." << std::endl;
+            } else {
+                decision = 0;
+                std::cout << " WARNING: 'device=auto' specified, but no GPU was found. Falling back to CPU." << std::endl;
+                std::cout << "          To suppress this warning, please explicitly set 'device=cpu' in your input." << std::endl;
+            }
+        } else { // device == "cpu"
+            decision = 0;
+            std::cout << " INFO: 'device=cpu' specified. CPU will be used." << std::endl;
+        }
+    }
+    
+    // Rank 0 broadcasts the final decision to all other ranks
+    MPI_Bcast(&decision, 1, MPI_INT, 0, MPI_COMM_WORLD);
+#else
+    // Non-MPI case: single process makes the decision
+    if (device == "gpu") {
+        if (probe_gpu_availability()) {
+            decision = 1;
+            std::cout << " INFO: 'device=gpu' specified. GPU will be used." << std::endl;
+        } else {
+            ModuleBase::WARNING_QUIT("device", "Device is set to 'gpu', but no available GPU was found. Please check your hardware/drivers or set 'device=cpu'.");
+        }
+    } else if (device == "auto") {
+        if (probe_gpu_availability()) {
+            decision = 1;
+            std::cout << " INFO: 'device=auto' specified. GPU detected and will be used." << std::endl;
+        } else {
+            decision = 0;
+            std::cout << " WARNING: 'device=auto' specified, but no GPU was found. Falling back to CPU." << std::endl;
+            std::cout << "          To suppress this warning, please explicitly set 'device=cpu' in your input." << std::endl;
+        }
+    } else { // device == "cpu"
+        decision = 0;
+        std::cout << " INFO: 'device=cpu' specified. CPU will be used." << std::endl;
+    }
 #endif
 
-if (basis_type == "lcao_in_pw") {
-  error_message +=
-      "The GPU currently does not support the basis type \"lcao_in_pw\"!";
-}
-if(error_message.empty())
-{
-  return "gpu"; // possibly automatically set to GPU
-}
-else if (device == "gpu")
-{
-  ModuleBase::WARNING_QUIT("device", error_message);
-}
-else { return "cpu";
-}
+    // 2. Final check for incompatible basis type
+    if (decision == 1 && basis_type == "lcao_in_pw") {
+        ModuleBase::WARNING_QUIT("device", "The GPU currently does not support the basis type \"lcao_in_pw\"!");
+    }
+
+    // 3. Return the final decision
+    return (decision == 1) ? "gpu" : "cpu";
 }
 
 int get_device_kpar(const int& kpar, const int& bndpar)
 {
 #if __MPI && (__CUDA || __ROCM)
+    // This function should only be called when GPU mode is active
+    // We use probe_gpu_availability to ensure GPU is actually available
+    if (!probe_gpu_availability()) {
+        // If no GPU available, return kpar unchanged
+        return kpar;
+    }
+
     int temp_nproc = 0;
     int new_kpar = kpar;
     MPI_Comm_size(MPI_COMM_WORLD, &temp_nproc);
@@ -213,15 +265,15 @@ int get_device_kpar(const int& kpar, const int& bndpar)
 
     int device_num = -1;
 #if defined(__CUDA)
-  cudaGetDeviceCount(&device_num); // get the number of GPU devices of current node
-  cudaSetDevice(node_rank % device_num); // band the CPU processor to the devices
+    cudaErrcheck(cudaGetDeviceCount(&device_num)); // get the number of GPU devices of current node
+    cudaErrcheck(cudaSetDevice(node_rank % device_num)); // bind the CPU processor to the devices
 #elif defined(__ROCM)
-  hipGetDeviceCount(&device_num);
-  hipSetDevice(node_rank % device_num);
+    hipErrcheck(hipGetDeviceCount(&device_num));
+    hipErrcheck(hipSetDevice(node_rank % device_num));
 #endif
-  return new_kpar;
+    return new_kpar;
 #endif
-  return kpar;
+    return kpar;
 }
 
 } // end of namespace information

source/source_base/module_device/device.h

@@ -44,6 +44,12 @@ void output_device_info(std::ostream& output);
  */
 int get_device_kpar(const int& kpar, const int& bndpar);
 
+/**
+ * @brief Safely probes for GPU availability without exiting on error.
+ * @return True if at least one GPU is found and usable, false otherwise.
+ */
+bool probe_gpu_availability();
+
 /**
  * @brief Get the device flag object
  * for source_io PARAM.inp.device