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