2 * This file is part of the GROMACS molecular simulation package.
4 * Copyright (c) 2019,2020,2021, 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 LINCS using CUDA
39 * This file contains implementation of LINCS constraints algorithm
40 * using CUDA, including class initialization, data-structures management
43 * \author Artem Zhmurov <zhmurov@gmail.com>
44 * \author Alan Gray <alang@nvidia.com>
46 * \ingroup module_mdlib
50 #include "lincs_gpu.cuh"
59 #include "gromacs/gpu_utils/cuda_arch_utils.cuh"
60 #include "gromacs/gpu_utils/cudautils.cuh"
61 #include "gromacs/gpu_utils/devicebuffer.cuh"
62 #include "gromacs/gpu_utils/gputraits.cuh"
63 #include "gromacs/gpu_utils/vectype_ops.cuh"
64 #include "gromacs/math/functions.h"
65 #include "gromacs/math/vec.h"
66 #include "gromacs/mdlib/constr.h"
67 #include "gromacs/pbcutil/pbc.h"
68 #include "gromacs/pbcutil/pbc_aiuc_cuda.cuh"
69 #include "gromacs/topology/forcefieldparameters.h"
70 #include "gromacs/topology/ifunc.h"
71 #include "gromacs/topology/topology.h"
76 //! Number of CUDA threads in a block
77 constexpr static int c_threadsPerBlock = 256;
78 //! Maximum number of threads in a block (for __launch_bounds__)
79 constexpr static int c_maxThreadsPerBlock = c_threadsPerBlock;
81 /*! \brief Main kernel for LINCS constraints.
83 * See Hess et al., J. Comput. Chem. 18: 1463-1472 (1997) for the description of the algorithm.
85 * In CUDA version, one thread is responsible for all computations for one constraint. The blocks are
86 * filled in a way that no constraint is coupled to the constraint from the next block. This is achieved
87 * by moving active threads to the next block, if the correspondent group of coupled constraints is to big
88 * to fit the current thread block. This may leave some 'dummy' threads in the end of the thread block, i.e.
89 * threads that are not required to do actual work. Since constraints from different blocks are not coupled,
90 * there is no need to synchronize across the device. However, extensive communication in a thread block
93 * \todo Reduce synchronization overhead. Some ideas are:
94 * 1. Consider going to warp-level synchronization for the coupled constraints.
95 * 2. Move more data to local/shared memory and try to get rid of atomic operations (at least on
97 * 3. Use analytical solution for matrix A inversion.
98 * 4. Introduce mapping of thread id to both single constraint and single atom, thus designating
99 * Nth threads to deal with Nat <= Nth coupled atoms and Nc <= Nth coupled constraints.
100 * See Issue #2885 for details (https://gitlab.com/gromacs/gromacs/-/issues/2885)
101 * \todo The use of __restrict__ for gm_xp and gm_v causes failure, probably because of the atomic
102 operations. Investigate this issue further.
104 * \param[in,out] kernelParams All parameters and pointers for the kernel condensed in single struct.
105 * \param[in] invdt Inverse timestep (needed to update velocities).
107 template<bool updateVelocities, bool computeVirial>
108 __launch_bounds__(c_maxThreadsPerBlock) __global__
109 void lincs_kernel(LincsGpuKernelParameters kernelParams,
110 const float3* __restrict__ gm_x,
115 const PbcAiuc pbcAiuc = kernelParams.pbcAiuc;
116 const int numConstraintsThreads = kernelParams.numConstraintsThreads;
117 const int numIterations = kernelParams.numIterations;
118 const int expansionOrder = kernelParams.expansionOrder;
119 const int2* __restrict__ gm_constraints = kernelParams.d_constraints;
120 const float* __restrict__ gm_constraintsTargetLengths = kernelParams.d_constraintsTargetLengths;
121 const int* __restrict__ gm_coupledConstraintsCounts = kernelParams.d_coupledConstraintsCounts;
122 const int* __restrict__ gm_coupledConstraintsIdxes = kernelParams.d_coupledConstraintsIndices;
123 const float* __restrict__ gm_massFactors = kernelParams.d_massFactors;
124 float* __restrict__ gm_matrixA = kernelParams.d_matrixA;
125 const float* __restrict__ gm_inverseMasses = kernelParams.d_inverseMasses;
126 float* __restrict__ gm_virialScaled = kernelParams.d_virialScaled;
128 int threadIndex = blockIdx.x * blockDim.x + threadIdx.x;
130 // numConstraintsThreads should be a integer multiple of blockSize (numConstraintsThreads = numBlocks*blockSize).
131 // This is to ensure proper synchronizations and reduction. All array are padded to the required size.
132 assert(threadIndex < numConstraintsThreads);
134 // Vectors connecting constrained atoms before algorithm was applied.
135 // Needed to construct constrain matrix A
136 extern __shared__ float3 sm_r[];
138 int2 pair = gm_constraints[threadIndex];
142 // Mass-scaled Lagrange multiplier
143 float lagrangeScaled = 0.0f;
148 float sqrtReducedMass;
154 // i == -1 indicates dummy constraint at the end of the thread block.
155 bool isDummyThread = (i == -1);
157 // Everything computed for these dummies will be equal to zero
163 sqrtReducedMass = 0.0f;
165 xi = make_float3(0.0f, 0.0f, 0.0f);
166 xj = make_float3(0.0f, 0.0f, 0.0f);
167 rc = make_float3(0.0f, 0.0f, 0.0f);
172 targetLength = gm_constraintsTargetLengths[threadIndex];
173 inverseMassi = gm_inverseMasses[i];
174 inverseMassj = gm_inverseMasses[j];
175 sqrtReducedMass = rsqrt(inverseMassi + inverseMassj);
180 float3 dx = pbcDxAiuc(pbcAiuc, xi, xj);
182 float rlen = rsqrtf(dx.x * dx.x + dx.y * dx.y + dx.z * dx.z);
186 sm_r[threadIdx.x] = rc;
187 // Make sure that all r's are saved into shared memory
188 // before they are accessed in the loop below
192 * Constructing LINCS matrix (A)
195 // Only non-zero values are saved (for coupled constraints)
196 int coupledConstraintsCount = gm_coupledConstraintsCounts[threadIndex];
197 for (int n = 0; n < coupledConstraintsCount; n++)
199 int index = n * numConstraintsThreads + threadIndex;
200 int c1 = gm_coupledConstraintsIdxes[index];
202 float3 rc1 = sm_r[c1 - blockIdx.x * blockDim.x];
203 gm_matrixA[index] = gm_massFactors[index] * (rc.x * rc1.x + rc.y * rc1.y + rc.z * rc1.z);
206 // Skipping in dummy threads
213 float3 dx = pbcDxAiuc(pbcAiuc, xi, xj);
215 float sol = sqrtReducedMass * ((rc.x * dx.x + rc.y * dx.y + rc.z * dx.z) - targetLength);
218 * Inverse matrix using a set of expansionOrder matrix multiplications
221 // This will use the same memory space as sm_r, which is no longer needed.
222 extern __shared__ float sm_rhs[];
223 // Save current right-hand-side vector in the shared memory
224 sm_rhs[threadIdx.x] = sol;
226 for (int rec = 0; rec < expansionOrder; rec++)
228 // Making sure that all sm_rhs are saved before they are accessed in a loop below
232 for (int n = 0; n < coupledConstraintsCount; n++)
234 int index = n * numConstraintsThreads + threadIndex;
235 int c1 = gm_coupledConstraintsIdxes[index];
236 // Convolute current right-hand-side with A
237 // Different, non overlapping parts of sm_rhs[..] are read during odd and even iterations
238 mvb = mvb + gm_matrixA[index] * sm_rhs[c1 - blockIdx.x * blockDim.x + blockDim.x * (rec % 2)];
240 // 'Switch' rhs vectors, save current result
241 // These values will be accessed in the loop above during the next iteration.
242 sm_rhs[threadIdx.x + blockDim.x * ((rec + 1) % 2)] = mvb;
246 // Current mass-scaled Lagrange multipliers
247 lagrangeScaled = sqrtReducedMass * sol;
249 // Save updated coordinates before correction for the rotational lengthening
250 float3 tmp = rc * lagrangeScaled;
252 // Writing for all but dummy constraints
255 atomicAdd(&gm_xp[i], -tmp * inverseMassi);
256 atomicAdd(&gm_xp[j], tmp * inverseMassj);
260 * Correction for centripetal effects
262 for (int iter = 0; iter < numIterations; iter++)
264 // Make sure that all xp's are saved: atomic operation calls before are
265 // communicating current xp[..] values across thread block.
274 float3 dx = pbcDxAiuc(pbcAiuc, xi, xj);
276 float len2 = targetLength * targetLength;
277 float dlen2 = 2.0f * len2 - norm2(dx);
279 // TODO A little bit more effective but slightly less readable version of the below would be:
280 // float proj = sqrtReducedMass*(targetLength - (dlen2 > 0.0f ? 1.0f : 0.0f)*dlen2*rsqrt(dlen2));
284 proj = sqrtReducedMass * (targetLength - dlen2 * rsqrt(dlen2));
288 proj = sqrtReducedMass * targetLength;
291 sm_rhs[threadIdx.x] = proj;
295 * Same matrix inversion as above is used for updated data
297 for (int rec = 0; rec < expansionOrder; rec++)
299 // Make sure that all elements of rhs are saved into shared memory
303 for (int n = 0; n < coupledConstraintsCount; n++)
305 int index = n * numConstraintsThreads + threadIndex;
306 int c1 = gm_coupledConstraintsIdxes[index];
308 mvb = mvb + gm_matrixA[index] * sm_rhs[c1 - blockIdx.x * blockDim.x + blockDim.x * (rec % 2)];
310 sm_rhs[threadIdx.x + blockDim.x * ((rec + 1) % 2)] = mvb;
314 // Add corrections to Lagrange multipliers
315 float sqrtmu_sol = sqrtReducedMass * sol;
316 lagrangeScaled += sqrtmu_sol;
318 // Save updated coordinates for the next iteration
319 // Dummy constraints are skipped
322 float3 tmp = rc * sqrtmu_sol;
323 atomicAdd(&gm_xp[i], -tmp * inverseMassi);
324 atomicAdd(&gm_xp[j], tmp * inverseMassj);
328 // Updating particle velocities for all but dummy threads
329 if (updateVelocities && !isDummyThread)
331 float3 tmp = rc * invdt * lagrangeScaled;
332 atomicAdd(&gm_v[i], -tmp * inverseMassi);
333 atomicAdd(&gm_v[j], tmp * inverseMassj);
339 // Virial is computed from Lagrange multiplier (lagrangeScaled), target constrain length
340 // (targetLength) and the normalized vector connecting constrained atoms before
341 // the algorithm was applied (rc). The evaluation of virial in each thread is
342 // followed by basic reduction for the values inside single thread block.
343 // Then, the values are reduced across grid by atomicAdd(...).
345 // TODO Shuffle reduction.
346 // TODO Should be unified and/or done once when virial is actually needed.
347 // TODO Recursive version that removes atomicAdd(...)'s entirely is needed. Ideally,
348 // one that works for any datatype.
350 // Save virial for each thread into the shared memory. Tensor is symmetrical, hence only
351 // 6 values are saved. Dummy threads will have zeroes in their virial: targetLength,
352 // lagrangeScaled and rc are all set to zero for them in the beginning of the kernel.
353 // The sm_threadVirial[..] will overlap with the sm_r[..] and sm_rhs[..], but the latter
354 // two are no longer in use.
355 extern __shared__ float sm_threadVirial[];
356 float mult = targetLength * lagrangeScaled;
357 sm_threadVirial[0 * blockDim.x + threadIdx.x] = mult * rc.x * rc.x;
358 sm_threadVirial[1 * blockDim.x + threadIdx.x] = mult * rc.x * rc.y;
359 sm_threadVirial[2 * blockDim.x + threadIdx.x] = mult * rc.x * rc.z;
360 sm_threadVirial[3 * blockDim.x + threadIdx.x] = mult * rc.y * rc.y;
361 sm_threadVirial[4 * blockDim.x + threadIdx.x] = mult * rc.y * rc.z;
362 sm_threadVirial[5 * blockDim.x + threadIdx.x] = mult * rc.z * rc.z;
366 // Reduce up to one virial per thread block. All blocks are divided by half, the first
367 // half of threads sums two virials. Then the first half is divided by two and the first
368 // half of it sums two values. This procedure is repeated until only one thread is left.
369 // Only works if the threads per blocks is a power of two (hence static_assert
370 // in the beginning of the kernel).
371 for (int divideBy = 2; divideBy <= static_cast<int>(blockDim.x); divideBy *= 2)
373 int dividedAt = blockDim.x / divideBy;
374 if (static_cast<int>(threadIdx.x) < dividedAt)
376 for (int d = 0; d < 6; d++)
378 sm_threadVirial[d * blockDim.x + threadIdx.x] +=
379 sm_threadVirial[d * blockDim.x + (threadIdx.x + dividedAt)];
382 // Syncronize if not within one warp
383 if (dividedAt > warpSize / 2)
388 // First 6 threads in the block add the results of 6 tensor components to the global memory address.
391 atomicAdd(&(gm_virialScaled[threadIdx.x]), sm_threadVirial[threadIdx.x * blockDim.x]);
398 /*! \brief Select templated kernel.
400 * Returns pointer to a CUDA kernel based on provided booleans.
402 * \param[in] updateVelocities If the velocities should be constrained.
403 * \param[in] computeVirial If virial should be updated.
405 * \return Pointer to CUDA kernel
407 inline auto getLincsKernelPtr(const bool updateVelocities, const bool computeVirial)
410 auto kernelPtr = lincs_kernel<true, true>;
411 if (updateVelocities && computeVirial)
413 kernelPtr = lincs_kernel<true, true>;
415 else if (updateVelocities && !computeVirial)
417 kernelPtr = lincs_kernel<true, false>;
419 else if (!updateVelocities && computeVirial)
421 kernelPtr = lincs_kernel<false, true>;
423 else if (!updateVelocities && !computeVirial)
425 kernelPtr = lincs_kernel<false, false>;
430 void LincsGpu::apply(const float3* d_x,
432 const bool updateVelocities,
435 const bool computeVirial,
437 const PbcAiuc pbcAiuc)
439 ensureNoPendingDeviceError("In CUDA version of LINCS");
441 // Early exit if no constraints
442 if (kernelParams_.numConstraintsThreads == 0)
449 // Fill with zeros so the values can be reduced to it
450 // Only 6 values are needed because virial is symmetrical
451 clearDeviceBufferAsync(&kernelParams_.d_virialScaled, 0, 6, deviceStream_);
454 auto kernelPtr = getLincsKernelPtr(updateVelocities, computeVirial);
456 KernelLaunchConfig config;
457 config.blockSize[0] = c_threadsPerBlock;
458 config.blockSize[1] = 1;
459 config.blockSize[2] = 1;
460 config.gridSize[0] = (kernelParams_.numConstraintsThreads + c_threadsPerBlock - 1) / c_threadsPerBlock;
461 config.gridSize[1] = 1;
462 config.gridSize[2] = 1;
464 // Shared memory is used to store:
465 // -- Current coordinates (3 floats per thread)
466 // -- Right-hand-sides for matrix inversion (2 floats per thread)
467 // -- Virial tensor components (6 floats per thread)
468 // Since none of these three are needed simultaneously, they can be saved at the same shared memory address
469 // (i.e. correspondent arrays are intentionally overlapped in address space). Consequently, only
470 // max{3, 2, 6} = 6 floats per thread are needed in case virial is computed, or max{3, 2} = 3 if not.
473 config.sharedMemorySize = c_threadsPerBlock * 6 * sizeof(float);
477 config.sharedMemorySize = c_threadsPerBlock * 3 * sizeof(float);
480 kernelParams_.pbcAiuc = pbcAiuc;
482 const auto kernelArgs =
483 prepareGpuKernelArguments(kernelPtr, config, &kernelParams_, &d_x, &d_xp, &d_v, &invdt);
485 launchGpuKernel(kernelPtr,
489 "lincs_kernel<updateVelocities, computeVirial>",
494 // Copy LINCS virial data and add it to the common virial
495 copyFromDeviceBuffer(h_virialScaled_.data(),
496 &kernelParams_.d_virialScaled,
500 GpuApiCallBehavior::Sync,
503 // Mapping [XX, XY, XZ, YY, YZ, ZZ] internal format to a tensor object
504 virialScaled[XX][XX] += h_virialScaled_[0];
505 virialScaled[XX][YY] += h_virialScaled_[1];
506 virialScaled[XX][ZZ] += h_virialScaled_[2];
508 virialScaled[YY][XX] += h_virialScaled_[1];
509 virialScaled[YY][YY] += h_virialScaled_[3];
510 virialScaled[YY][ZZ] += h_virialScaled_[4];
512 virialScaled[ZZ][XX] += h_virialScaled_[2];
513 virialScaled[ZZ][YY] += h_virialScaled_[4];
514 virialScaled[ZZ][ZZ] += h_virialScaled_[5];
520 LincsGpu::LincsGpu(int numIterations,
522 const DeviceContext& deviceContext,
523 const DeviceStream& deviceStream) :
524 deviceContext_(deviceContext),
525 deviceStream_(deviceStream)
527 kernelParams_.numIterations = numIterations;
528 kernelParams_.expansionOrder = expansionOrder;
530 static_assert(sizeof(real) == sizeof(float),
531 "Real numbers should be in single precision in GPU code.");
533 gmx::isPowerOfTwo(c_threadsPerBlock),
534 "Number of threads per block should be a power of two in order for reduction to work.");
536 allocateDeviceBuffer(&kernelParams_.d_virialScaled, 6, deviceContext_);
537 h_virialScaled_.resize(6);
539 // The data arrays should be expanded/reallocated on first call of set() function.
540 numConstraintsThreadsAlloc_ = 0;
544 LincsGpu::~LincsGpu()
546 freeDeviceBuffer(&kernelParams_.d_virialScaled);
548 if (numConstraintsThreadsAlloc_ > 0)
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);
558 if (numAtomsAlloc_ > 0)
560 freeDeviceBuffer(&kernelParams_.d_inverseMasses);
564 //! Helper type for discovering coupled constraints
565 struct AtomsAdjacencyListElement
567 AtomsAdjacencyListElement(const int indexOfSecondConstrainedAtom,
568 const int indexOfConstraint,
569 const int signFactor) :
570 indexOfSecondConstrainedAtom_(indexOfSecondConstrainedAtom),
571 indexOfConstraint_(indexOfConstraint),
572 signFactor_(signFactor)
575 //! The index of the other atom constrained to this atom.
576 int indexOfSecondConstrainedAtom_;
577 //! The index of this constraint in the container of constraints.
578 int indexOfConstraint_;
579 /*! \brief A multiplicative factor that indicates the relative
580 * order of the atoms in the atom list.
582 * Used for computing the mass factor of this constraint
583 * relative to any coupled constraints. */
586 //! Constructs and returns an atom constraint adjacency list
587 static std::vector<std::vector<AtomsAdjacencyListElement>>
588 constructAtomsAdjacencyList(const int numAtoms, ArrayRef<const int> iatoms)
590 const int stride = 1 + NRAL(F_CONSTR);
591 const int numConstraints = iatoms.ssize() / stride;
592 std::vector<std::vector<AtomsAdjacencyListElement>> atomsAdjacencyList(numAtoms);
593 for (int c = 0; c < numConstraints; c++)
595 int a1 = iatoms[stride * c + 1];
596 int a2 = iatoms[stride * c + 2];
598 // Each constraint will be represented as a tuple, containing index of the second
599 // constrained atom, index of the constraint and a sign that indicates the order of atoms in
600 // which they are listed. Sign is needed to compute the mass factors.
601 atomsAdjacencyList[a1].emplace_back(a2, c, +1);
602 atomsAdjacencyList[a2].emplace_back(a1, c, -1);
605 return atomsAdjacencyList;
608 /*! \brief Helper function to go through constraints recursively.
610 * For each constraint, counts the number of coupled constraints and stores the value in \p numCoupledConstraints array.
611 * This information is used to split the array of constraints between thread blocks on a GPU so there is no
612 * coupling between constraints from different thread blocks. After the \p numCoupledConstraints array is filled, the
613 * value \p numCoupledConstraints[c] should be equal to the number of constraints that are coupled to \p c and located
614 * after it in the constraints array.
616 * \param[in] a Atom index.
617 * \param[in,out] numCoupledConstraints Indicates if the constraint was already counted and stores
618 * the number of constraints (i) connected to it and (ii) located
619 * after it in memory. This array is filled by this recursive function.
620 * For a set of coupled constraints, only for the first one in this list
621 * the number of consecutive coupled constraints is needed: if there is
622 * not enough space for this set of constraints in the thread block,
623 * the group has to be moved to the next one.
624 * \param[in] atomsAdjacencyList Stores information about connections between atoms.
626 inline int countCoupled(int a,
627 ArrayRef<int> numCoupledConstraints,
628 ArrayRef<const std::vector<AtomsAdjacencyListElement>> atomsAdjacencyList)
632 for (const auto& adjacentAtom : atomsAdjacencyList[a])
634 const int c2 = adjacentAtom.indexOfConstraint_;
635 if (numCoupledConstraints[c2] == -1)
637 numCoupledConstraints[c2] = 0; // To indicate we've been here
639 + countCoupled(adjacentAtom.indexOfSecondConstrainedAtom_,
640 numCoupledConstraints,
647 /*! \brief Add constraint to \p splitMap with all constraints coupled to it.
649 * Adds the constraint \p c from the constrain list \p iatoms to the map \p splitMap
650 * if it was not yet added. Then goes through all the constraints coupled to \p c
651 * and calls itself recursively. This ensures that all the coupled constraints will
652 * be added to neighboring locations in the final data structures on the device,
653 * hence mapping all coupled constraints to the same thread block. A value of -1 in
654 * the \p splitMap is used to flag that constraint was not yet added to the \p splitMap.
656 * \param[in] iatoms The list of constraints.
657 * \param[in] stride Number of elements per constraint in \p iatoms.
658 * \param[in] atomsAdjacencyList Information about connections between atoms.
659 * \param[our] splitMap Map of sequential constraint indexes to indexes to be on the device
660 * \param[in] c Sequential index for constraint to consider adding.
661 * \param[in,out] currentMapIndex The rolling index for the constraints mapping.
663 inline void addWithCoupled(ArrayRef<const int> iatoms,
665 ArrayRef<const std::vector<AtomsAdjacencyListElement>> atomsAdjacencyList,
666 ArrayRef<int> splitMap,
668 int* currentMapIndex)
670 if (splitMap[c] == -1)
672 splitMap[c] = *currentMapIndex;
673 (*currentMapIndex)++;
675 // Constraints, coupled through both atoms.
676 for (int atomIndexInConstraint = 0; atomIndexInConstraint < 2; atomIndexInConstraint++)
678 const int a1 = iatoms[stride * c + 1 + atomIndexInConstraint];
679 for (const auto& adjacentAtom : atomsAdjacencyList[a1])
681 const int c2 = adjacentAtom.indexOfConstraint_;
684 addWithCoupled(iatoms, stride, atomsAdjacencyList, splitMap, c2, currentMapIndex);
691 /*! \brief Computes and returns how many constraints are coupled to each constraint
693 * Needed to introduce splits in data so that all coupled constraints will be computed in a
694 * single GPU block. The position \p c of the vector \p numCoupledConstraints should have the number
695 * of constraints that are coupled to a constraint \p c and are after \p c in the vector. Only
696 * first index of the connected group of the constraints is needed later in the code, hence the
697 * numCoupledConstraints vector is also used to keep track if the constrain was already counted.
699 static std::vector<int> countNumCoupledConstraints(ArrayRef<const int> iatoms,
700 ArrayRef<const std::vector<AtomsAdjacencyListElement>> atomsAdjacencyList)
702 const int stride = 1 + NRAL(F_CONSTR);
703 const int numConstraints = iatoms.ssize() / stride;
704 std::vector<int> numCoupledConstraints(numConstraints, -1);
705 for (int c = 0; c < numConstraints; c++)
707 const int a1 = iatoms[stride * c + 1];
708 const int a2 = iatoms[stride * c + 2];
709 if (numCoupledConstraints[c] == -1)
711 numCoupledConstraints[c] = countCoupled(a1, numCoupledConstraints, atomsAdjacencyList)
712 + countCoupled(a2, numCoupledConstraints, atomsAdjacencyList);
716 return numCoupledConstraints;
719 bool LincsGpu::isNumCoupledConstraintsSupported(const gmx_mtop_t& mtop)
721 for (const gmx_moltype_t& molType : mtop.moltype)
723 ArrayRef<const int> iatoms = molType.ilist[F_CONSTR].iatoms;
724 const auto atomsAdjacencyList = constructAtomsAdjacencyList(molType.atoms.nr, iatoms);
725 // Compute, how many constraints are coupled to each constraint
726 const auto numCoupledConstraints = countNumCoupledConstraints(iatoms, atomsAdjacencyList);
727 for (const int numCoupled : numCoupledConstraints)
729 if (numCoupled > c_threadsPerBlock)
739 void LincsGpu::set(const InteractionDefinitions& idef, const int numAtoms, const real* invmass)
741 // List of constrained atoms (CPU memory)
742 std::vector<int2> constraintsHost;
743 // Equilibrium distances for the constraints (CPU)
744 std::vector<float> constraintsTargetLengthsHost;
745 // Number of constraints, coupled with the current one (CPU)
746 std::vector<int> coupledConstraintsCountsHost;
747 // List of coupled with the current one (CPU)
748 std::vector<int> coupledConstraintsIndicesHost;
749 // Mass factors (CPU)
750 std::vector<float> massFactorsHost;
752 // List of constrained atoms in local topology
753 ArrayRef<const int> iatoms = idef.il[F_CONSTR].iatoms;
754 const int stride = NRAL(F_CONSTR) + 1;
755 const int numConstraints = idef.il[F_CONSTR].size() / stride;
757 // Early exit if no constraints
758 if (numConstraints == 0)
760 kernelParams_.numConstraintsThreads = 0;
764 // Construct the adjacency list, a useful intermediate structure
765 const auto atomsAdjacencyList = constructAtomsAdjacencyList(numAtoms, iatoms);
767 // Compute, how many constraints are coupled to each constraint
768 const auto numCoupledConstraints = countNumCoupledConstraints(iatoms, atomsAdjacencyList);
770 // Map of splits in the constraints data. For each 'old' constraint index gives 'new' which
771 // takes into account the empty spaces which might be needed in the end of each thread block.
772 std::vector<int> splitMap(numConstraints, -1);
773 int currentMapIndex = 0;
774 for (int c = 0; c < numConstraints; c++)
776 // Check if coupled constraints all fit in one block
777 if (numCoupledConstraints.at(c) > c_threadsPerBlock)
780 "Maximum number of coupled constraints (%d) exceeds the size of the CUDA "
781 "thread block (%d). Most likely, you are trying to use the GPU version of "
782 "LINCS with constraints on all-bonds, which is not supported for large "
783 "molecules. When compatible with the force field and integration settings, "
784 "using constraints on H-bonds only.",
785 numCoupledConstraints.at(c),
788 if (currentMapIndex / c_threadsPerBlock
789 != (currentMapIndex + numCoupledConstraints.at(c)) / c_threadsPerBlock)
791 currentMapIndex = ((currentMapIndex / c_threadsPerBlock) + 1) * c_threadsPerBlock;
793 addWithCoupled(iatoms, stride, atomsAdjacencyList, splitMap, c, ¤tMapIndex);
796 kernelParams_.numConstraintsThreads =
797 currentMapIndex + c_threadsPerBlock - currentMapIndex % c_threadsPerBlock;
798 GMX_RELEASE_ASSERT(kernelParams_.numConstraintsThreads % c_threadsPerBlock == 0,
799 "Number of threads should be a multiple of the block size");
801 // Initialize constraints and their target indexes taking into account the splits in the data arrays.
805 constraintsHost.resize(kernelParams_.numConstraintsThreads, pair);
806 std::fill(constraintsHost.begin(), constraintsHost.end(), pair);
807 constraintsTargetLengthsHost.resize(kernelParams_.numConstraintsThreads, 0.0);
808 std::fill(constraintsTargetLengthsHost.begin(), constraintsTargetLengthsHost.end(), 0.0);
809 for (int c = 0; c < numConstraints; c++)
811 int a1 = iatoms[stride * c + 1];
812 int a2 = iatoms[stride * c + 2];
813 int type = iatoms[stride * c];
818 constraintsHost.at(splitMap.at(c)) = pair;
819 constraintsTargetLengthsHost.at(splitMap.at(c)) = idef.iparams[type].constr.dA;
822 // The adjacency list of constraints (i.e. the list of coupled constraints for each constraint).
823 // We map a single thread to a single constraint, hence each thread 'c' will be using one
824 // element from coupledConstraintsCountsHost array, which is the number of constraints coupled
825 // to the constraint 'c'. The coupled constraints indexes are placed into the
826 // coupledConstraintsIndicesHost array. Latter is organized as a one-dimensional array to ensure
827 // good memory alignment. It is addressed as [c + i*numConstraintsThreads], where 'i' goes from
828 // zero to the number of constraints coupled to 'c'. 'numConstraintsThreads' is the width of the
829 // array --- a number, greater then total number of constraints, taking into account the splits
830 // in the constraints array due to the GPU block borders. This number can be adjusted to improve
831 // memory access pattern. Mass factors are saved in a similar data structure.
832 int maxCoupledConstraints = 0;
833 bool maxCoupledConstraintsHasIncreased = false;
834 for (int c = 0; c < numConstraints; c++)
836 int a1 = iatoms[stride * c + 1];
837 int a2 = iatoms[stride * c + 2];
839 // Constraint 'c' is counted twice, but it should be excluded altogether. Hence '-2'.
840 int nCoupledConstraints = atomsAdjacencyList.at(a1).size() + atomsAdjacencyList.at(a2).size() - 2;
842 if (nCoupledConstraints > maxCoupledConstraints)
844 maxCoupledConstraints = nCoupledConstraints;
845 maxCoupledConstraintsHasIncreased = true;
849 coupledConstraintsCountsHost.resize(kernelParams_.numConstraintsThreads, 0);
850 coupledConstraintsIndicesHost.resize(maxCoupledConstraints * kernelParams_.numConstraintsThreads, -1);
851 massFactorsHost.resize(maxCoupledConstraints * kernelParams_.numConstraintsThreads, -1);
853 for (int c1 = 0; c1 < numConstraints; c1++)
855 coupledConstraintsCountsHost.at(splitMap.at(c1)) = 0;
856 int c1a1 = iatoms[stride * c1 + 1];
857 int c1a2 = iatoms[stride * c1 + 2];
859 // Constraints, coupled through the first atom.
861 for (const auto& atomAdjacencyList : atomsAdjacencyList[c1a1])
863 int c2 = atomAdjacencyList.indexOfConstraint_;
867 int c2a2 = atomAdjacencyList.indexOfSecondConstrainedAtom_;
868 int sign = atomAdjacencyList.signFactor_;
869 int index = kernelParams_.numConstraintsThreads
870 * coupledConstraintsCountsHost.at(splitMap.at(c1))
873 coupledConstraintsIndicesHost.at(index) = splitMap.at(c2);
877 float sqrtmu1 = 1.0 / sqrt(invmass[c1a1] + invmass[c1a2]);
878 float sqrtmu2 = 1.0 / sqrt(invmass[c2a1] + invmass[c2a2]);
880 massFactorsHost.at(index) = -sign * invmass[center] * sqrtmu1 * sqrtmu2;
882 coupledConstraintsCountsHost.at(splitMap.at(c1))++;
886 // Constraints, coupled through the second atom.
888 for (const auto& atomAdjacencyList : atomsAdjacencyList[c1a2])
890 int c2 = atomAdjacencyList.indexOfConstraint_;
894 int c2a2 = atomAdjacencyList.indexOfSecondConstrainedAtom_;
895 int sign = atomAdjacencyList.signFactor_;
896 int index = kernelParams_.numConstraintsThreads
897 * coupledConstraintsCountsHost.at(splitMap.at(c1))
900 coupledConstraintsIndicesHost.at(index) = splitMap.at(c2);
904 float sqrtmu1 = 1.0 / sqrt(invmass[c1a1] + invmass[c1a2]);
905 float sqrtmu2 = 1.0 / sqrt(invmass[c2a1] + invmass[c2a2]);
907 massFactorsHost.at(index) = sign * invmass[center] * sqrtmu1 * sqrtmu2;
909 coupledConstraintsCountsHost.at(splitMap.at(c1))++;
914 // (Re)allocate the memory, if the number of constraints has increased.
915 if ((kernelParams_.numConstraintsThreads > numConstraintsThreadsAlloc_) || maxCoupledConstraintsHasIncreased)
917 // Free memory if it was allocated before (i.e. if not the first time here).
918 if (numConstraintsThreadsAlloc_ > 0)
920 freeDeviceBuffer(&kernelParams_.d_constraints);
921 freeDeviceBuffer(&kernelParams_.d_constraintsTargetLengths);
923 freeDeviceBuffer(&kernelParams_.d_coupledConstraintsCounts);
924 freeDeviceBuffer(&kernelParams_.d_coupledConstraintsIndices);
925 freeDeviceBuffer(&kernelParams_.d_massFactors);
926 freeDeviceBuffer(&kernelParams_.d_matrixA);
929 numConstraintsThreadsAlloc_ = kernelParams_.numConstraintsThreads;
931 allocateDeviceBuffer(
932 &kernelParams_.d_constraints, kernelParams_.numConstraintsThreads, deviceContext_);
933 allocateDeviceBuffer(&kernelParams_.d_constraintsTargetLengths,
934 kernelParams_.numConstraintsThreads,
937 allocateDeviceBuffer(&kernelParams_.d_coupledConstraintsCounts,
938 kernelParams_.numConstraintsThreads,
940 allocateDeviceBuffer(&kernelParams_.d_coupledConstraintsIndices,
941 maxCoupledConstraints * kernelParams_.numConstraintsThreads,
943 allocateDeviceBuffer(&kernelParams_.d_massFactors,
944 maxCoupledConstraints * kernelParams_.numConstraintsThreads,
946 allocateDeviceBuffer(&kernelParams_.d_matrixA,
947 maxCoupledConstraints * kernelParams_.numConstraintsThreads,
951 // (Re)allocate the memory, if the number of atoms has increased.
952 if (numAtoms > numAtomsAlloc_)
954 if (numAtomsAlloc_ > 0)
956 freeDeviceBuffer(&kernelParams_.d_inverseMasses);
958 numAtomsAlloc_ = numAtoms;
959 allocateDeviceBuffer(&kernelParams_.d_inverseMasses, numAtoms, deviceContext_);
963 copyToDeviceBuffer(&kernelParams_.d_constraints,
964 constraintsHost.data(),
966 kernelParams_.numConstraintsThreads,
968 GpuApiCallBehavior::Sync,
970 copyToDeviceBuffer(&kernelParams_.d_constraintsTargetLengths,
971 constraintsTargetLengthsHost.data(),
973 kernelParams_.numConstraintsThreads,
975 GpuApiCallBehavior::Sync,
977 copyToDeviceBuffer(&kernelParams_.d_coupledConstraintsCounts,
978 coupledConstraintsCountsHost.data(),
980 kernelParams_.numConstraintsThreads,
982 GpuApiCallBehavior::Sync,
984 copyToDeviceBuffer(&kernelParams_.d_coupledConstraintsIndices,
985 coupledConstraintsIndicesHost.data(),
987 maxCoupledConstraints * kernelParams_.numConstraintsThreads,
989 GpuApiCallBehavior::Sync,
991 copyToDeviceBuffer(&kernelParams_.d_massFactors,
992 massFactorsHost.data(),
994 maxCoupledConstraints * kernelParams_.numConstraintsThreads,
996 GpuApiCallBehavior::Sync,
999 GMX_RELEASE_ASSERT(invmass != nullptr, "Masses of atoms should be specified.\n");
1001 &kernelParams_.d_inverseMasses, invmass, 0, numAtoms, deviceStream_, GpuApiCallBehavior::Sync, nullptr);