src/gromacs/ewald/pme_only.cpp

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 1991-2000, University of Groningen, The Netherlands.
   5  * Copyright (c) 2001-2004, The GROMACS development team.
   6  * Copyright (c) 2013,2014,2015,2016,2017,2018,2019, by the GROMACS development team, led by
   7  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   8  * and including many others, as listed in the AUTHORS file in the
   9  * top-level source directory and at http://www.gromacs.org.
  10  *
  11  * GROMACS is free software; you can redistribute it and/or
  12  * modify it under the terms of the GNU Lesser General Public License
  13  * as published by the Free Software Foundation; either version 2.1
  14  * of the License, or (at your option) any later version.
  15  *
  16  * GROMACS is distributed in the hope that it will be useful,
  17  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  18  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  19  * Lesser General Public License for more details.
  20  *
  21  * You should have received a copy of the GNU Lesser General Public
  22  * License along with GROMACS; if not, see
  23  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
  24  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  25  *
  26  * If you want to redistribute modifications to GROMACS, please
  27  * consider that scientific software is very special. Version
  28  * control is crucial - bugs must be traceable. We will be happy to
  29  * consider code for inclusion in the official distribution, but
  30  * derived work must not be called official GROMACS. Details are found
  31  * in the README & COPYING files - if they are missing, get the
  32  * official version at http://www.gromacs.org.
  33  *
  34  * To help us fund GROMACS development, we humbly ask that you cite
  35  * the research papers on the package. Check out http://www.gromacs.org.
  36  */
  37 /* IMPORTANT FOR DEVELOPERS:
  38  *
  39  * Triclinic pme stuff isn't entirely trivial, and we've experienced
  40  * some bugs during development (many of them due to me). To avoid
  41  * this in the future, please check the following things if you make
  42  * changes in this file:
  43  *
  44  * 1. You should obtain identical (at least to the PME precision)
  45  *    energies, forces, and virial for
  46  *    a rectangular box and a triclinic one where the z (or y) axis is
  47  *    tilted a whole box side. For instance you could use these boxes:
  48  *
  49  *    rectangular       triclinic
  50  *     2  0  0           2  0  0
  51  *     0  2  0           0  2  0
  52  *     0  0  6           2  2  6
  53  *
  54  * 2. You should check the energy conservation in a triclinic box.
  55  *
  56  * It might seem an overkill, but better safe than sorry.
  57  * /Erik 001109
  58  */
  59
  60 #include "gmxpre.h"
  61
  62 #include "config.h"
  63
  64 #include <cassert>
  65 #include <cmath>
  66 #include <cstdio>
  67 #include <cstdlib>
  68 #include <cstring>
  69
  70 #include <memory>
  71 #include <numeric>
  72 #include <vector>
  73
  74 #include "gromacs/domdec/domdec.h"
  75 #include "gromacs/ewald/pme.h"
  76 #include "gromacs/fft/parallel_3dfft.h"
  77 #include "gromacs/fileio/pdbio.h"
  78 #include "gromacs/gmxlib/network.h"
  79 #include "gromacs/gmxlib/nrnb.h"
  80 #include "gromacs/gpu_utils/hostallocator.h"
  81 #include "gromacs/math/gmxcomplex.h"
  82 #include "gromacs/math/units.h"
  83 #include "gromacs/math/vec.h"
  84 #include "gromacs/mdtypes/commrec.h"
  85 #include "gromacs/mdtypes/forceoutput.h"
  86 #include "gromacs/mdtypes/inputrec.h"
  87 #include "gromacs/timing/cyclecounter.h"
  88 #include "gromacs/timing/wallcycle.h"
  89 #include "gromacs/utility/fatalerror.h"
  90 #include "gromacs/utility/futil.h"
  91 #include "gromacs/utility/gmxmpi.h"
  92 #include "gromacs/utility/gmxomp.h"
  93 #include "gromacs/utility/smalloc.h"
  94
  95 #include "pme_gpu_internal.h"
  96 #include "pme_internal.h"
  97 #include "pme_pp_communication.h"
  98
  99 //! Contains information about the PP ranks that partner this PME rank.
 100 struct PpRanks
 101 {
 102     //! The MPI rank ID of this partner PP rank.
 103     int rankId;
 104     //! The number of atoms to communicate with this partner PP rank.
 105     int numAtoms;
 106 };
 107
 108 /*! \brief Master PP-PME communication data structure */
 109 struct gmx_pme_pp {
 110     MPI_Comm             mpi_comm_mysim; /**< MPI communicator for this simulation */
 111     std::vector<PpRanks> ppRanks;        /**< The PP partner ranks                 */
 112     int                  peerRankId;     /**< The peer PP rank id                  */
 113     //@{
 114     /**< Vectors of A- and B-state parameters used to transfer vectors to PME ranks  */
 115     gmx::PaddedHostVector<real> chargeA;
 116     std::vector<real>           chargeB;
 117     std::vector<real>           sqrt_c6A;
 118     std::vector<real>           sqrt_c6B;
 119     std::vector<real>           sigmaA;
 120     std::vector<real>           sigmaB;
 121     //@}
 122     gmx::HostVector<gmx::RVec>  x; /**< Vector of atom coordinates to transfer to PME ranks */
 123     std::vector<gmx::RVec>      f; /**< Vector of atom forces received from PME ranks */
 124     //@{
 125     /**< Vectors of MPI objects used in non-blocking communication between multiple PP ranks per PME rank */
 126     std::vector<MPI_Request> req;
 127     std::vector<MPI_Status>  stat;
 128     //@}
 129 };
 130
 131 /*! \brief Initialize the PME-only side of the PME <-> PP communication */
 132 static std::unique_ptr<gmx_pme_pp> gmx_pme_pp_init(const t_commrec *cr)
 133 {
 134     auto pme_pp = std::make_unique<gmx_pme_pp>();
 135
 136 #if GMX_MPI
 137     int rank;
 138
 139     pme_pp->mpi_comm_mysim = cr->mpi_comm_mysim;
 140     MPI_Comm_rank(cr->mpi_comm_mygroup, &rank);
 141     auto ppRanks = get_pme_ddranks(cr, rank);
 142     pme_pp->ppRanks.reserve(ppRanks.size());
 143     for (const auto &ppRankId : ppRanks)
 144     {
 145         pme_pp->ppRanks.push_back({ppRankId, 0});
 146     }
 147     // The peer PP rank is the last one.
 148     pme_pp->peerRankId = pme_pp->ppRanks.back().rankId;
 149     pme_pp->req.resize(eCommType_NR*pme_pp->ppRanks.size());
 150     pme_pp->stat.resize(eCommType_NR*pme_pp->ppRanks.size());
 151 #else
 152     GMX_UNUSED_VALUE(cr);
 153 #endif
 154
 155     return pme_pp;
 156 }
 157
 158 static void reset_pmeonly_counters(gmx_wallcycle_t           wcycle,
 159                                    gmx_walltime_accounting_t walltime_accounting,
 160                                    t_nrnb                   *nrnb,
 161                                    int64_t                   step,
 162                                    bool                      useGpuForPme)
 163 {
 164     /* Reset all the counters related to performance over the run */
 165     wallcycle_stop(wcycle, ewcRUN);
 166     wallcycle_reset_all(wcycle);
 167     init_nrnb(nrnb);
 168     wallcycle_start(wcycle, ewcRUN);
 169     walltime_accounting_reset_time(walltime_accounting, step);
 170
 171     if (useGpuForPme)
 172     {
 173         resetGpuProfiler();
 174     }
 175 }
 176
 177 static gmx_pme_t *gmx_pmeonly_switch(std::vector<gmx_pme_t *> *pmedata,
 178                                      const ivec grid_size,
 179                                      real ewaldcoeff_q, real ewaldcoeff_lj,
 180                                      const t_commrec *cr, const t_inputrec *ir)
 181 {
 182     GMX_ASSERT(pmedata, "Bad PME tuning list pointer");
 183     for (auto &pme : *pmedata)
 184     {
 185         GMX_ASSERT(pme, "Bad PME tuning list element pointer");
 186         if (pme->nkx == grid_size[XX] &&
 187             pme->nky == grid_size[YY] &&
 188             pme->nkz == grid_size[ZZ])
 189         {
 190             /* Here we have found an existing PME data structure that suits us.
 191              * However, in the GPU case, we have to reinitialize it - there's only one GPU structure.
 192              * This should not cause actual GPU reallocations, at least (the allocated buffers are never shrunk).
 193              * So, just some grid size updates in the GPU kernel parameters.
 194              * TODO: this should be something like gmx_pme_update_split_params()
 195              */
 196             gmx_pme_reinit(&pme, cr, pme, ir, grid_size, ewaldcoeff_q, ewaldcoeff_lj);
 197             return pme;
 198         }
 199     }
 200
 201     const auto &pme          = pmedata->back();
 202     gmx_pme_t  *newStructure = nullptr;
 203     // Copy last structure with new grid params
 204     gmx_pme_reinit(&newStructure, cr, pme, ir, grid_size, ewaldcoeff_q, ewaldcoeff_lj);
 205     pmedata->push_back(newStructure);
 206     return newStructure;
 207 }
 208
 209 /*! \brief Called by PME-only ranks to receive coefficients and coordinates
 210  *
 211  * \param[in,out] pme_pp    PME-PP communication structure.
 212  * \param[out] natoms       Number of received atoms.
 213  * \param[out] box        System box, if received.
 214  * \param[out] maxshift_x        Maximum shift in X direction, if received.
 215  * \param[out] maxshift_y        Maximum shift in Y direction, if received.
 216  * \param[out] lambda_q         Free-energy lambda for electrostatics, if received.
 217  * \param[out] lambda_lj         Free-energy lambda for Lennard-Jones, if received.
 218  * \param[out] bEnerVir          Set to true if this is an energy/virial calculation step, otherwise set to false.
 219  * \param[out] step              MD integration step number.
 220  * \param[out] grid_size         PME grid size, if received.
 221  * \param[out] ewaldcoeff_q         Ewald cut-off parameter for electrostatics, if received.
 222  * \param[out] ewaldcoeff_lj         Ewald cut-off parameter for Lennard-Jones, if received.
 223  * \param[out] atomSetChanged    Set to true only if the local domain atom data (charges/coefficients)
 224  *                               has been received (after DD) and should be reinitialized. Otherwise not changed.
 225  *
 226  * \retval pmerecvqxX             All parameters were set, chargeA and chargeB can be NULL.
 227  * \retval pmerecvqxFINISH        No parameters were set.
 228  * \retval pmerecvqxSWITCHGRID    Only grid_size and *ewaldcoeff were set.
 229  * \retval pmerecvqxRESETCOUNTERS *step was set.
 230  */
 231 static int gmx_pme_recv_coeffs_coords(gmx_pme_pp        *pme_pp,
 232                                       int               *natoms,
 233                                       matrix             box,
 234                                       int               *maxshift_x,
 235                                       int               *maxshift_y,
 236                                       real              *lambda_q,
 237                                       real              *lambda_lj,
 238                                       gmx_bool          *bEnerVir,
 239                                       int64_t           *step,
 240                                       ivec              *grid_size,
 241                                       real              *ewaldcoeff_q,
 242                                       real              *ewaldcoeff_lj,
 243                                       bool              *atomSetChanged)
 244 {
 245     int status = -1;
 246     int nat    = 0;
 247
 248 #if GMX_MPI
 249     unsigned int flags    = 0;
 250     int          messages = 0;
 251
 252     do
 253     {
 254         gmx_pme_comm_n_box_t cnb;
 255         cnb.flags = 0;
 256
 257         /* Receive the send count, box and time step from the peer PP node */
 258         MPI_Recv(&cnb, sizeof(cnb), MPI_BYTE,
 259                  pme_pp->peerRankId, eCommType_CNB,
 260                  pme_pp->mpi_comm_mysim, MPI_STATUS_IGNORE);
 261
 262         /* We accumulate all received flags */
 263         flags |= cnb.flags;
 264
 265         *step  = cnb.step;
 266
 267         if (debug)
 268         {
 269             fprintf(debug, "PME only rank receiving:%s%s%s%s%s\n",
 270                     (cnb.flags & PP_PME_CHARGE)        ? " charges" : "",
 271                     (cnb.flags & PP_PME_COORD )        ? " coordinates" : "",
 272                     (cnb.flags & PP_PME_FINISH)        ? " finish" : "",
 273                     (cnb.flags & PP_PME_SWITCHGRID)    ? " switch grid" : "",
 274                     (cnb.flags & PP_PME_RESETCOUNTERS) ? " reset counters" : "");
 275         }
 276
 277         if (cnb.flags & PP_PME_FINISH)
 278         {
 279             status = pmerecvqxFINISH;
 280         }
 281
 282         if (cnb.flags & PP_PME_SWITCHGRID)
 283         {
 284             /* Special case, receive the new parameters and return */
 285             copy_ivec(cnb.grid_size, *grid_size);
 286             *ewaldcoeff_q  = cnb.ewaldcoeff_q;
 287             *ewaldcoeff_lj = cnb.ewaldcoeff_lj;
 288
 289             status         = pmerecvqxSWITCHGRID;
 290         }
 291
 292         if (cnb.flags & PP_PME_RESETCOUNTERS)
 293         {
 294             /* Special case, receive the step (set above) and return */
 295             status = pmerecvqxRESETCOUNTERS;
 296         }
 297
 298         if (cnb.flags & (PP_PME_CHARGE | PP_PME_SQRTC6 | PP_PME_SIGMA))
 299         {
 300             *atomSetChanged = true;
 301
 302             /* Receive the send counts from the other PP nodes */
 303             for (auto &sender : pme_pp->ppRanks)
 304             {
 305                 if (sender.rankId == pme_pp->peerRankId)
 306                 {
 307                     sender.numAtoms = cnb.natoms;
 308                 }
 309                 else
 310                 {
 311                     MPI_Irecv(&sender.numAtoms, sizeof(sender.numAtoms),
 312                               MPI_BYTE,
 313                               sender.rankId, eCommType_CNB,
 314                               pme_pp->mpi_comm_mysim, &pme_pp->req[messages++]);
 315                 }
 316             }
 317             MPI_Waitall(messages, pme_pp->req.data(), pme_pp->stat.data());
 318             messages = 0;
 319
 320             nat = 0;
 321             for (const auto &sender : pme_pp->ppRanks)
 322             {
 323                 nat += sender.numAtoms;
 324             }
 325
 326             if (cnb.flags & PP_PME_CHARGE)
 327             {
 328                 pme_pp->chargeA.resizeWithPadding(nat);
 329             }
 330             if (cnb.flags & PP_PME_CHARGEB)
 331             {
 332                 pme_pp->chargeB.resize(nat);
 333             }
 334             if (cnb.flags & PP_PME_SQRTC6)
 335             {
 336                 pme_pp->sqrt_c6A.resize(nat);
 337             }
 338             if (cnb.flags & PP_PME_SQRTC6B)
 339             {
 340                 pme_pp->sqrt_c6B.resize(nat);
 341             }
 342             if (cnb.flags & PP_PME_SIGMA)
 343             {
 344                 pme_pp->sigmaA.resize(nat);
 345             }
 346             if (cnb.flags & PP_PME_SIGMAB)
 347             {
 348                 pme_pp->sigmaB.resize(nat);
 349             }
 350             pme_pp->x.resize(nat);
 351             pme_pp->f.resize(nat);
 352
 353             /* maxshift is sent when the charges are sent */
 354             *maxshift_x = cnb.maxshift_x;
 355             *maxshift_y = cnb.maxshift_y;
 356
 357             /* Receive the charges in place */
 358             for (int q = 0; q < eCommType_NR; q++)
 359             {
 360                 real *bufferPtr;
 361
 362                 if (!(cnb.flags & (PP_PME_CHARGE<<q)))
 363                 {
 364                     continue;
 365                 }
 366                 switch (q)
 367                 {
 368                     case eCommType_ChargeA: bufferPtr = pme_pp->chargeA.data();  break;
 369                     case eCommType_ChargeB: bufferPtr = pme_pp->chargeB.data();  break;
 370                     case eCommType_SQRTC6A: bufferPtr = pme_pp->sqrt_c6A.data(); break;
 371                     case eCommType_SQRTC6B: bufferPtr = pme_pp->sqrt_c6B.data(); break;
 372                     case eCommType_SigmaA:  bufferPtr = pme_pp->sigmaA.data();   break;
 373                     case eCommType_SigmaB:  bufferPtr = pme_pp->sigmaB.data();   break;
 374                     default: gmx_incons("Wrong eCommType");
 375                 }
 376                 nat = 0;
 377                 for (const auto &sender : pme_pp->ppRanks)
 378                 {
 379                     if (sender.numAtoms > 0)
 380                     {
 381                         MPI_Irecv(bufferPtr+nat,
 382                                   sender.numAtoms*sizeof(real),
 383                                   MPI_BYTE,
 384                                   sender.rankId, q,
 385                                   pme_pp->mpi_comm_mysim,
 386                                   &pme_pp->req[messages++]);
 387                         nat += sender.numAtoms;
 388                         if (debug)
 389                         {
 390                             fprintf(debug, "Received from PP rank %d: %d %s\n",
 391                                     sender.rankId, sender.numAtoms,
 392                                     (q == eCommType_ChargeA ||
 393                                      q == eCommType_ChargeB) ? "charges" : "params");
 394                         }
 395                     }
 396                 }
 397             }
 398         }
 399
 400         if (cnb.flags & PP_PME_COORD)
 401         {
 402             /* The box, FE flag and lambda are sent along with the coordinates
 403              *  */
 404             copy_mat(cnb.box, box);
 405             *lambda_q       = cnb.lambda_q;
 406             *lambda_lj      = cnb.lambda_lj;
 407             *bEnerVir       = ((cnb.flags & PP_PME_ENER_VIR) != 0u);
 408             *step           = cnb.step;
 409
 410             /* Receive the coordinates in place */
 411             nat = 0;
 412             for (const auto &sender : pme_pp->ppRanks)
 413             {
 414                 if (sender.numAtoms > 0)
 415                 {
 416                     MPI_Irecv(pme_pp->x[nat],
 417                               sender.numAtoms*sizeof(rvec),
 418                               MPI_BYTE,
 419                               sender.rankId, eCommType_COORD,
 420                               pme_pp->mpi_comm_mysim, &pme_pp->req[messages++]);
 421                     nat += sender.numAtoms;
 422                     if (debug)
 423                     {
 424                         fprintf(debug, "Received from PP rank %d: %d "
 425                                 "coordinates\n",
 426                                 sender.rankId, sender.numAtoms);
 427                     }
 428                 }
 429             }
 430
 431             status = pmerecvqxX;
 432         }
 433
 434         /* Wait for the coordinates and/or charges to arrive */
 435         MPI_Waitall(messages, pme_pp->req.data(), pme_pp->stat.data());
 436         messages = 0;
 437     }
 438     while (status == -1);
 439 #else
 440     GMX_UNUSED_VALUE(pme_pp);
 441     GMX_UNUSED_VALUE(box);
 442     GMX_UNUSED_VALUE(maxshift_x);
 443     GMX_UNUSED_VALUE(maxshift_y);
 444     GMX_UNUSED_VALUE(lambda_q);
 445     GMX_UNUSED_VALUE(lambda_lj);
 446     GMX_UNUSED_VALUE(bEnerVir);
 447     GMX_UNUSED_VALUE(step);
 448     GMX_UNUSED_VALUE(grid_size);
 449     GMX_UNUSED_VALUE(ewaldcoeff_q);
 450     GMX_UNUSED_VALUE(ewaldcoeff_lj);
 451     GMX_UNUSED_VALUE(atomSetChanged);
 452
 453     status = pmerecvqxX;
 454 #endif
 455
 456     if (status == pmerecvqxX)
 457     {
 458         *natoms   = nat;
 459     }
 460
 461     return status;
 462 }
 463
 464 /*! \brief Send the PME mesh force, virial and energy to the PP-only ranks. */
 465 static void gmx_pme_send_force_vir_ener(gmx_pme_pp *pme_pp,
 466                                         const PmeOutput &output,
 467                                         real dvdlambda_q, real dvdlambda_lj,
 468                                         float cycles)
 469 {
 470 #if GMX_MPI
 471     gmx_pme_comm_vir_ene_t cve;
 472     int                    messages, ind_start, ind_end;
 473     cve.cycles = cycles;
 474
 475     /* Now the evaluated forces have to be transferred to the PP nodes */
 476     messages = 0;
 477     ind_end  = 0;
 478     for (const auto &receiver : pme_pp->ppRanks)
 479     {
 480         ind_start = ind_end;
 481         ind_end   = ind_start + receiver.numAtoms;
 482         if (MPI_Isend(const_cast<void *>(static_cast<const void *>(output.forces_[ind_start])),
 483                       (ind_end-ind_start)*sizeof(rvec), MPI_BYTE,
 484                       receiver.rankId, 0,
 485                       pme_pp->mpi_comm_mysim, &pme_pp->req[messages++]) != 0)
 486         {
 487             gmx_comm("MPI_Isend failed in do_pmeonly");
 488         }
 489     }
 490
 491     /* send virial and energy to our last PP node */
 492     copy_mat(output.coulombVirial_, cve.vir_q);
 493     copy_mat(output.lennardJonesVirial_, cve.vir_lj);
 494     cve.energy_q     = output.coulombEnergy_;
 495     cve.energy_lj    = output.lennardJonesEnergy_;
 496     cve.dvdlambda_q  = dvdlambda_q;
 497     cve.dvdlambda_lj = dvdlambda_lj;
 498     /* check for the signals to send back to a PP node */
 499     cve.stop_cond = gmx_get_stop_condition();
 500
 501     cve.cycles = cycles;
 502
 503     if (debug)
 504     {
 505         fprintf(debug, "PME rank sending to PP rank %d: virial and energy\n",
 506                 pme_pp->peerRankId);
 507     }
 508     MPI_Isend(&cve, sizeof(cve), MPI_BYTE,
 509               pme_pp->peerRankId, 1,
 510               pme_pp->mpi_comm_mysim, &pme_pp->req[messages++]);
 511
 512     /* Wait for the forces to arrive */
 513     MPI_Waitall(messages, pme_pp->req.data(), pme_pp->stat.data());
 514 #else
 515     gmx_call("MPI not enabled");
 516     GMX_UNUSED_VALUE(pme_pp);
 517     GMX_UNUSED_VALUE(output);
 518     GMX_UNUSED_VALUE(dvdlambda_q);
 519     GMX_UNUSED_VALUE(dvdlambda_lj);
 520     GMX_UNUSED_VALUE(cycles);
 521 #endif
 522 }
 523
 524 int gmx_pmeonly(struct gmx_pme_t *pme,
 525                 const t_commrec *cr, t_nrnb *mynrnb,
 526                 gmx_wallcycle  *wcycle,
 527                 gmx_walltime_accounting_t walltime_accounting,
 528                 t_inputrec *ir, PmeRunMode runMode)
 529 {
 530     int                ret;
 531     int                natoms = 0;
 532     matrix             box;
 533     real               lambda_q   = 0;
 534     real               lambda_lj  = 0;
 535     int                maxshift_x = 0, maxshift_y = 0;
 536     real               dvdlambda_q, dvdlambda_lj;
 537     float              cycles;
 538     int                count;
 539     gmx_bool           bEnerVir = FALSE;
 540     int64_t            step;
 541
 542     /* This data will only use with PME tuning, i.e. switching PME grids */
 543     std::vector<gmx_pme_t *> pmedata;
 544     pmedata.push_back(pme);
 545
 546     auto       pme_pp       = gmx_pme_pp_init(cr);
 547     //TODO the variable below should be queried from the task assignment info
 548     const bool useGpuForPme = (runMode == PmeRunMode::GPU) || (runMode == PmeRunMode::Mixed);
 549     if (useGpuForPme)
 550     {
 551         changePinningPolicy(&pme_pp->chargeA, pme_get_pinning_policy());
 552         changePinningPolicy(&pme_pp->x, pme_get_pinning_policy());
 553     }
 554
 555     init_nrnb(mynrnb);
 556
 557     count = 0;
 558     do /****** this is a quasi-loop over time steps! */
 559     {
 560         /* The reason for having a loop here is PME grid tuning/switching */
 561         do
 562         {
 563             /* Domain decomposition */
 564             ivec newGridSize;
 565             bool atomSetChanged = false;
 566             real ewaldcoeff_q   = 0, ewaldcoeff_lj = 0;
 567             ret = gmx_pme_recv_coeffs_coords(pme_pp.get(),
 568                                              &natoms,
 569                                              box,
 570                                              &maxshift_x, &maxshift_y,
 571                                              &lambda_q, &lambda_lj,
 572                                              &bEnerVir,
 573                                              &step,
 574                                              &newGridSize,
 575                                              &ewaldcoeff_q,
 576                                              &ewaldcoeff_lj,
 577                                              &atomSetChanged);
 578
 579             if (ret == pmerecvqxSWITCHGRID)
 580             {
 581                 /* Switch the PME grid to newGridSize */
 582                 pme = gmx_pmeonly_switch(&pmedata, newGridSize, ewaldcoeff_q, ewaldcoeff_lj, cr, ir);
 583             }
 584
 585             if (atomSetChanged)
 586             {
 587                 gmx_pme_reinit_atoms(pme, natoms, pme_pp->chargeA.data());
 588             }
 589
 590             if (ret == pmerecvqxRESETCOUNTERS)
 591             {
 592                 /* Reset the cycle and flop counters */
 593                 reset_pmeonly_counters(wcycle, walltime_accounting, mynrnb, step, useGpuForPme);
 594             }
 595         }
 596         while (ret == pmerecvqxSWITCHGRID || ret == pmerecvqxRESETCOUNTERS);
 597
 598         if (ret == pmerecvqxFINISH)
 599         {
 600             /* We should stop: break out of the loop */
 601             break;
 602         }
 603
 604         if (count == 0)
 605         {
 606             wallcycle_start(wcycle, ewcRUN);
 607             walltime_accounting_start_time(walltime_accounting);
 608         }
 609
 610         wallcycle_start(wcycle, ewcPMEMESH);
 611
 612         dvdlambda_q  = 0;
 613         dvdlambda_lj = 0;
 614
 615         // TODO Make a struct of array refs onto these per-atom fields
 616         // of pme_pp (maybe box, energy and virial, too; and likewise
 617         // from mdatoms for the other call to gmx_pme_do), so we have
 618         // fewer lines of code and less parameter passing.
 619         const int pmeFlags = GMX_PME_DO_ALL_F | (bEnerVir ? GMX_PME_CALC_ENER_VIR : 0);
 620         PmeOutput output   = {gmx::EmptyArrayRef {}, 0, {{0}}, 0, {{0}}};
 621         if (useGpuForPme)
 622         {
 623             const bool boxChanged = false;
 624             //TODO this should be set properly by gmx_pme_recv_coeffs_coords,
 625             // or maybe use inputrecDynamicBox(ir), at the very least - change this when this codepath is tested!
 626             pme_gpu_prepare_computation(pme, boxChanged, box, wcycle, pmeFlags);
 627             pme_gpu_launch_spread(pme, as_rvec_array(pme_pp->x.data()), wcycle);
 628             pme_gpu_launch_complex_transforms(pme, wcycle);
 629             pme_gpu_launch_gather(pme, wcycle, PmeForceOutputHandling::Set);
 630             output = pme_gpu_wait_finish_task(pme, pmeFlags, wcycle);
 631             pme_gpu_reinit_computation(pme, wcycle);
 632         }
 633         else
 634         {
 635             gmx_pme_do(pme, 0, natoms, as_rvec_array(pme_pp->x.data()), as_rvec_array(pme_pp->f.data()),
 636                        pme_pp->chargeA.data(), pme_pp->chargeB.data(),
 637                        pme_pp->sqrt_c6A.data(), pme_pp->sqrt_c6B.data(),
 638                        pme_pp->sigmaA.data(), pme_pp->sigmaB.data(), box,
 639                        cr, maxshift_x, maxshift_y, mynrnb, wcycle,
 640                        output.coulombVirial_, output.lennardJonesVirial_,
 641                        &output.coulombEnergy_, &output.lennardJonesEnergy_,
 642                        lambda_q, lambda_lj, &dvdlambda_q, &dvdlambda_lj,
 643                        pmeFlags);
 644             output.forces_ = pme_pp->f;
 645         }
 646
 647         cycles = wallcycle_stop(wcycle, ewcPMEMESH);
 648
 649         gmx_pme_send_force_vir_ener(pme_pp.get(), output,
 650                                     dvdlambda_q, dvdlambda_lj, cycles);
 651
 652         count++;
 653     } /***** end of quasi-loop, we stop with the break above */
 654     while (TRUE);
 655
 656     walltime_accounting_end_time(walltime_accounting);
 657
 658     return 0;
 659 }