src/gromacs/domdec/collect.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 functions to collect state data to the master rank.
  38  *
  39  * \author Berk Hess <hess@kth.se>
  40  * \ingroup module_domdec
  41  */
  42
  43 #include "gmxpre.h"
  44
  45 #include "collect.h"
  46
  47 #include "config.h"
  48
  49 #include "gromacs/domdec/domdec_network.h"
  50 #include "gromacs/math/vec.h"
  51 #include "gromacs/mdtypes/state.h"
  52 #include "gromacs/utility/fatalerror.h"
  53
  54 #include "atomdistribution.h"
  55 #include "distribute.h"
  56 #include "domdec_internal.h"
  57
  58 static void dd_collect_cg(gmx_domdec_t* dd, const t_state* state_local)
  59 {
  60     if (state_local->ddp_count == dd->comm->master_cg_ddp_count)
  61     {
  62         /* The master has the correct distribution */
  63         return;
  64     }
  65
  66     gmx::ArrayRef<const int> atomGroups;
  67     int                      nat_home = 0;
  68
  69     if (state_local->ddp_count == dd->ddp_count)
  70     {
  71         /* The local state and DD are in sync, use the DD indices */
  72         atomGroups = gmx::constArrayRefFromArray(dd->globalAtomGroupIndices.data(), dd->ncg_home);
  73         nat_home   = dd->comm->atomRanges.numHomeAtoms();
  74     }
  75     else if (state_local->ddp_count_cg_gl == state_local->ddp_count)
  76     {
  77         /* The DD is out of sync with the local state, but we have stored
  78          * the cg indices with the local state, so we can use those.
  79          */
  80         atomGroups = state_local->cg_gl;
  81         nat_home   = atomGroups.size();
  82     }
  83     else
  84     {
  85         gmx_incons(
  86                 "Attempted to collect a vector for a state for which the charge group distribution "
  87                 "is unknown");
  88     }
  89
  90     AtomDistribution* ma = dd->ma.get();
  91
  92     /* Collect the charge group and atom counts on the master */
  93     int localBuffer[2] = { static_cast<int>(atomGroups.size()), nat_home };
  94     dd_gather(dd, 2 * sizeof(int), localBuffer, DDMASTER(dd) ? ma->intBuffer.data() : nullptr);
  95
  96     if (DDMASTER(dd))
  97     {
  98         int groupOffset = 0;
  99         for (int rank = 0; rank < dd->nnodes; rank++)
 100         {
 101             auto& domainGroups = ma->domainGroups[rank];
 102             int   numGroups    = ma->intBuffer[2 * rank];
 103
 104             domainGroups.atomGroups =
 105                     gmx::constArrayRefFromArray(ma->atomGroups.data() + groupOffset, numGroups);
 106
 107             domainGroups.numAtoms = ma->intBuffer[2 * rank + 1];
 108
 109             groupOffset += numGroups;
 110         }
 111
 112         if (debug)
 113         {
 114             fprintf(debug, "Initial charge group distribution: ");
 115             for (int rank = 0; rank < dd->nnodes; rank++)
 116             {
 117                 fprintf(debug, " %td", ma->domainGroups[rank].atomGroups.ssize());
 118             }
 119             fprintf(debug, "\n");
 120         }
 121
 122         /* Make byte counts and indices */
 123         int offset = 0;
 124         for (int rank = 0; rank < dd->nnodes; rank++)
 125         {
 126             int numGroups                    = ma->domainGroups[rank].atomGroups.size();
 127             ma->intBuffer[rank]              = numGroups * sizeof(int);
 128             ma->intBuffer[dd->nnodes + rank] = offset * sizeof(int);
 129             offset += numGroups;
 130         }
 131     }
 132
 133     /* Collect the charge group indices on the master */
 134     dd_gatherv(dd, atomGroups.size() * sizeof(int), atomGroups.data(),
 135                DDMASTER(dd) ? ma->intBuffer.data() : nullptr,
 136                DDMASTER(dd) ? ma->intBuffer.data() + dd->nnodes : nullptr,
 137                DDMASTER(dd) ? ma->atomGroups.data() : nullptr);
 138
 139     dd->comm->master_cg_ddp_count = state_local->ddp_count;
 140 }
 141
 142 static void dd_collect_vec_sendrecv(gmx_domdec_t*                  dd,
 143                                     gmx::ArrayRef<const gmx::RVec> lv,
 144                                     gmx::ArrayRef<gmx::RVec>       v)
 145 {
 146     if (!DDMASTER(dd))
 147     {
 148 #if GMX_MPI
 149         const int numHomeAtoms = dd->comm->atomRanges.numHomeAtoms();
 150         MPI_Send(const_cast<void*>(static_cast<const void*>(lv.data())),
 151                  numHomeAtoms * sizeof(rvec), MPI_BYTE, dd->masterrank, dd->rank, dd->mpi_comm_all);
 152 #endif
 153     }
 154     else
 155     {
 156         AtomDistribution& ma = *dd->ma;
 157
 158         int rank      = dd->masterrank;
 159         int localAtom = 0;
 160         for (const int& globalAtom : ma.domainGroups[rank].atomGroups)
 161         {
 162             copy_rvec(lv[localAtom++], v[globalAtom]);
 163         }
 164
 165         for (int rank = 0; rank < dd->nnodes; rank++)
 166         {
 167             if (rank != dd->rank)
 168             {
 169                 const auto& domainGroups = ma.domainGroups[rank];
 170
 171                 GMX_RELEASE_ASSERT(v.data() != ma.rvecBuffer.data(),
 172                                    "We need different communication and return buffers");
 173
 174                 /* When we send/recv instead of scatter/gather, we might need
 175                  * to increase the communication buffer size here.
 176                  */
 177                 if (static_cast<size_t>(domainGroups.numAtoms) > ma.rvecBuffer.size())
 178                 {
 179                     ma.rvecBuffer.resize(domainGroups.numAtoms);
 180                 }
 181
 182 #if GMX_MPI
 183                 MPI_Recv(ma.rvecBuffer.data(), domainGroups.numAtoms * sizeof(rvec), MPI_BYTE, rank,
 184                          rank, dd->mpi_comm_all, MPI_STATUS_IGNORE);
 185 #endif
 186                 int localAtom = 0;
 187                 for (const int& globalAtom : domainGroups.atomGroups)
 188                 {
 189                     copy_rvec(ma.rvecBuffer[localAtom++], v[globalAtom]);
 190                 }
 191             }
 192         }
 193     }
 194 }
 195
 196 static void dd_collect_vec_gatherv(gmx_domdec_t*                  dd,
 197                                    gmx::ArrayRef<const gmx::RVec> lv,
 198                                    gmx::ArrayRef<gmx::RVec>       v)
 199 {
 200     int* recvCounts    = nullptr;
 201     int* displacements = nullptr;
 202
 203     if (DDMASTER(dd))
 204     {
 205         get_commbuffer_counts(dd->ma.get(), &recvCounts, &displacements);
 206     }
 207
 208     const int numHomeAtoms = dd->comm->atomRanges.numHomeAtoms();
 209     dd_gatherv(dd, numHomeAtoms * sizeof(rvec), lv.data(), recvCounts, displacements,
 210                DDMASTER(dd) ? dd->ma->rvecBuffer.data() : nullptr);
 211
 212     if (DDMASTER(dd))
 213     {
 214         const AtomDistribution& ma = *dd->ma;
 215
 216         int bufferAtom = 0;
 217         for (int rank = 0; rank < dd->nnodes; rank++)
 218         {
 219             const auto& domainGroups = ma.domainGroups[rank];
 220             for (const int& globalAtom : domainGroups.atomGroups)
 221             {
 222                 copy_rvec(ma.rvecBuffer[bufferAtom++], v[globalAtom]);
 223             }
 224         }
 225     }
 226 }
 227
 228 void dd_collect_vec(gmx_domdec_t*                  dd,
 229                     const t_state*                 state_local,
 230                     gmx::ArrayRef<const gmx::RVec> lv,
 231                     gmx::ArrayRef<gmx::RVec>       v)
 232 {
 233     dd_collect_cg(dd, state_local);
 234
 235     if (dd->nnodes <= c_maxNumRanksUseSendRecvForScatterAndGather)
 236     {
 237         dd_collect_vec_sendrecv(dd, lv, v);
 238     }
 239     else
 240     {
 241         dd_collect_vec_gatherv(dd, lv, v);
 242     }
 243 }
 244
 245
 246 void dd_collect_state(gmx_domdec_t* dd, const t_state* state_local, t_state* state)
 247 {
 248     int nh = state_local->nhchainlength;
 249
 250     if (DDMASTER(dd))
 251     {
 252         GMX_RELEASE_ASSERT(state->nhchainlength == nh,
 253                            "The global and local Nose-Hoover chain lengths should match");
 254
 255         for (int i = 0; i < efptNR; i++)
 256         {
 257             state->lambda[i] = state_local->lambda[i];
 258         }
 259         state->fep_state = state_local->fep_state;
 260         state->veta      = state_local->veta;
 261         state->vol0      = state_local->vol0;
 262         copy_mat(state_local->box, state->box);
 263         copy_mat(state_local->boxv, state->boxv);
 264         copy_mat(state_local->svir_prev, state->svir_prev);
 265         copy_mat(state_local->fvir_prev, state->fvir_prev);
 266         copy_mat(state_local->pres_prev, state->pres_prev);
 267
 268         for (int i = 0; i < state_local->ngtc; i++)
 269         {
 270             for (int j = 0; j < nh; j++)
 271             {
 272                 state->nosehoover_xi[i * nh + j]  = state_local->nosehoover_xi[i * nh + j];
 273                 state->nosehoover_vxi[i * nh + j] = state_local->nosehoover_vxi[i * nh + j];
 274             }
 275             state->therm_integral[i] = state_local->therm_integral[i];
 276         }
 277         for (int i = 0; i < state_local->nnhpres; i++)
 278         {
 279             for (int j = 0; j < nh; j++)
 280             {
 281                 state->nhpres_xi[i * nh + j]  = state_local->nhpres_xi[i * nh + j];
 282                 state->nhpres_vxi[i * nh + j] = state_local->nhpres_vxi[i * nh + j];
 283             }
 284         }
 285         state->baros_integral     = state_local->baros_integral;
 286         state->pull_com_prev_step = state_local->pull_com_prev_step;
 287     }
 288     if (state_local->flags & (1 << estX))
 289     {
 290         auto globalXRef = state ? state->x : gmx::ArrayRef<gmx::RVec>();
 291         dd_collect_vec(dd, state_local, state_local->x, globalXRef);
 292     }
 293     if (state_local->flags & (1 << estV))
 294     {
 295         auto globalVRef = state ? state->v : gmx::ArrayRef<gmx::RVec>();
 296         dd_collect_vec(dd, state_local, state_local->v, globalVRef);
 297     }
 298     if (state_local->flags & (1 << estCGP))
 299     {
 300         auto globalCgpRef = state ? state->cg_p : gmx::ArrayRef<gmx::RVec>();
 301         dd_collect_vec(dd, state_local, state_local->cg_p, globalCgpRef);
 302     }
 303 }