src/gromacs/ewald/pme_gather.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,2017,2018 by the GROMACS development team.
   7  * Copyright (c) 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
  39 #include "gmxpre.h"
  40
  41 #include "pme_gather.h"
  42
  43 #include "gromacs/math/vec.h"
  44 #include "gromacs/simd/simd.h"
  45 #include "gromacs/utility/basedefinitions.h"
  46 #include "gromacs/utility/gmxassert.h"
  47 #include "gromacs/utility/smalloc.h"
  48 #include "gromacs/utility/typetraits.h"
  49
  50 #include "pme_internal.h"
  51 #include "pme_simd.h"
  52 #include "pme_spline_work.h"
  53
  54 using namespace gmx; // TODO: Remove when this file is moved into gmx namespace
  55
  56 /* Spline function. Goals: 1) Force compiler to instantiate function separately
  57    for each compile-time value of order and once more for any possible (runtime)
  58    value. 2) Allow overloading for specific compile-time values.
  59    The Int template argument can be either int (runtime) or an object of type
  60    integral_constant<int, N> (compile-time). Common runtime values can be
  61    converted to compile-time values with a switch statement. For the compile
  62    value the compiler is required to instantiate the function separately for
  63    each value. The function can be overloaded for specific compile-time values
  64    using integral_constant<int, N> where N is either a specific value or an
  65    enable_if constrained non-type template parameter. The most specific overload
  66    (specific value > int template parameter > general function) is called. Inside
  67    the function the order argument can be used as regular int because
  68    integral_constant has a proper conversion.
  69
  70    SIMD do_fspline() template funtions will be used for PME order 4 and 5
  71    when the SIMD module has support for SIMD4 for the architecture used.
  72    For SIMD4 without unaligned load/store support:
  73      order 4 and 5 use the order 4+5 aligned SIMD template
  74    For SIMD4 with unaligned load/store support:
  75      order 4 uses the order 4 unaligned SIMD template
  76      order 5 uses the order 4+5 aligned SIMD template
  77  */
  78 struct do_fspline
  79 {
  80     do_fspline(const gmx_pme_t*                 pme,
  81                const real* gmx_restrict         grid,
  82                const PmeAtomComm* gmx_restrict  atc,
  83                const splinedata_t* gmx_restrict spline,
  84                int                              nn) :
  85         pme(pme), grid(grid), atc(atc), spline(spline), nn(nn)
  86     {
  87     }
  88
  89     template<typename Int>
  90     RVec operator()(Int order) const
  91     {
  92         static_assert(isIntegralConstant<Int, int>::value || std::is_same_v<Int, int>,
  93                       "'order' needs to be either of type integral_constant<int,N> or int.");
  94
  95         const int norder = nn * order;
  96
  97         /* Pointer arithmetic alert, next six statements */
  98         const real* const gmx_restrict thx  = spline->theta.coefficients[XX] + norder;
  99         const real* const gmx_restrict thy  = spline->theta.coefficients[YY] + norder;
 100         const real* const gmx_restrict thz  = spline->theta.coefficients[ZZ] + norder;
 101         const real* const gmx_restrict dthx = spline->dtheta.coefficients[XX] + norder;
 102         const real* const gmx_restrict dthy = spline->dtheta.coefficients[YY] + norder;
 103         const real* const gmx_restrict dthz = spline->dtheta.coefficients[ZZ] + norder;
 104
 105         RVec f(0, 0, 0);
 106
 107         for (int ithx = 0; (ithx < order); ithx++)
 108         {
 109             const int  index_x = (idxX + ithx) * gridNY * gridNZ;
 110             const real tx      = thx[ithx];
 111             const real dx      = dthx[ithx];
 112
 113             for (int ithy = 0; (ithy < order); ithy++)
 114             {
 115                 const int  index_xy = index_x + (idxY + ithy) * gridNZ;
 116                 const real ty       = thy[ithy];
 117                 const real dy       = dthy[ithy];
 118                 real       fxy1 = 0, fz1 = 0;
 119
 120                 for (int ithz = 0; (ithz < order); ithz++)
 121                 {
 122                     const real gval = grid[index_xy + (idxZ + ithz)];
 123                     fxy1 += thz[ithz] * gval;
 124                     fz1 += dthz[ithz] * gval;
 125                 }
 126                 f[XX] += dx * ty * fxy1;
 127                 f[YY] += tx * dy * fxy1;
 128                 f[ZZ] += tx * ty * fz1;
 129             }
 130         }
 131
 132         return f;
 133     }
 134
 135 // TODO: Consider always have at least a dummy implementation of Simd (enough for first phase of two-phase lookup) and then use enable_if instead of #ifdef
 136 #if PME_4NSIMD_GATHER
 137     /* Gather for one charge with pme_order=4 with unaligned SIMD4 load+store.
 138      * Uses 4N SIMD where N is SIMD_WIDTH/4 to operate on all of z and N of y.
 139      * This code does not assume any memory alignment for the grid.
 140      */
 141     RVec operator()(std::integral_constant<int, 4> /*unused*/) const
 142     {
 143         const int norder = nn * 4;
 144         /* Pointer arithmetic alert, next six statements */
 145         const real* const gmx_restrict thx  = spline->theta.coefficients[XX] + norder;
 146         const real* const gmx_restrict thy  = spline->theta.coefficients[YY] + norder;
 147         const real* const gmx_restrict thz  = spline->theta.coefficients[ZZ] + norder;
 148         const real* const gmx_restrict dthx = spline->dtheta.coefficients[XX] + norder;
 149         const real* const gmx_restrict dthy = spline->dtheta.coefficients[YY] + norder;
 150         const real* const gmx_restrict dthz = spline->dtheta.coefficients[ZZ] + norder;
 151
 152         Simd4NReal fx_S = setZero();
 153         Simd4NReal fy_S = setZero();
 154         Simd4NReal fz_S = setZero();
 155
 156         /* With order 4 the z-spline is actually aligned */
 157         const Simd4NReal tz_S = load4DuplicateN(thz);
 158         const Simd4NReal dz_S = load4DuplicateN(dthz);
 159
 160         for (int ithx = 0; ithx < 4; ithx++)
 161         {
 162             const int        index_x = (idxX + ithx) * gridNY * gridNZ;
 163             const Simd4NReal tx_S    = Simd4NReal(thx[ithx]);
 164             const Simd4NReal dx_S    = Simd4NReal(dthx[ithx]);
 165
 166             for (int ithy = 0; ithy < 4; ithy += GMX_SIMD4N_REAL_WIDTH / 4)
 167             {
 168                 const int index_xy = index_x + (idxY + ithy) * gridNZ;
 169
 170                 const Simd4NReal ty_S = loadUNDuplicate4(thy + ithy);
 171                 const Simd4NReal dy_S = loadUNDuplicate4(dthy + ithy);
 172
 173                 const Simd4NReal gval_S = loadU4NOffset(grid + index_xy + idxZ, gridNZ);
 174
 175
 176                 const Simd4NReal fxy1_S = tz_S * gval_S;
 177                 const Simd4NReal fz1_S  = dz_S * gval_S;
 178
 179                 fx_S = fma(dx_S * ty_S, fxy1_S, fx_S);
 180                 fy_S = fma(tx_S * dy_S, fxy1_S, fy_S);
 181                 fz_S = fma(tx_S * ty_S, fz1_S, fz_S);
 182             }
 183         }
 184
 185         return { reduce(fx_S), reduce(fy_S), reduce(fz_S) };
 186     }
 187 #endif
 188
 189 #ifdef PME_SIMD4_SPREAD_GATHER
 190     /* Load order elements from unaligned memory into two 4-wide SIMD */
 191     template<int order>
 192     static inline void loadOrderU(const real* data,
 193                                   std::integral_constant<int, order> /*unused*/,
 194                                   int        offset,
 195                                   Simd4Real* S0,
 196                                   Simd4Real* S1)
 197     {
 198 #    ifdef PME_SIMD4_UNALIGNED
 199         *S0 = load4U(data - offset);
 200         *S1 = load4U(data - offset + 4);
 201 #    else
 202         alignas(GMX_SIMD_ALIGNMENT) real buf_aligned[GMX_SIMD4_WIDTH * 2];
 203         /* Copy data to an aligned buffer */
 204         for (int i = 0; i < order; i++)
 205         {
 206             buf_aligned[offset + i] = data[i];
 207         }
 208         *S0 = load4(buf_aligned);
 209         *S1 = load4(buf_aligned + 4);
 210 #    endif
 211     }
 212 #endif
 213
 214 #ifdef PME_SIMD4_SPREAD_GATHER
 215     /* This code assumes that the grid is allocated 4-real aligned
 216      * and that pme->pmegrid_nz is a multiple of 4.
 217      * This code supports pme_order <= 5.
 218      */
 219     template<int Order>
 220     std::enable_if_t<Order == 4 || Order == 5, RVec> operator()(std::integral_constant<int, Order> order) const
 221     {
 222         const int norder = nn * order;
 223         GMX_ASSERT(gridNZ % 4 == 0,
 224                    "For aligned SIMD4 operations the grid size has to be padded up to a multiple "
 225                    "of 4");
 226         /* Pointer arithmetic alert, next six statements */
 227         const real* const gmx_restrict thx  = spline->theta.coefficients[XX] + norder;
 228         const real* const gmx_restrict thy  = spline->theta.coefficients[YY] + norder;
 229         const real* const gmx_restrict thz  = spline->theta.coefficients[ZZ] + norder;
 230         const real* const gmx_restrict dthx = spline->dtheta.coefficients[XX] + norder;
 231         const real* const gmx_restrict dthy = spline->dtheta.coefficients[YY] + norder;
 232         const real* const gmx_restrict dthz = spline->dtheta.coefficients[ZZ] + norder;
 233
 234         struct pme_spline_work* const work = pme->spline_work;
 235
 236         const int offset = idxZ & 3;
 237
 238         Simd4Real fx_S = setZero();
 239         Simd4Real fy_S = setZero();
 240         Simd4Real fz_S = setZero();
 241
 242         Simd4Real tz_S0, tz_S1, dz_S0, dz_S1;
 243         loadOrderU(thz, order, offset, &tz_S0, &tz_S1);
 244         loadOrderU(dthz, order, offset, &dz_S0, &dz_S1);
 245
 246         tz_S0 = selectByMask(tz_S0, work->mask_S0[offset]);
 247         dz_S0 = selectByMask(dz_S0, work->mask_S0[offset]);
 248         tz_S1 = selectByMask(tz_S1, work->mask_S1[offset]);
 249         dz_S1 = selectByMask(dz_S1, work->mask_S1[offset]);
 250
 251         for (int ithx = 0; (ithx < order); ithx++)
 252         {
 253             const int       index_x = (idxX + ithx) * gridNY * gridNZ;
 254             const Simd4Real tx_S    = Simd4Real(thx[ithx]);
 255             const Simd4Real dx_S    = Simd4Real(dthx[ithx]);
 256
 257             for (int ithy = 0; (ithy < order); ithy++)
 258             {
 259                 const int       index_xy = index_x + (idxY + ithy) * gridNZ;
 260                 const Simd4Real ty_S     = Simd4Real(thy[ithy]);
 261                 const Simd4Real dy_S     = Simd4Real(dthy[ithy]);
 262
 263                 const Simd4Real gval_S0 = load4(grid + index_xy + idxZ - offset);
 264                 const Simd4Real gval_S1 = load4(grid + index_xy + idxZ - offset + 4);
 265
 266                 const Simd4Real fxy1_S0 = tz_S0 * gval_S0;
 267                 const Simd4Real fz1_S0  = dz_S0 * gval_S0;
 268                 const Simd4Real fxy1_S1 = tz_S1 * gval_S1;
 269                 const Simd4Real fz1_S1  = dz_S1 * gval_S1;
 270
 271                 const Simd4Real fxy1_S = fxy1_S0 + fxy1_S1;
 272                 const Simd4Real fz1_S  = fz1_S0 + fz1_S1;
 273
 274                 fx_S = fma(dx_S * ty_S, fxy1_S, fx_S);
 275                 fy_S = fma(tx_S * dy_S, fxy1_S, fy_S);
 276                 fz_S = fma(tx_S * ty_S, fz1_S, fz_S);
 277             }
 278         }
 279
 280         return { reduce(fx_S), reduce(fy_S), reduce(fz_S) };
 281     }
 282 #endif
 283 private:
 284     const gmx_pme_t* const                 pme;
 285     const real* const gmx_restrict         grid;
 286     const PmeAtomComm* const gmx_restrict  atc;
 287     const splinedata_t* const gmx_restrict spline;
 288     const int                              nn;
 289
 290     const int gridNY = pme->pmegrid_ny;
 291     const int gridNZ = pme->pmegrid_nz;
 292
 293     const int* const idxptr = atc->idx[spline->ind[nn]];
 294     const int        idxX   = idxptr[XX];
 295     const int        idxY   = idxptr[YY];
 296     const int        idxZ   = idxptr[ZZ];
 297 };
 298
 299
 300 void gather_f_bsplines(const gmx_pme_t*    pme,
 301                        const real*         grid,
 302                        gmx_bool            bClearF,
 303                        const PmeAtomComm*  atc,
 304                        const splinedata_t* spline,
 305                        real                scale)
 306 {
 307     /* sum forces for local particles */
 308
 309     const int order = pme->pme_order;
 310     const int nx    = pme->nkx;
 311     const int ny    = pme->nky;
 312     const int nz    = pme->nkz;
 313
 314     const real rxx = pme->recipbox[XX][XX];
 315     const real ryx = pme->recipbox[YY][XX];
 316     const real ryy = pme->recipbox[YY][YY];
 317     const real rzx = pme->recipbox[ZZ][XX];
 318     const real rzy = pme->recipbox[ZZ][YY];
 319     const real rzz = pme->recipbox[ZZ][ZZ];
 320
 321     /* Extract the buffer for force output */
 322     rvec* gmx_restrict force = as_rvec_array(atc->f.data());
 323
 324     /* Note that unrolling this loop by templating this function on order
 325      * deteriorates performance significantly with gcc5/6/7.
 326      */
 327     for (int nn = 0; nn < spline->n; nn++)
 328     {
 329         const int  n           = spline->ind[nn];
 330         const real coefficient = scale * atc->coefficient[n];
 331
 332         if (bClearF)
 333         {
 334             force[n][XX] = 0;
 335             force[n][YY] = 0;
 336             force[n][ZZ] = 0;
 337         }
 338         if (coefficient != 0)
 339         {
 340             RVec       f;
 341             const auto spline_func = do_fspline(pme, grid, atc, spline, nn);
 342
 343             switch (order)
 344             {
 345                 case 4: f = spline_func(std::integral_constant<int, 4>()); break;
 346                 case 5: f = spline_func(std::integral_constant<int, 5>()); break;
 347                 default: f = spline_func(order); break;
 348             }
 349
 350             force[n][XX] += -coefficient * (f[XX] * nx * rxx);
 351             force[n][YY] += -coefficient * (f[XX] * nx * ryx + f[YY] * ny * ryy);
 352             force[n][ZZ] += -coefficient * (f[XX] * nx * rzx + f[YY] * ny * rzy + f[ZZ] * nz * rzz);
 353         }
 354     }
 355     /* Since the energy and not forces are interpolated
 356      * the net force might not be exactly zero.
 357      * This can be solved by also interpolating F, but
 358      * that comes at a cost.
 359      * A better hack is to remove the net force every
 360      * step, but that must be done at a higher level
 361      * since this routine doesn't see all atoms if running
 362      * in parallel. Don't know how important it is?  EL 990726
 363      */
 364 }
 365
 366
 367 real gather_energy_bsplines(gmx_pme_t* pme, const real* grid, PmeAtomComm* atc)
 368 {
 369     splinedata_t* spline;
 370     int           ithx, ithy, ithz, i0, j0, k0;
 371     int           index_x, index_xy;
 372     int*          idxptr;
 373     real          energy, pot, tx, ty, coefficient, gval;
 374     real *        thx, *thy, *thz;
 375     int           norder;
 376     int           order;
 377
 378     spline = &atc->spline[0];
 379
 380     order = pme->pme_order;
 381
 382     energy = 0;
 383     for (int n = 0; n < atc->numAtoms(); n++)
 384     {
 385         coefficient = atc->coefficient[n];
 386
 387         if (coefficient != 0)
 388         {
 389             idxptr = atc->idx[n];
 390             norder = n * order;
 391
 392             i0 = idxptr[XX];
 393             j0 = idxptr[YY];
 394             k0 = idxptr[ZZ];
 395
 396             /* Pointer arithmetic alert, next three statements */
 397             thx = spline->theta.coefficients[XX] + norder;
 398             thy = spline->theta.coefficients[YY] + norder;
 399             thz = spline->theta.coefficients[ZZ] + norder;
 400
 401             pot = 0;
 402             for (ithx = 0; (ithx < order); ithx++)
 403             {
 404                 index_x = (i0 + ithx) * pme->pmegrid_ny * pme->pmegrid_nz;
 405                 tx      = thx[ithx];
 406
 407                 for (ithy = 0; (ithy < order); ithy++)
 408                 {
 409                     index_xy = index_x + (j0 + ithy) * pme->pmegrid_nz;
 410                     ty       = thy[ithy];
 411
 412                     for (ithz = 0; (ithz < order); ithz++)
 413                     {
 414                         gval = grid[index_xy + (k0 + ithz)];
 415                         pot += tx * ty * thz[ithz] * gval;
 416                     }
 417                 }
 418             }
 419
 420             energy += pot * coefficient;
 421         }
 422     }
 423
 424     return energy;
 425 }