[alexxy/gromacs.git] / src / gromacs / ewald / pme_gather.cpp

/*
 * This file is part of the GROMACS molecular simulation package.
 *
 * Copyright (c) 1991-2000, University of Groningen, The Netherlands.
 * Copyright (c) 2001-2004, The GROMACS development team.
 * Copyright (c) 2013,2014,2015,2017,2018 by the GROMACS development team.
 * Copyright (c) 2019,2020, by the GROMACS development team, led by
 * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
 * and including many others, as listed in the AUTHORS file in the
 * top-level source directory and at http://www.gromacs.org.
 *
 * GROMACS is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public License
 * as published by the Free Software Foundation; either version 2.1
 * of the License, or (at your option) any later version.
 *
 * GROMACS is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with GROMACS; if not, see
 * http://www.gnu.org/licenses, or write to the Free Software Foundation,
 * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
 *
 * If you want to redistribute modifications to GROMACS, please
 * consider that scientific software is very special. Version
 * control is crucial - bugs must be traceable. We will be happy to
 * consider code for inclusion in the official distribution, but
 * derived work must not be called official GROMACS. Details are found
 * in the README & COPYING files - if they are missing, get the
 * official version at http://www.gromacs.org.
 *
 * To help us fund GROMACS development, we humbly ask that you cite
 * the research papers on the package. Check out http://www.gromacs.org.
 */

#include "gmxpre.h"

#include "pme_gather.h"

#include "gromacs/math/vec.h"
#include "gromacs/simd/simd.h"
#include "gromacs/utility/basedefinitions.h"
#include "gromacs/utility/gmxassert.h"
#include "gromacs/utility/smalloc.h"
#include "gromacs/utility/typetraits.h"

#include "pme_internal.h"
#include "pme_simd.h"
#include "pme_spline_work.h"

using namespace gmx; // TODO: Remove when this file is moved into gmx namespace

/* Spline function. Goals: 1) Force compiler to instantiate function separately
   for each compile-time value of order and once more for any possible (runtime)
   value. 2) Allow overloading for specific compile-time values.
   The Int template argument can be either int (runtime) or an object of type
   integral_constant<int, N> (compile-time). Common runtime values can be
   converted to compile-time values with a switch statement. For the compile
   value the compiler is required to instantiate the function separately for
   each value. The function can be overloaded for specific compile-time values
   using integral_constant<int, N> where N is either a specific value or an
   enable_if constrained non-type template parameter. The most specific overload
   (specific value > int template parameter > general function) is called. Inside
   the function the order argument can be used as regular int because
   integral_constant has a proper conversion.

   SIMD do_fspline() template funtions will be used for PME order 4 and 5
   when the SIMD module has support for SIMD4 for the architecture used.
   For SIMD4 without unaligned load/store support:
     order 4 and 5 use the order 4+5 aligned SIMD template
   For SIMD4 with unaligned load/store support:
     order 4 uses the order 4 unaligned SIMD template
     order 5 uses the order 4+5 aligned SIMD template
 */
struct do_fspline
{
    do_fspline(const gmx_pme_t* pme,
               const real* gmx_restrict grid,
               const PmeAtomComm* gmx_restrict atc,
               const splinedata_t* gmx_restrict spline,
               int                              nn) :
        pme(pme),
        grid(grid),
        atc(atc),
        spline(spline),
        nn(nn)
    {
    }

