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