src/gromacs/ewald/pme_pp.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 /*! \internal \file
  39  *
  40  * \brief This file contains function definitions necessary for
  41  * managing the offload of long-ranged PME work to separate MPI rank,
  42  * for computing energies and forces (Coulomb and LJ).
  43  *
  44  * \author Berk Hess <hess@kth.se>
  45  * \ingroup module_ewald
  46  */
  47
  48 #include "gmxpre.h"
  49
  50 #include "pme_pp.h"
  51
  52 #include "config.h"
  53
  54 #include <cstdio>
  55 #include <cstring>
  56
  57 #include "gromacs/domdec/domdec.h"
  58 #include "gromacs/domdec/domdec_struct.h"
  59 #include "gromacs/ewald/pme.h"
  60 #include "gromacs/ewald/pme_pp_comm_gpu.h"
  61 #include "gromacs/gmxlib/network.h"
  62 #include "gromacs/math/vec.h"
  63 #include "gromacs/mdlib/gmx_omp_nthreads.h"
  64 #include "gromacs/mdtypes/commrec.h"
  65 #include "gromacs/mdtypes/forceoutput.h"
  66 #include "gromacs/mdtypes/forcerec.h"
  67 #include "gromacs/mdtypes/interaction_const.h"
  68 #include "gromacs/mdtypes/md_enums.h"
  69 #include "gromacs/mdtypes/state_propagator_data_gpu.h"
  70 #include "gromacs/nbnxm/nbnxm.h"
  71 #include "gromacs/timing/wallcycle.h"
  72 #include "gromacs/utility/fatalerror.h"
  73 #include "gromacs/utility/gmxmpi.h"
  74 #include "gromacs/utility/smalloc.h"
  75
  76 #include "pme_pp_communication.h"
  77
  78 /*! \brief Block to wait for communication to PME ranks to complete
  79  *
  80  * This should be faster with a real non-blocking MPI implementation
  81  */
  82 static constexpr bool c_useDelayedWait = false;
  83
  84 /*! \brief Wait for the pending data send requests to PME ranks to complete */
  85 static void gmx_pme_send_coeffs_coords_wait(gmx_domdec_t* dd)
  86 {
  87     if (dd->nreq_pme)
  88     {
  89 #if GMX_MPI
  90         MPI_Waitall(dd->nreq_pme, dd->req_pme, MPI_STATUSES_IGNORE);
  91 #endif
  92         dd->nreq_pme = 0;
  93     }
  94 }
  95
  96 /*! \brief Send data to PME ranks */
  97 static void gmx_pme_send_coeffs_coords(t_forcerec*         fr,
  98                                        const t_commrec*    cr,
  99                                        unsigned int        flags,
 100                                        gmx::ArrayRef<real> chargeA,
 101                                        gmx::ArrayRef<real> chargeB,
 102                                        gmx::ArrayRef<real> c6A,
 103                                        gmx::ArrayRef<real> c6B,
 104                                        gmx::ArrayRef<real> sigmaA,
 105                                        gmx::ArrayRef<real> sigmaB,
 106                                        const matrix        box,
 107                                        const rvec gmx_unused* x,
 108                                        real                   lambda_q,
 109                                        real                   lambda_lj,
 110                                        int                    maxshift_x,
 111                                        int                    maxshift_y,
 112                                        int64_t                step,
 113                                        bool                   useGpuPmePpComms,
 114                                        bool                   reinitGpuPmePpComms,
 115                                        bool                   sendCoordinatesFromGpu,
 116                                        GpuEventSynchronizer*  coordinatesReadyOnDeviceEvent)
 117 {
 118     gmx_domdec_t*         dd;
 119     gmx_pme_comm_n_box_t* cnb;
 120     int                   n;
 121
 122     dd = cr->dd;
 123     n  = dd_numHomeAtoms(*dd);
 124
 125     if (debug)
 126     {
 127         fprintf(debug,
 128                 "PP rank %d sending to PME rank %d: %d%s%s%s%s\n",
 129                 cr->sim_nodeid,
 130                 dd->pme_nodeid,
 131                 n,
 132                 (flags & PP_PME_CHARGE) ? " charges" : "",
 133                 (flags & PP_PME_SQRTC6) ? " sqrtC6" : "",
 134                 (flags & PP_PME_SIGMA) ? " sigma" : "",
 135                 (flags & PP_PME_COORD) ? " coordinates" : "");
 136     }
 137
 138     if (useGpuPmePpComms)
 139     {
 140         flags |= PP_PME_GPUCOMMS;
 141     }
 142
 143     if (c_useDelayedWait)
 144     {
 145         /* We can not use cnb until pending communication has finished */
 146         gmx_pme_send_coeffs_coords_wait(dd);
 147     }
 148
 149     if (dd->pme_receive_vir_ener)
 150     {
 151         /* Peer PP node: communicate all data */
 152         if (dd->cnb == nullptr)
 153         {
 154             snew(dd->cnb, 1);
 155         }
 156         cnb = dd->cnb;
 157
 158         cnb->flags      = flags;
 159         cnb->natoms     = n;
 160         cnb->maxshift_x = maxshift_x;
 161         cnb->maxshift_y = maxshift_y;
 162         cnb->lambda_q   = lambda_q;
 163         cnb->lambda_lj  = lambda_lj;
 164         cnb->step       = step;
 165         if (flags & PP_PME_COORD)
 166         {
 167             copy_mat(box, cnb->box);
 168         }
 169 #if GMX_MPI
 170         MPI_Isend(cnb,
 171                   sizeof(*cnb),
 172                   MPI_BYTE,
 173                   dd->pme_nodeid,
 174                   eCommType_CNB,
 175                   cr->mpi_comm_mysim,
 176                   &dd->req_pme[dd->nreq_pme++]);
 177 #endif
 178     }
 179     else if (flags & (PP_PME_CHARGE | PP_PME_SQRTC6 | PP_PME_SIGMA))
 180     {
 181 #if GMX_MPI
 182         /* Communicate only the number of atoms */
 183         MPI_Isend(&n,
 184                   sizeof(n),
 185                   MPI_BYTE,
 186                   dd->pme_nodeid,
 187                   eCommType_CNB,
 188                   cr->mpi_comm_mysim,
 189                   &dd->req_pme[dd->nreq_pme++]);
 190 #endif
 191     }
 192
 193 #if GMX_MPI
 194     if (n > 0)
 195     {
 196         if (flags & PP_PME_CHARGE)
 197         {
 198             MPI_Isend(chargeA.data(),
 199                       n * sizeof(real),
 200                       MPI_BYTE,
 201                       dd->pme_nodeid,
 202                       eCommType_ChargeA,
 203                       cr->mpi_comm_mysim,
 204                       &dd->req_pme[dd->nreq_pme++]);
 205         }
 206         if (flags & PP_PME_CHARGEB)
 207         {
 208             MPI_Isend(chargeB.data(),
 209                       n * sizeof(real),
 210                       MPI_BYTE,
 211                       dd->pme_nodeid,
 212                       eCommType_ChargeB,
 213                       cr->mpi_comm_mysim,
 214                       &dd->req_pme[dd->nreq_pme++]);
 215         }
 216         if (flags & PP_PME_SQRTC6)
 217         {
 218             MPI_Isend(c6A.data(),
 219                       n * sizeof(real),
 220                       MPI_BYTE,
 221                       dd->pme_nodeid,
 222                       eCommType_SQRTC6A,
 223                       cr->mpi_comm_mysim,
 224                       &dd->req_pme[dd->nreq_pme++]);
 225         }
 226         if (flags & PP_PME_SQRTC6B)
 227         {
 228             MPI_Isend(c6B.data(),
 229                       n * sizeof(real),
 230                       MPI_BYTE,
 231                       dd->pme_nodeid,
 232                       eCommType_SQRTC6B,
 233                       cr->mpi_comm_mysim,
 234                       &dd->req_pme[dd->nreq_pme++]);
 235         }
 236         if (flags & PP_PME_SIGMA)
 237         {
 238             MPI_Isend(sigmaA.data(),
 239                       n * sizeof(real),
 240                       MPI_BYTE,
 241                       dd->pme_nodeid,
 242                       eCommType_SigmaA,
 243                       cr->mpi_comm_mysim,
 244                       &dd->req_pme[dd->nreq_pme++]);
 245         }
 246         if (flags & PP_PME_SIGMAB)
 247         {
 248             MPI_Isend(sigmaB.data(),
 249                       n * sizeof(real),
 250                       MPI_BYTE,
 251                       dd->pme_nodeid,
 252                       eCommType_SigmaB,
 253                       cr->mpi_comm_mysim,
 254                       &dd->req_pme[dd->nreq_pme++]);
 255         }
 256         if (flags & PP_PME_COORD)
 257         {
 258             if (reinitGpuPmePpComms)
 259             {
 260                 fr->pmePpCommGpu->reinit(n);
 261             }
 262
 263
 264             /* MPI_Isend does not accept a const buffer pointer */
 265             real* xRealPtr = const_cast<real*>(x[0]);
 266             if (useGpuPmePpComms && (fr != nullptr))
 267             {
 268                 if (sendCoordinatesFromGpu)
 269                 {
 270                     fr->pmePpCommGpu->sendCoordinatesToPmeFromGpu(
 271                             fr->stateGpu->getCoordinates(), n, coordinatesReadyOnDeviceEvent);
 272                 }
 273                 else
 274                 {
 275                     fr->pmePpCommGpu->sendCoordinatesToPmeFromCpu(
 276                             reinterpret_cast<gmx::RVec*>(xRealPtr), n, coordinatesReadyOnDeviceEvent);
 277                 }
 278             }
 279             else
 280             {
 281                 MPI_Isend(xRealPtr,
 282                           n * sizeof(rvec),
 283                           MPI_BYTE,
 284                           dd->pme_nodeid,
 285                           eCommType_COORD,
 286                           cr->mpi_comm_mysim,
 287                           &dd->req_pme[dd->nreq_pme++]);
 288             }
 289         }
 290     }
 291 #else
 292     GMX_UNUSED_VALUE(fr);
 293     GMX_UNUSED_VALUE(chargeA);
 294     GMX_UNUSED_VALUE(chargeB);
 295     GMX_UNUSED_VALUE(c6A);
 296     GMX_UNUSED_VALUE(c6B);
 297     GMX_UNUSED_VALUE(sigmaA);
 298     GMX_UNUSED_VALUE(sigmaB);
 299     GMX_UNUSED_VALUE(reinitGpuPmePpComms);
 300     GMX_UNUSED_VALUE(sendCoordinatesFromGpu);
 301     GMX_UNUSED_VALUE(coordinatesReadyOnDeviceEvent);
 302 #endif
 303     if (!c_useDelayedWait)
 304     {
 305         /* Wait for the data to arrive */
 306         /* We can skip this wait as we are sure x and q will not be modified
 307          * before the next call to gmx_pme_send_x_q or gmx_pme_receive_f.
 308          */
 309         gmx_pme_send_coeffs_coords_wait(dd);
 310     }
 311 }
 312
 313 void gmx_pme_send_parameters(const t_commrec*           cr,
 314                              const interaction_const_t& interactionConst,
 315                              bool                       bFreeEnergy_q,
 316                              bool                       bFreeEnergy_lj,
 317                              gmx::ArrayRef<real>        chargeA,
 318                              gmx::ArrayRef<real>        chargeB,
 319                              gmx::ArrayRef<real>        sqrt_c6A,
 320                              gmx::ArrayRef<real>        sqrt_c6B,
 321                              gmx::ArrayRef<real>        sigmaA,
 322                              gmx::ArrayRef<real>        sigmaB,
 323                              int                        maxshift_x,
 324                              int                        maxshift_y)
 325 {
 326     unsigned int flags = 0;
 327
 328     if (EEL_PME(interactionConst.eeltype))
 329     {
 330         flags |= PP_PME_CHARGE;
 331     }
 332     if (EVDW_PME(interactionConst.vdwtype))
 333     {
 334         flags |= (PP_PME_SQRTC6 | PP_PME_SIGMA);
 335     }
 336     if (bFreeEnergy_q || bFreeEnergy_lj)
 337     {
 338         /* Assumes that the B state flags are in the bits just above
 339          * the ones for the A state. */
 340         flags |= (flags << 1);
 341     }
 342
 343     gmx_pme_send_coeffs_coords(nullptr,
 344                                cr,
 345                                flags,
 346                                chargeA,
 347                                chargeB,
 348                                sqrt_c6A,
 349                                sqrt_c6B,
 350                                sigmaA,
 351                                sigmaB,
 352                                nullptr,
 353                                nullptr,
 354                                0,
 355                                0,
 356                                maxshift_x,
 357                                maxshift_y,
 358                                -1,
 359                                false,
 360                                false,
 361                                false,
 362                                nullptr);
 363 }
 364
 365 void gmx_pme_send_coordinates(t_forcerec*           fr,
 366                               const t_commrec*      cr,
 367                               const matrix          box,
 368                               const rvec*           x,
 369                               real                  lambda_q,
 370                               real                  lambda_lj,
 371                               bool                  computeEnergyAndVirial,
 372                               int64_t               step,
 373                               bool                  useGpuPmePpComms,
 374                               bool                  receiveCoordinateAddressFromPme,
 375                               bool                  sendCoordinatesFromGpu,
 376                               GpuEventSynchronizer* coordinatesReadyOnDeviceEvent,
 377                               gmx_wallcycle*        wcycle)
 378 {
 379     wallcycle_start(wcycle, WallCycleCounter::PpPmeSendX);
 380
 381     unsigned int flags = PP_PME_COORD;
 382     if (computeEnergyAndVirial)
 383     {
 384         flags |= PP_PME_ENER_VIR;
 385     }
 386     gmx_pme_send_coeffs_coords(fr,
 387                                cr,
 388                                flags,
 389                                {},
 390                                {},
 391                                {},
 392                                {},
 393                                {},
 394                                {},
 395                                box,
 396                                x,
 397                                lambda_q,
 398                                lambda_lj,
 399                                0,
 400                                0,
 401                                step,
 402                                useGpuPmePpComms,
 403                                receiveCoordinateAddressFromPme,
 404                                sendCoordinatesFromGpu,
 405                                coordinatesReadyOnDeviceEvent);
 406
 407     wallcycle_stop(wcycle, WallCycleCounter::PpPmeSendX);
 408 }
 409
 410 void gmx_pme_send_finish(const t_commrec* cr)
 411 {
 412     unsigned int flags = PP_PME_FINISH;
 413
 414     gmx_pme_send_coeffs_coords(
 415             nullptr, cr, flags, {}, {}, {}, {}, {}, {}, nullptr, nullptr, 0, 0, 0, 0, -1, false, false, false, nullptr);
 416 }
 417
 418 void gmx_pme_send_switchgrid(const t_commrec* cr, ivec grid_size, real ewaldcoeff_q, real ewaldcoeff_lj)
 419 {
 420 #if GMX_MPI
 421     gmx_pme_comm_n_box_t cnb;
 422
 423     /* Only let one PP node signal each PME node */
 424     if (cr->dd->pme_receive_vir_ener)
 425     {
 426         cnb.flags = PP_PME_SWITCHGRID;
 427         copy_ivec(grid_size, cnb.grid_size);
 428         cnb.ewaldcoeff_q  = ewaldcoeff_q;
 429         cnb.ewaldcoeff_lj = ewaldcoeff_lj;
 430
 431         /* We send this, uncommon, message blocking to simplify the code */
 432         MPI_Send(&cnb, sizeof(cnb), MPI_BYTE, cr->dd->pme_nodeid, eCommType_CNB, cr->mpi_comm_mysim);
 433     }
 434 #else
 435     GMX_UNUSED_VALUE(cr);
 436     GMX_UNUSED_VALUE(grid_size);
 437     GMX_UNUSED_VALUE(ewaldcoeff_q);
 438     GMX_UNUSED_VALUE(ewaldcoeff_lj);
 439 #endif
 440 }
 441
 442 void gmx_pme_send_resetcounters(const t_commrec gmx_unused* cr, int64_t gmx_unused step)
 443 {
 444 #if GMX_MPI
 445     gmx_pme_comm_n_box_t cnb;
 446
 447     /* Only let one PP node signal each PME node */
 448     if (cr->dd->pme_receive_vir_ener)
 449     {
 450         cnb.flags = PP_PME_RESETCOUNTERS;
 451         cnb.step  = step;
 452
 453         /* We send this, uncommon, message blocking to simplify the code */
 454         MPI_Send(&cnb, sizeof(cnb), MPI_BYTE, cr->dd->pme_nodeid, eCommType_CNB, cr->mpi_comm_mysim);
 455     }
 456 #endif
 457 }
 458
 459 /*! \brief Receive virial and energy from PME rank */
 460 static void receive_virial_energy(const t_commrec*      cr,
 461                                   gmx::ForceWithVirial* forceWithVirial,
 462                                   real*                 energy_q,
 463                                   real*                 energy_lj,
 464                                   real*                 dvdlambda_q,
 465                                   real*                 dvdlambda_lj,
 466                                   float*                pme_cycles)
 467 {
 468     gmx_pme_comm_vir_ene_t cve;
 469
 470     if (cr->dd->pme_receive_vir_ener)
 471     {
 472         if (debug)
 473         {
 474             fprintf(debug,
 475                     "PP rank %d receiving from PME rank %d: virial and energy\n",
 476                     cr->sim_nodeid,
 477                     cr->dd->pme_nodeid);
 478         }
 479 #if GMX_MPI
 480         MPI_Recv(&cve, sizeof(cve), MPI_BYTE, cr->dd->pme_nodeid, 1, cr->mpi_comm_mysim, MPI_STATUS_IGNORE);
 481 #else
 482         memset(&cve, 0, sizeof(cve));
 483 #endif
 484
 485         forceWithVirial->addVirialContribution(cve.vir_q);
 486         forceWithVirial->addVirialContribution(cve.vir_lj);
 487         *energy_q  = cve.energy_q;
 488         *energy_lj = cve.energy_lj;
 489         *dvdlambda_q += cve.dvdlambda_q;
 490         *dvdlambda_lj += cve.dvdlambda_lj;
 491         *pme_cycles = cve.cycles;
 492
 493         if (cve.stop_cond != gmx_stop_cond_none)
 494         {
 495             gmx_set_stop_condition(cve.stop_cond);
 496         }
 497     }
 498     else
 499     {
 500         *energy_q   = 0;
 501         *energy_lj  = 0;
 502         *pme_cycles = 0;
 503     }
 504 }
 505
 506 /*! \brief Recieve force data from PME ranks */
 507 static void recvFFromPme(gmx::PmePpCommGpu* pmePpCommGpu,
 508                          void*              recvptr,
 509                          int                n,
 510                          const t_commrec*   cr,
 511                          bool               useGpuPmePpComms,
 512                          bool               receivePmeForceToGpu)
 513 {
 514     if (useGpuPmePpComms)
 515     {
 516         GMX_ASSERT(pmePpCommGpu != nullptr, "Need valid pmePpCommGpu");
 517         // Receive forces from PME rank
 518         pmePpCommGpu->receiveForceFromPme(static_cast<gmx::RVec*>(recvptr), n, receivePmeForceToGpu);
 519     }
 520     else
 521     {
 522         // Receive data using MPI
 523 #if GMX_MPI
 524         MPI_Recv(recvptr, n * sizeof(rvec), MPI_BYTE, cr->dd->pme_nodeid, 0, cr->mpi_comm_mysim, MPI_STATUS_IGNORE);
 525 #else
 526         GMX_UNUSED_VALUE(cr);
 527 #endif
 528     }
 529 }
 530
 531
 532 void gmx_pme_receive_f(gmx::PmePpCommGpu*    pmePpCommGpu,
 533                        const t_commrec*      cr,
 534                        gmx::ForceWithVirial* forceWithVirial,
 535                        real*                 energy_q,
 536                        real*                 energy_lj,
 537                        real*                 dvdlambda_q,
 538                        real*                 dvdlambda_lj,
 539                        bool                  useGpuPmePpComms,
 540                        bool                  receivePmeForceToGpu,
 541                        float*                pme_cycles)
 542 {
 543     if (c_useDelayedWait)
 544     {
 545         /* Wait for the x request to finish */
 546         gmx_pme_send_coeffs_coords_wait(cr->dd);
 547     }
 548
 549     const int               natoms = dd_numHomeAtoms(*cr->dd);
 550     std::vector<gmx::RVec>& buffer = cr->dd->pmeForceReceiveBuffer;
 551     buffer.resize(natoms);
 552
 553     void* recvptr = reinterpret_cast<void*>(buffer.data());
 554     recvFFromPme(pmePpCommGpu, recvptr, natoms, cr, useGpuPmePpComms, receivePmeForceToGpu);
 555
 556     int nt = gmx_omp_nthreads_get_simple_rvec_task(ModuleMultiThread::Default, natoms);
 557
 558     gmx::ArrayRef<gmx::RVec> f = forceWithVirial->force_;
 559
 560     if (!receivePmeForceToGpu)
 561     {
 562         /* Note that we would like to avoid this conditional by putting it
 563          * into the omp pragma instead, but then we still take the full
 564          * omp parallel for overhead (at least with gcc5).
 565          */
 566         if (nt == 1)
 567         {
 568             for (int i = 0; i < natoms; i++)
 569             {
 570                 f[i] += buffer[i];
 571             }
 572         }
 573         else
 574         {
 575 #pragma omp parallel for num_threads(nt) schedule(static)
 576             for (int i = 0; i < natoms; i++)
 577             {
 578                 f[i] += buffer[i];
 579             }
 580         }
 581     }
 582
 583     receive_virial_energy(cr, forceWithVirial, energy_q, energy_lj, dvdlambda_q, dvdlambda_lj, pme_cycles);
 584 }