|
1 | 1 | #include "parallel_grid.h" |
| 2 | + |
2 | 3 | #include "source_base/global_function.h" |
3 | 4 | #include "source_base/global_variable.h" |
4 | 5 |
|
5 | 6 | #ifdef __MPI |
6 | 7 | #include "source_base/parallel_comm.h" // use POOL_WORLD |
| 8 | + |
7 | 9 | #include <mpi.h> |
8 | 10 | #endif |
9 | 11 |
|
@@ -167,6 +169,95 @@ void Parallel_Grid::z_distribution() |
167 | 169 | return; |
168 | 170 | } |
169 | 171 |
|
| 172 | +void Parallel_Grid::reduce_across_pools(double* data) const |
| 173 | +{ |
| 174 | +#ifdef __MPI |
| 175 | + if (GlobalV::KPAR <= 1) |
| 176 | + { |
| 177 | + return; |
| 178 | + } |
| 179 | + |
| 180 | + assert(data != nullptr); |
| 181 | + if (KP_WORLD != MPI_COMM_NULL) |
| 182 | + { |
| 183 | + // Equal-sized pools give corresponding ranks identical z-slab layouts, |
| 184 | + // so their local buffers can be summed directly without redistribution. |
| 185 | + MPI_Allreduce(MPI_IN_PLACE, data, this->nrxx, MPI_DOUBLE, MPI_SUM, KP_WORLD); |
| 186 | + return; |
| 187 | + } |
| 188 | + |
| 189 | + // Uneven pool sizes have no KP_WORLD and may assign different z-slabs to |
| 190 | + // corresponding ranks. Validate the local distribution before rebuilding |
| 191 | + // a common global layout for the cross-pool reduction. |
| 192 | + assert(!this->numz.empty()); |
| 193 | + assert(GlobalV::MY_POOL >= 0 && GlobalV::MY_POOL < static_cast<int>(this->numz.size())); |
| 194 | + assert(GlobalV::RANK_IN_POOL >= 0 && GlobalV::RANK_IN_POOL < static_cast<int>(this->numz[GlobalV::MY_POOL].size())); |
| 195 | + assert(this->nczp == this->numz[GlobalV::MY_POOL][GlobalV::RANK_IN_POOL]); |
| 196 | + assert(this->nrxx == this->ncxy * this->nczp); |
| 197 | + |
| 198 | + const int pool_size = this->nproc_in_pool[GlobalV::MY_POOL]; |
| 199 | + std::vector<int> receive_counts(pool_size); |
| 200 | + std::vector<int> displacements(pool_size); |
| 201 | + for (int ip = 0; ip < pool_size; ++ip) |
| 202 | + { |
| 203 | + receive_counts[ip] = this->numz[GlobalV::MY_POOL][ip] * this->ncxy; |
| 204 | + displacements[ip] = this->startz[GlobalV::MY_POOL][ip] * this->ncxy; |
| 205 | + } |
| 206 | + |
| 207 | + std::vector<double> local_data(this->nrxx); |
| 208 | + // The allgather below replicates one complete pool grid on every rank in |
| 209 | + // that pool. INT_BGROUP then sums all of those replicas, so divide each |
| 210 | + // local slab by the pool size to make each pool contribute exactly once. |
| 211 | + const double pool_normalization = 1.0 / static_cast<double>(pool_size); |
| 212 | + for (int ir = 0; ir < this->nrxx; ++ir) |
| 213 | + { |
| 214 | + local_data[ir] = data[ir] * pool_normalization; |
| 215 | + } |
| 216 | + |
| 217 | + std::vector<double> pool_data(this->ncxyz); |
| 218 | + // Collect the rank-local [xy][local_z] slabs into rank-contiguous blocks. |
| 219 | + MPI_Allgatherv(local_data.data(), |
| 220 | + this->nrxx, |
| 221 | + MPI_DOUBLE, |
| 222 | + pool_data.data(), |
| 223 | + receive_counts.data(), |
| 224 | + displacements.data(), |
| 225 | + MPI_DOUBLE, |
| 226 | + POOL_WORLD); |
| 227 | + |
| 228 | + std::vector<double> global_layout(this->ncxyz); |
| 229 | + // Convert the rank-contiguous allgather result to the canonical |
| 230 | + // [xy][global_z] order required for element-wise reduction across pools. |
| 231 | + for (int ip = 0; ip < pool_size; ++ip) |
| 232 | + { |
| 233 | + const int local_nz = this->numz[GlobalV::MY_POOL][ip]; |
| 234 | + const int global_z_start = this->startz[GlobalV::MY_POOL][ip]; |
| 235 | + const int gathered_start = global_z_start * this->ncxy; |
| 236 | + for (int ixy = 0; ixy < this->ncxy; ++ixy) |
| 237 | + { |
| 238 | + for (int iz = 0; iz < local_nz; ++iz) |
| 239 | + { |
| 240 | + global_layout[ixy * this->ncz + global_z_start + iz] = pool_data[gathered_start + ixy * local_nz + iz]; |
| 241 | + } |
| 242 | + } |
| 243 | + } |
| 244 | + |
| 245 | + MPI_Allreduce(MPI_IN_PLACE, global_layout.data(), this->ncxyz, MPI_DOUBLE, MPI_SUM, INT_BGROUP); |
| 246 | + |
| 247 | + // Return only the z-slab owned by this rank under its pool's distribution. |
| 248 | + const int local_z_start = this->startz[GlobalV::MY_POOL][GlobalV::RANK_IN_POOL]; |
| 249 | + for (int ixy = 0; ixy < this->ncxy; ++ixy) |
| 250 | + { |
| 251 | + for (int iz = 0; iz < this->nczp; ++iz) |
| 252 | + { |
| 253 | + data[ixy * this->nczp + iz] = global_layout[ixy * this->ncz + local_z_start + iz]; |
| 254 | + } |
| 255 | + } |
| 256 | +#else |
| 257 | + (void)data; |
| 258 | +#endif |
| 259 | +} |
| 260 | + |
170 | 261 | #ifdef __MPI |
171 | 262 | void Parallel_Grid::bcast(const double* const data_global, double* data_local, const int& rank, const bool is_sdft) const |
172 | 263 | { |
|
0 commit comments