Skip to content

Commit 4ef3500

Browse files
Flying-dragon-boxingFlying-dragon-boxing
authored
Refactor: EXX PW now uses the new Parallel_device communication interface (#7312)
* Harden GPU MPI staging helpers * Add NCCL collectives for parallel_device * Fix NCCL headers in parallel_device * Route PGemm collectives through device wrappers * Tighten NCCL collective correctness * Relax NCCL discovery for existing environments * Decouple NCCL parallel_device from CUDA-aware MPI * Propagate NCCL headers to subdirectory targets * Fix: narrow CPU staging guards in para_gemm to respect NCCL collectives isend_dev has no NCCL path — keep guard as #ifndef __CUDA_MPI. reduce_dev / gatherv_dev have NCCL early-returns — exclude CPU staging when __NCCL_PARALLEL_DEVICE is defined (&& !defined). * Refactor(exx_pw): unify GPU/CPU bcast via Parallel_Common::bcast_dev Add configurable root parameter to bcast_data/nccl_bcast_data/bcast_dev (default root=0 for backward compat). Replace manual MPI_Bcast with GPU/CPU branching in EXX PW operator with unified bcast_dev/reduce_dev. * Fix(sdft): update bcast_dev call with root parameter * Fix: add missing include for base_device namespace in parallel_device.h * Fix: ELF on GPU and MPI-disabled compile errors - ELF GPU: use static_cast with Device-typed psi to bypass virtual dispatch mismatch where DEVICE_GPU cal_tau() did not override base class DEVICE_CPU cal_tau(). Meta-GGA path was unaffected because tau is computed in psiToRho during SCF. - ACE EXX: guard Parallel_Common::reduce_dev calls with __MPI since POOL_WORLD and reduce_dev are only declared when MPI is enabled. * Docs: fix bcast_dev doxygen to match actual signature
1 parent 8dfe3bc commit 4ef3500

5 files changed

Lines changed: 49 additions & 85 deletions

File tree

‎source/source_base/parallel_device.cpp‎

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -86,14 +86,14 @@ NcclCommRegistry& get_nccl_registry()
8686
}
8787

8888
template <typename T>
89-
void nccl_bcast_impl(T* object, const int n, MPI_Comm& comm, ncclDataType_t datatype, const int count_scale = 1)
89+
void nccl_bcast_impl(T* object, const int n, MPI_Comm& comm, ncclDataType_t datatype, int root = 0, const int count_scale = 1)
9090
{
9191
NcclCommContext& ctx = get_nccl_registry().get(comm);
9292
if (ctx.size <= 1 || n <= 0)
9393
{
9494
return;
9595
}
96-
CHECK_NCCL(ncclBroadcast(object, object, static_cast<size_t>(n) * count_scale, datatype, 0, ctx.comm, ctx.stream));
96+
CHECK_NCCL(ncclBroadcast(object, object, static_cast<size_t>(n) * count_scale, datatype, root, ctx.comm, ctx.stream));
9797
CHECK_CUDA(cudaStreamSynchronize(ctx.stream));
9898
}
9999

@@ -183,24 +183,24 @@ void nccl_gatherv_impl(const T* sendbuf,
183183
}
184184
} // namespace
185185

186-
void nccl_bcast_data(double* object, const int& n, MPI_Comm& comm)
186+
void nccl_bcast_data(double* object, const int& n, MPI_Comm& comm, int root)
187187
{
188-
nccl_bcast_impl(object, n, comm, ncclDouble);
188+
nccl_bcast_impl(object, n, comm, ncclDouble, root);
189189
}
190190

191-
void nccl_bcast_data(std::complex<double>* object, const int& n, MPI_Comm& comm)
191+
void nccl_bcast_data(std::complex<double>* object, const int& n, MPI_Comm& comm, int root)
192192
{
193-
nccl_bcast_impl(reinterpret_cast<double*>(object), n, comm, ncclDouble, 2);
193+
nccl_bcast_impl(reinterpret_cast<double*>(object), n, comm, ncclDouble, root, 2);
194194
}
195195

196-
void nccl_bcast_data(float* object, const int& n, MPI_Comm& comm)
196+
void nccl_bcast_data(float* object, const int& n, MPI_Comm& comm, int root)
197197
{
198-
nccl_bcast_impl(object, n, comm, ncclFloat);
198+
nccl_bcast_impl(object, n, comm, ncclFloat, root);
199199
}
200200

201-
void nccl_bcast_data(std::complex<float>* object, const int& n, MPI_Comm& comm)
201+
void nccl_bcast_data(std::complex<float>* object, const int& n, MPI_Comm& comm, int root)
202202
{
203-
nccl_bcast_impl(reinterpret_cast<float*>(object), n, comm, ncclFloat, 2);
203+
nccl_bcast_impl(reinterpret_cast<float*>(object), n, comm, ncclFloat, root, 2);
204204
}
205205

