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