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,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.
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.
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.
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.
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.
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.
37 /* IMPORTANT FOR DEVELOPERS:
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:
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:
49 * rectangular triclinic
54 * 2. You should check the energy conservation in a triclinic box.
56 * It might seem an overkill, but better safe than sorry.
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"
98 #include "pme_gpu_internal.h"
99 #include "pme_internal.h"
100 #include "pme_pp_communication.h"
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);
106 /*! \brief Master PP-PME communication data structure */
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 */
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;
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 */
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;
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;
134 /*! \brief whether GPU direct communications are active for PME-PP transfers */
135 bool useGpuDirectComm = false;
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)
141 auto pme_pp = std::make_unique<gmx_pme_pp>();
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)
152 pme_pp->ppRanks.push_back({ ppRankId, 0 });
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());
159 GMX_UNUSED_VALUE(cr);
165 static void reset_pmeonly_counters(gmx_wallcycle_t wcycle,
166 gmx_walltime_accounting_t walltime_accounting,
171 /* Reset all the counters related to performance over the run */
172 wallcycle_stop(wcycle, ewcRUN);
173 wallcycle_reset_all(wcycle);
175 wallcycle_start(wcycle, ewcRUN);
176 walltime_accounting_reset_time(walltime_accounting, step);
184 static gmx_pme_t* gmx_pmeonly_switch(std::vector<gmx_pme_t*>* pmedata,
185 const ivec grid_size,
189 const t_inputrec* ir)
191 GMX_ASSERT(pmedata, "Bad PME tuning list pointer");
192 for (auto& pme : *pmedata)
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])
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()
203 gmx_pme_reinit(&pme, cr, pme, ir, grid_size, ewaldcoeff_q, ewaldcoeff_lj);
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);
216 /*! \brief Called by PME-only ranks to receive coefficients and coordinates
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
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.
238 static int gmx_pme_recv_coeffs_coords(struct gmx_pme_t* pme,
252 gmx::StatePropagatorDataGpu* stateGpu,
253 PmeRunMode gmx_unused runMode)
259 unsigned int flags = 0;
261 bool atomSetChanged = false;
265 gmx_pme_comm_n_box_t cnb;
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);
272 /* We accumulate all received flags */
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" : "");
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");
292 if (cnb.flags & PP_PME_FINISH)
294 status = pmerecvqxFINISH;
297 if (cnb.flags & PP_PME_SWITCHGRID)
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;
304 status = pmerecvqxSWITCHGRID;
307 if (cnb.flags & PP_PME_RESETCOUNTERS)
309 /* Special case, receive the step (set above) and return */
310 status = pmerecvqxRESETCOUNTERS;
313 if (cnb.flags & (PP_PME_CHARGE | PP_PME_SQRTC6 | PP_PME_SIGMA))
315 atomSetChanged = true;
317 /* Receive the send counts from the other PP nodes */
318 for (auto& sender : pme_pp->ppRanks)
320 if (sender.rankId == pme_pp->peerRankId)
322 sender.numAtoms = cnb.natoms;
326 MPI_Irecv(&sender.numAtoms, sizeof(sender.numAtoms), MPI_BYTE, sender.rankId,
327 eCommType_CNB, pme_pp->mpi_comm_mysim, &pme_pp->req[messages++]);
330 MPI_Waitall(messages, pme_pp->req.data(), pme_pp->stat.data());
334 for (const auto& sender : pme_pp->ppRanks)
336 nat += sender.numAtoms;
339 if (cnb.flags & PP_PME_CHARGE)
341 pme_pp->chargeA.resizeWithPadding(nat);
343 if (cnb.flags & PP_PME_CHARGEB)
345 pme_pp->chargeB.resize(nat);
347 if (cnb.flags & PP_PME_SQRTC6)
349 pme_pp->sqrt_c6A.resize(nat);
351 if (cnb.flags & PP_PME_SQRTC6B)
353 pme_pp->sqrt_c6B.resize(nat);
355 if (cnb.flags & PP_PME_SIGMA)
357 pme_pp->sigmaA.resize(nat);
359 if (cnb.flags & PP_PME_SIGMAB)
361 pme_pp->sigmaB.resize(nat);
363 pme_pp->x.resize(nat);
364 pme_pp->f.resize(nat);
366 /* maxshift is sent when the charges are sent */
367 *maxshift_x = cnb.maxshift_x;
368 *maxshift_y = cnb.maxshift_y;
370 /* Receive the charges in place */
371 for (int q = 0; q < eCommType_NR; q++)
375 if (!(cnb.flags & (PP_PME_CHARGE << q)))
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");
390 for (const auto& sender : pme_pp->ppRanks)
392 if (sender.numAtoms > 0)
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;
399 fprintf(debug, "Received from PP rank %d: %d %s\n", sender.rankId,
401 (q == eCommType_ChargeA || q == eCommType_ChargeB) ? "charges"
409 if (cnb.flags & PP_PME_COORD)
413 gmx_pme_reinit_atoms(pme, nat, pme_pp->chargeA.data());
416 stateGpu->reinit(nat, nat);
417 pme_gpu_set_device_x(pme, stateGpu->getCoordinates());
419 if (pme_pp->useGpuDirectComm)
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");
425 // This rank will have its data accessed directly by PP rank, so needs to send the remote addresses.
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));
432 pme_pp->pmeCoordinateReceiverGpu->sendCoordinateBufferAddressToPpRanks(d_x);
433 pme_pp->pmeForceSenderGpu->sendForceBufferAddressToPpRanks(d_f);
438 /* The box, FE flag and lambda are sent along with the coordinates
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);
446 /* Receive the coordinates in place */
448 for (const auto& sender : pme_pp->ppRanks)
450 if (sender.numAtoms > 0)
452 if (pme_pp->useGpuDirectComm)
454 pme_pp->pmeCoordinateReceiverGpu->launchReceiveCoordinatesFromPpCudaDirect(
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++]);
462 nat += sender.numAtoms;
466 "Received from PP rank %d: %d "
468 sender.rankId, sender.numAtoms);
473 if (pme_pp->useGpuDirectComm)
475 pme_pp->pmeCoordinateReceiverGpu->enqueueWaitReceiveCoordinatesFromPpCudaDirect();
481 /* Wait for the coordinates and/or charges to arrive */
482 MPI_Waitall(messages, pme_pp->req.data(), pme_pp->stat.data());
484 } while (status == -1);
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);
504 if (status == pmerecvqxX)
513 /*! \brief Send force data to PP ranks */
514 static void sendFToPP(void* sendbuf, PpRanks receiver, gmx_pme_pp* pme_pp, int* messages)
517 if (pme_pp->useGpuDirectComm)
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");
523 pme_pp->pmeForceSenderGpu->sendFToPpCudaDirect(receiver.rankId);
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;
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,
538 const PmeOutput& output,
544 gmx_pme_comm_vir_ene_t cve;
545 int messages, ind_start, ind_end;
548 /* Now the evaluated forces have to be transferred to the PP nodes */
551 for (const auto& receiver : pme_pp->ppRanks)
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)
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]);
562 sendFToPP(sendbuf, receiver, pme_pp, &messages);
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();
579 fprintf(debug, "PME rank sending to PP rank %d: virial and energy\n", pme_pp->peerRankId);
581 MPI_Isend(&cve, sizeof(cve), MPI_BYTE, pme_pp->peerRankId, 1, pme_pp->mpi_comm_mysim,
582 &pme_pp->req[messages++]);
584 /* Wait for the forces to arrive */
585 MPI_Waitall(messages, pme_pp->req.data(), pme_pp->stat.data());
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);
597 int gmx_pmeonly(struct gmx_pme_t* pme,
600 gmx_wallcycle* wcycle,
601 gmx_walltime_accounting_t walltime_accounting,
610 int maxshift_x = 0, maxshift_y = 0;
611 real dvdlambda_q, dvdlambda_lj;
614 gmx_bool bEnerVir = FALSE;
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);
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);
629 changePinningPolicy(&pme_pp->chargeA, pme_get_pinning_policy());
630 changePinningPolicy(&pme_pp->x, pme_get_pinning_policy());
631 if (c_enableGpuPmePpComms)
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);
640 std::unique_ptr<gmx::StatePropagatorDataGpu> stateGpu;
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, wcycle);
653 do /****** this is a quasi-loop over time steps! */
655 /* The reason for having a loop here is PME grid tuning/switching */
658 /* Domain decomposition */
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);
666 if (ret == pmerecvqxSWITCHGRID)
668 /* Switch the PME grid to newGridSize */
669 pme = gmx_pmeonly_switch(&pmedata, newGridSize, ewaldcoeff_q, ewaldcoeff_lj, cr, ir);
672 if (ret == pmerecvqxRESETCOUNTERS)
674 /* Reset the cycle and flop counters */
675 reset_pmeonly_counters(wcycle, walltime_accounting, mynrnb, step, useGpuForPme);
677 } while (ret == pmerecvqxSWITCHGRID || ret == pmerecvqxRESETCOUNTERS);
679 if (ret == pmerecvqxFINISH)
681 /* We should stop: break out of the loop */
687 wallcycle_start(wcycle, ewcRUN);
688 walltime_accounting_start_time(walltime_accounting);
691 wallcycle_start(wcycle, ewcPMEMESH);
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);
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)
711 stateGpu->copyCoordinatesToGpu(gmx::ArrayRef<gmx::RVec>(pme_pp->x), gmx::AtomLocality::All);
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;
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);
725 GMX_ASSERT(pme_pp->x.size() == static_cast<size_t>(natoms),
726 "The coordinate buffer should have size natoms");
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;
737 cycles = wallcycle_stop(wcycle, ewcPMEMESH);
738 gmx_pme_send_force_vir_ener(*pme, pme_pp.get(), output, dvdlambda_q, dvdlambda_lj, cycles);
741 } /***** end of quasi-loop, we stop with the break above */
744 walltime_accounting_end_time(walltime_accounting);