f3aa6660fec8dbfb0b8a50691870bf2d67b413f5
[alexxy/gromacs.git] / 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, 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-gpu-utils.h"
53 #include "pme-grid.h"
54
55 /*
56  * This define affects the spline calculation behaviour in the kernel.
57  * 0: a single GPU thread handles a single dimension of a single particle (calculating and storing (order) spline values and derivatives).
58  * 1: (order) threads do redundant work on this same task, each one stores only a single theta and single dtheta into global arrays.
59  * The only efficiency difference is less global store operations, countered by more redundant spline computation.
60  *
61  * TODO: estimate if this should be a boolean parameter (and add it to the unit test if so).
62  */
63 #define PME_GPU_PARALLEL_SPLINE 0
64
65
66 /*! \brief
67  * General purpose function for loading atom-related data from global to shared memory.
68  *
69  * \tparam[in] T                 Data type (float/int/...)
70  * \tparam[in] atomsPerBlock     Number of atoms processed by a block - should be accounted for in the size of the shared memory array.
71  * \tparam[in] dataCountPerAtom  Number of data elements per single atom (e.g. DIM for an rvec coordinates array).
72  * \param[in]  kernelParams      Input PME CUDA data in constant memory.
73  * \param[out] sm_destination    Shared memory array for output.
74  * \param[in]  gm_source         Global memory array for input.
75  */
76 template<typename T,
77          const int atomsPerBlock,
78          const int dataCountPerAtom>
79 __device__  __forceinline__
80 void pme_gpu_stage_atom_data(const PmeGpuCudaKernelParams       kernelParams,
81                              T * __restrict__                   sm_destination,
82                              const T * __restrict__             gm_source)
83 {
84     static_assert(c_usePadding, "With padding disabled, index checking should be fixed to account for spline theta/dtheta per-warp alignment");
85     const int blockIndex       = blockIdx.y * gridDim.x + blockIdx.x;
86     const int threadLocalIndex = ((threadIdx.z * blockDim.y + threadIdx.y) * blockDim.x) + threadIdx.x;
87     const int localIndex       = threadLocalIndex;
88     const int globalIndexBase  = blockIndex * atomsPerBlock * dataCountPerAtom;
89     const int globalIndex      = globalIndexBase + localIndex;
90     const int globalCheck      = pme_gpu_check_atom_data_index(globalIndex, kernelParams.atoms.nAtoms * dataCountPerAtom);
91     if ((localIndex < atomsPerBlock * dataCountPerAtom) & globalCheck)
92     {
93         assert(isfinite(float(gm_source[globalIndex])));
94         sm_destination[localIndex] = gm_source[globalIndex];
95     }
96 }
97
98 /*! \brief
99  * PME GPU spline parameter and gridline indices calculation.
100  * This corresponds to the CPU functions calc_interpolation_idx() and make_bsplines().
101  * First stage of the whole kernel.
102  *
103  * \tparam[in] order                PME interpolation order.
104  * \tparam[in] atomsPerBlock        Number of atoms processed by a block - should be accounted for
105  *                                  in the sizes of the shared memory arrays.
106  * \param[in]  kernelParams         Input PME CUDA data in constant memory.
107  * \param[in]  atomIndexOffset      Starting atom index for the execution block w.r.t. global memory.
108  * \param[in]  sm_coordinates       Atom coordinates in the shared memory.
109  * \param[in]  sm_coefficients      Atom charges/coefficients in the shared memory.
110  * \param[out] sm_theta             Atom spline values in the shared memory.
111  * \param[out] sm_gridlineIndices   Atom gridline indices in the shared memory.
112  */
113 template <const int order,
114           const int atomsPerBlock>
115 __device__ __forceinline__ void calculate_splines(const PmeGpuCudaKernelParams           kernelParams,
116                                                   const int                              atomIndexOffset,
117                                                   const float3 * __restrict__            sm_coordinates,
118                                                   const float * __restrict__             sm_coefficients,
119                                                   float * __restrict__                   sm_theta,
120                                                   int * __restrict__                     sm_gridlineIndices)
121 {
122     /* Global memory pointers for output */
123     float * __restrict__ gm_theta           = kernelParams.atoms.d_theta;
124     float * __restrict__ gm_dtheta          = kernelParams.atoms.d_dtheta;
125     int * __restrict__   gm_gridlineIndices = kernelParams.atoms.d_gridlineIndices;
126
127     /* Fractional coordinates */
128     __shared__ float sm_fractCoords[atomsPerBlock * DIM];
129
130     /* Thread index w.r.t. block */
131     const int threadLocalId = (threadIdx.z * (blockDim.x * blockDim.y))
132         + (threadIdx.y * blockDim.x) + threadIdx.x;
133     /* Warp index w.r.t. block - could probably be obtained easier? */
134     const int warpIndex = threadLocalId / warp_size;
135     /* Thread index w.r.t. warp */
136     const int threadWarpIndex = threadLocalId % warp_size;
137     /* Atom index w.r.t. warp - alternating 0 1 0 1 .. */
138     const int atomWarpIndex = threadWarpIndex % PME_SPREADGATHER_ATOMS_PER_WARP;
139     /* Atom index w.r.t. block/shared memory */
140     const int atomIndexLocal = warpIndex * PME_SPREADGATHER_ATOMS_PER_WARP + atomWarpIndex;
141
142     /* Atom index w.r.t. global memory */
143     const int atomIndexGlobal = atomIndexOffset + atomIndexLocal;
144     /* Spline contribution index in one dimension */
145     const int orderIndex = threadWarpIndex / (PME_SPREADGATHER_ATOMS_PER_WARP * DIM);
146     /* Dimension index */
147     const int dimIndex = (threadWarpIndex - orderIndex * (PME_SPREADGATHER_ATOMS_PER_WARP * DIM)) / PME_SPREADGATHER_ATOMS_PER_WARP;
148
149     /* Multi-purpose index of rvec/ivec atom data */
150     const int sharedMemoryIndex = atomIndexLocal * DIM + dimIndex;
151
152     /* Spline parameters need a working buffer.
153      * With PME_GPU_PARALLEL_SPLINE == 0 it is just a local array of (order) values for some of the threads, which is fine;
154      * With PME_GPU_PARALLEL_SPLINE == 1 (order) times more threads are involved, so the shared memory is used to avoid overhead.
155      * The buffer's size, striding and indexing are adjusted accordingly.
156      * The buffer is accessed with SPLINE_DATA_PTR and SPLINE_DATA macros.
157      */
158 #if PME_GPU_PARALLEL_SPLINE
159     const int        splineDataStride  = atomsPerBlock * DIM;
160     const int        splineDataIndex   = sharedMemoryIndex;
161     __shared__ float sm_splineData[splineDataStride * order];
162     float           *splineDataPtr = sm_splineData;
163 #else
164     const int        splineDataStride = 1;
165     const int        splineDataIndex  = 0;
166     float            splineData[splineDataStride * order];
167     float           *splineDataPtr = splineData;
168 #endif
169
170 #define SPLINE_DATA_PTR(i) (splineDataPtr + ((i) * splineDataStride + splineDataIndex))
171 #define SPLINE_DATA(i) (*SPLINE_DATA_PTR(i))
172
173     const int localCheck  = (dimIndex < DIM) && (orderIndex < (PME_GPU_PARALLEL_SPLINE ? order : 1));
174     const int globalCheck = pme_gpu_check_atom_data_index(atomIndexGlobal, kernelParams.atoms.nAtoms);
175
176     if (localCheck && globalCheck)
177     {
178         /* Indices interpolation */
179
180         if (orderIndex == 0)
181         {
182             int           tableIndex, tInt;
183             float         n, t;
184             const float3  x = sm_coordinates[atomIndexLocal];
185             /* Accessing fields in fshOffset/nXYZ/recipbox/... with dimIndex offset
186              * puts them into local memory(!) instead of accessing the constant memory directly.
187              * That's the reason for the switch, to unroll explicitly.
188              * The commented parts correspond to the 0 components of the recipbox.
189              */
190             switch (dimIndex)
191             {
192                 case XX:
193                     tableIndex  = kernelParams.grid.tablesOffsets[XX];
194                     n           = kernelParams.grid.realGridSizeFP[XX];
195                     t           = x.x * kernelParams.current.recipBox[dimIndex][XX] + x.y * kernelParams.current.recipBox[dimIndex][YY] + x.z * kernelParams.current.recipBox[dimIndex][ZZ];
196                     break;
197
198                 case YY:
199                     tableIndex  = kernelParams.grid.tablesOffsets[YY];
200                     n           = kernelParams.grid.realGridSizeFP[YY];
201                     t           = /*x.x * kernelParams.current.recipbox[dimIndex][XX] + */ x.y * kernelParams.current.recipBox[dimIndex][YY] + x.z * kernelParams.current.recipBox[dimIndex][ZZ];
202                     break;
203
204                 case ZZ:
205                     tableIndex  = kernelParams.grid.tablesOffsets[ZZ];
206                     n           = kernelParams.grid.realGridSizeFP[ZZ];
207                     t           = /*x.x * kernelParams.current.recipbox[dimIndex][XX] + x.y * kernelParams.current.recipbox[dimIndex][YY] + */ x.z * kernelParams.current.recipBox[dimIndex][ZZ];
208                     break;
209             }
210             const float shift = c_pmeMaxUnitcellShift;
211             /* Fractional coordinates along box vectors, adding a positive shift to ensure t is positive for triclinic boxes */
212             t    = (t + shift) * n;
213             tInt = (int)t;
214             sm_fractCoords[sharedMemoryIndex] = t - tInt;
215             tableIndex                       += tInt;
216             assert(tInt >= 0);
217             assert(tInt < c_pmeNeighborUnitcellCount * n);
218
219             // TODO have shared table for both parameters to share the fetch, as index is always same?
220             // TODO compare texture/LDG performance
221             sm_fractCoords[sharedMemoryIndex] +=
222                 fetchFromParamLookupTable(kernelParams.grid.d_fractShiftsTable,
223                                           kernelParams.fractShiftsTableTexture,
224                                           tableIndex);
225             sm_gridlineIndices[sharedMemoryIndex] =
226                 fetchFromParamLookupTable(kernelParams.grid.d_gridlineIndicesTable,
227                                           kernelParams.gridlineIndicesTableTexture,
228                                           tableIndex);
229             gm_gridlineIndices[atomIndexOffset * DIM + sharedMemoryIndex] = sm_gridlineIndices[sharedMemoryIndex];
230         }
231
232         /* B-spline calculation */
233
234         const int chargeCheck = pme_gpu_check_atom_charge(sm_coefficients[atomIndexLocal]);
235         if (chargeCheck)
236         {
237             float       div;
238             int         o = orderIndex; // This is an index that is set once for PME_GPU_PARALLEL_SPLINE == 1
239
240             const float dr = sm_fractCoords[sharedMemoryIndex];
241             assert(isfinite(dr));
242
243             /* dr is relative offset from lower cell limit */
244             *SPLINE_DATA_PTR(order - 1) = 0.0f;
245             *SPLINE_DATA_PTR(1)         = dr;
246             *SPLINE_DATA_PTR(0)         = 1.0f - dr;
247
248 #pragma unroll
249             for (int k = 3; k < order; k++)
250             {
251                 div                     = 1.0f / (k - 1.0f);
252                 *SPLINE_DATA_PTR(k - 1) = div * dr * SPLINE_DATA(k - 2);
253 #pragma unroll
254                 for (int l = 1; l < (k - 1); l++)
255                 {
256                     *SPLINE_DATA_PTR(k - l - 1) = div * ((dr + l) * SPLINE_DATA(k - l - 2) + (k - l - dr) * SPLINE_DATA(k - l - 1));
257                 }
258                 *SPLINE_DATA_PTR(0) = div * (1.0f - dr) * SPLINE_DATA(0);
259             }
260
261             const int thetaIndexBase        = getSplineParamIndexBase<order>(warpIndex, atomWarpIndex);
262             const int thetaGlobalOffsetBase = atomIndexOffset * DIM * order;
263
264             /* Differentiation and storing the spline derivatives (dtheta) */
265 #if !PME_GPU_PARALLEL_SPLINE
266             // With PME_GPU_PARALLEL_SPLINE == 1, o is already set to orderIndex;
267             // With PME_GPU_PARALLEL_SPLINE == 0, we loop o over range(order).
268 #pragma unroll
269             for (o = 0; o < order; o++)
270 #endif
271             {
272                 const int   thetaIndex       = getSplineParamIndex<order>(thetaIndexBase, dimIndex, o);
273                 const int   thetaGlobalIndex = thetaGlobalOffsetBase + thetaIndex;
274
275                 const float dtheta = ((o > 0) ? SPLINE_DATA(o - 1) : 0.0f) - SPLINE_DATA(o);
276                 assert(isfinite(dtheta));
277                 gm_dtheta[thetaGlobalIndex] = dtheta;
278             }
279
280             div  = 1.0f / (order - 1.0f);
281             *SPLINE_DATA_PTR(order - 1) = div * dr * SPLINE_DATA(order - 2);
282 #pragma unroll
283             for (int k = 1; k < (order - 1); k++)
284             {
285                 *SPLINE_DATA_PTR(order - k - 1) = div * ((dr + k) * SPLINE_DATA(order - k - 2) + (order - k - dr) * SPLINE_DATA(order - k - 1));
286             }
287             *SPLINE_DATA_PTR(0) = div * (1.0f - dr) * SPLINE_DATA(0);
288
289             /* Storing the spline values (theta) */
290 #if !PME_GPU_PARALLEL_SPLINE
291             // See comment for the loop above
292 #pragma unroll
293             for (o = 0; o < order; o++)
294 #endif
295             {
296                 const int thetaIndex       = getSplineParamIndex<order>(thetaIndexBase, dimIndex, o);
297                 const int thetaGlobalIndex = thetaGlobalOffsetBase + thetaIndex;
298
299                 sm_theta[thetaIndex]       = SPLINE_DATA(o);
300                 assert(isfinite(sm_theta[thetaIndex]));
301                 gm_theta[thetaGlobalIndex] = SPLINE_DATA(o);
302             }
303         }
304     }
305 }
306
307 /*! \brief
308  * Charge spreading onto the grid.
309  * This corresponds to the CPU function spread_coefficients_bsplines_thread().
310  * Second stage of the whole kernel.
311  *
312  * \tparam[in] order                PME interpolation order.
313  * \tparam[in] wrapX                A boolean which tells if the grid overlap in dimension X should be wrapped.
314  * \tparam[in] wrapY                A boolean which tells if the grid overlap in dimension Y should be wrapped.
315  * \param[in]  kernelParams         Input PME CUDA data in constant memory.
316  * \param[in]  atomIndexOffset      Starting atom index for the execution block w.r.t. global memory.
317  * \param[in]  sm_coefficients      Atom coefficents/charges in the shared memory.
318  * \param[in]  sm_gridlineIndices   Atom gridline indices in the shared memory.
319  * \param[in]  sm_theta             Atom spline values in the shared memory.
320  */
321 template <
322     const int order, const bool wrapX, const bool wrapY>
323 __device__ __forceinline__ void spread_charges(const PmeGpuCudaKernelParams           kernelParams,
324                                                int                                    atomIndexOffset,
325                                                const float * __restrict__             sm_coefficients,
326                                                const int * __restrict__               sm_gridlineIndices,
327                                                const float * __restrict__             sm_theta)
328 {
329     /* Global memory pointer to the output grid */
330     float * __restrict__ gm_grid = kernelParams.grid.d_realGrid;
331
332
333     const int nx  = kernelParams.grid.realGridSize[XX];
334     const int ny  = kernelParams.grid.realGridSize[YY];
335     const int nz  = kernelParams.grid.realGridSize[ZZ];
336     const int pny = kernelParams.grid.realGridSizePadded[YY];
337     const int pnz = kernelParams.grid.realGridSizePadded[ZZ];
338
339     const int offx = 0, offy = 0, offz = 0; // unused for now
340
341     const int atomIndexLocal  = threadIdx.z;
342     const int atomIndexGlobal = atomIndexOffset + atomIndexLocal;
343
344     const int globalCheck = pme_gpu_check_atom_data_index(atomIndexGlobal, kernelParams.atoms.nAtoms);
345     const int chargeCheck = pme_gpu_check_atom_charge(sm_coefficients[atomIndexLocal]);
346     if (chargeCheck & globalCheck)
347     {
348         // Spline Y/Z coordinates
349         const int ithy   = threadIdx.y;
350         const int ithz   = threadIdx.x;
351         const int ixBase = sm_gridlineIndices[atomIndexLocal * DIM + XX] - offx;
352         int       iy     = sm_gridlineIndices[atomIndexLocal * DIM + YY] - offy + ithy;
353         if (wrapY & (iy >= ny))
354         {
355             iy -= ny;
356         }
357         int iz  = sm_gridlineIndices[atomIndexLocal * DIM + ZZ] - offz + ithz;
358         if (iz >= nz)
359         {
360             iz -= nz;
361         }
362
363         /* Atom index w.r.t. warp - alternating 0 1 0 1 .. */
364         const int    atomWarpIndex   = atomIndexLocal % PME_SPREADGATHER_ATOMS_PER_WARP;
365         /* Warp index w.r.t. block - could probably be obtained easier? */
366         const int    warpIndex       = atomIndexLocal / PME_SPREADGATHER_ATOMS_PER_WARP;
367
368         const int    splineIndexBase = getSplineParamIndexBase<order>(warpIndex, atomWarpIndex);
369         const int    splineIndexZ    = getSplineParamIndex<order>(splineIndexBase, ZZ, ithz);
370         const float  thetaZ          = sm_theta[splineIndexZ];
371         const int    splineIndexY    = getSplineParamIndex<order>(splineIndexBase, YY, ithy);
372         const float  thetaY          = sm_theta[splineIndexY];
373         const float  constVal        = thetaZ * thetaY * sm_coefficients[atomIndexLocal];
374         assert(isfinite(constVal));
375         const int    constOffset     = iy * pnz + iz;
376
377 #pragma unroll
378         for (int ithx = 0; (ithx < order); ithx++)
379         {
380             int ix = ixBase + ithx;
381             if (wrapX & (ix >= nx))
382             {
383                 ix -= nx;
384             }
385             const int   gridIndexGlobal = ix * pny * pnz + constOffset;
386             const int   splineIndexX    = getSplineParamIndex<order>(splineIndexBase, XX, ithx);
387             const float thetaX          = sm_theta[splineIndexX];
388             assert(isfinite(thetaX));
389             assert(isfinite(gm_grid[gridIndexGlobal]));
390             atomicAdd(gm_grid + gridIndexGlobal, thetaX * constVal);
391         }
392     }
393 }
394
395 /*! \brief
396  * A spline computation and charge spreading kernel function.
397  *
398  * \tparam[in] order                PME interpolation order.
399  * \tparam[in] computeSplines       A boolean which tells if the spline parameter and
400  *                                  gridline indices' computation should be performed.
401  * \tparam[in] spreadCharges        A boolean which tells if the charge spreading should be performed.
402  * \tparam[in] wrapX                A boolean which tells if the grid overlap in dimension X should be wrapped.
403  * \tparam[in] wrapY                A boolean which tells if the grid overlap in dimension Y should be wrapped.
404  * \param[in]  kernelParams         Input PME CUDA data in constant memory.
405  */
406 template <
407     const int order,
408     const bool computeSplines,
409     const bool spreadCharges,
410     const bool wrapX,
411     const bool wrapY
412     >
413 __launch_bounds__(c_spreadMaxThreadsPerBlock)
414 __global__ void pme_spline_and_spread_kernel(const PmeGpuCudaKernelParams kernelParams)
415 {
416     const int        atomsPerBlock = c_spreadMaxThreadsPerBlock / PME_SPREADGATHER_THREADS_PER_ATOM;
417     // Gridline indices, ivec
418     __shared__ int   sm_gridlineIndices[atomsPerBlock * DIM];
419     // Charges
420     __shared__ float sm_coefficients[atomsPerBlock];
421     // Spline values
422     __shared__ float sm_theta[atomsPerBlock * DIM * order];
423
424     const int        blockIndex      = blockIdx.y * gridDim.x + blockIdx.x;
425     const int        atomIndexOffset = blockIndex * atomsPerBlock;
426
427     /* Early return for fully empty blocks at the end
428      * (should only happen on Fermi or billions of input atoms)
429      */
430     if (atomIndexOffset >= kernelParams.atoms.nAtoms)
431     {
432         return;
433     }
434
435     /* Staging coefficients/charges for both spline and spread */
436     pme_gpu_stage_atom_data<float, atomsPerBlock, 1>(kernelParams, sm_coefficients, kernelParams.atoms.d_coefficients);
437
438     if (computeSplines)
439     {
440         /* Staging coordinates */
441         __shared__ float sm_coordinates[DIM * atomsPerBlock];
442         pme_gpu_stage_atom_data<float, atomsPerBlock, DIM>(kernelParams, sm_coordinates, kernelParams.atoms.d_coordinates);
443
444         __syncthreads();
445         calculate_splines<order, atomsPerBlock>(kernelParams, atomIndexOffset, (const float3 *)sm_coordinates,
446                                                 sm_coefficients, sm_theta, sm_gridlineIndices);
447         gmx_syncwarp();
448     }
449     else
450     {
451         /* Staging the data for spread
452          * (the data is assumed to be in GPU global memory with proper layout already,
453          * as in after running the spline kernel)
454          */
455         /* Spline data - only thetas (dthetas will only be needed in gather) */
456         pme_gpu_stage_atom_data<float, atomsPerBlock, DIM * order>(kernelParams, sm_theta, kernelParams.atoms.d_theta);
457         /* Gridline indices */
458         pme_gpu_stage_atom_data<int, atomsPerBlock, DIM>(kernelParams, sm_gridlineIndices, kernelParams.atoms.d_gridlineIndices);
459
460         __syncthreads();
461     }
462
463     /* Spreading */
464     if (spreadCharges)
465     {
466         spread_charges<order, wrapX, wrapY>(kernelParams, atomIndexOffset, sm_coefficients, sm_gridlineIndices, sm_theta);
467     }
468 }
469
470 //! Kernel instantiations
471 template __global__ void pme_spline_and_spread_kernel<4, true, true, true, true>(const PmeGpuCudaKernelParams);
472 template __global__ void pme_spline_and_spread_kernel<4, true, false, true, true>(const PmeGpuCudaKernelParams);
473 template __global__ void pme_spline_and_spread_kernel<4, false, true, true, true>(const PmeGpuCudaKernelParams);