src/gromacs/domdec/distribute.cpp

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2018,2019,2020, by the GROMACS development team, led by
   5  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   6  * and including many others, as listed in the AUTHORS file in the
   7  * top-level source directory and at http://www.gromacs.org.
   8  *
   9  * GROMACS is free software; you can redistribute it and/or
  10  * modify it under the terms of the GNU Lesser General Public License
  11  * as published by the Free Software Foundation; either version 2.1
  12  * of the License, or (at your option) any later version.
  13  *
  14  * GROMACS is distributed in the hope that it will be useful,
  15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  17  * Lesser General Public License for more details.
  18  *
  19  * You should have received a copy of the GNU Lesser General Public
  20  * License along with GROMACS; if not, see
  21  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
  22  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  23  *
  24  * If you want to redistribute modifications to GROMACS, please
  25  * consider that scientific software is very special. Version
  26  * control is crucial - bugs must be traceable. We will be happy to
  27  * consider code for inclusion in the official distribution, but
  28  * derived work must not be called official GROMACS. Details are found
  29  * in the README & COPYING files - if they are missing, get the
  30  * official version at http://www.gromacs.org.
  31  *
  32  * To help us fund GROMACS development, we humbly ask that you cite
  33  * the research papers on the package. Check out http://www.gromacs.org.
  34  */
  35 /* \internal \file
  36  *
  37  * \brief Implements atom distribution functions.
  38  *
  39  * \author Berk Hess <hess@kth.se>
  40  * \ingroup module_domdec
  41  */
  42
  43 #include "gmxpre.h"
  44
  45 #include "distribute.h"
  46
  47 #include "config.h"
  48
  49 #include <vector>
  50
  51 #include "gromacs/domdec/domdec_network.h"
  52 #include "gromacs/math/vec.h"
  53 #include "gromacs/mdtypes/commrec.h"
  54 #include "gromacs/mdtypes/df_history.h"
  55 #include "gromacs/mdtypes/state.h"
  56 #include "gromacs/topology/topology.h"
  57 #include "gromacs/utility/fatalerror.h"
  58 #include "gromacs/utility/logger.h"
  59
  60 #include "atomdistribution.h"
  61 #include "cellsizes.h"
  62 #include "domdec_internal.h"
  63 #include "utility.h"
  64
  65 static void distributeVecSendrecv(gmx_domdec_t*                  dd,
  66                                   gmx::ArrayRef<const gmx::RVec> globalVec,
  67                                   gmx::ArrayRef<gmx::RVec>       localVec)
  68 {
  69     if (DDMASTER(dd))
  70     {
  71         std::vector<gmx::RVec> buffer;
  72
  73         for (int rank = 0; rank < dd->nnodes; rank++)
  74         {
  75             if (rank != dd->rank)
  76             {
  77                 const auto& domainGroups = dd->ma->domainGroups[rank];
  78
  79                 buffer.resize(domainGroups.numAtoms);
  80
  81                 int localAtom = 0;
  82                 for (const int& globalAtom : domainGroups.atomGroups)
  83                 {
  84                     buffer[localAtom++] = globalVec[globalAtom];
  85                 }
  86                 GMX_RELEASE_ASSERT(localAtom == domainGroups.numAtoms,
  87                                    "The index count and number of indices should match");
  88
  89 #if GMX_MPI
  90                 MPI_Send(buffer.data(), domainGroups.numAtoms * sizeof(gmx::RVec), MPI_BYTE, rank,
  91                          rank, dd->mpi_comm_all);
  92 #endif
  93             }
  94         }
  95
  96         const auto& domainGroups = dd->ma->domainGroups[dd->masterrank];
  97         int         localAtom    = 0;
  98         for (const int& globalAtom : domainGroups.atomGroups)
  99         {
 100             localVec[localAtom++] = globalVec[globalAtom];
 101         }
 102     }
 103     else
 104     {
 105 #if GMX_MPI
 106         int numHomeAtoms = dd->comm->atomRanges.numHomeAtoms();
 107         MPI_Recv(localVec.data(), numHomeAtoms * sizeof(gmx::RVec), MPI_BYTE, dd->masterrank,
 108                  MPI_ANY_TAG, dd->mpi_comm_all, MPI_STATUS_IGNORE);
 109 #endif
 110     }
 111 }
 112
 113 static void distributeVecScatterv(gmx_domdec_t*                  dd,
 114                                   gmx::ArrayRef<const gmx::RVec> globalVec,
 115                                   gmx::ArrayRef<gmx::RVec>       localVec)
 116 {
 117     int* sendCounts    = nullptr;
 118     int* displacements = nullptr;
 119
 120     if (DDMASTER(dd))
 121     {
 122         AtomDistribution& ma = *dd->ma;
 123
 124         get_commbuffer_counts(&ma, &sendCounts, &displacements);
 125
 126         gmx::ArrayRef<gmx::RVec> buffer    = ma.rvecBuffer;
 127         int                      localAtom = 0;
 128         for (int rank = 0; rank < dd->nnodes; rank++)
 129         {
 130             const auto& domainGroups = ma.domainGroups[rank];
 131             for (const int& globalAtom : domainGroups.atomGroups)
 132             {
 133                 buffer[localAtom++] = globalVec[globalAtom];
 134             }
 135         }
 136     }
 137
 138     int numHomeAtoms = dd->comm->atomRanges.numHomeAtoms();
 139     dd_scatterv(dd, sendCounts, displacements, DDMASTER(dd) ? dd->ma->rvecBuffer.data() : nullptr,
 140                 numHomeAtoms * sizeof(gmx::RVec), localVec.data());
 141 }
 142
 143 static void distributeVec(gmx_domdec_t*                  dd,
 144                           gmx::ArrayRef<const gmx::RVec> globalVec,
 145                           gmx::ArrayRef<gmx::RVec>       localVec)
 146 {
 147     if (dd->nnodes <= c_maxNumRanksUseSendRecvForScatterAndGather)
 148     {
 149         distributeVecSendrecv(dd, globalVec, localVec);
 150     }
 151     else
 152     {
 153         distributeVecScatterv(dd, globalVec, localVec);
 154     }
 155 }
 156
 157 static void dd_distribute_dfhist(gmx_domdec_t* dd, df_history_t* dfhist)
 158 {
 159     if (dfhist == nullptr)
 160     {
 161         return;
 162     }
 163
 164     dd_bcast(dd, sizeof(int), &dfhist->bEquil);
 165     dd_bcast(dd, sizeof(int), &dfhist->nlambda);
 166     dd_bcast(dd, sizeof(real), &dfhist->wl_delta);
 167
 168     if (dfhist->nlambda > 0)
 169     {
 170         int nlam = dfhist->nlambda;
 171         dd_bcast(dd, sizeof(int) * nlam, dfhist->n_at_lam);
 172         dd_bcast(dd, sizeof(real) * nlam, dfhist->wl_histo);
 173         dd_bcast(dd, sizeof(real) * nlam, dfhist->sum_weights);
 174         dd_bcast(dd, sizeof(real) * nlam, dfhist->sum_dg);
 175         dd_bcast(dd, sizeof(real) * nlam, dfhist->sum_minvar);
 176         dd_bcast(dd, sizeof(real) * nlam, dfhist->sum_variance);
 177
 178         for (int i = 0; i < nlam; i++)
 179         {
 180             dd_bcast(dd, sizeof(real) * nlam, dfhist->accum_p[i]);
 181             dd_bcast(dd, sizeof(real) * nlam, dfhist->accum_m[i]);
 182             dd_bcast(dd, sizeof(real) * nlam, dfhist->accum_p2[i]);
 183             dd_bcast(dd, sizeof(real) * nlam, dfhist->accum_m2[i]);
 184             dd_bcast(dd, sizeof(real) * nlam, dfhist->Tij[i]);
 185             dd_bcast(dd, sizeof(real) * nlam, dfhist->Tij_empirical[i]);
 186         }
 187     }
 188 }
 189
 190 static void dd_distribute_state(gmx_domdec_t* dd, const t_state* state, t_state* state_local)
 191 {
 192     int nh = state_local->nhchainlength;
 193
 194     if (DDMASTER(dd))
 195     {
 196         GMX_RELEASE_ASSERT(state->nhchainlength == nh,
 197                            "The global and local Nose-Hoover chain lengths should match");
 198
 199         for (int i = 0; i < efptNR; i++)
 200         {
 201             state_local->lambda[i] = state->lambda[i];
 202         }
 203         state_local->fep_state = state->fep_state;
 204         state_local->veta      = state->veta;
 205         state_local->vol0      = state->vol0;
 206         copy_mat(state->box, state_local->box);
 207         copy_mat(state->box_rel, state_local->box_rel);
 208         copy_mat(state->boxv, state_local->boxv);
 209         copy_mat(state->svir_prev, state_local->svir_prev);
 210         copy_mat(state->fvir_prev, state_local->fvir_prev);
 211         if (state->dfhist != nullptr)
 212         {
 213             copy_df_history(state_local->dfhist, state->dfhist);
 214         }
 215         for (int i = 0; i < state_local->ngtc; i++)
 216         {
 217             for (int j = 0; j < nh; j++)
 218             {
 219                 state_local->nosehoover_xi[i * nh + j]  = state->nosehoover_xi[i * nh + j];
 220                 state_local->nosehoover_vxi[i * nh + j] = state->nosehoover_vxi[i * nh + j];
 221             }
 222             state_local->therm_integral[i] = state->therm_integral[i];
 223         }
 224         for (int i = 0; i < state_local->nnhpres; i++)
 225         {
 226             for (int j = 0; j < nh; j++)
 227             {
 228                 state_local->nhpres_xi[i * nh + j]  = state->nhpres_xi[i * nh + j];
 229                 state_local->nhpres_vxi[i * nh + j] = state->nhpres_vxi[i * nh + j];
 230             }
 231         }
 232         state_local->baros_integral = state->baros_integral;
 233     }
 234     dd_bcast(dd, ((efptNR) * sizeof(real)), state_local->lambda.data());
 235     dd_bcast(dd, sizeof(int), &state_local->fep_state);
 236     dd_bcast(dd, sizeof(real), &state_local->veta);
 237     dd_bcast(dd, sizeof(real), &state_local->vol0);
 238     dd_bcast(dd, sizeof(state_local->box), state_local->box);
 239     dd_bcast(dd, sizeof(state_local->box_rel), state_local->box_rel);
 240     dd_bcast(dd, sizeof(state_local->boxv), state_local->boxv);
 241     dd_bcast(dd, sizeof(state_local->svir_prev), state_local->svir_prev);
 242     dd_bcast(dd, sizeof(state_local->fvir_prev), state_local->fvir_prev);
 243     dd_bcast(dd, ((state_local->ngtc * nh) * sizeof(double)), state_local->nosehoover_xi.data());
 244     dd_bcast(dd, ((state_local->ngtc * nh) * sizeof(double)), state_local->nosehoover_vxi.data());
 245     dd_bcast(dd, state_local->ngtc * sizeof(double), state_local->therm_integral.data());
 246     dd_bcast(dd, ((state_local->nnhpres * nh) * sizeof(double)), state_local->nhpres_xi.data());
 247     dd_bcast(dd, ((state_local->nnhpres * nh) * sizeof(double)), state_local->nhpres_vxi.data());
 248
 249     /* communicate df_history -- required for restarting from checkpoint */
 250     dd_distribute_dfhist(dd, state_local->dfhist);
 251
 252     state_change_natoms(state_local, dd->comm->atomRanges.numHomeAtoms());
 253
 254     if (state_local->flags & (1 << estX))
 255     {
 256         distributeVec(dd, DDMASTER(dd) ? state->x : gmx::ArrayRef<const gmx::RVec>(), state_local->x);
 257     }
 258     if (state_local->flags & (1 << estV))
 259     {
 260         distributeVec(dd, DDMASTER(dd) ? state->v : gmx::ArrayRef<const gmx::RVec>(), state_local->v);
 261     }
 262     if (state_local->flags & (1 << estCGP))
 263     {
 264         distributeVec(dd, DDMASTER(dd) ? state->cg_p : gmx::ArrayRef<const gmx::RVec>(), state_local->cg_p);
 265     }
 266 }
 267
 268 /* Computes and returns the domain index for the given atom group.
 269  *
 270  * Also updates the coordinates in pos for PBC, when necessary.
 271  */
 272 static inline int computeAtomGroupDomainIndex(const gmx_domdec_t& dd,
 273                                               const gmx_ddbox_t&  ddbox,
 274                                               const matrix&       triclinicCorrectionMatrix,
 275                                               gmx::ArrayRef<const std::vector<real>> cellBoundaries,
 276                                               int                                    atomBegin,
 277                                               int                                    atomEnd,
 278                                               const matrix                           box,
 279                                               rvec*                                  pos)
 280 {
 281     /* Set the reference location cg_cm for assigning the group */
 282     rvec cog;
 283     int  numAtoms = atomEnd - atomBegin;
 284     if (numAtoms == 1)
 285     {
 286         copy_rvec(pos[atomBegin], cog);
 287     }
 288     else
 289     {
 290         real invNumAtoms = 1 / static_cast<real>(numAtoms);
 291
 292         clear_rvec(cog);
 293         for (int a = atomBegin; a < atomEnd; a++)
 294         {
 295             rvec_inc(cog, pos[a]);
 296         }
 297         for (int d = 0; d < DIM; d++)
 298         {
 299             cog[d] *= invNumAtoms;
 300         }
 301     }
 302     /* Put the charge group in the box and determine the cell index ind */
 303     ivec ind;
 304     for (int d = DIM - 1; d >= 0; d--)
 305     {
 306         real pos_d = cog[d];
 307         if (d < dd.unitCellInfo.npbcdim)
 308         {
 309             bool bScrew = (dd.unitCellInfo.haveScrewPBC && d == XX);
 310             if (ddbox.tric_dir[d] && dd.numCells[d] > 1)
 311             {
 312                 /* Use triclinic coordinates for this dimension */
 313                 for (int j = d + 1; j < DIM; j++)
 314                 {
 315                     pos_d += cog[j] * triclinicCorrectionMatrix[j][d];
 316                 }
 317             }
 318             while (pos_d >= box[d][d])
 319             {
 320                 pos_d -= box[d][d];
 321                 rvec_dec(cog, box[d]);
 322                 if (bScrew)
 323                 {
 324                     cog[YY] = box[YY][YY] - cog[YY];
 325                     cog[ZZ] = box[ZZ][ZZ] - cog[ZZ];
 326                 }
 327                 for (int a = atomBegin; a < atomEnd; a++)
 328                 {
 329                     rvec_dec(pos[a], box[d]);
 330                     if (bScrew)
 331                     {
 332                         pos[a][YY] = box[YY][YY] - pos[a][YY];
 333                         pos[a][ZZ] = box[ZZ][ZZ] - pos[a][ZZ];
 334                     }
 335                 }
 336             }
 337             while (pos_d < 0)
 338             {
 339                 pos_d += box[d][d];
 340                 rvec_inc(cog, box[d]);
 341                 if (bScrew)
 342                 {
 343                     cog[YY] = box[YY][YY] - cog[YY];
 344                     cog[ZZ] = box[ZZ][ZZ] - cog[ZZ];
 345                 }
 346                 for (int a = atomBegin; a < atomEnd; a++)
 347                 {
 348                     rvec_inc(pos[a], box[d]);
 349                     if (bScrew)
 350                     {
 351                         pos[a][YY] = box[YY][YY] - pos[a][YY];
 352                         pos[a][ZZ] = box[ZZ][ZZ] - pos[a][ZZ];
 353                     }
 354                 }
 355             }
 356         }
 357         /* This could be done more efficiently */
 358         ind[d] = 0;
 359         while (ind[d] + 1 < dd.numCells[d] && pos_d >= cellBoundaries[d][ind[d] + 1])
 360         {
 361             ind[d]++;
 362         }
 363     }
 364
 365     return dd_index(dd.numCells, ind);
 366 }
 367
 368
 369 static std::vector<std::vector<int>> getAtomGroupDistribution(const gmx::MDLogger& mdlog,
 370                                                               const gmx_mtop_t&    mtop,
 371                                                               const matrix         box,
 372                                                               const gmx_ddbox_t&   ddbox,
 373                                                               rvec                 pos[],
 374                                                               gmx_domdec_t*        dd)
 375 {
 376     AtomDistribution& ma = *dd->ma;
 377
 378     /* Clear the count */
 379     for (int rank = 0; rank < dd->nnodes; rank++)
 380     {
 381         ma.domainGroups[rank].numAtoms = 0;
 382     }
 383
 384     matrix triclinicCorrectionMatrix;
 385     make_tric_corr_matrix(dd->unitCellInfo.npbcdim, box, triclinicCorrectionMatrix);
 386
 387     ivec       npulse;
 388     const auto cellBoundaries = set_dd_cell_sizes_slb(dd, &ddbox, setcellsizeslbMASTER, npulse);
 389
 390     std::vector<std::vector<int>> indices(dd->nnodes);
 391
 392     if (dd->comm->systemInfo.useUpdateGroups)
 393     {
 394         int atomOffset = 0;
 395         for (const gmx_molblock_t& molblock : mtop.molblock)
 396         {
 397             const auto& updateGrouping =
 398                     dd->comm->systemInfo.updateGroupingPerMoleculetype[molblock.type];
 399
 400             for (int mol = 0; mol < molblock.nmol; mol++)
 401             {
 402                 for (int g = 0; g < updateGrouping.numBlocks(); g++)
 403                 {
 404                     const auto& block     = updateGrouping.block(g);
 405                     const int   atomBegin = atomOffset + block.begin();
 406                     const int   atomEnd   = atomOffset + block.end();
 407                     const int   domainIndex =
 408                             computeAtomGroupDomainIndex(*dd, ddbox, triclinicCorrectionMatrix,
 409                                                         cellBoundaries, atomBegin, atomEnd, box, pos);
 410
 411                     for (int atomIndex : block)
 412                     {
 413                         indices[domainIndex].push_back(atomOffset + atomIndex);
 414                     }
 415                     ma.domainGroups[domainIndex].numAtoms += block.size();
 416                 }
 417
 418                 atomOffset += updateGrouping.fullRange().end();
 419             }
 420         }
 421
 422         GMX_RELEASE_ASSERT(atomOffset == mtop.natoms, "Should distribute all atoms");
 423     }
 424     else
 425     {
 426         /* Compute the center of geometry for all atoms */
 427         for (int atom = 0; atom < mtop.natoms; atom++)
 428         {
 429             int domainIndex = computeAtomGroupDomainIndex(*dd, ddbox, triclinicCorrectionMatrix,
 430                                                           cellBoundaries, atom, atom + 1, box, pos);
 431
 432             indices[domainIndex].push_back(atom);
 433             ma.domainGroups[domainIndex].numAtoms += 1;
 434         }
 435     }
 436
 437     {
 438         // Use double for the sums to avoid natoms^2 overflowing
 439         // (65537^2 > 2^32)
 440         int    nat_sum, nat_min, nat_max;
 441         double nat2_sum;
 442
 443         nat_sum  = 0;
 444         nat2_sum = 0;
 445         nat_min  = ma.domainGroups[0].numAtoms;
 446         nat_max  = ma.domainGroups[0].numAtoms;
 447         for (int rank = 0; rank < dd->nnodes; rank++)
 448         {
 449             int numAtoms = ma.domainGroups[rank].numAtoms;
 450             nat_sum += numAtoms;
 451             // convert to double to avoid integer overflows when squaring
 452             nat2_sum += gmx::square(double(numAtoms));
 453             nat_min = std::min(nat_min, numAtoms);
 454             nat_max = std::max(nat_max, numAtoms);
 455         }
 456         nat_sum /= dd->nnodes;
 457         nat2_sum /= dd->nnodes;
 458
 459         GMX_LOG(mdlog.info)
 460                 .appendTextFormatted(
 461                         "Atom distribution over %d domains: av %d stddev %d min %d max %d",
 462                         dd->nnodes, nat_sum,
 463                         gmx::roundToInt(std::sqrt(nat2_sum - gmx::square(static_cast<double>(nat_sum)))),
 464                         nat_min, nat_max);
 465     }
 466
 467     return indices;
 468 }
 469
 470 static void distributeAtomGroups(const gmx::MDLogger& mdlog,
 471                                  gmx_domdec_t*        dd,
 472                                  const gmx_mtop_t&    mtop,
 473                                  const matrix         box,
 474                                  const gmx_ddbox_t*   ddbox,
 475                                  rvec                 pos[])
 476 {
 477     AtomDistribution* ma = dd->ma.get();
 478     int *             ibuf, buf2[2] = { 0, 0 };
 479     gmx_bool          bMaster = DDMASTER(dd);
 480
 481     std::vector<std::vector<int>> groupIndices;
 482
 483     if (bMaster)
 484     {
 485         GMX_ASSERT(box && pos, "box or pos not set on master");
 486
 487         if (dd->unitCellInfo.haveScrewPBC)
 488         {
 489             check_screw_box(box);
 490         }
 491
 492         groupIndices = getAtomGroupDistribution(mdlog, mtop, box, *ddbox, pos, dd);
 493
 494         for (int rank = 0; rank < dd->nnodes; rank++)
 495         {
 496             ma->intBuffer[rank * 2]     = groupIndices[rank].size();
 497             ma->intBuffer[rank * 2 + 1] = ma->domainGroups[rank].numAtoms;
 498         }
 499         ibuf = ma->intBuffer.data();
 500     }
 501     else
 502     {
 503         ibuf = nullptr;
 504     }
 505     dd_scatter(dd, 2 * sizeof(int), ibuf, buf2);
 506
 507     dd->ncg_home = buf2[0];
 508     dd->comm->atomRanges.setEnd(DDAtomRanges::Type::Home, buf2[1]);
 509     dd->globalAtomGroupIndices.resize(dd->ncg_home);
 510     dd->globalAtomIndices.resize(dd->comm->atomRanges.numHomeAtoms());
 511
 512     if (bMaster)
 513     {
 514         ma->atomGroups.clear();
 515
 516         int groupOffset = 0;
 517         for (int rank = 0; rank < dd->nnodes; rank++)
 518         {
 519             ma->intBuffer[rank]              = groupIndices[rank].size() * sizeof(int);
 520             ma->intBuffer[dd->nnodes + rank] = groupOffset * sizeof(int);
 521
 522             ma->atomGroups.insert(ma->atomGroups.end(), groupIndices[rank].begin(),
 523                                   groupIndices[rank].end());
 524
 525             ma->domainGroups[rank].atomGroups = gmx::constArrayRefFromArray(
 526                     ma->atomGroups.data() + groupOffset, groupIndices[rank].size());
 527
 528             groupOffset += groupIndices[rank].size();
 529         }
 530     }
 531
 532     dd_scatterv(dd, bMaster ? ma->intBuffer.data() : nullptr,
 533                 bMaster ? ma->intBuffer.data() + dd->nnodes : nullptr,
 534                 bMaster ? ma->atomGroups.data() : nullptr, dd->ncg_home * sizeof(int),
 535                 dd->globalAtomGroupIndices.data());
 536
 537     if (debug)
 538     {
 539         fprintf(debug, "Home charge groups:\n");
 540         for (int i = 0; i < dd->ncg_home; i++)
 541         {
 542             fprintf(debug, " %d", dd->globalAtomGroupIndices[i]);
 543             if (i % 10 == 9)
 544             {
 545                 fprintf(debug, "\n");
 546             }
 547         }
 548         fprintf(debug, "\n");
 549     }
 550 }
 551
 552 void distributeState(const gmx::MDLogger& mdlog,
 553                      gmx_domdec_t*        dd,
 554                      const gmx_mtop_t&    mtop,
 555                      t_state*             state_global,
 556                      const gmx_ddbox_t&   ddbox,
 557                      t_state*             state_local)
 558 {
 559     rvec* xGlobal = (DDMASTER(dd) ? state_global->x.rvec_array() : nullptr);
 560
 561     distributeAtomGroups(mdlog, dd, mtop, DDMASTER(dd) ? state_global->box : nullptr, &ddbox, xGlobal);
 562
 563     dd_distribute_state(dd, state_global, state_local);
 564 }