    template<typename Int>
    RVec operator()(Int order) const
    {
        static_assert(isIntegralConstant<Int, int>::value || std::is_same_v<Int, int>,
                      "'order' needs to be either of type integral_constant<int,N> or int.");

        const int norder = nn * order;

        /* Pointer arithmetic alert, next six statements */
        const real* const gmx_restrict thx  = spline->theta.coefficients[XX] + norder;
        const real* const gmx_restrict thy  = spline->theta.coefficients[YY] + norder;
        const real* const gmx_restrict thz  = spline->theta.coefficients[ZZ] + norder;
        const real* const gmx_restrict dthx = spline->dtheta.coefficients[XX] + norder;
        const real* const gmx_restrict dthy = spline->dtheta.coefficients[YY] + norder;
        const real* const gmx_restrict dthz = spline->dtheta.coefficients[ZZ] + norder;

        RVec f(0, 0, 0);

        for (int ithx = 0; (ithx < order); ithx++)
        {
            const int  index_x = (idxX + ithx) * gridNY * gridNZ;
            const real tx      = thx[ithx];
            const real dx      = dthx[ithx];

            for (int ithy = 0; (ithy < order); ithy++)
            {
                const int  index_xy = index_x + (idxY + ithy) * gridNZ;
                const real ty       = thy[ithy];
                const real dy       = dthy[ithy];
                real       fxy1 = 0, fz1 = 0;

                for (int ithz = 0; (ithz < order); ithz++)
                {
                    const real gval = grid[index_xy + (idxZ + ithz)];
                    fxy1 += thz[ithz] * gval;
                    fz1 += dthz[ithz] * gval;
                }
                f[XX] += dx * ty * fxy1;
                f[YY] += tx * dy * fxy1;
                f[ZZ] += tx * ty * fz1;
            }
        }

        return f;
    }

// 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
#if PME_4NSIMD_GATHER
    /* Gather for one charge with pme_order=4 with unaligned SIMD4 load+store.
     * Uses 4N SIMD where N is SIMD_WIDTH/4 to operate on all of z and N of y.
     * This code does not assume any memory alignment for the grid.
     */
    RVec operator()(std::integral_constant<int, 4> /*unused*/) const
    {
        const int norder = nn * 4;
        /* Pointer arithmetic alert, next six statements */
        const real* const gmx_restrict thx  = spline->theta.coefficients[XX] + norder;
        const real* const gmx_restrict thy  = spline->theta.coefficients[YY] + norder;
        const real* const gmx_restrict thz  = spline->theta.coefficients[ZZ] + norder;
        const real* const gmx_restrict dthx = spline->dtheta.coefficients[XX] + norder;
        const real* const gmx_restrict dthy = spline->dtheta.coefficients[YY] + norder;
        const real* const gmx_restrict dthz = spline->dtheta.coefficients[ZZ] + norder;

        Simd4NReal fx_S = setZero();
        Simd4NReal fy_S = setZero();
        Simd4NReal fz_S = setZero();

        /* With order 4 the z-spline is actually aligned */
        const Simd4NReal tz_S = load4DuplicateN(thz);
        const Simd4NReal dz_S = load4DuplicateN(dthz);

        for (int ithx = 0; ithx < 4; ithx++)
        {
            const int        index_x = (idxX + ithx) * gridNY * gridNZ;
            const Simd4NReal tx_S    = Simd4NReal(thx[ithx]);
            const Simd4NReal dx_S    = Simd4NReal(dthx[ithx]);

            for (int ithy = 0; ithy < 4; ithy += GMX_SIMD4N_REAL_WIDTH / 4)
            {
                const int index_xy = index_x + (idxY + ithy) * gridNZ;

                const Simd4NReal ty_S = loadUNDuplicate4(thy + ithy);
                const Simd4NReal dy_S = loadUNDuplicate4(dthy + ithy);

                const Simd4NReal gval_S = loadU4NOffset(grid + index_xy + idxZ, gridNZ);


                const Simd4NReal fxy1_S = tz_S * gval_S;
                const Simd4NReal fz1_S  = dz_S * gval_S;

                fx_S = fma(dx_S * ty_S, fxy1_S, fx_S);
                fy_S = fma(tx_S * dy_S, fxy1_S, fy_S);
                fz_S = fma(tx_S * ty_S, fz1_S, fz_S);
            }
        }

        return { reduce(fx_S), reduce(fy_S), reduce(fz_S) };
    }
#endif

#ifdef PME_SIMD4_SPREAD_GATHER
    /* Load order elements from unaligned memory into two 4-wide SIMD */
    template<int order>
    static inline void loadOrderU(const real* data,
                                  std::integral_constant<int, order> /*unused*/,
                                  int        offset,
                                  Simd4Real* S0,
                                  Simd4Real* S1)
    {
#    ifdef PME_SIMD4_UNALIGNED
        *S0 = load4U(data - offset);
        *S1 = load4U(data - offset + 4);
#    else
        alignas(GMX_SIMD_ALIGNMENT) real buf_aligned[GMX_SIMD4_WIDTH * 2];
        /* Copy data to an aligned buffer */
        for (int i = 0; i < order; i++)
        {
            buf_aligned[offset + i] = data[i];
        }
        *S0 = load4(buf_aligned);
        *S1 = load4(buf_aligned + 4);
#    endif
    }
#endif

#ifdef PME_SIMD4_SPREAD_GATHER
    /* This code assumes that the grid is allocated 4-real aligned
     * and that pme->pmegrid_nz is a multiple of 4.
     * This code supports pme_order <= 5.
     */
    template<int Order>
    std::enable_if_t<Order == 4 || Order == 5, RVec> operator()(std::integral_constant<int, Order> order) const
    {
        const int norder = nn * order;
        GMX_ASSERT(gridNZ % 4 == 0,
                   "For aligned SIMD4 operations the grid size has to be padded up to a multiple "
                   "of 4");
        /* Pointer arithmetic alert, next six statements */
        const real* const gmx_restrict thx  = spline->theta.coefficients[XX] + norder;
        const real* const gmx_restrict thy  = spline->theta.coefficients[YY] + norder;
        const real* const gmx_restrict thz  = spline->theta.coefficients[ZZ] + norder;
        const real* const gmx_restrict dthx = spline->dtheta.coefficients[XX] + norder;
        const real* const gmx_restrict dthy = spline->dtheta.coefficients[YY] + norder;
        const real* const gmx_restrict dthz = spline->dtheta.coefficients[ZZ] + norder;

        struct pme_spline_work* const work = pme->spline_work;

        const int offset = idxZ & 3;

        Simd4Real fx_S = setZero();
        Simd4Real fy_S = setZero();
        Simd4Real fz_S = setZero();

        Simd4Real tz_S0, tz_S1, dz_S0, dz_S1;
        loadOrderU(thz, order, offset, &tz_S0, &tz_S1);
        loadOrderU(dthz, order, offset, &dz_S0, &dz_S1);

        tz_S0 = selectByMask(tz_S0, work->mask_S0[offset]);
        dz_S0 = selectByMask(dz_S0, work->mask_S0[offset]);
        tz_S1 = selectByMask(tz_S1, work->mask_S1[offset]);
        dz_S1 = selectByMask(dz_S1, work->mask_S1[offset]);

        for (int ithx = 0; (ithx < order); ithx++)
        {
            const int       index_x = (idxX + ithx) * gridNY * gridNZ;
            const Simd4Real tx_S    = Simd4Real(thx[ithx]);
            const Simd4Real dx_S    = Simd4Real(dthx[ithx]);

            for (int ithy = 0; (ithy < order); ithy++)
            {
                const int       index_xy = index_x + (idxY + ithy) * gridNZ;
                const Simd4Real ty_S     = Simd4Real(thy[ithy]);
                const Simd4Real dy_S     = Simd4Real(dthy[ithy]);

                const Simd4Real gval_S0 = load4(grid + index_xy + idxZ - offset);
                const Simd4Real gval_S1 = load4(grid + index_xy + idxZ - offset + 4);

                const Simd4Real fxy1_S0 = tz_S0 * gval_S0;
                const Simd4Real fz1_S0  = dz_S0 * gval_S0;
                const Simd4Real fxy1_S1 = tz_S1 * gval_S1;
                const Simd4Real fz1_S1  = dz_S1 * gval_S1;

                const Simd4Real fxy1_S = fxy1_S0 + fxy1_S1;
                const Simd4Real fz1_S  = fz1_S0 + fz1_S1;

                fx_S = fma(dx_S * ty_S, fxy1_S, fx_S);
                fy_S = fma(tx_S * dy_S, fxy1_S, fy_S);
                fz_S = fma(tx_S * ty_S, fz1_S, fz_S);
            }
        }

        return { reduce(fx_S), reduce(fy_S), reduce(fz_S) };
    }
#endif
private:
    const gmx_pme_t* const pme;
    const real* const gmx_restrict grid;
    const PmeAtomComm* const gmx_restrict atc;
    const splinedata_t* const gmx_restrict spline;
    const int                              nn;

