|
1 | | - |
2 | 1 | #include "device.h" |
3 | 2 |
|
4 | 3 | #include "source_base/tool_quit.h" |
@@ -176,69 +175,40 @@ std::string get_device_flag(const std::string &device, |
176 | 175 | ModuleBase::WARNING_QUIT("device", "Parameter \"device\" can only be set to \"cpu\", \"gpu\", or \"auto\"!"); |
177 | 176 | } |
178 | 177 |
|
179 | | - int decision = 0; // 0 for CPU, 1 for GPU |
180 | | - |
181 | | -#ifdef __MPI |
182 | | - int world_rank = 0; |
183 | | - MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); |
| 178 | + // NOTE: This function is called only on rank 0 during input parsing. |
| 179 | + // The result will be broadcast to other ranks via the standard bcast mechanism. |
| 180 | + // DO NOT use MPI_Bcast here as other ranks are not in this code path. |
184 | 181 |
|
185 | | - if (world_rank == 0) { |
186 | | - // Rank 0 makes the decision |
187 | | - if (device == "gpu") { |
188 | | - if (probe_gpu_availability()) { |
189 | | - decision = 1; |
190 | | - std::cout << " INFO: 'device=gpu' specified. GPU will be used." << std::endl; |
191 | | - } else { |
192 | | - ModuleBase::WARNING_QUIT("device", "Device is set to 'gpu', but no available GPU was found. Please check your hardware/drivers or set 'device=cpu'."); |
193 | | - } |
194 | | - } else if (device == "auto") { |
195 | | - if (probe_gpu_availability()) { |
196 | | - decision = 1; |
197 | | - std::cout << " INFO: 'device=auto' specified. GPU detected and will be used." << std::endl; |
198 | | - } else { |
199 | | - decision = 0; |
200 | | - std::cout << " WARNING: 'device=auto' specified, but no GPU was found. Falling back to CPU." << std::endl; |
201 | | - std::cout << " To suppress this warning, please explicitly set 'device=cpu' in your input." << std::endl; |
202 | | - } |
203 | | - } else { // device == "cpu" |
204 | | - decision = 0; |
205 | | - std::cout << " INFO: 'device=cpu' specified. CPU will be used." << std::endl; |
206 | | - } |
207 | | - } |
| 182 | + std::string result = "cpu"; |
208 | 183 |
|
209 | | - // Rank 0 broadcasts the final decision to all other ranks |
210 | | - MPI_Bcast(&decision, 1, MPI_INT, 0, MPI_COMM_WORLD); |
211 | | -#else |
212 | | - // Non-MPI case: single process makes the decision |
213 | 184 | if (device == "gpu") { |
214 | 185 | if (probe_gpu_availability()) { |
215 | | - decision = 1; |
216 | | - std::cout << " INFO: 'device=gpu' specified. GPU will be used." << std::endl; |
| 186 | + result = "gpu"; |
| 187 | + // std::cout << " INFO: 'device=gpu' specified. GPU will be used." << std::endl; |
217 | 188 | } else { |
218 | 189 | ModuleBase::WARNING_QUIT("device", "Device is set to 'gpu', but no available GPU was found. Please check your hardware/drivers or set 'device=cpu'."); |
219 | 190 | } |
220 | 191 | } else if (device == "auto") { |
221 | 192 | if (probe_gpu_availability()) { |
222 | | - decision = 1; |
223 | | - std::cout << " INFO: 'device=auto' specified. GPU detected and will be used." << std::endl; |
| 193 | + result = "gpu"; |
| 194 | + // std::cout << " INFO: 'device=auto' specified. GPU detected and will be used." << std::endl; |
224 | 195 | } else { |
225 | | - decision = 0; |
226 | | - std::cout << " WARNING: 'device=auto' specified, but no GPU was found. Falling back to CPU." << std::endl; |
227 | | - std::cout << " To suppress this warning, please explicitly set 'device=cpu' in your input." << std::endl; |
| 196 | + result = "cpu"; |
| 197 | + // std::cout << " WARNING: 'device=auto' specified, but no GPU was found. Falling back to CPU." << std::endl; |
| 198 | + // std::cout << " To suppress this warning, please explicitly set 'device=cpu' in your input." << std::endl; |
228 | 199 | } |
229 | 200 | } else { // device == "cpu" |
230 | | - decision = 0; |
231 | | - std::cout << " INFO: 'device=cpu' specified. CPU will be used." << std::endl; |
| 201 | + result = "cpu"; |
| 202 | + // std::cout << " INFO: 'device=cpu' specified. CPU will be used." << std::endl; |
232 | 203 | } |
233 | | -#endif |
234 | 204 |
|
235 | 205 | // 2. Final check for incompatible basis type |
236 | | - if (decision == 1 && basis_type == "lcao_in_pw") { |
| 206 | + if (result == "gpu" && basis_type == "lcao_in_pw") { |
237 | 207 | ModuleBase::WARNING_QUIT("device", "The GPU currently does not support the basis type \"lcao_in_pw\"!"); |
238 | 208 | } |
239 | 209 |
|
240 | 210 | // 3. Return the final decision |
241 | | - return (decision == 1) ? "gpu" : "cpu"; |
| 211 | + return result; |
242 | 212 | } |
243 | 213 |
|
244 | 214 | int get_device_kpar(const int& kpar, const int& bndpar) |
|
0 commit comments