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