2 * This file is part of the GROMACS molecular simulation package.
4 * Copyright (c) 2016,2017,2018,2019,2020, by the GROMACS development team, led by
5 * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
6 * and including many others, as listed in the AUTHORS file in the
7 * top-level source directory and at http://www.gromacs.org.
9 * GROMACS is free software; you can redistribute it and/or
10 * modify it under the terms of the GNU Lesser General Public License
11 * as published by the Free Software Foundation; either version 2.1
12 * of the License, or (at your option) any later version.
14 * GROMACS is distributed in the hope that it will be useful,
15 * but WITHOUT ANY WARRANTY; without even the implied warranty of
16 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
17 * Lesser General Public License for more details.
19 * You should have received a copy of the GNU Lesser General Public
20 * License along with GROMACS; if not, see
21 * http://www.gnu.org/licenses, or write to the Free Software Foundation,
22 * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
24 * If you want to redistribute modifications to GROMACS, please
25 * consider that scientific software is very special. Version
26 * control is crucial - bugs must be traceable. We will be happy to
27 * consider code for inclusion in the official distribution, but
28 * derived work must not be called official GROMACS. Details are found
29 * in the README & COPYING files - if they are missing, get the
30 * official version at http://www.gromacs.org.
32 * To help us fund GROMACS development, we humbly ask that you cite
33 * the research papers on the package. Check out http://www.gromacs.org.
37 * \brief Implements helper routines for PME gather and spline routines.
39 * \author Aleksei Iupinov <a.yupinov@gmail.com>
46 #include "gromacs/gpu_utils/cuda_kernel_utils.cuh"
52 * Gets a base of the unique index to an element in a spline parameter buffer (theta/dtheta),
53 * which is laid out for GPU spread/gather kernels. The base only corresponds to the atom index within the execution block.
54 * Feed the result into getSplineParamIndex() to get a full index.
55 * TODO: it's likely that both parameters can be just replaced with a single atom index, as they are derived from it.
56 * Do that, verifying that the generated code is not bloated, and/or revise the spline indexing scheme.
57 * Removing warp dependency would also be nice (and would probably coincide with removing c_pmeSpreadGatherAtomsPerWarp).
59 * \tparam order PME order
60 * \tparam atomsPerWarp Number of atoms processed by a warp
61 * \param[in] warpIndex Warp index wrt the block.
62 * \param[in] atomWarpIndex Atom index wrt the warp (from 0 to atomsPerWarp - 1).
64 * \returns Index into theta or dtheta array using GPU layout.
66 template<int order, int atomsPerWarp>
67 int __device__ __forceinline__ getSplineParamIndexBase(int warpIndex, int atomWarpIndex)
69 assert((atomWarpIndex >= 0) && (atomWarpIndex < atomsPerWarp));
70 const int dimIndex = 0;
71 const int splineIndex = 0;
72 // The zeroes are here to preserve the full index formula for reference
73 return (((splineIndex + order * warpIndex) * DIM + dimIndex) * atomsPerWarp + atomWarpIndex);
77 * Gets a unique index to an element in a spline parameter buffer (theta/dtheta),
78 * which is laid out for GPU spread/gather kernels. The index is wrt to the execution block,
79 * in range(0, atomsPerBlock * order * DIM).
80 * This function consumes result of getSplineParamIndexBase() and adjusts it for \p dimIndex and \p splineIndex.
82 * \tparam order PME order
83 * \tparam atomsPerWarp Number of atoms processed by a warp
84 * \param[in] paramIndexBase Must be result of getSplineParamIndexBase().
85 * \param[in] dimIndex Dimension index (from 0 to 2)
86 * \param[in] splineIndex Spline contribution index (from 0 to \p order - 1)
88 * \returns Index into theta or dtheta array using GPU layout.
90 template<int order, int atomsPerWarp>
91 int __device__ __forceinline__ getSplineParamIndex(int paramIndexBase, int dimIndex, int splineIndex)
93 assert((dimIndex >= XX) && (dimIndex < DIM));
94 assert((splineIndex >= 0) && (splineIndex < order));
95 return (paramIndexBase + (splineIndex * DIM + dimIndex) * atomsPerWarp);
99 * An inline CUDA function for checking the global atom data indices against the atom data array sizes.
101 * \param[in] atomDataIndex The atom data index.
102 * \param[in] nAtomData The atom data array element count.
103 * \returns Non-0 if index is within bounds (or PME data padding is enabled), 0 otherwise.
105 * This is called from the spline_and_spread and gather PME kernels.
106 * The goal is to isolate the global range checks, and allow avoiding them with c_usePadding enabled.
108 int __device__ __forceinline__ pme_gpu_check_atom_data_index(const int atomDataIndex, const int nAtomData)
110 return c_usePadding ? 1 : (atomDataIndex < nAtomData);
114 * An inline CUDA function for skipping the zero-charge atoms.
116 * \returns Non-0 if atom should be processed, 0 otherwise.
117 * \param[in] coefficient The atom charge.
119 * This is called from the spline_and_spread and gather PME kernels.
121 int __device__ __forceinline__ pme_gpu_check_atom_charge(const float coefficient)
123 assert(isfinite(coefficient));
124 return c_skipNeutralAtoms ? (coefficient != 0.0f) : 1;
127 //! Controls if the atom and charge data is prefeched into shared memory or loaded per thread from global
128 static const bool c_useAtomDataPrefetch = true;
130 /*! \brief Asserts if the argument is finite.
132 * The function works for any data type, that can be casted to float. Note that there is also
133 * a specialized implementation for float3 data type.
135 * \param[in] arg Argument to check.
138 __device__ inline void assertIsFinite(T arg);
141 __device__ inline void assertIsFinite(float3 arg)
143 assert(isfinite(float(arg.x)));
144 assert(isfinite(float(arg.y)));
145 assert(isfinite(float(arg.z)));
149 __device__ inline void assertIsFinite(T arg)
151 assert(isfinite(float(arg)));
155 * General purpose function for loading atom-related data from global to shared memory.
157 * \tparam[in] T Data type (float/int/...)
158 * \tparam[in] atomsPerBlock Number of atoms processed by a block - should be accounted for in the size of the shared memory array.
159 * \tparam[in] dataCountPerAtom Number of data elements per single atom (e.g. DIM for an rvec coordinates array).
160 * \param[in] kernelParams Input PME CUDA data in constant memory.
161 * \param[out] sm_destination Shared memory array for output.
162 * \param[in] gm_source Global memory array for input.
164 template<typename T, const int atomsPerBlock, const int dataCountPerAtom>
165 __device__ __forceinline__ void pme_gpu_stage_atom_data(const PmeGpuCudaKernelParams kernelParams,
166 T* __restrict__ sm_destination,
167 const T* __restrict__ gm_source)
169 static_assert(c_usePadding,
170 "With padding disabled, index checking should be fixed to account for spline "
171 "theta/dtheta pr-warp alignment");
172 const int blockIndex = blockIdx.y * gridDim.x + blockIdx.x;
173 const int threadLocalIndex = ((threadIdx.z * blockDim.y + threadIdx.y) * blockDim.x) + threadIdx.x;
174 const int localIndex = threadLocalIndex;
175 const int globalIndexBase = blockIndex * atomsPerBlock * dataCountPerAtom;
176 const int globalIndex = globalIndexBase + localIndex;
177 const int globalCheck =
178 pme_gpu_check_atom_data_index(globalIndex, kernelParams.atoms.nAtoms * dataCountPerAtom);
179 if ((localIndex < atomsPerBlock * dataCountPerAtom) & globalCheck)
181 assertIsFinite(gm_source[globalIndex]);
182 sm_destination[localIndex] = gm_source[globalIndex];
187 * PME GPU spline parameter and gridline indices calculation.
188 * This corresponds to the CPU functions calc_interpolation_idx() and make_bsplines().
189 * First stage of the whole kernel.
191 * \tparam[in] order PME interpolation order.
192 * \tparam[in] atomsPerBlock Number of atoms processed by a block - should be accounted for
193 * in the sizes of the shared memory arrays.
194 * \tparam[in] atomsPerWarp Number of atoms processed by a warp
195 * \tparam[in] writeSmDtheta Bool controling if the theta derivative should be written to shared memory. Enables calculation of dtheta if set.
196 * \tparam[in] writeGlobal A boolean which tells if the theta values and gridlines should be written to global memory. Enables calculation of dtheta if set.
197 * \param[in] kernelParams Input PME CUDA data in constant memory.
198 * \param[in] atomIndexOffset Starting atom index for the execution block w.r.t. global memory.
199 * \param[in] atomX Atom coordinate of atom processed by thread.
200 * \param[in] atomCharge Atom charge/coefficient of atom processed by thread.
201 * \param[out] sm_theta Atom spline values in the shared memory.
202 * \param[out] sm_dtheta Derivative of atom spline values in shared memory.
203 * \param[out] sm_gridlineIndices Atom gridline indices in the shared memory.
206 template<const int order, const int atomsPerBlock, const int atomsPerWarp, const bool writeSmDtheta, const bool writeGlobal>
207 __device__ __forceinline__ void calculate_splines(const PmeGpuCudaKernelParams kernelParams,
208 const int atomIndexOffset,
210 const float atomCharge,
211 float* __restrict__ sm_theta,
212 float* __restrict__ sm_dtheta,
213 int* __restrict__ sm_gridlineIndices)
215 /* Global memory pointers for output */
216 float* __restrict__ gm_theta = kernelParams.atoms.d_theta;
217 float* __restrict__ gm_dtheta = kernelParams.atoms.d_dtheta;
218 int* __restrict__ gm_gridlineIndices = kernelParams.atoms.d_gridlineIndices;
220 /* Fractional coordinates */
221 __shared__ float sm_fractCoords[atomsPerBlock * DIM];
223 /* Thread index w.r.t. block */
224 const int threadLocalId =
225 (threadIdx.z * (blockDim.x * blockDim.y)) + (threadIdx.y * blockDim.x) + threadIdx.x;
226 /* Warp index w.r.t. block - could probably be obtained easier? */
227 const int warpIndex = threadLocalId / warp_size;
228 /* Atom index w.r.t. warp - alternating 0 1 0 1 .. */
229 const int atomWarpIndex = threadIdx.z % atomsPerWarp;
230 /* Atom index w.r.t. block/shared memory */
231 const int atomIndexLocal = warpIndex * atomsPerWarp + atomWarpIndex;
233 /* Atom index w.r.t. global memory */
234 const int atomIndexGlobal = atomIndexOffset + atomIndexLocal;
235 /* Spline contribution index in one dimension */
236 const int threadLocalIdXY = (threadIdx.y * blockDim.x) + threadIdx.x;
237 const int orderIndex = threadLocalIdXY / DIM;
238 /* Dimension index */
239 const int dimIndex = threadLocalIdXY % DIM;
241 /* Multi-purpose index of rvec/ivec atom data */
242 const int sharedMemoryIndex = atomIndexLocal * DIM + dimIndex;
244 float splineData[order];
246 const int localCheck = (dimIndex < DIM) && (orderIndex < 1);
247 const int globalCheck = pme_gpu_check_atom_data_index(atomIndexGlobal, kernelParams.atoms.nAtoms);
249 /* we have 4 threads per atom, but can only use 3 here for the dimensions */
250 if (localCheck && globalCheck)
252 /* Indices interpolation */
256 int tableIndex, tInt;
258 assert(atomIndexLocal < DIM * atomsPerBlock);
259 /* Accessing fields in fshOffset/nXYZ/recipbox/... with dimIndex offset
260 * puts them into local memory(!) instead of accessing the constant memory directly.
261 * That's the reason for the switch, to unroll explicitly.
262 * The commented parts correspond to the 0 components of the recipbox.
267 tableIndex = kernelParams.grid.tablesOffsets[XX];
268 n = kernelParams.grid.realGridSizeFP[XX];
269 t = atomX.x * kernelParams.current.recipBox[dimIndex][XX]
270 + atomX.y * kernelParams.current.recipBox[dimIndex][YY]
271 + atomX.z * kernelParams.current.recipBox[dimIndex][ZZ];
275 tableIndex = kernelParams.grid.tablesOffsets[YY];
276 n = kernelParams.grid.realGridSizeFP[YY];
277 t = /*atomX.x * kernelParams.current.recipBox[dimIndex][XX] + */ atomX.y
278 * kernelParams.current.recipBox[dimIndex][YY]
279 + atomX.z * kernelParams.current.recipBox[dimIndex][ZZ];
283 tableIndex = kernelParams.grid.tablesOffsets[ZZ];
284 n = kernelParams.grid.realGridSizeFP[ZZ];
285 t = /*atomX.x * kernelParams.current.recipBox[dimIndex][XX] + atomX.y * kernelParams.current.recipBox[dimIndex][YY] + */ atomX
287 * kernelParams.current.recipBox[dimIndex][ZZ];
290 const float shift = c_pmeMaxUnitcellShift;
291 /* Fractional coordinates along box vectors, adding a positive shift to ensure t is positive for triclinic boxes */
294 assert(sharedMemoryIndex < atomsPerBlock * DIM);
295 sm_fractCoords[sharedMemoryIndex] = t - tInt;
298 assert(tInt < c_pmeNeighborUnitcellCount * n);
300 // TODO have shared table for both parameters to share the fetch, as index is always same?
301 // TODO compare texture/LDG performance
302 sm_fractCoords[sharedMemoryIndex] +=
303 fetchFromParamLookupTable(kernelParams.grid.d_fractShiftsTable,
304 kernelParams.fractShiftsTableTexture, tableIndex);
305 sm_gridlineIndices[sharedMemoryIndex] =
306 fetchFromParamLookupTable(kernelParams.grid.d_gridlineIndicesTable,
307 kernelParams.gridlineIndicesTableTexture, tableIndex);
310 gm_gridlineIndices[atomIndexOffset * DIM + sharedMemoryIndex] =
311 sm_gridlineIndices[sharedMemoryIndex];
315 /* B-spline calculation */
317 const int chargeCheck = pme_gpu_check_atom_charge(atomCharge);
321 int o = orderIndex; // This is an index that is set once for PME_GPU_PARALLEL_SPLINE == 1
323 const float dr = sm_fractCoords[sharedMemoryIndex];
324 assert(isfinite(dr));
326 /* dr is relative offset from lower cell limit */
327 splineData[order - 1] = 0.0f;
329 splineData[0] = 1.0f - dr;
332 for (int k = 3; k < order; k++)
334 div = 1.0f / (k - 1.0f);
335 splineData[k - 1] = div * dr * splineData[k - 2];
337 for (int l = 1; l < (k - 1); l++)
339 splineData[k - l - 1] =
340 div * ((dr + l) * splineData[k - l - 2] + (k - l - dr) * splineData[k - l - 1]);
342 splineData[0] = div * (1.0f - dr) * splineData[0];
345 const int thetaIndexBase =
346 getSplineParamIndexBase<order, atomsPerWarp>(warpIndex, atomWarpIndex);
347 const int thetaGlobalOffsetBase = atomIndexOffset * DIM * order;
348 /* only calculate dtheta if we are saving it to shared or global memory */
349 if (writeSmDtheta || writeGlobal)
351 /* Differentiation and storing the spline derivatives (dtheta) */
353 for (o = 0; o < order; o++)
355 const int thetaIndex =
356 getSplineParamIndex<order, atomsPerWarp>(thetaIndexBase, dimIndex, o);
358 const float dtheta = ((o > 0) ? splineData[o - 1] : 0.0f) - splineData[o];
359 assert(isfinite(dtheta));
360 assert(thetaIndex < order * DIM * atomsPerBlock);
363 sm_dtheta[thetaIndex] = dtheta;
367 const int thetaGlobalIndex = thetaGlobalOffsetBase + thetaIndex;
368 gm_dtheta[thetaGlobalIndex] = dtheta;
373 div = 1.0f / (order - 1.0f);
374 splineData[order - 1] = div * dr * splineData[order - 2];
376 for (int k = 1; k < (order - 1); k++)
378 splineData[order - k - 1] = div
379 * ((dr + k) * splineData[order - k - 2]
380 + (order - k - dr) * splineData[order - k - 1]);
382 splineData[0] = div * (1.0f - dr) * splineData[0];
384 /* Storing the spline values (theta) */
386 for (o = 0; o < order; o++)
388 const int thetaIndex =
389 getSplineParamIndex<order, atomsPerWarp>(thetaIndexBase, dimIndex, o);
390 assert(thetaIndex < order * DIM * atomsPerBlock);
391 sm_theta[thetaIndex] = splineData[o];
392 assert(isfinite(sm_theta[thetaIndex]));
395 const int thetaGlobalIndex = thetaGlobalOffsetBase + thetaIndex;
396 gm_theta[thetaGlobalIndex] = splineData[o];