src/gromacs/ewald/pme_internal.h

   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, 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 declarations necessary for
  41  * computing energies and forces for the PME long-ranged part (Coulomb
  42  * and LJ).
  43  *
  44  * \author Berk Hess <hess@kth.se>
  45  * \author Mark Abraham <mark.j.abraham@gmail.com>
  46  * \ingroup module_ewald
  47  */
  48
  49 /* TODO This file is a temporary holding area for stuff local to the
  50  * PME code, before it acquires some more normal ewald/file.c and
  51  * ewald/file.h structure.  In future clean up, get rid of this file,
  52  * to build more normal. */
  53
  54 #ifndef GMX_EWALD_PME_INTERNAL_H
  55 #define GMX_EWALD_PME_INTERNAL_H
  56
  57 #include "config.h"
  58
  59 #include "gromacs/math/gmxcomplex.h"
  60 #include "gromacs/utility/basedefinitions.h"
  61 #include "gromacs/utility/defaultinitializationallocator.h"
  62 #include "gromacs/utility/gmxmpi.h"
  63 #include "gromacs/utility/smalloc.h"
  64
  65 #include "pme_gpu_types_host.h"
  66
  67 //! A repeat of typedef from parallel_3dfft.h
  68 typedef struct gmx_parallel_3dfft* gmx_parallel_3dfft_t;
  69
  70 struct t_commrec;
  71 struct t_inputrec;
  72 struct PmeGpu;
  73
  74 //@{
  75 //! Grid indices for A state for charge and Lennard-Jones C6
  76 #define PME_GRID_QA 0
  77 #define PME_GRID_C6A 2
  78 //@}
  79
  80 //@{
  81 /*! \brief Flags that indicate the number of PME grids in use */
  82 #define DO_Q 2           /* Electrostatic grids have index q<2 */
  83 #define DO_Q_AND_LJ 4    /* non-LB LJ grids have index 2 <= q < 4 */
  84 #define DO_Q_AND_LJ_LB 9 /* With LB rules we need a total of 2+7 grids */
  85 //@}
  86
  87 /*! \brief Pascal triangle coefficients scaled with (1/2)^6 for LJ-PME with LB-rules */
  88 static const real lb_scale_factor[] = { 1.0 / 64,  6.0 / 64, 15.0 / 64, 20.0 / 64,
  89                                         15.0 / 64, 6.0 / 64, 1.0 / 64 };
  90
  91 /*! \brief Pascal triangle coefficients used in solve_pme_lj_yzx, only need to do 4 calculations due to symmetry */
  92 static const real lb_scale_factor_symm[] = { 2.0 / 64, 12.0 / 64, 30.0 / 64, 20.0 / 64 };
  93
  94 /*! \brief We only define a maximum to be able to use local arrays without allocation.
  95  * An order larger than 12 should never be needed, even for test cases.
  96  * If needed it can be changed here.
  97  */
  98 #define PME_ORDER_MAX 12
  99
 100 /*! \brief As gmx_pme_init, but takes most settings, except the grid/Ewald coefficients, from
 101  * pme_src. This is only called when the PME cut-off/grid size changes.
 102  */
 103 void gmx_pme_reinit(struct gmx_pme_t** pmedata,
 104                     const t_commrec*   cr,
 105                     struct gmx_pme_t*  pme_src,
 106                     const t_inputrec*  ir,
 107                     const ivec         grid_size,
 108                     real               ewaldcoeff_q,
 109                     real               ewaldcoeff_lj);
 110
 111
 112 /* Temporary suppression until these structs become opaque and don't live in
 113  * a header that is included by other headers. Also, until then I have no
 114  * idea what some of the names mean. */
 115
 116 //! @cond Doxygen_Suppress
 117
 118 /*! \brief Data structure for grid communication */
 119 struct pme_grid_comm_t
 120 {
 121     int send_id; //!< Source rank id
 122     int send_index0;
 123     int send_nindex;
 124     int recv_id; //!< Destination rank id
 125     int recv_index0;
 126     int recv_nindex;
 127     int recv_size = 0; //!< Receive buffer width, used with OpenMP
 128 };
 129
 130 /*! \brief Data structure for grid overlap communication in a single dimension */
 131 struct pme_overlap_t
 132 {
 133     MPI_Comm                     mpi_comm;  //!< MPI communcator
 134     int                          nnodes;    //!< Number of ranks
 135     int                          nodeid;    //!< Unique rank identifcator
 136     std::vector<int>             s2g0;      //!< The local interpolation grid start
 137     std::vector<int>             s2g1;      //!< The local interpolation grid end
 138     int                          send_size; //!< Send buffer width, used with OpenMP
 139     std::vector<pme_grid_comm_t> comm_data; //!< All the individual communication data for each rank
 140     std::vector<real>            sendbuf;   //!< Shared buffer for sending
 141     std::vector<real>            recvbuf;   //!< Shared buffer for receiving
 142 };
 143
 144 template<typename T>
 145 using AlignedVector = std::vector<T, gmx::AlignedAllocator<T>>;
 146
 147 template<typename T>
 148 using FastVector = std::vector<T, gmx::DefaultInitializationAllocator<T>>;
 149
 150 /*! \brief Data structure for organizing particle allocation to threads */
 151 struct AtomToThreadMap
 152 {
 153     //! Cumulative counts of the number of particles per thread
 154     int* n = nullptr;
 155     //! Storage buffer for n
 156     std::vector<int> nBuffer;
 157     //! Particle indices ordered on thread index (n)
 158     FastVector<int> i;
 159 };
 160
 161 /*! \brief Helper typedef for spline vectors */
 162 typedef real* splinevec[DIM];
 163
 164 /*! \internal
 165  * \brief Coefficients for theta or dtheta
 166  */
 167 class SplineCoefficients
 168 {
 169 public:
 170     //! Reallocate for use with up to nalloc coefficients
 171     void realloc(int nalloc);
 172
 173     //! Pointers to the coefficient buffer for x, y, z
 174     splinevec coefficients = { nullptr };
 175
 176 private:
 177     //! Storage for x coefficients
 178     std::vector<real> bufferX_;
 179     //! Storage for y coefficients
 180     std::vector<real> bufferY_;
 181     //! Storage for z coefficients, aligned for SIMD load
 182     AlignedVector<real> bufferZ_;
 183 };
 184
 185 /*! \brief Data structure for beta-spline interpolation */
 186 struct splinedata_t
 187 {
 188     int                n = 0;
 189     FastVector<int>    ind;
 190     SplineCoefficients theta;
 191     SplineCoefficients dtheta;
 192     int                nalloc = 0;
 193 };
 194
 195 /*! \brief PME slab MPI communication setup */
 196 struct SlabCommSetup
 197 {
 198     //! The nodes to send x and q to with DD
 199     int node_dest;
 200     //! The nodes to receive x and q from with DD
 201     int node_src;
 202     //! Index for commnode into the buffers
 203     int buf_index;
 204     //! The number of atoms to receive
 205     int rcount;
 206 };
 207
 208 /*! \internal
 209  * \brief Data structure for coordinating transfers between PME ranks along one dimension
 210  *
 211  * Also used for passing coordinates, coefficients and forces to and from PME routines.
 212  */
 213 class PmeAtomComm
 214 {
 215 public:
 216     //! Constructor, \p PmeMpiCommunicator is the communicator for this dimension
 217     PmeAtomComm(MPI_Comm PmeMpiCommunicator, int numThreads, int pmeOrder, int dimIndex, bool doSpread);
 218
 219     //! Set the atom count and when necessary resizes atom buffers
 220     void setNumAtoms(int numAtoms);
 221
 222     //! Returns the atom count
 223     int numAtoms() const { return numAtoms_; }
 224
 225     //! Returns the number of atoms to send to each rank
 226     gmx::ArrayRef<int> sendCount()
 227     {
 228         GMX_ASSERT(!count_thread.empty(), "Need at least one thread_count");
 229         return count_thread[0];
 230     }
 231
 232     //! The index of the dimension, 0=x, 1=y
 233     int dimind = 0;
 234     //! The number of slabs and ranks this dimension is decomposed over
 235     int nslab = 1;
 236     //! Our MPI rank index
 237     int nodeid = 0;
 238     //! Communicator for this dimension
 239     MPI_Comm mpi_comm;
 240
 241     //! Communication setup for each slab, only present with nslab > 1
 242     std::vector<SlabCommSetup> slabCommSetup;
 243     //! The maximum communication distance counted in MPI ranks
 244     int maxshift = 0;
 245
 246     //! The target slab index for each particle
 247     FastVector<int> pd;
 248     //! Target particle counts for each slab, for each thread
 249     std::vector<std::vector<int>> count_thread;
 250
 251 private:
 252     //! The number of atoms
 253     int numAtoms_ = 0;
 254
 255 public:
 256     //! The coordinates
 257     gmx::ArrayRef<const gmx::RVec> x;
 258     //! The coefficient, charges or LJ C6
 259     gmx::ArrayRef<const real> coefficient;
 260     //! The forces
 261     gmx::ArrayRef<gmx::RVec> f;
 262     //! Coordinate buffer, used only with nslab > 1
 263     FastVector<gmx::RVec> xBuffer;
 264     //! Coefficient buffer, used only with nslab > 1
 265     FastVector<real> coefficientBuffer;
 266     //! Force buffer, used only with nslab > 1
 267     FastVector<gmx::RVec> fBuffer;
 268     //! Tells whether these coordinates are used for spreading
 269     bool bSpread;
 270     //! The PME order
 271     int pme_order;
 272     //! The grid index per atom
 273     FastVector<gmx::IVec> idx;
 274     //! Fractional atom coordinates relative to the lower cell boundary
 275     FastVector<gmx::RVec> fractx;
 276
 277     //! The number of threads to use in PME
 278     int nthread;
 279     //! Thread index for each atom
 280     FastVector<int>              thread_idx;
 281     std::vector<AtomToThreadMap> threadMap;
 282     std::vector<splinedata_t>    spline;
 283 };
 284
 285 /*! \brief Data structure for a single PME grid */
 286 struct pmegrid_t
 287 {
 288     ivec  ci;     /* The spatial location of this grid         */
 289     ivec  n;      /* The used size of *grid, including order-1 */
 290     ivec  offset; /* The grid offset from the full node grid   */
 291     int   order;  /* PME spreading order                       */
 292     ivec  s;      /* The allocated size of *grid, s >= n       */
 293     real* grid;   /* The grid local thread, size n             */
 294 };
 295
 296 /*! \brief Data structures for PME grids */
 297 struct pmegrids_t
 298 {
 299     pmegrid_t  grid;         /* The full node grid (non thread-local)            */
 300     int        nthread;      /* The number of threads operating on this grid     */
 301     ivec       nc;           /* The local spatial decomposition over the threads */
 302     pmegrid_t* grid_th;      /* Array of grids for each thread                   */
 303     real*      grid_all;     /* Allocated array for the grids in *grid_th        */
 304     int*       g2t[DIM];     /* The grid to thread index                         */
 305     ivec       nthread_comm; /* The number of threads to communicate with        */
 306 };
 307
 308 /*! \brief Data structure for spline-interpolation working buffers */
 309 struct pme_spline_work;
 310
 311 /*! \brief Data structure for working buffers */
 312 struct pme_solve_work_t;
 313
 314 /*! \brief Master PME data structure */
 315 struct gmx_pme_t
 316 {                   //NOLINT(clang-analyzer-optin.performance.Padding)
 317     int ndecompdim; /* The number of decomposition dimensions */
 318     int nodeid;     /* Our nodeid in mpi->mpi_comm */
 319     int nodeid_major;
 320     int nodeid_minor;
 321     int nnodes; /* The number of nodes doing PME */
 322     int nnodes_major;
 323     int nnodes_minor;
 324
 325     MPI_Comm mpi_comm;
 326     MPI_Comm mpi_comm_d[2]; /* Indexed on dimension, 0=x, 1=y */
 327 #if GMX_MPI
 328     MPI_Datatype rvec_mpi; /* the pme vector's MPI type */
 329 #endif
 330
 331     gmx_bool bUseThreads; /* Does any of the PME ranks have nthread>1 ?  */
 332     int      nthread;     /* The number of threads doing PME on our rank */
 333
 334     gmx_bool bPPnode;   /* Node also does particle-particle forces */
 335     bool     doCoulomb; /* Apply PME to electrostatics */
 336     bool     doLJ;      /* Apply PME to Lennard-Jones r^-6 interactions */
 337     gmx_bool bFEP;      /* Compute Free energy contribution */
 338     gmx_bool bFEP_q;
 339     gmx_bool bFEP_lj;
 340     int      nkx, nky, nkz; /* Grid dimensions */
 341     gmx_bool bP3M;          /* Do P3M: optimize the influence function */
 342     int      pme_order;
 343     real     ewaldcoeff_q;  /* Ewald splitting coefficient for Coulomb */
 344     real     ewaldcoeff_lj; /* Ewald splitting coefficient for r^-6 */
 345     real     epsilon_r;
 346
 347
 348     enum PmeRunMode runMode; /* Which codepath is the PME runner taking - CPU, GPU, mixed;
 349                               * TODO: this is the information that should be owned by the task
 350                               * scheduler, and ideally not be duplicated here.
 351                               */
 352
 353     PmeGpu* gpu; /* A pointer to the GPU data.
 354                   * TODO: this should be unique or a shared pointer.
 355                   * Currently in practice there is a single gmx_pme_t instance while a code
 356                   * is partially set up for many of them. The PME tuning calls gmx_pme_reinit()
 357                   * which fully reinitializes the one and only PME structure anew while maybe
 358                   * keeping the old grid buffers if they were already large enough.
 359                   * This small choice should be made clear in the later refactoring -
 360                   * do we store many PME objects for different grid sizes,
 361                   * or a single PME object that handles different grid sizes gracefully.
 362                   */
 363
 364
 365     class EwaldBoxZScaler* boxScaler; /**< The scaling data Ewald uses with walls (set at pme_init constant for the entire run) */
 366
 367
 368     int ljpme_combination_rule; /* Type of combination rule in LJ-PME */
 369
 370     int ngrids; /* number of grids we maintain for pmegrid, (c)fftgrid and pfft_setups*/
 371
 372     pmegrids_t pmegrid[DO_Q_AND_LJ_LB]; /* Grids on which we do spreading/interpolation,
 373                                          * includes overlap Grid indices are ordered as
 374                                          * follows:
 375                                          * 0: Coloumb PME, state A
 376                                          * 1: Coloumb PME, state B
 377                                          * 2-8: LJ-PME
 378                                          * This can probably be done in a better way
 379                                          * but this simple hack works for now
 380                                          */
 381
 382     /* The PME coefficient spreading grid sizes/strides, includes pme_order-1 */
 383     int pmegrid_nx, pmegrid_ny, pmegrid_nz;
 384     /* pmegrid_nz might be larger than strictly necessary to ensure
 385      * memory alignment, pmegrid_nz_base gives the real base size.
 386      */
 387     int pmegrid_nz_base;
 388     /* The local PME grid starting indices */
 389     int pmegrid_start_ix, pmegrid_start_iy, pmegrid_start_iz;
 390
 391     /* Work data for spreading and gathering */
 392     pme_spline_work* spline_work;
 393
 394     real** fftgrid; /* Grids for FFT. With 1D FFT decomposition this can be a pointer */
 395     /* inside the interpolation grid, but separate for 2D PME decomp. */
 396     int fftgrid_nx, fftgrid_ny, fftgrid_nz;
 397
 398     t_complex** cfftgrid; /* Grids for complex FFT data */
 399
 400     int cfftgrid_nx, cfftgrid_ny, cfftgrid_nz;
 401
 402     gmx_parallel_3dfft_t* pfft_setup;
 403
 404     int * nnx, *nny, *nnz;
 405     real *fshx, *fshy, *fshz;
 406
 407     std::vector<PmeAtomComm> atc; /* Indexed on decomposition index */
 408     matrix                   recipbox;
 409     real                     boxVolume;
 410     splinevec                bsp_mod;
 411     /* Buffers to store data for local atoms for L-B combination rule
 412      * calculations in LJ-PME. lb_buf1 stores either the coefficients
 413      * for spreading/gathering (in serial), or the C6 coefficient for
 414      * local atoms (in parallel).  lb_buf2 is only used in parallel,
 415      * and stores the sigma values for local atoms. */
 416     FastVector<real> lb_buf1;
 417     FastVector<real> lb_buf2;
 418
 419     pme_overlap_t overlap[2]; /* Indexed on dimension, 0=x, 1=y */
 420
 421     /* Atom step for energy only calculation in gmx_pme_calc_energy() */
 422     std::unique_ptr<PmeAtomComm> atc_energy;
 423
 424     /* Communication buffers */
 425     rvec* bufv;       /* Communication buffer */
 426     real* bufr;       /* Communication buffer */
 427     int   buf_nalloc; /* The communication buffer size */
 428
 429     /* thread local work data for solve_pme */
 430     struct pme_solve_work_t* solve_work;
 431
 432     /* Work data for sum_qgrid */
 433     real* sum_qgrid_tmp;
 434     real* sum_qgrid_dd_tmp;
 435 };
 436
 437 //! @endcond
 438
 439 /*! \brief
 440  * Finds out if PME is currently running on GPU.
 441  * TODO: should this be removed eventually?
 442  *
 443  * \param[in] pme  The PME structure.
 444  * \returns        True if PME runs on GPU currently, false otherwise.
 445  */
 446 inline bool pme_gpu_active(const gmx_pme_t* pme)
 447 {
 448     return (pme != nullptr) && (pme->runMode != PmeRunMode::CPU);
 449 }
 450
 451 /*! \brief Tell our PME-only node to switch to a new grid size */
 452 void gmx_pme_send_switchgrid(const t_commrec* cr, ivec grid_size, real ewaldcoeff_q, real ewaldcoeff_lj);
 453
 454 #endif