206206
void nccl_reduce_data(double* object, const int& n, MPI_Comm& comm)
@@ -302,21 +302,21 @@ void recv_data(std::complex<float>* buf, int count, int source, int tag, MPI_Com
302302
{
303303
MPI_Recv(buf, count, MPI_COMPLEX, source, tag, comm, status);
304304
}
305-
void bcast_data(std::complex<double>* object, const int& n, const MPI_Comm& comm)
305+
void bcast_data(std::complex<double>* object, const int& n, const MPI_Comm& comm, int root)
306306
{
307-
MPI_Bcast(object, n * 2, MPI_DOUBLE, 0, comm);
307+
MPI_Bcast(object, n * 2, MPI_DOUBLE, root, comm);
308308
}
309-
void bcast_data(std::complex<float>* object, const int& n, const MPI_Comm& comm)
309+
void bcast_data(std::complex<float>* object, const int& n, const MPI_Comm& comm, int root)
310310
{
311-
MPI_Bcast(object, n * 2, MPI_FLOAT, 0, comm);
311+
MPI_Bcast(object, n * 2, MPI_FLOAT, root, comm);
312312
}
313-
void bcast_data(double* object, const int& n, const MPI_Comm& comm)
313+
void bcast_data(double* object, const int& n, const MPI_Comm& comm, int root)
314314
{
315-
MPI_Bcast(object, n, MPI_DOUBLE, 0, comm);
315+
MPI_Bcast(object, n, MPI_DOUBLE, root, comm);
316316
}
317-
void bcast_data(float* object, const int& n, const MPI_Comm& comm)
317+
void bcast_data(float* object, const int& n, const MPI_Comm& comm, int root)
318318
{
319-
MPI_Bcast(object, n, MPI_FLOAT, 0, comm);
319+
MPI_Bcast(object, n, MPI_FLOAT, root, comm);
320320
}
321321
void reduce_data(std::complex<double>* object, const int& n, const MPI_Comm& comm)
322322
{

‎source/source_base/parallel_device.h‎

Lines changed: 20 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
#include "mpi.h"
55
#include <complex>
66
#include <type_traits>
7+
#include "source_base/module_device/types.h"
78
namespace Parallel_Common
89
{
910
void isend_data(const double* buf, int count, int dest, int tag, MPI_Comm& comm, MPI_Request* request);
@@ -18,10 +19,10 @@ void recv_data(double* buf, int count, int source, int tag, MPI_Comm& comm, MPI_
1819
void recv_data(std::complex<double>* buf, int count, int source, int tag, MPI_Comm& comm, MPI_Status* status);
1920
void recv_data(float* buf, int count, int source, int tag, MPI_Comm& comm, MPI_Status* status);
2021
void recv_data(std::complex<float>* buf, int count, int source, int tag, MPI_Comm& comm, MPI_Status* status);
21-
void bcast_data(std::complex<double>* object, const int& n, const MPI_Comm& comm);
22-
void bcast_data(std::complex<float>* object, const int& n, const MPI_Comm& comm);
23-
void bcast_data(double* object, const int& n, const MPI_Comm& comm);
24-
void bcast_data(float* object, const int& n, const MPI_Comm& comm);
22+
void bcast_data(std::complex<double>* object, const int& n, const MPI_Comm& comm, int root = 0);
23+
void bcast_data(std::complex<float>* object, const int& n, const MPI_Comm& comm, int root = 0);
24+
void bcast_data(double* object, const int& n, const MPI_Comm& comm, int root = 0);
25+
void bcast_data(float* object, const int& n, const MPI_Comm& comm, int root = 0);
2526
void reduce_data(std::complex<double>* object, const int& n, const MPI_Comm& comm);
2627
void reduce_data(std::complex<float>* object, const int& n, const MPI_Comm& comm);
2728
void reduce_data(double* object, const int& n, const MPI_Comm& comm);
@@ -32,10 +33,10 @@ void gatherv_data(const float* sendbuf, int sendcount, float* recvbuf, const int
3233
void gatherv_data(const std::complex<float>* sendbuf, int sendcount, std::complex<float>* recvbuf, const int* recvcounts, const int* displs, MPI_Comm& comm);
3334

3435
#if defined(__NCCL_PARALLEL_DEVICE)
35-
void nccl_bcast_data(double* object, const int& n, MPI_Comm& comm);
36-
void nccl_bcast_data(std::complex<double>* object, const int& n, MPI_Comm& comm);
37-
void nccl_bcast_data(float* object, const int& n, MPI_Comm& comm);
38-
void nccl_bcast_data(std::complex<float>* object, const int& n, MPI_Comm& comm);
36+
void nccl_bcast_data(double* object, const int& n, MPI_Comm& comm, int root = 0);
37+
void nccl_bcast_data(std::complex<double>* object, const int& n, MPI_Comm& comm, int root = 0);
38+
void nccl_bcast_data(float* object, const int& n, MPI_Comm& comm, int root = 0);
39+
void nccl_bcast_data(std::complex<float>* object, const int& n, MPI_Comm& comm, int root = 0);
3940
void nccl_reduce_data(double* object, const int& n, MPI_Comm& comm);
4041
void nccl_reduce_data(std::complex<double>* object, const int& n, MPI_Comm& comm);
4142
void nccl_reduce_data(float* object, const int& n, MPI_Comm& comm);
@@ -116,35 +117,35 @@ void recv_dev(T* object, int count, int source, int tag, MPI_Comm& comm, MPI_Sta
116117
}
117118

118119
/**
119-
* @brief bcast data in Device
120+
* @brief broadcast data in Device
120121
*
121122
* @tparam T: float, double, std::complex<float>, std::complex<double>
122123
* @tparam Device
123-
* @param ctx Device ctx
124-
* @param object complex arrays in Device
125-
* @param n the size of complex arrays
124+
* @param object arrays in Device
125+
* @param n the size of array
126126
* @param comm MPI_Comm
127-
* @param tmp_space tmp space in CPU
127+
* @param root root rank (default 0)
128+
* @param tmp_space optional tmp space in CPU (default nullptr)
128129
*/
129130
template <typename T, typename Device>
130-
void bcast_dev(T* object, const int& n, const MPI_Comm& comm, T* tmp_space = nullptr)
131+
void bcast_dev(T* object, const int& n, const MPI_Comm& comm, int root = 0, T* tmp_space = nullptr)
131132
{
132133
#if defined(__NCCL_PARALLEL_DEVICE)
133134
if (std::is_same<Device, base_device::DEVICE_GPU>::value)
134135
{
135-
nccl_bcast_data(object, n, const_cast<MPI_Comm&>(comm));
136+
nccl_bcast_data(object, n, const_cast<MPI_Comm&>(comm), root);
136137
return;
137138
}
138139
#endif
139140
#ifdef __CUDA_MPI
140-
bcast_data(object, n, comm);
141+
bcast_data(object, n, comm, root);
141142
#else
142143
object_cpu_point<T,Device> o;
143144
int rank = 0;
144145
MPI_Comm_rank(comm, &rank);
145-
T* object_cpu = rank == 0 ? o.get(object, n, tmp_space) : o.get_buffer(object, n, tmp_space);
146-
bcast_data(object_cpu, n, comm);
147-
if (rank != 0)
146+
T* object_cpu = rank == root ? o.get(object, n, tmp_space) : o.get_buffer(object, n, tmp_space);
147+
bcast_data(object_cpu, n, comm, root);
148+
if (rank != root)
148149
{
149150
o.sync_h2d(object, object_cpu, n);
150151
}

‎source/source_hsolver/hsolver_pw_sdft.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ void HSolverPW_SDFT<T, Device>::solve(const UnitCell& ucell,
6161
#ifdef __MPI
6262
if (nbands > 0 && !PARAM.globalv.all_ks_run)
6363
{
64-
Parallel_Common::bcast_dev<T,Device>(&psi(ik, 0, 0), npwx * nbands, BP_WORLD, &psi_cpu(ik, 0, 0));
64+
Parallel_Common::bcast_dev<T,Device>(&psi(ik, 0, 0), npwx * nbands, BP_WORLD, 0, &psi_cpu(ik, 0, 0));
6565
MPI_Bcast(&pes->ekb(ik, 0), nbands, MPI_DOUBLE, 0, BP_WORLD);
6666
}
6767
#endif

‎source/source_pw/module_pwdft/op_pw_exx.cpp‎

Lines changed: 2 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
#include "source_base/constants.h"
44
#include "source_base/global_variable.h"
55
#include "source_base/parallel_common.h"
6+
#include "source_base/parallel_device.h"
67
#include "source_base/parallel_comm.h" // use KP_WORLD
78
#include "source_base/parallel_reduce.h"
89
#include "source_base/module_external/lapack_connector.h"
@@ -350,27 +351,7 @@ void OperatorEXXPW<T, Device>::act_op_kpar(const int nbands,
350351
// send
351352
}
352353
#ifdef __MPI
353-
#ifdef __CUDA_MPI
354-
MPI_Bcast(psi_mq_real, wfcpw->nrxx, MPI_DOUBLE_COMPLEX, iq_pool, KP_WORLD);
355-
#else
356-
if (PARAM.inp.device == "cpu")
357-
{
358-
MPI_Bcast(psi_mq_real, wfcpw->nrxx, MPI_DOUBLE_COMPLEX, iq_pool, KP_WORLD);
359-
}
360-
else if (PARAM.inp.device == "gpu")
361-
{
362-
// need to copy to cpu first
363-
T* psi_mq_real_cpu = new T[wfcpw->nrxx];
364-
syncmem_complex_d2c_op()(psi_mq_real_cpu, psi_mq_real, wfcpw->nrxx);
365-
MPI_Bcast(psi_mq_real_cpu, wfcpw->nrxx, MPI_DOUBLE_COMPLEX, iq_pool, KP_WORLD);
366-
syncmem_complex_c2d_op()(psi_mq_real, psi_mq_real_cpu, wfcpw->nrxx);
367-
delete[] psi_mq_real_cpu;
368-
}
369-
else
370-
{
371-
ModuleBase::WARNING_QUIT("OperatorEXXPW", "construct_ace: unknown device");
372-
}
373-
#endif
354+
Parallel_Common::bcast_dev<T, Device>(psi_mq_real, wfcpw->nrxx, KP_WORLD, iq_pool);
374355
#endif
375356
for (int n_iband = 0; n_iband < nbands; n_iband++)
376357
{

‎source/source_pw/module_pwdft/op_pw_exx_ace.cpp‎

Lines changed: 8 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#include "op_pw_exx.h"
22
#include "source_base/parallel_comm.h"
3+
#include "source_base/parallel_device.h"
34
#include "source_base/parallel_reduce.h"
45
#include "source_io/module_parameter/parameter.h"
56
#include "source_hamilt/module_xc/exx_info.h"
@@ -46,7 +47,9 @@ void OperatorEXXPW<T, Device>::act_op_ace(const int nbands,
4647
nbands_tot
4748
);
4849

49-
Parallel_Reduce::reduce_pool(Xi_psi, nbands_tot * nbands);
50+
#ifdef __MPI
51+
Parallel_Common::reduce_dev<T, Device>(Xi_psi, nbands_tot * nbands, POOL_WORLD);
52+
#endif
5053

5154
// Xi^\dagger * (Xi * psi)
5255
gemm_complex_op()(trans_C,
@@ -179,32 +182,9 @@ void OperatorEXXPW<T, Device>::construct_ace() const
179182
{
180183
const T* psi_mq = get_pw(m_iband, iq_loc);
181184
wfcpw->recip_to_real(ctx, psi_mq, psi_mq_real, iq_loc);
182-
// send
183185
}
184-
// if (iq == 0)
185-
// std::cout << "Bcast psi_mq_real" << std::endl;
186186
#ifdef __MPI
187-
#ifdef __CUDA_MPI
188-
MPI_Bcast(psi_mq_real, wfcpw->nrxx, MPI_DOUBLE_COMPLEX, iq_pool, KP_WORLD);
189-
#else
190-
if (PARAM.inp.device == "cpu")
191-
{
192-
MPI_Bcast(psi_mq_real, wfcpw->nrxx, MPI_DOUBLE_COMPLEX, iq_pool, KP_WORLD);
193-
}
194-
else if (PARAM.inp.device == "gpu")
195-
{
196-
// need to copy to cpu first
197-
T* psi_mq_real_cpu = new T[wfcpw->nrxx];
198-
syncmem_complex_d2c_op()(psi_mq_real_cpu, psi_mq_real, wfcpw->nrxx);
199-
MPI_Bcast(psi_mq_real_cpu, wfcpw->nrxx, MPI_DOUBLE_COMPLEX, iq_pool, KP_WORLD);
200-
syncmem_complex_c2d_op()(psi_mq_real, psi_mq_real_cpu, wfcpw->nrxx);
201-
delete[] psi_mq_real_cpu;
202-
}
203-
else
204-
{
205-
ModuleBase::WARNING_QUIT("OperatorEXXPW", "construct_ace: unknown device");
206-
}
207-
#endif
187+
Parallel_Common::bcast_dev<T, Device>(psi_mq_real, wfcpw->nrxx, KP_WORLD, iq_pool);
208188
#endif
209189

210190
} // end of iq
@@ -232,7 +212,9 @@ void OperatorEXXPW<T, Device>::construct_ace() const
232212
nbands);
233213

234214
// reduction of psi_h_psi_ace, due to distributed memory
235-
Parallel_Reduce::reduce_pool(psi_h_psi_ace, nbands * nbands);
215+
#ifdef __MPI
216+
Parallel_Common::reduce_dev<T, Device>(psi_h_psi_ace, nbands * nbands, POOL_WORLD);
217+
#endif
236218

237219
T intermediate_minus_one = -1.0;
238220
axpy_complex_op()(nbands * nbands,

0 commit comments

Comments
 (0)