src/gromacs/domdec/distribute.cpp

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