src/gromacs/ewald/pme_spread.cu

   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-2016,2017,2018,2019,2020, 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.
  10  *
  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.
  15  *
  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.
  20  *
  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.
  25  *
  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.
  33  *
  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.
  36  */
  37
  38 /*! \internal \file
  39  *  \brief Implements PME GPU spline calculation and charge spreading in CUDA.
  40  *  TODO: consider always pre-sorting particles (as in DD case).
  41  *
  42  *  \author Aleksei Iupinov <a.yupinov@gmail.com>
  43  */
  44
  45 #include "gmxpre.h"
  46
  47 #include <cassert>
  48
  49 #include "gromacs/gpu_utils/cuda_kernel_utils.cuh"
  50 #include "gromacs/gpu_utils/typecasts.cuh"
  51
  52 #include "pme.cuh"
  53 #include "pme_gpu_calculate_splines.cuh"
  54 #include "pme_grid.h"
  55
  56 /*! \brief
  57  * Charge spreading onto the grid.
  58  * This corresponds to the CPU function spread_coefficients_bsplines_thread().
  59  * Optional second stage of the spline_and_spread_kernel.
  60  *
  61  * \tparam[in] order                PME interpolation order.
  62  * \tparam[in] wrapX                Whether the grid overlap in dimension X should be wrapped.
  63  * \tparam[in] wrapY                Whether the grid overlap in dimension Y should be wrapped.
  64  * \tparam[in] threadsPerAtom       How many threads work on each atom
  65  *
  66  * \param[in]  kernelParams         Input PME CUDA data in constant memory.
  67  * \param[in]  atomCharge           Atom charge/coefficient of atom processed by thread.
  68  * \param[in]  sm_gridlineIndices   Atom gridline indices in the shared memory.
  69  * \param[in]  sm_theta             Atom spline values in the shared memory.
  70  */
  71 template<int order, bool wrapX, bool wrapY, ThreadsPerAtom threadsPerAtom>
  72 __device__ __forceinline__ void spread_charges(const PmeGpuCudaKernelParams kernelParams,
  73                                                const float*                 atomCharge,
  74                                                const int* __restrict__ sm_gridlineIndices,
  75                                                const float* __restrict__ sm_theta)
  76 {
  77     /* Global memory pointer to the output grid */
  78     float* __restrict__ gm_grid = kernelParams.grid.d_realGrid;
  79
  80     // Number of atoms processed by a single warp in spread and gather
  81     const int threadsPerAtomValue = (threadsPerAtom == ThreadsPerAtom::Order) ? order : order * order;
  82     const int atomsPerWarp        = warp_size / threadsPerAtomValue;
  83
  84     const int nx  = kernelParams.grid.realGridSize[XX];
  85     const int ny  = kernelParams.grid.realGridSize[YY];
  86     const int nz  = kernelParams.grid.realGridSize[ZZ];
  87     const int pny = kernelParams.grid.realGridSizePadded[YY];
  88     const int pnz = kernelParams.grid.realGridSizePadded[ZZ];
  89
  90     const int offx = 0, offy = 0, offz = 0; // unused for now
  91
  92     const int atomIndexLocal = threadIdx.z;
  93
  94     const int chargeCheck = pme_gpu_check_atom_charge(*atomCharge);
  95     if (chargeCheck)
  96     {
  97         // Spline Z coordinates
  98         const int ithz = threadIdx.x;
  99
 100         const int ixBase = sm_gridlineIndices[atomIndexLocal * DIM + XX] - offx;
 101         const int iyBase = sm_gridlineIndices[atomIndexLocal * DIM + YY] - offy;
 102         int       iz     = sm_gridlineIndices[atomIndexLocal * DIM + ZZ] - offz + ithz;
 103         if (iz >= nz)
 104         {
 105             iz -= nz;
 106         }
 107         /* Atom index w.r.t. warp - alternating 0 1 0 1 .. */
 108         const int atomWarpIndex = atomIndexLocal % atomsPerWarp;
 109         /* Warp index w.r.t. block - could probably be obtained easier? */
 110         const int warpIndex = atomIndexLocal / atomsPerWarp;
 111
 112         const int splineIndexBase = getSplineParamIndexBase<order, atomsPerWarp>(warpIndex, atomWarpIndex);
 113         const int splineIndexZ = getSplineParamIndex<order, atomsPerWarp>(splineIndexBase, ZZ, ithz);
 114         const float thetaZ     = sm_theta[splineIndexZ];
 115
 116         /* loop not used if order*order threads per atom */
 117         const int ithyMin = (threadsPerAtom == ThreadsPerAtom::Order) ? 0 : threadIdx.y;
 118         const int ithyMax = (threadsPerAtom == ThreadsPerAtom::Order) ? order : threadIdx.y + 1;
 119         for (int ithy = ithyMin; ithy < ithyMax; ithy++)
 120         {
 121             int iy = iyBase + ithy;
 122             if (wrapY & (iy >= ny))
 123             {
 124                 iy -= ny;
 125             }
 126
 127             const int splineIndexY = getSplineParamIndex<order, atomsPerWarp>(splineIndexBase, YY, ithy);
 128             float       thetaY = sm_theta[splineIndexY];
 129             const float Val    = thetaZ * thetaY * (*atomCharge);
 130             assert(isfinite(Val));
 131             const int offset = iy * pnz + iz;
 132
 133 #pragma unroll
 134             for (int ithx = 0; (ithx < order); ithx++)
 135             {
 136                 int ix = ixBase + ithx;
 137                 if (wrapX & (ix >= nx))
 138                 {
 139                     ix -= nx;
 140                 }
 141                 const int gridIndexGlobal = ix * pny * pnz + offset;
 142                 const int splineIndexX =
 143                         getSplineParamIndex<order, atomsPerWarp>(splineIndexBase, XX, ithx);
 144                 const float thetaX = sm_theta[splineIndexX];
 145                 assert(isfinite(thetaX));
 146                 assert(isfinite(gm_grid[gridIndexGlobal]));
 147                 atomicAdd(gm_grid + gridIndexGlobal, thetaX * Val);
 148             }
 149         }
 150     }
 151 }
 152
 153 /*! \brief
 154  * A spline computation and charge spreading kernel function.
 155  *
 156  * Two tuning parameters can be used for additional performance. For small systems and for debugging
 157  * writeGlobal should be used removing the need to recalculate the theta values in the gather kernel.
 158  * Similarly for useOrderThreads large systems order threads per atom gives higher performance than order*order threads
 159  *
 160  * \tparam[in] order                PME interpolation order.
 161  * \tparam[in] computeSplines       A boolean which tells if the spline parameter and
 162  *                                  gridline indices' computation should be performed.
 163  * \tparam[in] spreadCharges        A boolean which tells if the charge spreading should be performed.
 164  * \tparam[in] wrapX                A boolean which tells if the grid overlap in dimension X should be wrapped.
 165  * \tparam[in] wrapY                A boolean which tells if the grid overlap in dimension Y should be wrapped.
 166  * \tparam[in] writeGlobal          A boolean which tells if the theta values and gridlines should be written to global memory.
 167  * \tparam[in] threadsPerAtom       How many threads work on each atom
 168  * \param[in]  kernelParams         Input PME CUDA data in constant memory.
 169  */
 170 template<int order, bool computeSplines, bool spreadCharges, bool wrapX, bool wrapY, bool writeGlobal, ThreadsPerAtom threadsPerAtom>
 171 __launch_bounds__(c_spreadMaxThreadsPerBlock) CLANG_DISABLE_OPTIMIZATION_ATTRIBUTE __global__
 172         void pme_spline_and_spread_kernel(const PmeGpuCudaKernelParams kernelParams)
 173 {
 174     const int threadsPerAtomValue = (threadsPerAtom == ThreadsPerAtom::Order) ? order : order * order;
 175     const int atomsPerBlock       = c_spreadMaxThreadsPerBlock / threadsPerAtomValue;
 176     // Number of atoms processed by a single warp in spread and gather
 177     const int atomsPerWarp = warp_size / threadsPerAtomValue;
 178     // Gridline indices, ivec
 179     __shared__ int sm_gridlineIndices[atomsPerBlock * DIM];
 180     // Spline values
 181     __shared__ float sm_theta[atomsPerBlock * DIM * order];
 182     float            dtheta;
 183
 184     float3 atomX;
 185     float  atomCharge;
 186
 187     const int blockIndex      = blockIdx.y * gridDim.x + blockIdx.x;
 188     const int atomIndexOffset = blockIndex * atomsPerBlock;
 189
 190     /* Thread index w.r.t. block */
 191     const int threadLocalId =
 192             (threadIdx.z * (blockDim.x * blockDim.y)) + (threadIdx.y * blockDim.x) + threadIdx.x;
 193     /* Warp index w.r.t. block - could probably be obtained easier? */
 194     const int warpIndex = threadLocalId / warp_size;
 195
 196     /* Atom index w.r.t. warp */
 197     const int atomWarpIndex = threadIdx.z % atomsPerWarp;
 198     /* Atom index w.r.t. block/shared memory */
 199     const int atomIndexLocal = warpIndex * atomsPerWarp + atomWarpIndex;
 200     /* Atom index w.r.t. global memory */
 201     const int atomIndexGlobal = atomIndexOffset + atomIndexLocal;
 202
 203     /* Early return for fully empty blocks at the end
 204      * (should only happen for billions of input atoms)
 205      */
 206     if (atomIndexOffset >= kernelParams.atoms.nAtoms)
 207     {
 208         return;
 209     }
 210     /* Charges, required for both spline and spread */
 211     if (c_useAtomDataPrefetch)
 212     {
 213         __shared__ float sm_coefficients[atomsPerBlock];
 214         pme_gpu_stage_atom_data<float, atomsPerBlock, 1>(sm_coefficients, kernelParams.atoms.d_coefficients);
 215         __syncthreads();
 216         atomCharge = sm_coefficients[atomIndexLocal];
 217     }
 218     else
 219     {
 220         atomCharge = kernelParams.atoms.d_coefficients[atomIndexGlobal];
 221     }
 222
 223     if (computeSplines)
 224     {
 225         const float3* __restrict__ gm_coordinates = asFloat3(kernelParams.atoms.d_coordinates);
 226         if (c_useAtomDataPrefetch)
 227         {
 228             // Coordinates
 229             __shared__ float3 sm_coordinates[atomsPerBlock];
 230
 231             /* Staging coordinates */
 232             pme_gpu_stage_atom_data<float3, atomsPerBlock, 1>(sm_coordinates, gm_coordinates);
 233             __syncthreads();
 234             atomX = sm_coordinates[atomIndexLocal];
 235         }
 236         else
 237         {
 238             atomX = gm_coordinates[atomIndexGlobal];
 239         }
 240         calculate_splines<order, atomsPerBlock, atomsPerWarp, false, writeGlobal>(
 241                 kernelParams, atomIndexOffset, atomX, atomCharge, sm_theta, &dtheta, sm_gridlineIndices);
 242         __syncwarp();
 243     }
 244     else
 245     {
 246         /* Staging the data for spread
 247          * (the data is assumed to be in GPU global memory with proper layout already,
 248          * as in after running the spline kernel)
 249          */
 250         /* Spline data - only thetas (dthetas will only be needed in gather) */
 251         pme_gpu_stage_atom_data<float, atomsPerBlock, DIM * order>(sm_theta, kernelParams.atoms.d_theta);
 252         /* Gridline indices */
 253         pme_gpu_stage_atom_data<int, atomsPerBlock, DIM>(sm_gridlineIndices,
 254                                                          kernelParams.atoms.d_gridlineIndices);
 255
 256         __syncthreads();
 257     }
 258
 259     /* Spreading */
 260     if (spreadCharges)
 261     {
 262         spread_charges<order, wrapX, wrapY, threadsPerAtom>(kernelParams, &atomCharge,
 263                                                             sm_gridlineIndices, sm_theta);
 264     }
 265 }
 266
 267 //! Kernel instantiations
 268 // clang-format off
 269 template __global__ void pme_spline_and_spread_kernel<4, true,  true,  true, true, true,  ThreadsPerAtom::Order>       (const PmeGpuCudaKernelParams);
 270 template __global__ void pme_spline_and_spread_kernel<4, true,  false, true, true, true,  ThreadsPerAtom::Order>       (const PmeGpuCudaKernelParams);
 271 template __global__ void pme_spline_and_spread_kernel<4, false, true,  true, true, true,  ThreadsPerAtom::Order>       (const PmeGpuCudaKernelParams);
 272 template __global__ void pme_spline_and_spread_kernel<4, true,  true,  true, true, false, ThreadsPerAtom::Order>       (const PmeGpuCudaKernelParams);
 273 template __global__ void pme_spline_and_spread_kernel<4, true,  true,  true, true, true,  ThreadsPerAtom::OrderSquared>(const PmeGpuCudaKernelParams);
 274 template __global__ void pme_spline_and_spread_kernel<4, true,  false, true, true, true,  ThreadsPerAtom::OrderSquared>(const PmeGpuCudaKernelParams);
 275 template __global__ void pme_spline_and_spread_kernel<4, false, true,  true, true, true,  ThreadsPerAtom::OrderSquared>(const PmeGpuCudaKernelParams);
 276 template __global__ void pme_spline_and_spread_kernel<4, true,  true,  true, true, false, ThreadsPerAtom::OrderSquared>(const PmeGpuCudaKernelParams);
 277 // clang-format on