    const int gridNY = pme->pmegrid_ny;
    const int gridNZ = pme->pmegrid_nz;

    const int* const idxptr = atc->idx[spline->ind[nn]];
    const int        idxX   = idxptr[XX];
    const int        idxY   = idxptr[YY];
    const int        idxZ   = idxptr[ZZ];
};


void gather_f_bsplines(const gmx_pme_t*    pme,
                       const real*         grid,
                       gmx_bool            bClearF,
                       const PmeAtomComm*  atc,
                       const splinedata_t* spline,
                       real                scale)
{
    /* sum forces for local particles */

    const int order = pme->pme_order;
    const int nx    = pme->nkx;
    const int ny    = pme->nky;
    const int nz    = pme->nkz;

    const real rxx = pme->recipbox[XX][XX];
    const real ryx = pme->recipbox[YY][XX];
    const real ryy = pme->recipbox[YY][YY];
    const real rzx = pme->recipbox[ZZ][XX];
    const real rzy = pme->recipbox[ZZ][YY];
    const real rzz = pme->recipbox[ZZ][ZZ];

    /* Extract the buffer for force output */
    rvec* gmx_restrict force = as_rvec_array(atc->f.data());

    /* Note that unrolling this loop by templating this function on order
     * deteriorates performance significantly with gcc5/6/7.
     */
    for (int nn = 0; nn < spline->n; nn++)
    {
        const int  n           = spline->ind[nn];
        const real coefficient = scale * atc->coefficient[n];

        if (bClearF)
        {
            force[n][XX] = 0;
            force[n][YY] = 0;
            force[n][ZZ] = 0;
        }
        if (coefficient != 0)
        {
            RVec       f;
            const auto spline_func = do_fspline(pme, grid, atc, spline, nn);

            switch (order)
            {
                case 4: f = spline_func(std::integral_constant<int, 4>()); break;
                case 5: f = spline_func(std::integral_constant<int, 5>()); break;
                default: f = spline_func(order); break;
            }

            force[n][XX] += -coefficient * (f[XX] * nx * rxx);
            force[n][YY] += -coefficient * (f[XX] * nx * ryx + f[YY] * ny * ryy);
            force[n][ZZ] += -coefficient * (f[XX] * nx * rzx + f[YY] * ny * rzy + f[ZZ] * nz * rzz);
        }
    }
    /* Since the energy and not forces are interpolated
     * the net force might not be exactly zero.
     * This can be solved by also interpolating F, but
     * that comes at a cost.
     * A better hack is to remove the net force every
     * step, but that must be done at a higher level
     * since this routine doesn't see all atoms if running
     * in parallel. Don't know how important it is?  EL 990726
     */
}


