Remove PImpl scaffolding from CUDA version of LINCS
[gromacs.git] / src / gromacs / mdlib / lincs_cuda.cu
blobdf018e9e918b658ad823ec86d157c8b24774855b
1 /*
2  * This file is part of the GROMACS molecular simulation package.
3  *
4  * Copyright (c) 2019, 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.
8  *
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.
13  *
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.
18  *
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.
23  *
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.
31  *
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.
34  */
35 /*! \internal \file
36  *
37  * \brief Implements LINCS using CUDA
38  *
39  * This file contains implementation of LINCS constraints algorithm
40  * using CUDA, including class initialization, data-structures management
41  * and GPU kernel.
42  *
43  * \note Management of CUDA stream and periodic boundary exists here as a temporary
44  *       scaffolding to allow this class to be used as a stand-alone unit. The scaffolding
45  *        is intended to be removed once constraints are fully integrated with update module.
46  * \todo Reconsider naming, i.e. "cuda" suffics should be changed to "gpu".
47  *
48  * \author Artem Zhmurov <zhmurov@gmail.com>
49  * \author Alan Gray <alang@nvidia.com>
50  *
51  * \ingroup module_mdlib
52  */
53 #include "gmxpre.h"
55 #include "lincs_cuda.cuh"
57 #include <assert.h>
58 #include <stdio.h>
60 #include <cmath>
62 #include <algorithm>
64 #include "gromacs/gpu_utils/cuda_arch_utils.cuh"
65 #include "gromacs/gpu_utils/cudautils.cuh"
66 #include "gromacs/gpu_utils/devicebuffer.cuh"
67 #include "gromacs/gpu_utils/gputraits.cuh"
68 #include "gromacs/gpu_utils/vectype_ops.cuh"
69 #include "gromacs/math/vec.h"
70 #include "gromacs/mdlib/constr.h"
71 #include "gromacs/pbcutil/pbc.h"
72 #include "gromacs/pbcutil/pbc_aiuc_cuda.cuh"
73 #include "gromacs/topology/ifunc.h"
75 namespace gmx
78 //! Number of CUDA threads in a block
79 constexpr static int c_threadsPerBlock = 256;
80 //! Maximum number of threads in a block (for __launch_bounds__)
81 constexpr static int c_maxThreadsPerBlock = c_threadsPerBlock;
83 /*! \brief Main kernel for LINCS constraints.
84  *
85  * See Hess et al., J. Comput. Chem. 18: 1463-1472 (1997) for the description of the algorithm.
86  *
87  * In CUDA version, one thread is responsible for all computations for one constraint. The blocks are
88  * filled in a way that no constraint is coupled to the constraint from the next block. This is achieved
89  * by moving active threads to the next block, if the correspondent group of coupled constraints is to big
90  * to fit the current thread block. This may leave some 'dummy' threads in the end of the thread block, i.e.
91  * threads that are not required to do actual work. Since constraints from different blocks are not coupled,
92  * there is no need to synchronize across the device. However, extensive communication in a thread block
93  * are still needed.
94  *
95  * \todo Reduce synchronization overhead. Some ideas are:
96  *        1. Consider going to warp-level synchronization for the coupled constraints.
97  *        2. Move more data to local/shared memory and try to get rid of atomic operations (at least on
98  *           the device level).
99  *        3. Use analytical solution for matrix A inversion.
100  *        4. Introduce mapping of thread id to both single constraint and single atom, thus designating
101  *           Nth threads to deal with Nat <= Nth coupled atoms and Nc <= Nth coupled constraints.
102  *       See Redmine issue #2885 for details (https://redmine.gromacs.org/issues/2885)
103  * \todo The use of __restrict__  for gm_xp and gm_v causes failure, probably because of the atomic
104          operations. Investigate this issue further.
106  * \param[in,out] kernelParams  All parameters and pointers for the kernel condensed in single struct.
107  * \param[in]     invdt         Inverse timestep (needed to update velocities).
108  */
109 template <bool updateVelocities, bool computeVirial>
110 __launch_bounds__(c_maxThreadsPerBlock)
111 __global__ void lincs_kernel(LincsCudaKernelParameters   kernelParams,
112                              const float3* __restrict__  gm_x,
113                              float3*                     gm_xp,
114                              float3*                     gm_v,
115                              const float                 invdt)
117     const PbcAiuc               pbcAiuc                     = kernelParams.pbcAiuc;
118     const int                   numConstraintsThreads       = kernelParams.numConstraintsThreads;
119     const int                   numIterations               = kernelParams.numIterations;
120     const int                   expansionOrder              = kernelParams.expansionOrder;
121     const int2*   __restrict__  gm_constraints              = kernelParams.d_constraints;
122     const float*  __restrict__  gm_constraintsTargetLengths = kernelParams.d_constraintsTargetLengths;
123     const int*    __restrict__  gm_coupledConstraintsCounts = kernelParams.d_coupledConstraintsCounts;
124     const int*    __restrict__  gm_coupledConstraintsIdxes  = kernelParams.d_coupledConstraintsIndices;
125     const float*  __restrict__  gm_massFactors              = kernelParams.d_massFactors;
126     float*  __restrict__        gm_matrixA                  = kernelParams.d_matrixA;
127     const float*  __restrict__  gm_inverseMasses            = kernelParams.d_inverseMasses;
128     float*  __restrict__        gm_virialScaled             = kernelParams.d_virialScaled;
130     int threadIndex                                         = blockIdx.x*blockDim.x+threadIdx.x;
132     // numConstraintsThreads should be a integer multiple of blockSize (numConstraintsThreads = numBlocks*blockSize).
133     // This is to ensure proper synchronizations and reduction. All array are padded to the required size.
134     assert(threadIndex < numConstraintsThreads);
136     // Vectors connecting constrained atoms before algorithm was applied.
137     // Needed to construct constrain matrix A
138     extern __shared__ float3 sm_r[];
140     int2                     pair = gm_constraints[threadIndex];
141     int                      i    = pair.x;
142     int                      j    = pair.y;
144     // Mass-scaled Lagrange multiplier
145     float  lagrangeScaled = 0.0f;
147     float  targetLength;
148     float  inverseMassi;
149     float  inverseMassj;
150     float  sqrtReducedMass;
152     float3 xi;
153     float3 xj;
154     float3 rc;
156     // i == -1 indicates dummy constraint at the end of the thread block.
157     bool isDummyThread = (i == -1);
159     // Everything computed for these dummies will be equal to zero
160     if (isDummyThread)
161     {
162         targetLength    = 0.0f;
163         inverseMassi    = 0.0f;
164         inverseMassj    = 0.0f;
165         sqrtReducedMass = 0.0f;
167         xi = make_float3(0.0f, 0.0f, 0.0f);
168         xj = make_float3(0.0f, 0.0f, 0.0f);
169         rc = make_float3(0.0f, 0.0f, 0.0f);
170     }
171     else
172     {
173         // Collecting data
174         targetLength    = gm_constraintsTargetLengths[threadIndex];
175         inverseMassi    = gm_inverseMasses[i];
176         inverseMassj    = gm_inverseMasses[j];
177         sqrtReducedMass = rsqrt(inverseMassi + inverseMassj);
179         xi = gm_x[i];
180         xj = gm_x[j];
182         float3 dx   = pbcDxAiuc(pbcAiuc, xi, xj);
184         float  rlen = rsqrtf(dx.x*dx.x + dx.y*dx.y + dx.z*dx.z);
185         rc   = rlen*dx;
186     }
188     sm_r[threadIdx.x] = rc;
189     // Make sure that all r's are saved into shared memory
190     // before they are accessed in the loop below
191     __syncthreads();
193     /*
194      * Constructing LINCS matrix (A)
195      */
197     // Only non-zero values are saved (for coupled constraints)
198     int coupledConstraintsCount = gm_coupledConstraintsCounts[threadIndex];
199     for (int n = 0; n < coupledConstraintsCount; n++)
200     {
201         int    index = n*numConstraintsThreads + threadIndex;
202         int    c1    = gm_coupledConstraintsIdxes[index];
204         float3 rc1 = sm_r[c1 - blockIdx.x*blockDim.x];
205         gm_matrixA[index] = gm_massFactors[index]*(rc.x*rc1.x + rc.y*rc1.y + rc.z*rc1.z);
206     }
208     // Skipping in dummy threads
209     if (!isDummyThread)
210     {
211         xi = gm_xp[i];
212         xj = gm_xp[j];
213     }
215     float3 dx = pbcDxAiuc(pbcAiuc, xi, xj);
217     float  sol = sqrtReducedMass*((rc.x*dx.x + rc.y*dx.y + rc.z*dx.z) - targetLength);
219     /*
220      *  Inverse matrix using a set of expansionOrder matrix multiplications
221      */
223     // This will use the same memory space as sm_r, which is no longer needed.
224     extern __shared__ float  sm_rhs[];
225     // Save current right-hand-side vector in the shared memory
226     sm_rhs[threadIdx.x] = sol;
228     for (int rec = 0; rec < expansionOrder; rec++)
229     {
230         // Making sure that all sm_rhs are saved before they are accessed in a loop below
231         __syncthreads();
232         float mvb = 0.0f;
234         for (int n = 0; n < coupledConstraintsCount; n++)
235         {
236             int index = n*numConstraintsThreads + threadIndex;
237             int c1    = gm_coupledConstraintsIdxes[index];
238             // Convolute current right-hand-side with A
239             // Different, non overlapping parts of sm_rhs[..] are read during odd and even iterations
240             mvb = mvb + gm_matrixA[index]*sm_rhs[c1 - blockIdx.x*blockDim.x + blockDim.x*(rec % 2)];
241         }
242         // 'Switch' rhs vectors, save current result
243         // These values will be accessed in the loop above during the next iteration.
244         sm_rhs[threadIdx.x + blockDim.x*((rec + 1) % 2)] = mvb;
245         sol  = sol + mvb;
246     }
248     // Current mass-scaled Lagrange multipliers
249     lagrangeScaled = sqrtReducedMass*sol;
251     // Save updated coordinates before correction for the rotational lengthening
252     float3 tmp     = rc*lagrangeScaled;
254     // Writing for all but dummy constraints
255     if (!isDummyThread)
256     {
257         atomicAdd(&gm_xp[i], -tmp*inverseMassi);
258         atomicAdd(&gm_xp[j],  tmp*inverseMassj);
259     }
261     /*
262      *  Correction for centripetal effects
263      */
264     for (int iter = 0; iter < numIterations; iter++)
265     {
266         // Make sure that all xp's are saved: atomic operation calls before are
267         // communicating current xp[..] values across thread block.
268         __syncthreads();
270         if (!isDummyThread)
271         {
272             xi = gm_xp[i];
273             xj = gm_xp[j];
274         }
276         float3 dx = pbcDxAiuc(pbcAiuc, xi, xj);
278         float  len2  = targetLength*targetLength;
279         float  dlen2 = 2.0f*len2 - norm2(dx);
281         // TODO A little bit more effective but slightly less readable version of the below would be:
282         //      float proj = sqrtReducedMass*(targetLength - (dlen2 > 0.0f ? 1.0f : 0.0f)*dlen2*rsqrt(dlen2));
283         float  proj;
284         if (dlen2 > 0.0f)
285         {
286             proj = sqrtReducedMass*(targetLength - dlen2*rsqrt(dlen2));
287         }
288         else
289         {
290             proj = sqrtReducedMass*targetLength;
291         }
293         sm_rhs[threadIdx.x]   = proj;
294         float sol = proj;
296         /*
297          * Same matrix inversion as above is used for updated data
298          */
299         for (int rec = 0; rec < expansionOrder; rec++)
300         {
301             // Make sure that all elements of rhs are saved into shared memory
302             __syncthreads();
303             float mvb = 0;
305             for (int n = 0; n < coupledConstraintsCount; n++)
306             {
307                 int index = n*numConstraintsThreads + threadIndex;
308                 int c1    = gm_coupledConstraintsIdxes[index];
310                 mvb = mvb + gm_matrixA[index]*sm_rhs[c1 - blockIdx.x*blockDim.x + blockDim.x*(rec % 2)];
312             }
313             sm_rhs[threadIdx.x + blockDim.x*((rec + 1) % 2)] = mvb;
314             sol  = sol + mvb;
315         }
317         // Add corrections to Lagrange multipliers
318         float sqrtmu_sol  = sqrtReducedMass*sol;
319         lagrangeScaled += sqrtmu_sol;
321         // Save updated coordinates for the next iteration
322         // Dummy constraints are skipped
323         if (!isDummyThread)
324         {
325             float3 tmp = rc*sqrtmu_sol;
326             atomicAdd(&gm_xp[i], -tmp*inverseMassi);
327             atomicAdd(&gm_xp[j],  tmp*inverseMassj);
328         }
329     }
331     // Updating particle velocities for all but dummy threads
332     if (updateVelocities && !isDummyThread)
333     {
334         float3 tmp     = rc*invdt*lagrangeScaled;
335         atomicAdd(&gm_v[i], -tmp*inverseMassi);
336         atomicAdd(&gm_v[j],  tmp*inverseMassj);
337     }
340     if (computeVirial)
341     {
342         // Virial is computed from Lagrange multiplier (lagrangeScaled), target constrain length
343         // (targetLength) and the normalized vector connecting constrained atoms before
344         // the algorithm was applied (rc). The evaluation of virial in each thread is
345         // followed by basic reduction for the values inside single thread block.
346         // Then, the values are reduced across grid by atomicAdd(...).
347         //
348         // TODO Shuffle reduction.
349         // TODO Should be unified and/or done once when virial is actually needed.
350         // TODO Recursive version that removes atomicAdd(...)'s entirely is needed. Ideally,
351         //      one that works for any datatype.
353         // Save virial for each thread into the shared memory. Tensor is symmetrical, hence only
354         // 6 values are saved. Dummy threads will have zeroes in their virial: targetLength,
355         // lagrangeScaled and rc are all set to zero for them in the beginning of the kernel.
356         // The sm_threadVirial[..] will overlap with the sm_r[..] and sm_rhs[..], but the latter
357         // two are no longer in use.
358         extern __shared__ float  sm_threadVirial[];
359         float                    mult = targetLength*lagrangeScaled;
360         sm_threadVirial[0*blockDim.x + threadIdx.x] = mult*rc.x*rc.x;
361         sm_threadVirial[1*blockDim.x + threadIdx.x] = mult*rc.x*rc.y;
362         sm_threadVirial[2*blockDim.x + threadIdx.x] = mult*rc.x*rc.z;
363         sm_threadVirial[3*blockDim.x + threadIdx.x] = mult*rc.y*rc.y;
364         sm_threadVirial[4*blockDim.x + threadIdx.x] = mult*rc.y*rc.z;
365         sm_threadVirial[5*blockDim.x + threadIdx.x] = mult*rc.z*rc.z;
367         __syncthreads();
369         // Reduce up to one virial per thread block. All blocks are divided by half, the first
370         // half of threads sums two virials. Then the first half is divided by two and the first
371         // half of it sums two values. This procedure is repeated until only one thread is left.
372         // Only works if the threads per blocks is a power of two (hence static_assert
373         // in the beginning of the kernel).
374         for (int divideBy = 2; divideBy <= static_cast<int>(blockDim.x); divideBy *= 2)
375         {
376             int dividedAt = blockDim.x/divideBy;
377             if (static_cast<int>(threadIdx.x) < dividedAt)
378             {
379                 for (int d = 0; d < 6; d++)
380                 {
381                     sm_threadVirial[d*blockDim.x + threadIdx.x] += sm_threadVirial[d*blockDim.x + (threadIdx.x + dividedAt)];
382                 }
383             }
384             // Syncronize if not within one warp
385             if (dividedAt > warpSize/2)
386             {
387                 __syncthreads();
388             }
389         }
390         // First 6 threads in the block add the results of 6 tensor components to the global memory address.
391         if (threadIdx.x < 6)
392         {
393             atomicAdd(&(gm_virialScaled[threadIdx.x]), sm_threadVirial[threadIdx.x*blockDim.x]);
394         }
395     }
397     return;
400 /*! \brief Select templated kernel.
402  * Returns pointer to a CUDA kernel based on provided booleans.
404  * \param[in] updateVelocities  If the velocities should be constrained.
405  * \param[in] computeVirial     If virial should be updated.
407  * \return                      Pointer to CUDA kernel
408  */
409 inline auto getLincsKernelPtr(const bool  updateVelocities,
410                               const bool  computeVirial)
413     auto kernelPtr = lincs_kernel<true, true>;
414     if (updateVelocities && computeVirial)
415     {
416         kernelPtr = lincs_kernel<true, true>;
417     }
418     else if (updateVelocities && !computeVirial)
419     {
420         kernelPtr = lincs_kernel<true, false>;
421     }
422     else if (!updateVelocities && computeVirial)
423     {
424         kernelPtr = lincs_kernel<false, true>;
425     }
426     else if (!updateVelocities && !computeVirial)
427     {
428         kernelPtr = lincs_kernel<false, false>;
429     }
430     return kernelPtr;
433 void LincsCuda::apply(const float3 *d_x,
434                       float3       *d_xp,
435                       const bool    updateVelocities,
436                       float3       *d_v,
437                       const real    invdt,
438                       const bool    computeVirial,
439                       tensor        virialScaled)
441     ensureNoPendingCudaError("In CUDA version of LINCS");
443     // Early exit if no constraints
444     if (kernelParams_.numConstraintsThreads == 0)
445     {
446         return;
447     }
449     if (computeVirial)
450     {
451         // Fill with zeros so the values can be reduced to it
452         // Only 6 values are needed because virial is symmetrical
453         clearDeviceBufferAsync(&kernelParams_.d_virialScaled, 0, 6, stream_);
454     }
456     auto               kernelPtr = getLincsKernelPtr(updateVelocities, computeVirial);
458     KernelLaunchConfig config;
459     config.blockSize[0]     = c_threadsPerBlock;
460     config.blockSize[1]     = 1;
461     config.blockSize[2]     = 1;
462     config.gridSize[0]      = (kernelParams_.numConstraintsThreads + c_threadsPerBlock - 1)/c_threadsPerBlock;
463     config.gridSize[1]      = 1;
464     config.gridSize[2]      = 1;
466     // Shared memory is used to store:
467     // -- Current coordinates (3 floats per thread)
468     // -- Right-hand-sides for matrix inversion (2 floats per thread)
469     // -- Virial tensor components (6 floats per thread)
470     // Since none of these three are needed simultaneously, they can be saved at the same shared memory address
471     // (i.e. correspondent arrays are intentionally overlapped in address space). Consequently, only
472     // max{3, 2, 6} = 6 floats per thread are needed in case virial is computed, or max{3, 2} = 3 if not.
473     if (computeVirial)
474     {
475         config.sharedMemorySize = c_threadsPerBlock*6*sizeof(float);
476     }
477     else
478     {
479         config.sharedMemorySize = c_threadsPerBlock*3*sizeof(float);
480     }
481     config.stream           = stream_;
483     // This is to satisfy prepareGpuKernelArguments(...)
484     // It there a better way?
485     const float3 * gm_x  = d_x;
486     float3       * gm_xp = d_xp;
487     float3       * gm_v  = d_v;
489     const auto     kernelArgs = prepareGpuKernelArguments(kernelPtr, config,
490                                                           &kernelParams_,
491                                                           &gm_x, &gm_xp,
492                                                           &gm_v, &invdt);
494     launchGpuKernel(kernelPtr, config, nullptr,
495                     "lincs_kernel<updateVelocities, computeVirial>", kernelArgs);
497     if (computeVirial)
498     {
499         // Copy LINCS virial data and add it to the common virial
500         copyFromDeviceBuffer(h_virialScaled_.data(), &kernelParams_.d_virialScaled,
501                              0, 6,
502                              stream_, GpuApiCallBehavior::Sync, nullptr);
504         // Mapping [XX, XY, XZ, YY, YZ, ZZ] internal format to a tensor object
505         virialScaled[XX][XX] += h_virialScaled_[0];
506         virialScaled[XX][YY] += h_virialScaled_[1];
507         virialScaled[XX][ZZ] += h_virialScaled_[2];
509         virialScaled[YY][XX] += h_virialScaled_[1];
510         virialScaled[YY][YY] += h_virialScaled_[3];
511         virialScaled[YY][ZZ] += h_virialScaled_[4];
513         virialScaled[ZZ][XX] += h_virialScaled_[2];
514         virialScaled[ZZ][YY] += h_virialScaled_[4];
515         virialScaled[ZZ][ZZ] += h_virialScaled_[5];
516     }
518     return;
521 LincsCuda::LincsCuda(int numIterations,
522                      int expansionOrder)
524     kernelParams_.numIterations              = numIterations;
525     kernelParams_.expansionOrder             = expansionOrder;
527     static_assert(sizeof(real) == sizeof(float),
528                   "Real numbers should be in single precision in GPU code.");
529     static_assert(c_threadsPerBlock > 0 && ((c_threadsPerBlock & (c_threadsPerBlock - 1)) == 0),
530                   "Number of threads per block should be a power of two in order for reduction to work.");
532     allocateDeviceBuffer(&kernelParams_.d_virialScaled, 6, nullptr);
533     h_virialScaled_.resize(6);
535     // The data arrays should be expanded/reallocated on first call of set() function.
536     numConstraintsThreadsAlloc_ = 0;
537     numAtomsAlloc_              = 0;
538     // Use default stream.
539     // TODO The stream should/can be assigned by the GPU schedule when the code will be integrated.
540     stream_ = nullptr;
544 LincsCuda::~LincsCuda()
546     freeDeviceBuffer(&kernelParams_.d_virialScaled);
548     if (numConstraintsThreadsAlloc_ > 0)
549     {
550         freeDeviceBuffer(&kernelParams_.d_constraints);
551         freeDeviceBuffer(&kernelParams_.d_constraintsTargetLengths);
553         freeDeviceBuffer(&kernelParams_.d_coupledConstraintsCounts);
554         freeDeviceBuffer(&kernelParams_.d_coupledConstraintsIndices);
555         freeDeviceBuffer(&kernelParams_.d_massFactors);
556         freeDeviceBuffer(&kernelParams_.d_matrixA);
557     }
558     if (numAtomsAlloc_ > 0)
559     {
560         freeDeviceBuffer(&kernelParams_.d_inverseMasses);
561     }
564 /*! \brief Helper function to go through constraints recursively.
566  *  For each constraint, counts the number of coupled constraints and stores the value in spaceNeeded array.
567  *  This information is used to split the array of constraints between thread blocks on a GPU so there is no
568  *  coupling between constraints from different thread blocks. After the 'spaceNeeded' array is filled, the
569  *  value spaceNeeded[c] should be equal to the number of constraints that are coupled to 'c' and located
570  *  after it in the constraints array.
572  * \param[in]     a                  Atom index.
573  * \param[in,out] spaceNeeded        Indicates if the constraint was already counted and stores
574  *                                   the number of constraints (i) connected to it and (ii) located
575  *                                   after it in memory. This array is filled by this recursive function.
576  *                                   For a set of coupled constraints, only for the first one in this list
577  *                                   the number of consecutive coupled constraints is needed: if there is
578  *                                   not enough space for this set of constraints in the thread block,
579  *                                   the group has to be moved to the next one.
580  * \param[in]     atomAdjacencyList  Stores information about connections between atoms.
581  */
582 inline int countCoupled(int a, std::vector<int> *spaceNeeded,
583                         std::vector<std::vector<std::tuple<int, int, int> > > *atomsAdjacencyList)
586     int c2, a2, sign;
587     int counted = 0;
588     for (unsigned i = 0; i < atomsAdjacencyList->at(a).size(); i++)
589     {
590         std::tie(a2, c2, sign) = atomsAdjacencyList->at(a).at(i);
591         if (spaceNeeded->at(c2) == -1)
592         {
593             spaceNeeded->at(c2) = 0; // To indicate we've been here
594             counted            += 1 + countCoupled(a2, spaceNeeded, atomsAdjacencyList);
595         }
596     }
597     return counted;
600 void LincsCuda::set(const t_idef    &idef,
601                     const t_mdatoms &md)
603     int                 numAtoms = md.nr;
604     // List of constrained atoms (CPU memory)
605     std::vector<int2>   constraintsHost;
606     // Equilibrium distances for the constraints (CPU)
607     std::vector<float>  constraintsTargetLengthsHost;
608     // Number of constraints, coupled with the current one (CPU)
609     std::vector<int>    coupledConstraintsCountsHost;
610     // List of coupled with the current one (CPU)
611     std::vector<int>    coupledConstraintsIndicesHost;
612     // Mass factors (CPU)
613     std::vector<float>  massFactorsHost;
615     // List of constrained atoms in local topology
616     t_iatom  *iatoms         = idef.il[F_CONSTR].iatoms;
617     const int stride         = NRAL(F_CONSTR) + 1;
618     const int numConstraints = idef.il[F_CONSTR].nr/stride;
620     // Early exit if no constraints
621     if (numConstraints == 0)
622     {
623         kernelParams_.numConstraintsThreads = 0;
624         return;
625     }
627     // Constructing adjacency list --- usefull intermediate structure
628     std::vector<std::vector<std::tuple<int, int, int> > > atomsAdjacencyList(numAtoms);
629     for (int c = 0; c < numConstraints; c++)
630     {
631         int a1     = iatoms[stride*c + 1];
632         int a2     = iatoms[stride*c + 2];
634         // Each constraint will be represented as a tuple, containing index of the second constrained atom,
635         // index of the constraint and a sign that indicates the order of atoms in which they are listed.
636         // Sign is needed to compute the mass factors.
637         atomsAdjacencyList.at(a1).push_back(std::make_tuple(a2, c, +1));
638         atomsAdjacencyList.at(a2).push_back(std::make_tuple(a1, c, -1));
639     }
641     // Compute, how many coupled constraints are in front of each constraint.
642     // Needed to introduce splits in data so that all coupled constraints will be computed in a single GPU block.
643     // The position 'c' of the vector spaceNeeded should have the number of constraints that are coupled to a constraint
644     // 'c' and are after 'c' in the vector. Only first index of the connected group of the constraints is needed later in the
645     // code, hence the spaceNeeded vector is also used to keep track if the constrain was already counted.
646     std::vector<int> spaceNeeded;
647     spaceNeeded.resize(numConstraints, -1);
648     std::fill(spaceNeeded.begin(), spaceNeeded.end(), -1);
649     for (int c = 0; c < numConstraints; c++)
650     {
651         int a1     = iatoms[stride*c + 1];
652         int a2     = iatoms[stride*c + 2];
653         if (spaceNeeded.at(c) == -1)
654         {
655             spaceNeeded.at(c) = countCoupled(a1, &spaceNeeded, &atomsAdjacencyList) +
656                 countCoupled(a2, &spaceNeeded, &atomsAdjacencyList);
657         }
658     }
660     // Map of splits in the constraints data. For each 'old' constraint index gives 'new' which
661     // takes into account the empty spaces which might be needed in the end of each thread block.
662     std::vector<int> splitMap;
663     splitMap.resize(numConstraints, -1);
664     int              currentMapIndex = 0;
665     for (int c = 0; c < numConstraints; c++)
666     {
667         // Check if coupled constraints all fit in one block
668         GMX_RELEASE_ASSERT(spaceNeeded.at(c) < c_threadsPerBlock, "Maximum number of coupled constraints exceedes the size of the CUDA thread block. "
669                            "Most likely, you are trying to use GPU version of LINCS with constraints on all-bonds, "
670                            "which is not supported. Try using H-bonds constraints only.");
671         if (currentMapIndex / c_threadsPerBlock != (currentMapIndex + spaceNeeded.at(c)) / c_threadsPerBlock)
672         {
673             currentMapIndex = ((currentMapIndex/c_threadsPerBlock) + 1) * c_threadsPerBlock;
674         }
675         splitMap.at(c) = currentMapIndex;
676         currentMapIndex++;
677     }
678     kernelParams_.numConstraintsThreads = currentMapIndex + c_threadsPerBlock - currentMapIndex % c_threadsPerBlock;
679     GMX_RELEASE_ASSERT(kernelParams_.numConstraintsThreads % c_threadsPerBlock == 0, "Number of threads should be a multiple of the block size");
681     // Initialize constraints and their target indexes taking into account the splits in the data arrays.
682     int2 pair;
683     pair.x = -1;
684     pair.y = -1;
685     constraintsHost.resize(kernelParams_.numConstraintsThreads, pair);
686     std::fill(constraintsHost.begin(), constraintsHost.end(), pair);
687     constraintsTargetLengthsHost.resize(kernelParams_.numConstraintsThreads, 0.0);
688     std::fill(constraintsTargetLengthsHost.begin(), constraintsTargetLengthsHost.end(), 0.0);
689     for (int c = 0; c < numConstraints; c++)
690     {
691         int  a1     = iatoms[stride*c + 1];
692         int  a2     = iatoms[stride*c + 2];
693         int  type   = iatoms[stride*c];
695         int2 pair;
696         pair.x = a1;
697         pair.y = a2;
698         constraintsHost.at(splitMap.at(c))              = pair;
699         constraintsTargetLengthsHost.at(splitMap.at(c)) = idef.iparams[type].constr.dA;
701     }
703     // The adjacency list of constraints (i.e. the list of coupled constraints for each constraint).
704     // We map a single thread to a single constraint, hence each thread 'c' will be using one element from
705     // coupledConstraintsCountsHost array, which is the number of constraints coupled to the constraint 'c'.
706     // The coupled constraints indexes are placed into the coupledConstraintsIndicesHost array. Latter is organized
707     // as a one-dimensional array to ensure good memory alignment. It is addressed as [c + i*numConstraintsThreads],
708     // where 'i' goes from zero to the number of constraints coupled to 'c'. 'numConstraintsThreads' is the width of
709     // the array --- a number, greater then total number of constraints, taking into account the splits in the
710     // constraints array due to the GPU block borders. This number can be adjusted to improve memory access pattern.
711     // Mass factors are saved in a similar data structure.
712     int              maxCoupledConstraints = 0;
713     for (int c = 0; c < numConstraints; c++)
714     {
715         int a1     = iatoms[stride*c + 1];
716         int a2     = iatoms[stride*c + 2];
718         // Constraint 'c' is counted twice, but it should be excluded altogether. Hence '-2'.
719         int nCoupedConstraints = atomsAdjacencyList.at(a1).size() + atomsAdjacencyList.at(a2).size() - 2;
721         if (nCoupedConstraints > maxCoupledConstraints)
722         {
723             maxCoupledConstraints = nCoupedConstraints;
724         }
725     }
727     coupledConstraintsCountsHost.resize(kernelParams_.numConstraintsThreads, 0);
728     coupledConstraintsIndicesHost.resize(maxCoupledConstraints*kernelParams_.numConstraintsThreads, -1);
729     massFactorsHost.resize(maxCoupledConstraints*kernelParams_.numConstraintsThreads, -1);
731     for (int c1 = 0; c1 < numConstraints; c1++)
732     {
733         coupledConstraintsCountsHost.at(splitMap.at(c1))  = 0;
734         int c1a1     = iatoms[stride*c1 + 1];
735         int c1a2     = iatoms[stride*c1 + 2];
736         int c2;
737         int c2a1;
738         int c2a2;
740         int sign;
742         // Constraints, coupled trough the first atom.
743         c2a1 = c1a1;
744         for (unsigned j = 0; j < atomsAdjacencyList.at(c1a1).size(); j++)
745         {
747             std::tie(c2a2, c2, sign) = atomsAdjacencyList.at(c1a1).at(j);
749             if (c1 != c2)
750             {
751                 int index = kernelParams_.numConstraintsThreads*coupledConstraintsCountsHost.at(splitMap.at(c1)) + splitMap.at(c1);
753                 coupledConstraintsIndicesHost.at(index) = splitMap.at(c2);
755                 int   center = c1a1;
757                 float sqrtmu1 = 1.0/sqrt(md.invmass[c1a1] + md.invmass[c1a2]);
758                 float sqrtmu2 = 1.0/sqrt(md.invmass[c2a1] + md.invmass[c2a2]);
760                 massFactorsHost.at(index) = -sign*md.invmass[center]*sqrtmu1*sqrtmu2;
762                 coupledConstraintsCountsHost.at(splitMap.at(c1))++;
764             }
765         }
767         // Constraints, coupled through the second atom.
768         c2a1 = c1a2;
769         for (unsigned j = 0; j < atomsAdjacencyList.at(c1a2).size(); j++)
770         {
772             std::tie(c2a2, c2, sign) = atomsAdjacencyList.at(c1a2).at(j);
774             if (c1 != c2)
775             {
776                 int index = kernelParams_.numConstraintsThreads*coupledConstraintsCountsHost.at(splitMap.at(c1)) + splitMap.at(c1);
778                 coupledConstraintsIndicesHost.at(index) = splitMap.at(c2);
780                 int   center = c1a2;
782                 float sqrtmu1 = 1.0/sqrt(md.invmass[c1a1] + md.invmass[c1a2]);
783                 float sqrtmu2 = 1.0/sqrt(md.invmass[c2a1] + md.invmass[c2a2]);
785                 massFactorsHost.at(index) = sign*md.invmass[center]*sqrtmu1*sqrtmu2;
787                 coupledConstraintsCountsHost.at(splitMap.at(c1))++;
789             }
790         }
791     }
793     // (Re)allocate the memory, if the number of constraints has increased.
794     if (kernelParams_.numConstraintsThreads > numConstraintsThreadsAlloc_)
795     {
796         // Free memory if it was allocated before (i.e. if not the first time here).
797         if (numConstraintsThreadsAlloc_ > 0)
798         {
799             freeDeviceBuffer(&kernelParams_.d_constraints);
800             freeDeviceBuffer(&kernelParams_.d_constraintsTargetLengths);
802             freeDeviceBuffer(&kernelParams_.d_coupledConstraintsCounts);
803             freeDeviceBuffer(&kernelParams_.d_coupledConstraintsIndices);
804             freeDeviceBuffer(&kernelParams_.d_massFactors);
805             freeDeviceBuffer(&kernelParams_.d_matrixA);
807         }
809         numConstraintsThreadsAlloc_ = kernelParams_.numConstraintsThreads;
811         allocateDeviceBuffer(&kernelParams_.d_constraints, kernelParams_.numConstraintsThreads, nullptr);
812         allocateDeviceBuffer(&kernelParams_.d_constraintsTargetLengths, kernelParams_.numConstraintsThreads, nullptr);
814         allocateDeviceBuffer(&kernelParams_.d_coupledConstraintsCounts, kernelParams_.numConstraintsThreads, nullptr);
815         allocateDeviceBuffer(&kernelParams_.d_coupledConstraintsIndices, maxCoupledConstraints*kernelParams_.numConstraintsThreads, nullptr);
816         allocateDeviceBuffer(&kernelParams_.d_massFactors, maxCoupledConstraints*kernelParams_.numConstraintsThreads, nullptr);
817         allocateDeviceBuffer(&kernelParams_.d_matrixA, maxCoupledConstraints*kernelParams_.numConstraintsThreads, nullptr);
819     }
821     // (Re)allocate the memory, if the number of atoms has increased.
822     if (numAtoms > numAtomsAlloc_)
823     {
824         if (numAtomsAlloc_ > 0)
825         {
826             freeDeviceBuffer(&kernelParams_.d_inverseMasses);
827         }
828         numAtomsAlloc_ = numAtoms;
829         allocateDeviceBuffer(&kernelParams_.d_inverseMasses, numAtoms, nullptr);
830     }
832     // Copy data to GPU.
833     copyToDeviceBuffer(&kernelParams_.d_constraints, constraintsHost.data(),
834                        0, kernelParams_.numConstraintsThreads,
835                        stream_, GpuApiCallBehavior::Sync, nullptr);
836     copyToDeviceBuffer(&kernelParams_.d_constraintsTargetLengths, constraintsTargetLengthsHost.data(),
837                        0, kernelParams_.numConstraintsThreads,
838                        stream_, GpuApiCallBehavior::Sync, nullptr);
839     copyToDeviceBuffer(&kernelParams_.d_coupledConstraintsCounts, coupledConstraintsCountsHost.data(),
840                        0, kernelParams_.numConstraintsThreads,
841                        stream_, GpuApiCallBehavior::Sync, nullptr);
842     copyToDeviceBuffer(&kernelParams_.d_coupledConstraintsIndices, coupledConstraintsIndicesHost.data(),
843                        0, maxCoupledConstraints*kernelParams_.numConstraintsThreads,
844                        stream_, GpuApiCallBehavior::Sync, nullptr);
845     copyToDeviceBuffer(&kernelParams_.d_massFactors, massFactorsHost.data(),
846                        0, maxCoupledConstraints*kernelParams_.numConstraintsThreads,
847                        stream_, GpuApiCallBehavior::Sync, nullptr);
849     GMX_RELEASE_ASSERT(md.invmass != nullptr, "Masses of attoms should be specified.\n");
850     copyToDeviceBuffer(&kernelParams_.d_inverseMasses, md.invmass,
851                        0, numAtoms,
852                        stream_, GpuApiCallBehavior::Sync, nullptr);
856 void LincsCuda::setPbc(const t_pbc *pbc)
858     setPbcAiuc(pbc->ndim_ePBC, pbc->box, &kernelParams_.pbcAiuc);
861 } // namespace gmx