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