2 * This file is part of the GROMACS molecular simulation package.
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.
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.
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.
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.
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.
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.
38 /* IMPORTANT FOR DEVELOPERS:
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:
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:
50 * rectangular triclinic
55 * 2. You should check the energy conservation in a triclinic box.
57 * It might seem an overkill, but better safe than sorry.
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"
103 #include "pme_gpu_internal.h"
104 #include "pme_internal.h"
105 #include "pme_output.h"
106 #include "pme_pp_communication.h"
108 /*! \brief Master PP-PME communication data structure */
111 MPI_Comm mpi_comm_mysim; /**< MPI communicator for this simulation */
112 std::vector<PpRanks> ppRanks; /**< The PP partner ranks */
113 int peerRankId; /**< The peer PP rank id */
115 /**< Vectors of A- and B-state parameters used to transfer vectors to PME ranks */
116 gmx::PaddedHostVector<real> chargeA;
117 std::vector<real> chargeB;
118 std::vector<real> sqrt_c6A;
119 std::vector<real> sqrt_c6B;
120 std::vector<real> sigmaA;
121 std::vector<real> sigmaB;
123 gmx::HostVector<gmx::RVec> x; /**< Vector of atom coordinates to transfer to PME ranks */
124 std::vector<gmx::RVec> f; /**< Vector of atom forces received from PME ranks */
126 /**< Vectors of MPI objects used in non-blocking communication between multiple PP ranks per PME rank */
127 std::vector<MPI_Request> req;
128 std::vector<MPI_Status> stat;
131 /*! \brief object for receiving coordinates using communications operating on GPU memory space */
132 std::unique_ptr<gmx::PmeCoordinateReceiverGpu> pmeCoordinateReceiverGpu;
133 /*! \brief object for sending PME force using communications operating on GPU memory space */
134 std::unique_ptr<gmx::PmeForceSenderGpu> pmeForceSenderGpu;
136 /*! \brief whether GPU direct communications are active for PME-PP transfers */
137 bool useGpuDirectComm = false;
138 /*! \brief whether GPU direct communications should send forces directly to remote GPU memory */
139 bool sendForcesDirectToPpGpu = false;
142 /*! \brief Initialize the PME-only side of the PME <-> PP communication */
143 static std::unique_ptr<gmx_pme_pp> gmx_pme_pp_init(const t_commrec* cr)
145 auto pme_pp = std::make_unique<gmx_pme_pp>();
150 pme_pp->mpi_comm_mysim = cr->mpi_comm_mysim;
151 MPI_Comm_rank(cr->mpi_comm_mygroup, &rank);
152 auto ppRanks = get_pme_ddranks(cr, rank);
153 pme_pp->ppRanks.reserve(ppRanks.size());
154 for (const auto& ppRankId : ppRanks)
156 pme_pp->ppRanks.push_back({ ppRankId, 0 });
158 // The peer PP rank is the last one.
159 pme_pp->peerRankId = pme_pp->ppRanks.back().rankId;
160 pme_pp->req.resize(eCommType_NR * pme_pp->ppRanks.size());
161 pme_pp->stat.resize(eCommType_NR * pme_pp->ppRanks.size());
163 GMX_UNUSED_VALUE(cr);
169 static void reset_pmeonly_counters(gmx_wallcycle* wcycle,
170 gmx_walltime_accounting_t walltime_accounting,
175 /* Reset all the counters related to performance over the run */
176 wallcycle_stop(wcycle, WallCycleCounter::Run);
177 wallcycle_reset_all(wcycle);
179 wallcycle_start(wcycle, WallCycleCounter::Run);
180 walltime_accounting_reset_time(walltime_accounting, step);
188 static gmx_pme_t* gmx_pmeonly_switch(std::vector<gmx_pme_t*>* pmedata,
189 const ivec grid_size,
193 const t_inputrec* ir)
195 GMX_ASSERT(pmedata, "Bad PME tuning list pointer");
196 for (auto& pme : *pmedata)
198 GMX_ASSERT(pme, "Bad PME tuning list element pointer");
199 if (gmx_pme_grid_matches(*pme, grid_size))
201 /* Here we have found an existing PME data structure that suits us.
202 * However, in the GPU case, we have to reinitialize it - there's only one GPU structure.
203 * This should not cause actual GPU reallocations, at least (the allocated buffers are never shrunk).
204 * So, just some grid size updates in the GPU kernel parameters.
205 * TODO: this should be something like gmx_pme_update_split_params()
207 gmx_pme_reinit(&pme, cr, pme, ir, grid_size, ewaldcoeff_q, ewaldcoeff_lj);
212 const auto& pme = pmedata->back();
213 gmx_pme_t* newStructure = nullptr;
214 // Copy last structure with new grid params
215 gmx_pme_reinit(&newStructure, cr, pme, ir, grid_size, ewaldcoeff_q, ewaldcoeff_lj);
216 pmedata->push_back(newStructure);
220 /*! \brief Called by PME-only ranks to receive coefficients and coordinates
222 * Note that with GPU direct communication the transfer is only initiated, it is the responsibility
223 * of the caller to synchronize prior to launching spread.
225 * \param[in] pme PME data structure.
226 * \param[in,out] pme_pp PME-PP communication structure.
227 * \param[out] natoms Number of received atoms.
228 * \param[out] box System box, if received.
229 * \param[out] maxshift_x Maximum shift in X direction, if received.
230 * \param[out] maxshift_y Maximum shift in Y direction, if received.
231 * \param[out] lambda_q Free-energy lambda for electrostatics, if received.
232 * \param[out] lambda_lj Free-energy lambda for Lennard-Jones, if received.
233 * \param[out] computeEnergyAndVirial Set to true if this is an energy/virial calculation
234 * step, otherwise set to false.
235 * \param[out] step MD integration step number.
236 * \param[out] grid_size PME grid size, if received.
237 * \param[out] ewaldcoeff_q Ewald cut-off parameter for electrostatics, if received.
238 * \param[out] ewaldcoeff_lj Ewald cut-off parameter for Lennard-Jones, if received.
239 * \param[in] useGpuForPme Flag on whether PME is on GPU.
240 * \param[in] stateGpu GPU state propagator object.
241 * \param[in] runMode PME run mode.
243 * \retval pmerecvqxX All parameters were set, chargeA and chargeB can be NULL.
244 * \retval pmerecvqxFINISH No parameters were set.
245 * \retval pmerecvqxSWITCHGRID Only grid_size and *ewaldcoeff were set.
246 * \retval pmerecvqxRESETCOUNTERS *step was set.
248 static int gmx_pme_recv_coeffs_coords(struct gmx_pme_t* pme,
256 gmx_bool* computeEnergyAndVirial,
262 gmx::StatePropagatorDataGpu* stateGpu,
263 PmeRunMode gmx_unused runMode)
270 bool atomSetChanged = false;
274 gmx_pme_comm_n_box_t cnb;
277 /* Receive the send count, box and time step from the peer PP node */
278 MPI_Recv(&cnb, sizeof(cnb), MPI_BYTE, pme_pp->peerRankId, eCommType_CNB, pme_pp->mpi_comm_mysim, MPI_STATUS_IGNORE);
285 "PME only rank receiving:%s%s%s%s%s\n",
286 (cnb.flags & PP_PME_CHARGE) ? " charges" : "",
287 (cnb.flags & PP_PME_COORD) ? " coordinates" : "",
288 (cnb.flags & PP_PME_FINISH) ? " finish" : "",
289 (cnb.flags & PP_PME_SWITCHGRID) ? " switch grid" : "",
290 (cnb.flags & PP_PME_RESETCOUNTERS) ? " reset counters" : "");
293 pme_pp->useGpuDirectComm = ((cnb.flags & PP_PME_GPUCOMMS) != 0);
294 GMX_ASSERT(!pme_pp->useGpuDirectComm || (pme_pp->pmeForceSenderGpu != nullptr),
295 "The use of GPU direct communication for PME-PP is enabled, "
296 "but the PME GPU force reciever object does not exist");
297 pme_pp->sendForcesDirectToPpGpu = ((cnb.flags & PP_PME_RECVFTOGPU) != 0);
299 if (cnb.flags & PP_PME_FINISH)
301 status = pmerecvqxFINISH;
304 if (cnb.flags & PP_PME_SWITCHGRID)
306 /* Special case, receive the new parameters and return */
307 copy_ivec(cnb.grid_size, *grid_size);
308 *ewaldcoeff_q = cnb.ewaldcoeff_q;
309 *ewaldcoeff_lj = cnb.ewaldcoeff_lj;
311 status = pmerecvqxSWITCHGRID;
314 if (cnb.flags & PP_PME_RESETCOUNTERS)
316 /* Special case, receive the step (set above) and return */
317 status = pmerecvqxRESETCOUNTERS;
320 if (cnb.flags & (PP_PME_CHARGE | PP_PME_SQRTC6 | PP_PME_SIGMA))
322 atomSetChanged = true;
324 /* Receive the send counts from the other PP nodes */
325 for (auto& sender : pme_pp->ppRanks)
327 if (sender.rankId == pme_pp->peerRankId)
329 sender.numAtoms = cnb.natoms;
333 MPI_Irecv(&sender.numAtoms,
334 sizeof(sender.numAtoms),
338 pme_pp->mpi_comm_mysim,
339 &pme_pp->req[messages++]);
342 MPI_Waitall(messages, pme_pp->req.data(), pme_pp->stat.data());
346 for (const auto& sender : pme_pp->ppRanks)
348 nat += sender.numAtoms;
351 if (cnb.flags & PP_PME_CHARGE)
353 pme_pp->chargeA.resizeWithPadding(nat);
355 if (cnb.flags & PP_PME_CHARGEB)
357 pme_pp->chargeB.resize(nat);
359 if (cnb.flags & PP_PME_SQRTC6)
361 pme_pp->sqrt_c6A.resize(nat);
363 if (cnb.flags & PP_PME_SQRTC6B)
365 pme_pp->sqrt_c6B.resize(nat);
367 if (cnb.flags & PP_PME_SIGMA)
369 pme_pp->sigmaA.resize(nat);
371 if (cnb.flags & PP_PME_SIGMAB)
373 pme_pp->sigmaB.resize(nat);
375 pme_pp->x.resize(nat);
376 pme_pp->f.resize(nat);
378 /* maxshift is sent when the charges are sent */
379 *maxshift_x = cnb.maxshift_x;
380 *maxshift_y = cnb.maxshift_y;
382 /* Receive the charges in place */
383 for (int q = 0; q < eCommType_NR; q++)
387 if (!(cnb.flags & (PP_PME_CHARGE << q)))
393 case eCommType_ChargeA: bufferPtr = pme_pp->chargeA.data(); break;
394 case eCommType_ChargeB: bufferPtr = pme_pp->chargeB.data(); break;
395 case eCommType_SQRTC6A: bufferPtr = pme_pp->sqrt_c6A.data(); break;
396 case eCommType_SQRTC6B: bufferPtr = pme_pp->sqrt_c6B.data(); break;
397 case eCommType_SigmaA: bufferPtr = pme_pp->sigmaA.data(); break;
398 case eCommType_SigmaB: bufferPtr = pme_pp->sigmaB.data(); break;
399 default: gmx_incons("Wrong eCommType");
402 for (const auto& sender : pme_pp->ppRanks)
404 if (sender.numAtoms > 0)
406 MPI_Irecv(bufferPtr + nat,
407 sender.numAtoms * sizeof(real),
411 pme_pp->mpi_comm_mysim,
412 &pme_pp->req[messages++]);
413 nat += sender.numAtoms;
417 "Received from PP rank %d: %d %s\n",
420 (q == eCommType_ChargeA || q == eCommType_ChargeB) ? "charges"
428 if (cnb.flags & PP_PME_COORD)
432 gmx_pme_reinit_atoms(pme, nat, pme_pp->chargeA, pme_pp->chargeB);
435 stateGpu->reinit(nat, nat);
436 pme_gpu_set_device_x(pme, stateGpu->getCoordinates());
438 if (pme_pp->useGpuDirectComm)
440 GMX_ASSERT(runMode == PmeRunMode::GPU,
441 "GPU Direct PME-PP communication has been enabled, "
442 "but PME run mode is not PmeRunMode::GPU\n");
444 // This rank will have its data accessed directly by PP rank, so needs to send the remote addresses and re-set atom ranges associated with transfers.
445 pme_pp->pmeCoordinateReceiverGpu->reinitCoordinateReceiver(stateGpu->getCoordinates());
446 pme_pp->pmeForceSenderGpu->setForceSendBuffer(pme_gpu_get_device_f(pme));
451 /* The box, FE flag and lambda are sent along with the coordinates
453 copy_mat(cnb.box, box);
454 *lambda_q = cnb.lambda_q;
455 *lambda_lj = cnb.lambda_lj;
456 *computeEnergyAndVirial = ((cnb.flags & PP_PME_ENER_VIR) != 0U);
459 /* Receive the coordinates in place */
461 for (const auto& sender : pme_pp->ppRanks)
463 if (sender.numAtoms > 0)
465 if (pme_pp->useGpuDirectComm)
469 pme_pp->pmeCoordinateReceiverGpu->receiveCoordinatesSynchronizerFromPpCudaDirect(
474 pme_pp->pmeCoordinateReceiverGpu->launchReceiveCoordinatesFromPpCudaMpi(
475 stateGpu->getCoordinates(), nat, sender.numAtoms * sizeof(rvec), sender.rankId);
480 MPI_Irecv(pme_pp->x[nat],
481 sender.numAtoms * sizeof(rvec),
485 pme_pp->mpi_comm_mysim,
486 &pme_pp->req[messages++]);
488 nat += sender.numAtoms;
492 "Received from PP rank %d: %d "
503 /* Wait for the coordinates and/or charges to arrive */
504 MPI_Waitall(messages, pme_pp->req.data(), pme_pp->stat.data());
506 } while (status == -1);
508 GMX_UNUSED_VALUE(pme);
509 GMX_UNUSED_VALUE(pme_pp);
510 GMX_UNUSED_VALUE(box);
511 GMX_UNUSED_VALUE(maxshift_x);
512 GMX_UNUSED_VALUE(maxshift_y);
513 GMX_UNUSED_VALUE(lambda_q);
514 GMX_UNUSED_VALUE(lambda_lj);
515 GMX_UNUSED_VALUE(computeEnergyAndVirial);
516 GMX_UNUSED_VALUE(step);
517 GMX_UNUSED_VALUE(grid_size);
518 GMX_UNUSED_VALUE(ewaldcoeff_q);
519 GMX_UNUSED_VALUE(ewaldcoeff_lj);
520 GMX_UNUSED_VALUE(useGpuForPme);
521 GMX_UNUSED_VALUE(stateGpu);
526 if (status == pmerecvqxX)
534 /*! \brief Send the PME mesh force, virial and energy to the PP-only ranks. */
535 static void gmx_pme_send_force_vir_ener(const gmx_pme_t& pme,
537 const PmeOutput& output,
543 gmx_pme_comm_vir_ene_t cve;
544 int messages, ind_start, ind_end;
547 if (pme_pp->useGpuDirectComm)
549 GMX_ASSERT((pme_pp->pmeForceSenderGpu != nullptr),
550 "The use of GPU direct communication for PME-PP is enabled, "
551 "but the PME GPU force reciever object does not exist");
557 /* Now the evaluated forces have to be transferred to the PP ranks */
558 if (pme_pp->useGpuDirectComm && GMX_THREAD_MPI)
560 int numPpRanks = static_cast<int>(pme_pp->ppRanks.size());
561 # pragma omp parallel for num_threads(std::min(numPpRanks, pme.nthread)) schedule(static)
562 for (int i = 0; i < numPpRanks; i++)
564 auto& receiver = pme_pp->ppRanks[i];
565 pme_pp->pmeForceSenderGpu->sendFToPpCudaDirect(
566 receiver.rankId, receiver.numAtoms, pme_pp->sendForcesDirectToPpGpu);
571 for (const auto& receiver : pme_pp->ppRanks)
574 ind_end = ind_start + receiver.numAtoms;
575 if (pme_pp->useGpuDirectComm)
577 pme_pp->pmeForceSenderGpu->sendFToPpCudaMpi(pme_gpu_get_device_f(&pme),
579 receiver.numAtoms * sizeof(rvec),
581 &pme_pp->req[messages]);
585 void* sendbuf = const_cast<void*>(static_cast<const void*>(output.forces_[ind_start]));
588 receiver.numAtoms * sizeof(rvec),
592 pme_pp->mpi_comm_mysim,
593 &pme_pp->req[messages]);
599 /* send virial and energy to our last PP node */
600 copy_mat(output.coulombVirial_, cve.vir_q);
601 copy_mat(output.lennardJonesVirial_, cve.vir_lj);
602 cve.energy_q = output.coulombEnergy_;
603 cve.energy_lj = output.lennardJonesEnergy_;
604 cve.dvdlambda_q = dvdlambda_q;
605 cve.dvdlambda_lj = dvdlambda_lj;
606 /* check for the signals to send back to a PP node */
607 cve.stop_cond = gmx_get_stop_condition();
613 fprintf(debug, "PME rank sending to PP rank %d: virial and energy\n", pme_pp->peerRankId);
615 MPI_Isend(&cve, sizeof(cve), MPI_BYTE, pme_pp->peerRankId, 1, pme_pp->mpi_comm_mysim, &pme_pp->req[messages++]);
617 /* Wait for the forces to arrive */
618 MPI_Waitall(messages, pme_pp->req.data(), pme_pp->stat.data());
620 GMX_RELEASE_ASSERT(false, "Invalid call to gmx_pme_send_force_vir_ener");
621 GMX_UNUSED_VALUE(pme);
622 GMX_UNUSED_VALUE(pme_pp);
623 GMX_UNUSED_VALUE(output);
624 GMX_UNUSED_VALUE(dvdlambda_q);
625 GMX_UNUSED_VALUE(dvdlambda_lj);
626 GMX_UNUSED_VALUE(cycles);
630 int gmx_pmeonly(struct gmx_pme_t* pme,
633 gmx_wallcycle* wcycle,
634 gmx_walltime_accounting_t walltime_accounting,
637 bool useGpuPmePpCommunication,
638 const gmx::DeviceStreamManager* deviceStreamManager)
645 int maxshift_x = 0, maxshift_y = 0;
646 real dvdlambda_q, dvdlambda_lj;
649 bool computeEnergyAndVirial = false;
652 /* This data will only use with PME tuning, i.e. switching PME grids */
653 std::vector<gmx_pme_t*> pmedata;
654 pmedata.push_back(pme);
656 auto pme_pp = gmx_pme_pp_init(cr);
658 std::unique_ptr<gmx::StatePropagatorDataGpu> stateGpu;
659 // TODO the variable below should be queried from the task assignment info
660 const bool useGpuForPme = (runMode == PmeRunMode::GPU) || (runMode == PmeRunMode::Mixed);
664 deviceStreamManager != nullptr,
665 "Device stream manager can not be nullptr when using GPU in PME-only rank.");
666 GMX_RELEASE_ASSERT(deviceStreamManager->streamIsValid(gmx::DeviceStreamType::Pme),
667 "Device stream can not be nullptr when using GPU in PME-only rank");
668 changePinningPolicy(&pme_pp->chargeA, pme_get_pinning_policy());
669 changePinningPolicy(&pme_pp->x, pme_get_pinning_policy());
670 if (useGpuPmePpCommunication)
672 pme_pp->pmeCoordinateReceiverGpu = std::make_unique<gmx::PmeCoordinateReceiverGpu>(
673 pme_pp->mpi_comm_mysim, deviceStreamManager->context(), pme_pp->ppRanks);
674 pme_pp->pmeForceSenderGpu =
675 std::make_unique<gmx::PmeForceSenderGpu>(pme_gpu_get_f_ready_synchronizer(pme),
676 pme_pp->mpi_comm_mysim,
677 deviceStreamManager->context(),
680 // TODO: Special PME-only constructor is used here. There is no mechanism to prevent from using the other constructor here.
681 // This should be made safer.
682 stateGpu = std::make_unique<gmx::StatePropagatorDataGpu>(
683 &deviceStreamManager->stream(gmx::DeviceStreamType::Pme),
684 deviceStreamManager->context(),
685 GpuApiCallBehavior::Async,
686 pme_gpu_get_block_size(pme),
693 do /****** this is a quasi-loop over time steps! */
695 /* The reason for having a loop here is PME grid tuning/switching */
698 /* Domain decomposition */
700 real ewaldcoeff_q = 0, ewaldcoeff_lj = 0;
701 ret = gmx_pme_recv_coeffs_coords(pme,
709 &computeEnergyAndVirial,
718 if (ret == pmerecvqxSWITCHGRID)
720 /* Switch the PME grid to newGridSize */
721 pme = gmx_pmeonly_switch(&pmedata, newGridSize, ewaldcoeff_q, ewaldcoeff_lj, cr, ir);
724 if (ret == pmerecvqxRESETCOUNTERS)
726 /* Reset the cycle and flop counters */
727 reset_pmeonly_counters(wcycle, walltime_accounting, mynrnb, step, useGpuForPme);
729 } while (ret == pmerecvqxSWITCHGRID || ret == pmerecvqxRESETCOUNTERS);
731 if (ret == pmerecvqxFINISH)
733 /* We should stop: break out of the loop */
739 wallcycle_start(wcycle, WallCycleCounter::Run);
740 walltime_accounting_start_time(walltime_accounting);
743 wallcycle_start(wcycle, WallCycleCounter::PmeMesh);
748 // TODO Make a struct of array refs onto these per-atom fields
749 // of pme_pp (maybe box, energy and virial, too; and likewise
750 // from mdatoms for the other call to gmx_pme_do), so we have
751 // fewer lines of code and less parameter passing.
752 gmx::StepWorkload stepWork;
753 stepWork.computeVirial = computeEnergyAndVirial;
754 stepWork.computeEnergy = computeEnergyAndVirial;
755 stepWork.computeForces = true;
756 PmeOutput output = { {}, false, 0, { { 0 } }, 0, 0, { { 0 } }, 0 };
759 stepWork.haveDynamicBox = false;
760 stepWork.useGpuPmeFReduction = pme_pp->useGpuDirectComm;
761 // TODO this should be set properly by gmx_pme_recv_coeffs_coords,
762 // or maybe use inputrecDynamicBox(ir), at the very least - change this when this codepath is tested!
763 pme_gpu_prepare_computation(pme, box, wcycle, stepWork);
764 if (!pme_pp->useGpuDirectComm)
766 stateGpu->copyCoordinatesToGpu(gmx::ArrayRef<gmx::RVec>(pme_pp->x),
767 gmx::AtomLocality::Local);
769 // On the separate PME rank we do not need a synchronizer as we schedule everything in a single stream
770 // TODO: with pme on GPU the receive should make a list of synchronizers and pass it here #3157
771 auto xReadyOnDevice = nullptr;
773 pme_gpu_launch_spread(pme,
777 pme_pp->useGpuDirectComm,
778 pme_pp->pmeCoordinateReceiverGpu.get());
779 pme_gpu_launch_complex_transforms(pme, wcycle, stepWork);
780 pme_gpu_launch_gather(pme, wcycle, lambda_q);
781 output = pme_gpu_wait_finish_task(pme, computeEnergyAndVirial, lambda_q, wcycle);
782 pme_gpu_reinit_computation(pme, wcycle);
786 GMX_ASSERT(pme_pp->x.size() == static_cast<size_t>(natoms),
787 "The coordinate buffer should have size natoms");
804 output.coulombVirial_,
805 output.lennardJonesVirial_,
806 &output.coulombEnergy_,
807 &output.lennardJonesEnergy_,
813 output.forces_ = pme_pp->f;
816 cycles = wallcycle_stop(wcycle, WallCycleCounter::PmeMesh);
817 gmx_pme_send_force_vir_ener(*pme, pme_pp.get(), output, dvdlambda_q, dvdlambda_lj, cycles);
820 } /***** end of quasi-loop, we stop with the break above */
823 walltime_accounting_end_time(walltime_accounting);