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, rank, dd->mpi_comm_all);
  91 #endif
  92             }
  93         }
  94
  95         const auto& domainGroups = dd->ma->domainGroups[dd->masterrank];
  96         int         localAtom    = 0;
  97         for (const int& globalAtom : domainGroups.atomGroups)
  98         {
  99             localVec[localAtom++] = globalVec[globalAtom];
 100         }
 101     }
 102     else
 103     {
 104 #if GMX_MPI
 105         int numHomeAtoms = dd->comm->atomRanges.numHomeAtoms();
 106         MPI_Recv(localVec.data(),
 107                  numHomeAtoms * sizeof(gmx::RVec),
 108                  MPI_BYTE,
 109                  dd->masterrank,
 110                  MPI_ANY_TAG,
 111                  dd->mpi_comm_all,
 112                  MPI_STATUS_IGNORE);
 113 #endif
 114     }
 115 }
 116
 117 static void distributeVecScatterv(gmx_domdec_t*                  dd,
 118                                   gmx::ArrayRef<const gmx::RVec> globalVec,
 119                                   gmx::ArrayRef<gmx::RVec>       localVec)
 120 {
 121     int* sendCounts    = nullptr;
 122     int* displacements = nullptr;
 123
 124     if (DDMASTER(dd))
 125     {
 126         AtomDistribution& ma = *dd->ma;
 127
 128         get_commbuffer_counts(&ma, &sendCounts, &displacements);
 129
 130         gmx::ArrayRef<gmx::RVec> buffer    = ma.rvecBuffer;
 131         int                      localAtom = 0;
 132         for (int rank = 0; rank < dd->nnodes; rank++)
 133         {
 134             const auto& domainGroups = ma.domainGroups[rank];
 135             for (const int& globalAtom : domainGroups.atomGroups)
 136             {
 137                 buffer[localAtom++] = globalVec[globalAtom];
 138             }
 139         }
 140     }
 141
 142     int numHomeAtoms = dd->comm->atomRanges.numHomeAtoms();
 143     dd_scatterv(dd,
 144                 sendCounts,
 145                 displacements,
 146                 DDMASTER(dd) ? dd->ma->rvecBuffer.data() : nullptr,
 147                 numHomeAtoms * sizeof(gmx::RVec),
 148                 localVec.data());
 149 }
 150
 151 static void distributeVec(gmx_domdec_t*                  dd,
 152                           gmx::ArrayRef<const gmx::RVec> globalVec,
 153                           gmx::ArrayRef<gmx::RVec>       localVec)
 154 {
 155     if (dd->nnodes <= c_maxNumRanksUseSendRecvForScatterAndGather)
 156     {
 157         distributeVecSendrecv(dd, globalVec, localVec);
 158     }
 159     else
 160     {
 161         distributeVecScatterv(dd, globalVec, localVec);
 162     }
 163 }
 164
 165 static void dd_distribute_dfhist(gmx_domdec_t* dd, df_history_t* dfhist)
 166 {
 167     if (dfhist == nullptr)
 168     {
 169         return;
 170     }
 171
 172     dd_bcast(dd, sizeof(int), &dfhist->bEquil);
 173     dd_bcast(dd, sizeof(int), &dfhist->nlambda);
 174     dd_bcast(dd, sizeof(real), &dfhist->wl_delta);
 175
 176     if (dfhist->nlambda > 0)
 177     {
 178         int nlam = dfhist->nlambda;
 179         dd_bcast(dd, sizeof(int) * nlam, dfhist->n_at_lam);
 180         dd_bcast(dd, sizeof(real) * nlam, dfhist->wl_histo);
 181         dd_bcast(dd, sizeof(real) * nlam, dfhist->sum_weights);
 182         dd_bcast(dd, sizeof(real) * nlam, dfhist->sum_dg);
 183         dd_bcast(dd, sizeof(real) * nlam, dfhist->sum_minvar);
 184         dd_bcast(dd, sizeof(real) * nlam, dfhist->sum_variance);
 185
 186         for (int i = 0; i < nlam; i++)
 187         {
 188             dd_bcast(dd, sizeof(real) * nlam, dfhist->accum_p[i]);
 189             dd_bcast(dd, sizeof(real) * nlam, dfhist->accum_m[i]);
 190             dd_bcast(dd, sizeof(real) * nlam, dfhist->accum_p2[i]);
 191             dd_bcast(dd, sizeof(real) * nlam, dfhist->accum_m2[i]);
 192             dd_bcast(dd, sizeof(real) * nlam, dfhist->Tij[i]);
 193             dd_bcast(dd, sizeof(real) * nlam, dfhist->Tij_empirical[i]);
 194         }
 195     }
 196 }
 197
 198 static void dd_distribute_state(gmx_domdec_t* dd, const t_state* state, t_state* state_local)
 199 {
 200     int nh = state_local->nhchainlength;
 201
 202     if (DDMASTER(dd))
 203     {
 204         GMX_RELEASE_ASSERT(state->nhchainlength == nh,
 205                            "The global and local Nose-Hoover chain lengths should match");
 206
 207         for (int i = 0; i < efptNR; i++)
 208         {
 209             state_local->lambda[i] = state->lambda[i];
 210         }
 211         state_local->fep_state = state->fep_state;
 212         state_local->veta      = state->veta;
 213         state_local->vol0      = state->vol0;
 214         copy_mat(state->box, state_local->box);
 215         copy_mat(state->box_rel, state_local->box_rel);
 216         copy_mat(state->boxv, state_local->boxv);
 217         copy_mat(state->svir_prev, state_local->svir_prev);
 218         copy_mat(state->fvir_prev, state_local->fvir_prev);
 219         if (state->dfhist != nullptr)
 220         {
 221             copy_df_history(state_local->dfhist, state->dfhist);
 222         }
 223         for (int i = 0; i < state_local->ngtc; i++)
 224         {
 225             for (int j = 0; j < nh; j++)
 226             {
 227                 state_local->nosehoover_xi[i * nh + j]  = state->nosehoover_xi[i * nh + j];
 228                 state_local->nosehoover_vxi[i * nh + j] = state->nosehoover_vxi[i * nh + j];
 229             }
 230             state_local->therm_integral[i] = state->therm_integral[i];
 231         }
 232         for (int i = 0; i < state_local->nnhpres; i++)
 233         {
 234             for (int j = 0; j < nh; j++)
 235             {
 236                 state_local->nhpres_xi[i * nh + j]  = state->nhpres_xi[i * nh + j];
 237                 state_local->nhpres_vxi[i * nh + j] = state->nhpres_vxi[i * nh + j];
 238             }
 239         }
 240         state_local->baros_integral = state->baros_integral;
 241     }
 242     dd_bcast(dd, ((efptNR) * sizeof(real)), state_local->lambda.data());
 243     dd_bcast(dd, sizeof(int), &state_local->fep_state);
 244     dd_bcast(dd, sizeof(real), &state_local->veta);
 245     dd_bcast(dd, sizeof(real), &state_local->vol0);
 246     dd_bcast(dd, sizeof(state_local->box), state_local->box);
 247     dd_bcast(dd, sizeof(state_local->box_rel), state_local->box_rel);
 248     dd_bcast(dd, sizeof(state_local->boxv), state_local->boxv);
 249     dd_bcast(dd, sizeof(state_local->svir_prev), state_local->svir_prev);
 250     dd_bcast(dd, sizeof(state_local->fvir_prev), state_local->fvir_prev);
 251     dd_bcast(dd, ((state_local->ngtc * nh) * sizeof(double)), state_local->nosehoover_xi.data());
 252     dd_bcast(dd, ((state_local->ngtc * nh) * sizeof(double)), state_local->nosehoover_vxi.data());
 253     dd_bcast(dd, state_local->ngtc * sizeof(double), state_local->therm_integral.data());
 254     dd_bcast(dd, ((state_local->nnhpres * nh) * sizeof(double)), state_local->nhpres_xi.data());
 255     dd_bcast(dd, ((state_local->nnhpres * nh) * sizeof(double)), state_local->nhpres_vxi.data());
 256
 257     /* communicate df_history -- required for restarting from checkpoint */
 258     dd_distribute_dfhist(dd, state_local->dfhist);
 259
 260     state_change_natoms(state_local, dd->comm->atomRanges.numHomeAtoms());
 261
 262     if (state_local->flags & (1 << estX))
 263     {
 264         distributeVec(dd, DDMASTER(dd) ? state->x : gmx::ArrayRef<const gmx::RVec>(), state_local->x);
 265     }
 266     if (state_local->flags & (1 << estV))
 267     {
 268         distributeVec(dd, DDMASTER(dd) ? state->v : gmx::ArrayRef<const gmx::RVec>(), state_local->v);
 269     }
 270     if (state_local->flags & (1 << estCGP))
 271     {
 272         distributeVec(dd, DDMASTER(dd) ? state->cg_p : gmx::ArrayRef<const gmx::RVec>(), state_local->cg_p);
 273     }
 274 }
 275
 276 /* Computes and returns the domain index for the given atom group.
 277  *
 278  * Also updates the coordinates in pos for PBC, when necessary.
 279  */
 280 static inline int computeAtomGroupDomainIndex(const gmx_domdec_t& dd,
 281                                               const gmx_ddbox_t&  ddbox,
 282                                               const matrix&       triclinicCorrectionMatrix,
 283                                               gmx::ArrayRef<const std::vector<real>> cellBoundaries,
 284                                               int                                    atomBegin,
 285                                               int                                    atomEnd,
 286                                               const matrix                           box,
 287                                               rvec*                                  pos)
 288 {
 289     /* Set the reference location cg_cm for assigning the group */
 290     rvec cog;
 291     int  numAtoms = atomEnd - atomBegin;
 292     if (numAtoms == 1)
 293     {
 294         copy_rvec(pos[atomBegin], cog);
 295     }
 296     else
 297     {
 298         real invNumAtoms = 1 / static_cast<real>(numAtoms);
 299
 300         clear_rvec(cog);
 301         for (int a = atomBegin; a < atomEnd; a++)
 302         {
 303             rvec_inc(cog, pos[a]);
 304         }
 305         for (int d = 0; d < DIM; d++)
 306         {
 307             cog[d] *= invNumAtoms;
 308         }
 309     }
 310     /* Put the charge group in the box and determine the cell index ind */
 311     ivec ind;
 312     for (int d = DIM - 1; d >= 0; d--)
 313     {
 314         real pos_d = cog[d];
 315         if (d < dd.unitCellInfo.npbcdim)
 316         {
 317             bool bScrew = (dd.unitCellInfo.haveScrewPBC && d == XX);
 318             if (ddbox.tric_dir[d] && dd.numCells[d] > 1)
 319             {
 320                 /* Use triclinic coordinates for this dimension */
 321                 for (int j = d + 1; j < DIM; j++)
 322                 {
 323                     pos_d += cog[j] * triclinicCorrectionMatrix[j][d];
 324                 }
 325             }
 326             while (pos_d >= box[d][d])
 327             {
 328                 pos_d -= box[d][d];
 329                 rvec_dec(cog, box[d]);
 330                 if (bScrew)
 331                 {
 332                     cog[YY] = box[YY][YY] - cog[YY];
 333                     cog[ZZ] = box[ZZ][ZZ] - cog[ZZ];
 334                 }
 335                 for (int a = atomBegin; a < atomEnd; a++)
 336                 {
 337                     rvec_dec(pos[a], box[d]);
 338                     if (bScrew)
 339                     {
 340                         pos[a][YY] = box[YY][YY] - pos[a][YY];
 341                         pos[a][ZZ] = box[ZZ][ZZ] - pos[a][ZZ];
 342                     }
 343                 }
 344             }
 345             while (pos_d < 0)
 346             {
 347                 pos_d += box[d][d];
 348                 rvec_inc(cog, box[d]);
 349                 if (bScrew)
 350                 {
 351                     cog[YY] = box[YY][YY] - cog[YY];
 352                     cog[ZZ] = box[ZZ][ZZ] - cog[ZZ];
 353                 }
 354                 for (int a = atomBegin; a < atomEnd; a++)
 355                 {
 356                     rvec_inc(pos[a], box[d]);
 357                     if (bScrew)
 358                     {
 359                         pos[a][YY] = box[YY][YY] - pos[a][YY];
 360                         pos[a][ZZ] = box[ZZ][ZZ] - pos[a][ZZ];
 361                     }
 362                 }
 363             }
 364         }
 365         /* This could be done more efficiently */
 366         ind[d] = 0;
 367         while (ind[d] + 1 < dd.numCells[d] && pos_d >= cellBoundaries[d][ind[d] + 1])
 368         {
 369             ind[d]++;
 370         }
 371     }
 372
 373     return dd_index(dd.numCells, ind);
 374 }
 375
 376
 377 static std::vector<std::vector<int>> getAtomGroupDistribution(const gmx::MDLogger& mdlog,
 378                                                               const gmx_mtop_t&    mtop,
 379                                                               const matrix         box,
 380                                                               const gmx_ddbox_t&   ddbox,
 381                                                               rvec                 pos[],
 382                                                               gmx_domdec_t*        dd)
 383 {
 384     AtomDistribution& ma = *dd->ma;
 385
 386     /* Clear the count */
 387     for (int rank = 0; rank < dd->nnodes; rank++)
 388     {
 389         ma.domainGroups[rank].numAtoms = 0;
 390     }
 391
 392     matrix triclinicCorrectionMatrix;
 393     make_tric_corr_matrix(dd->unitCellInfo.npbcdim, box, triclinicCorrectionMatrix);
 394
 395     ivec       npulse;
 396     const auto cellBoundaries = set_dd_cell_sizes_slb(dd, &ddbox, setcellsizeslbMASTER, npulse);
 397
 398     std::vector<std::vector<int>> indices(dd->nnodes);
 399
 400     if (dd->comm->systemInfo.useUpdateGroups)
 401     {
 402         int atomOffset = 0;
 403         for (const gmx_molblock_t& molblock : mtop.molblock)
 404         {
 405             const auto& updateGrouping =
 406                     dd->comm->systemInfo.updateGroupingPerMoleculetype[molblock.type];
 407
 408             for (int mol = 0; mol < molblock.nmol; mol++)
 409             {
 410                 for (int g = 0; g < updateGrouping.numBlocks(); g++)
 411                 {
 412                     const auto& block       = updateGrouping.block(g);
 413                     const int   atomBegin   = atomOffset + block.begin();
 414                     const int   atomEnd     = atomOffset + block.end();
 415                     const int   domainIndex = computeAtomGroupDomainIndex(
 416                             *dd, ddbox, triclinicCorrectionMatrix, cellBoundaries, atomBegin, atomEnd, box, pos);
 417
 418                     for (int atomIndex : block)
 419                     {
 420                         indices[domainIndex].push_back(atomOffset + atomIndex);
 421                     }
 422                     ma.domainGroups[domainIndex].numAtoms += block.size();
 423                 }
 424
 425                 atomOffset += updateGrouping.fullRange().end();
 426             }
 427         }
 428
 429         GMX_RELEASE_ASSERT(atomOffset == mtop.natoms, "Should distribute all atoms");
 430     }
 431     else
 432     {
 433         /* Compute the center of geometry for all atoms */
 434         for (int atom = 0; atom < mtop.natoms; atom++)
 435         {
 436             int domainIndex = computeAtomGroupDomainIndex(
 437                     *dd, ddbox, triclinicCorrectionMatrix, cellBoundaries, atom, atom + 1, box, pos);
 438
 439             indices[domainIndex].push_back(atom);
 440             ma.domainGroups[domainIndex].numAtoms += 1;
 441         }
 442     }
 443
 444     {
 445         // Use double for the sums to avoid natoms^2 overflowing
 446         // (65537^2 > 2^32)
 447         int    nat_sum, nat_min, nat_max;
 448         double nat2_sum;
 449
 450         nat_sum  = 0;
 451         nat2_sum = 0;
 452         nat_min  = ma.domainGroups[0].numAtoms;
 453         nat_max  = ma.domainGroups[0].numAtoms;
 454         for (int rank = 0; rank < dd->nnodes; rank++)
 455         {
 456             int numAtoms = ma.domainGroups[rank].numAtoms;
 457             nat_sum += numAtoms;
 458             // convert to double to avoid integer overflows when squaring
 459             nat2_sum += gmx::square(double(numAtoms));
 460             nat_min = std::min(nat_min, numAtoms);
 461             nat_max = std::max(nat_max, numAtoms);
 462         }
 463         nat_sum /= dd->nnodes;
 464         nat2_sum /= dd->nnodes;
 465
 466         GMX_LOG(mdlog.info)
 467                 .appendTextFormatted(
 468                         "Atom distribution over %d domains: av %d stddev %d min %d max %d",
 469                         dd->nnodes,
 470                         nat_sum,
 471                         gmx::roundToInt(std::sqrt(nat2_sum - gmx::square(static_cast<double>(nat_sum)))),
 472                         nat_min,
 473                         nat_max);
 474     }
 475
 476     return indices;
 477 }
 478
 479 static void distributeAtomGroups(const gmx::MDLogger& mdlog,
 480                                  gmx_domdec_t*        dd,
 481                                  const gmx_mtop_t&    mtop,
 482                                  const matrix         box,
 483                                  const gmx_ddbox_t*   ddbox,
 484                                  rvec                 pos[])
 485 {
 486     AtomDistribution* ma = dd->ma.get();
 487     int *             ibuf, buf2[2] = { 0, 0 };
 488     gmx_bool          bMaster = DDMASTER(dd);
 489
 490     std::vector<std::vector<int>> groupIndices;
 491
 492     if (bMaster)
 493     {
 494         GMX_ASSERT(box && pos, "box or pos not set on master");
 495
 496         if (dd->unitCellInfo.haveScrewPBC)
 497         {
 498             check_screw_box(box);
 499         }
 500
 501         groupIndices = getAtomGroupDistribution(mdlog, mtop, box, *ddbox, pos, dd);
 502
 503         for (int rank = 0; rank < dd->nnodes; rank++)
 504         {
 505             ma->intBuffer[rank * 2]     = groupIndices[rank].size();
 506             ma->intBuffer[rank * 2 + 1] = ma->domainGroups[rank].numAtoms;
 507         }
 508         ibuf = ma->intBuffer.data();
 509     }
 510     else
 511     {
 512         ibuf = nullptr;
 513     }
 514     dd_scatter(dd, 2 * sizeof(int), ibuf, buf2);
 515
 516     dd->ncg_home = buf2[0];
 517     dd->comm->atomRanges.setEnd(DDAtomRanges::Type::Home, buf2[1]);
 518     dd->globalAtomGroupIndices.resize(dd->ncg_home);
 519     dd->globalAtomIndices.resize(dd->comm->atomRanges.numHomeAtoms());
 520
 521     if (bMaster)
 522     {
 523         ma->atomGroups.clear();
 524
 525         int groupOffset = 0;
 526         for (int rank = 0; rank < dd->nnodes; rank++)
 527         {
 528             ma->intBuffer[rank]              = groupIndices[rank].size() * sizeof(int);
 529             ma->intBuffer[dd->nnodes + rank] = groupOffset * sizeof(int);
 530
 531             ma->atomGroups.insert(
 532                     ma->atomGroups.end(), groupIndices[rank].begin(), groupIndices[rank].end());
 533
 534             ma->domainGroups[rank].atomGroups = gmx::constArrayRefFromArray(
 535                     ma->atomGroups.data() + groupOffset, groupIndices[rank].size());
 536
 537             groupOffset += groupIndices[rank].size();
 538         }
 539     }
 540
 541     dd_scatterv(dd,
 542                 bMaster ? ma->intBuffer.data() : nullptr,
 543                 bMaster ? ma->intBuffer.data() + dd->nnodes : nullptr,
 544                 bMaster ? ma->atomGroups.data() : nullptr,
 545                 dd->ncg_home * sizeof(int),
 546                 dd->globalAtomGroupIndices.data());
 547
 548     if (debug)
 549     {
 550         fprintf(debug, "Home charge groups:\n");
 551         for (int i = 0; i < dd->ncg_home; i++)
 552         {
 553             fprintf(debug, " %d", dd->globalAtomGroupIndices[i]);
 554             if (i % 10 == 9)
 555             {
 556                 fprintf(debug, "\n");
 557             }
 558         }
 559         fprintf(debug, "\n");
 560     }
 561 }
 562
 563 void distributeState(const gmx::MDLogger& mdlog,
 564                      gmx_domdec_t*        dd,
 565                      const gmx_mtop_t&    mtop,
 566                      t_state*             state_global,
 567                      const gmx_ddbox_t&   ddbox,
 568                      t_state*             state_local)
 569 {
 570     rvec* xGlobal = (DDMASTER(dd) ? state_global->x.rvec_array() : nullptr);
 571
 572     distributeAtomGroups(mdlog, dd, mtop, DDMASTER(dd) ? state_global->box : nullptr, &ddbox, xGlobal);
 573
 574     dd_distribute_state(dd, state_global, state_local);
 575 }