real gather_energy_bsplines(gmx_pme_t* pme, const real* grid, PmeAtomComm* atc)
{
    splinedata_t* spline;
    int           ithx, ithy, ithz, i0, j0, k0;
    int           index_x, index_xy;
    int*          idxptr;
    real          energy, pot, tx, ty, coefficient, gval;
    real *        thx, *thy, *thz;
    int           norder;
    int           order;

    spline = &atc->spline[0];

    order = pme->pme_order;

    energy = 0;
    for (int n = 0; n < atc->numAtoms(); n++)
    {
        coefficient = atc->coefficient[n];

        if (coefficient != 0)
        {
            idxptr = atc->idx[n];
            norder = n * order;

            i0 = idxptr[XX];
            j0 = idxptr[YY];
            k0 = idxptr[ZZ];

            /* Pointer arithmetic alert, next three statements */
            thx = spline->theta.coefficients[XX] + norder;
            thy = spline->theta.coefficients[YY] + norder;
            thz = spline->theta.coefficients[ZZ] + norder;

            pot = 0;
            for (ithx = 0; (ithx < order); ithx++)
            {
                index_x = (i0 + ithx) * pme->pmegrid_ny * pme->pmegrid_nz;
                tx      = thx[ithx];

                for (ithy = 0; (ithy < order); ithy++)
                {
                    index_xy = index_x + (j0 + ithy) * pme->pmegrid_nz;
                    ty       = thy[ithy];

                    for (ithz = 0; (ithz < order); ithz++)
                    {
                        gval = grid[index_xy + (k0 + ithz)];
                        pot += tx * ty * thz[ithz] * gval;
                    }
                }
            }

            energy += pot * coefficient;
        }
    }

    return energy;
}