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