2 * This file is part of the GROMACS molecular simulation package.
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.
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 * \note Management of periodic boundary should be unified with SETTLE and
45 * \todo Reconsider naming, i.e. "cuda" suffics should be changed to "gpu".
47 * \author Artem Zhmurov <zhmurov@gmail.com>
48 * \author Alan Gray <alang@nvidia.com>
50 * \ingroup module_mdlib
54 #include "lincs_cuda.cuh"
63 #include "gromacs/gpu_utils/cuda_arch_utils.cuh"
64 #include "gromacs/gpu_utils/cudautils.cuh"
65 #include "gromacs/gpu_utils/devicebuffer.cuh"
66 #include "gromacs/gpu_utils/gputraits.cuh"
67 #include "gromacs/gpu_utils/vectype_ops.cuh"
68 #include "gromacs/math/vec.h"
69 #include "gromacs/mdlib/constr.h"
70 #include "gromacs/pbcutil/pbc.h"
71 #include "gromacs/pbcutil/pbc_aiuc_cuda.cuh"
72 #include "gromacs/topology/ifunc.h"
77 //! Number of CUDA threads in a block
78 constexpr static int c_threadsPerBlock = 256;
79 //! Maximum number of threads in a block (for __launch_bounds__)
80 constexpr static int c_maxThreadsPerBlock = c_threadsPerBlock;
82 /*! \brief Main kernel for LINCS constraints.
84 * See Hess et al., J. Comput. Chem. 18: 1463-1472 (1997) for the description of the algorithm.
86 * In CUDA version, one thread is responsible for all computations for one constraint. The blocks are
87 * filled in a way that no constraint is coupled to the constraint from the next block. This is achieved
88 * by moving active threads to the next block, if the correspondent group of coupled constraints is to big
89 * to fit the current thread block. This may leave some 'dummy' threads in the end of the thread block, i.e.
90 * threads that are not required to do actual work. Since constraints from different blocks are not coupled,
91 * there is no need to synchronize across the device. However, extensive communication in a thread block
94 * \todo Reduce synchronization overhead. Some ideas are:
95 * 1. Consider going to warp-level synchronization for the coupled constraints.
96 * 2. Move more data to local/shared memory and try to get rid of atomic operations (at least on
98 * 3. Use analytical solution for matrix A inversion.
99 * 4. Introduce mapping of thread id to both single constraint and single atom, thus designating
100 * Nth threads to deal with Nat <= Nth coupled atoms and Nc <= Nth coupled constraints.
101 * See Redmine issue #2885 for details (https://redmine.gromacs.org/issues/2885)
102 * \todo The use of __restrict__ for gm_xp and gm_v causes failure, probably because of the atomic
103 operations. Investigate this issue further.
105 * \param[in,out] kernelParams All parameters and pointers for the kernel condensed in single struct.
106 * \param[in] invdt Inverse timestep (needed to update velocities).
108 template<bool updateVelocities, bool computeVirial>
109 __launch_bounds__(c_maxThreadsPerBlock) __global__
110 void lincs_kernel(LincsCudaKernelParameters kernelParams,
111 const float3* __restrict__ gm_x,
116 const PbcAiuc pbcAiuc = kernelParams.pbcAiuc;
117 const int numConstraintsThreads = kernelParams.numConstraintsThreads;
118 const int numIterations = kernelParams.numIterations;
119 const int expansionOrder = kernelParams.expansionOrder;
120 const int2* __restrict__ gm_constraints = kernelParams.d_constraints;
121 const float* __restrict__ gm_constraintsTargetLengths = kernelParams.d_constraintsTargetLengths;
122 const int* __restrict__ gm_coupledConstraintsCounts = kernelParams.d_coupledConstraintsCounts;
123 const int* __restrict__ gm_coupledConstraintsIdxes = kernelParams.d_coupledConstraintsIndices;
124 const float* __restrict__ gm_massFactors = kernelParams.d_massFactors;
125 float* __restrict__ gm_matrixA = kernelParams.d_matrixA;
126 const float* __restrict__ gm_inverseMasses = kernelParams.d_inverseMasses;
127 float* __restrict__ gm_virialScaled = kernelParams.d_virialScaled;
129 int threadIndex = blockIdx.x * blockDim.x + threadIdx.x;
131 // numConstraintsThreads should be a integer multiple of blockSize (numConstraintsThreads = numBlocks*blockSize).
132 // This is to ensure proper synchronizations and reduction. All array are padded to the required size.
133 assert(threadIndex < numConstraintsThreads);
135 // Vectors connecting constrained atoms before algorithm was applied.
136 // Needed to construct constrain matrix A
137 extern __shared__ float3 sm_r[];
139 int2 pair = gm_constraints[threadIndex];
143 // Mass-scaled Lagrange multiplier
144 float lagrangeScaled = 0.0f;
149 float sqrtReducedMass;
155 // i == -1 indicates dummy constraint at the end of the thread block.
156 bool isDummyThread = (i == -1);
158 // Everything computed for these dummies will be equal to zero
164 sqrtReducedMass = 0.0f;
166 xi = make_float3(0.0f, 0.0f, 0.0f);
167 xj = make_float3(0.0f, 0.0f, 0.0f);
168 rc = make_float3(0.0f, 0.0f, 0.0f);
173 targetLength = gm_constraintsTargetLengths[threadIndex];
174 inverseMassi = gm_inverseMasses[i];
175 inverseMassj = gm_inverseMasses[j];
176 sqrtReducedMass = rsqrt(inverseMassi + inverseMassj);
181 float3 dx = pbcDxAiuc(pbcAiuc, xi, xj);
183 float rlen = rsqrtf(dx.x * dx.x + dx.y * dx.y + dx.z * dx.z);
187 sm_r[threadIdx.x] = rc;
188 // Make sure that all r's are saved into shared memory
189 // before they are accessed in the loop below
193 * Constructing LINCS matrix (A)
196 // Only non-zero values are saved (for coupled constraints)
197 int coupledConstraintsCount = gm_coupledConstraintsCounts[threadIndex];
198 for (int n = 0; n < coupledConstraintsCount; n++)
200 int index = n * numConstraintsThreads + threadIndex;
201 int c1 = gm_coupledConstraintsIdxes[index];
203 float3 rc1 = sm_r[c1 - blockIdx.x * blockDim.x];
204 gm_matrixA[index] = gm_massFactors[index] * (rc.x * rc1.x + rc.y * rc1.y + rc.z * rc1.z);
207 // Skipping in dummy threads
214 float3 dx = pbcDxAiuc(pbcAiuc, xi, xj);
216 float sol = sqrtReducedMass * ((rc.x * dx.x + rc.y * dx.y + rc.z * dx.z) - targetLength);
219 * Inverse matrix using a set of expansionOrder matrix multiplications
222 // This will use the same memory space as sm_r, which is no longer needed.
223 extern __shared__ float sm_rhs[];
224 // Save current right-hand-side vector in the shared memory
225 sm_rhs[threadIdx.x] = sol;
227 for (int rec = 0; rec < expansionOrder; rec++)
229 // Making sure that all sm_rhs are saved before they are accessed in a loop below
233 for (int n = 0; n < coupledConstraintsCount; n++)
235 int index = n * numConstraintsThreads + threadIndex;
236 int c1 = gm_coupledConstraintsIdxes[index];
237 // Convolute current right-hand-side with A
238 // Different, non overlapping parts of sm_rhs[..] are read during odd and even iterations
239 mvb = mvb + gm_matrixA[index] * sm_rhs[c1 - blockIdx.x * blockDim.x + blockDim.x * (rec % 2)];
241 // 'Switch' rhs vectors, save current result
242 // These values will be accessed in the loop above during the next iteration.
243 sm_rhs[threadIdx.x + blockDim.x * ((rec + 1) % 2)] = mvb;
247 // Current mass-scaled Lagrange multipliers
248 lagrangeScaled = sqrtReducedMass * sol;
250 // Save updated coordinates before correction for the rotational lengthening
251 float3 tmp = rc * lagrangeScaled;
253 // Writing for all but dummy constraints
256 atomicAdd(&gm_xp[i], -tmp * inverseMassi);
257 atomicAdd(&gm_xp[j], tmp * inverseMassj);
261 * Correction for centripetal effects
263 for (int iter = 0; iter < numIterations; iter++)
265 // Make sure that all xp's are saved: atomic operation calls before are
266 // communicating current xp[..] values across thread block.
275 float3 dx = pbcDxAiuc(pbcAiuc, xi, xj);
277 float len2 = targetLength * targetLength;
278 float dlen2 = 2.0f * len2 - norm2(dx);
280 // TODO A little bit more effective but slightly less readable version of the below would be:
281 // float proj = sqrtReducedMass*(targetLength - (dlen2 > 0.0f ? 1.0f : 0.0f)*dlen2*rsqrt(dlen2));
285 proj = sqrtReducedMass * (targetLength - dlen2 * rsqrt(dlen2));
289 proj = sqrtReducedMass * targetLength;
292 sm_rhs[threadIdx.x] = proj;
296 * Same matrix inversion as above is used for updated data
298 for (int rec = 0; rec < expansionOrder; rec++)
300 // Make sure that all elements of rhs are saved into shared memory
304 for (int n = 0; n < coupledConstraintsCount; n++)
306 int index = n * numConstraintsThreads + threadIndex;
307 int c1 = gm_coupledConstraintsIdxes[index];
309 mvb = mvb + gm_matrixA[index] * sm_rhs[c1 - blockIdx.x * blockDim.x + blockDim.x * (rec % 2)];
311 sm_rhs[threadIdx.x + blockDim.x * ((rec + 1) % 2)] = mvb;
315 // Add corrections to Lagrange multipliers
316 float sqrtmu_sol = sqrtReducedMass * sol;
317 lagrangeScaled += sqrtmu_sol;
319 // Save updated coordinates for the next iteration
320 // Dummy constraints are skipped
323 float3 tmp = rc * sqrtmu_sol;
324 atomicAdd(&gm_xp[i], -tmp * inverseMassi);
325 atomicAdd(&gm_xp[j], tmp * inverseMassj);
329 // Updating particle velocities for all but dummy threads
330 if (updateVelocities && !isDummyThread)
332 float3 tmp = rc * invdt * lagrangeScaled;
333 atomicAdd(&gm_v[i], -tmp * inverseMassi);
334 atomicAdd(&gm_v[j], tmp * inverseMassj);
340 // Virial is computed from Lagrange multiplier (lagrangeScaled), target constrain length
341 // (targetLength) and the normalized vector connecting constrained atoms before
342 // the algorithm was applied (rc). The evaluation of virial in each thread is
343 // followed by basic reduction for the values inside single thread block.
344 // Then, the values are reduced across grid by atomicAdd(...).
346 // TODO Shuffle reduction.
347 // TODO Should be unified and/or done once when virial is actually needed.
348 // TODO Recursive version that removes atomicAdd(...)'s entirely is needed. Ideally,
349 // one that works for any datatype.
351 // Save virial for each thread into the shared memory. Tensor is symmetrical, hence only
352 // 6 values are saved. Dummy threads will have zeroes in their virial: targetLength,
353 // lagrangeScaled and rc are all set to zero for them in the beginning of the kernel.
354 // The sm_threadVirial[..] will overlap with the sm_r[..] and sm_rhs[..], but the latter
355 // two are no longer in use.
356 extern __shared__ float sm_threadVirial[];
357 float mult = targetLength * lagrangeScaled;
358 sm_threadVirial[0 * blockDim.x + threadIdx.x] = mult * rc.x * rc.x;
359 sm_threadVirial[1 * blockDim.x + threadIdx.x] = mult * rc.x * rc.y;
360 sm_threadVirial[2 * blockDim.x + threadIdx.x] = mult * rc.x * rc.z;
361 sm_threadVirial[3 * blockDim.x + threadIdx.x] = mult * rc.y * rc.y;
362 sm_threadVirial[4 * blockDim.x + threadIdx.x] = mult * rc.y * rc.z;
363 sm_threadVirial[5 * blockDim.x + threadIdx.x] = mult * rc.z * rc.z;
367 // Reduce up to one virial per thread block. All blocks are divided by half, the first
368 // half of threads sums two virials. Then the first half is divided by two and the first
369 // half of it sums two values. This procedure is repeated until only one thread is left.
370 // Only works if the threads per blocks is a power of two (hence static_assert
371 // in the beginning of the kernel).
372 for (int divideBy = 2; divideBy <= static_cast<int>(blockDim.x); divideBy *= 2)
374 int dividedAt = blockDim.x / divideBy;
375 if (static_cast<int>(threadIdx.x) < dividedAt)
377 for (int d = 0; d < 6; d++)
379 sm_threadVirial[d * blockDim.x + threadIdx.x] +=
380 sm_threadVirial[d * blockDim.x + (threadIdx.x + dividedAt)];
383 // Syncronize if not within one warp
384 if (dividedAt > warpSize / 2)
389 // First 6 threads in the block add the results of 6 tensor components to the global memory address.
392 atomicAdd(&(gm_virialScaled[threadIdx.x]), sm_threadVirial[threadIdx.x * blockDim.x]);
399 /*! \brief Select templated kernel.
401 * Returns pointer to a CUDA kernel based on provided booleans.
403 * \param[in] updateVelocities If the velocities should be constrained.
404 * \param[in] computeVirial If virial should be updated.
406 * \return Pointer to CUDA kernel
408 inline auto getLincsKernelPtr(const bool updateVelocities, const bool computeVirial)
411 auto kernelPtr = lincs_kernel<true, true>;
412 if (updateVelocities && computeVirial)
414 kernelPtr = lincs_kernel<true, true>;
416 else if (updateVelocities && !computeVirial)
418 kernelPtr = lincs_kernel<true, false>;
420 else if (!updateVelocities && computeVirial)
422 kernelPtr = lincs_kernel<false, true>;
424 else if (!updateVelocities && !computeVirial)
426 kernelPtr = lincs_kernel<false, false>;
431 void LincsCuda::apply(const float3* d_x,
433 const bool updateVelocities,
436 const bool computeVirial,
439 ensureNoPendingCudaError("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, commandStream_);
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);
479 config.stream = commandStream_;
481 const auto kernelArgs =
482 prepareGpuKernelArguments(kernelPtr, config, &kernelParams_, &d_x, &d_xp, &d_v, &invdt);
484 launchGpuKernel(kernelPtr, config, nullptr, "lincs_kernel<updateVelocities, computeVirial>", kernelArgs);
488 // Copy LINCS virial data and add it to the common virial
489 copyFromDeviceBuffer(h_virialScaled_.data(), &kernelParams_.d_virialScaled, 0, 6,
490 commandStream_, GpuApiCallBehavior::Sync, nullptr);
492 // Mapping [XX, XY, XZ, YY, YZ, ZZ] internal format to a tensor object
493 virialScaled[XX][XX] += h_virialScaled_[0];
494 virialScaled[XX][YY] += h_virialScaled_[1];
495 virialScaled[XX][ZZ] += h_virialScaled_[2];
497 virialScaled[YY][XX] += h_virialScaled_[1];
498 virialScaled[YY][YY] += h_virialScaled_[3];
499 virialScaled[YY][ZZ] += h_virialScaled_[4];
501 virialScaled[ZZ][XX] += h_virialScaled_[2];
502 virialScaled[ZZ][YY] += h_virialScaled_[4];
503 virialScaled[ZZ][ZZ] += h_virialScaled_[5];
509 LincsCuda::LincsCuda(int numIterations, int expansionOrder, CommandStream commandStream) :
510 commandStream_(commandStream)
512 kernelParams_.numIterations = numIterations;
513 kernelParams_.expansionOrder = expansionOrder;
515 static_assert(sizeof(real) == sizeof(float),
516 "Real numbers should be in single precision in GPU code.");
518 c_threadsPerBlock > 0 && ((c_threadsPerBlock & (c_threadsPerBlock - 1)) == 0),
519 "Number of threads per block should be a power of two in order for reduction to work.");
521 allocateDeviceBuffer(&kernelParams_.d_virialScaled, 6, nullptr);
522 h_virialScaled_.resize(6);
524 // The data arrays should be expanded/reallocated on first call of set() function.
525 numConstraintsThreadsAlloc_ = 0;
529 LincsCuda::~LincsCuda()
531 freeDeviceBuffer(&kernelParams_.d_virialScaled);
533 if (numConstraintsThreadsAlloc_ > 0)
535 freeDeviceBuffer(&kernelParams_.d_constraints);
536 freeDeviceBuffer(&kernelParams_.d_constraintsTargetLengths);
538 freeDeviceBuffer(&kernelParams_.d_coupledConstraintsCounts);
539 freeDeviceBuffer(&kernelParams_.d_coupledConstraintsIndices);
540 freeDeviceBuffer(&kernelParams_.d_massFactors);
541 freeDeviceBuffer(&kernelParams_.d_matrixA);
543 if (numAtomsAlloc_ > 0)
545 freeDeviceBuffer(&kernelParams_.d_inverseMasses);
549 /*! \brief Helper function to go through constraints recursively.
551 * For each constraint, counts the number of coupled constraints and stores the value in spaceNeeded array.
552 * This information is used to split the array of constraints between thread blocks on a GPU so there is no
553 * coupling between constraints from different thread blocks. After the 'spaceNeeded' array is filled, the
554 * value spaceNeeded[c] should be equal to the number of constraints that are coupled to 'c' and located
555 * after it in the constraints array.
557 * \param[in] a Atom index.
558 * \param[in,out] spaceNeeded Indicates if the constraint was already counted and stores
559 * the number of constraints (i) connected to it and (ii) located
560 * after it in memory. This array is filled by this recursive function.
561 * For a set of coupled constraints, only for the first one in this list
562 * the number of consecutive coupled constraints is needed: if there is
563 * not enough space for this set of constraints in the thread block,
564 * the group has to be moved to the next one.
565 * \param[in] atomAdjacencyList Stores information about connections between atoms.
567 inline int countCoupled(int a,
568 std::vector<int>* spaceNeeded,
569 std::vector<std::vector<std::tuple<int, int, int>>>* atomsAdjacencyList)
574 for (unsigned i = 0; i < atomsAdjacencyList->at(a).size(); i++)
576 std::tie(a2, c2, sign) = atomsAdjacencyList->at(a).at(i);
577 if (spaceNeeded->at(c2) == -1)
579 spaceNeeded->at(c2) = 0; // To indicate we've been here
580 counted += 1 + countCoupled(a2, spaceNeeded, atomsAdjacencyList);
586 void LincsCuda::set(const t_idef& idef, const t_mdatoms& md)
588 int numAtoms = md.nr;
589 // List of constrained atoms (CPU memory)
590 std::vector<int2> constraintsHost;
591 // Equilibrium distances for the constraints (CPU)
592 std::vector<float> constraintsTargetLengthsHost;
593 // Number of constraints, coupled with the current one (CPU)
594 std::vector<int> coupledConstraintsCountsHost;
595 // List of coupled with the current one (CPU)
596 std::vector<int> coupledConstraintsIndicesHost;
597 // Mass factors (CPU)
598 std::vector<float> massFactorsHost;
600 // List of constrained atoms in local topology
601 t_iatom* iatoms = idef.il[F_CONSTR].iatoms;
602 const int stride = NRAL(F_CONSTR) + 1;
603 const int numConstraints = idef.il[F_CONSTR].nr / stride;
605 // Early exit if no constraints
606 if (numConstraints == 0)
608 kernelParams_.numConstraintsThreads = 0;
612 // Constructing adjacency list --- usefull intermediate structure
613 std::vector<std::vector<std::tuple<int, int, int>>> atomsAdjacencyList(numAtoms);
614 for (int c = 0; c < numConstraints; c++)
616 int a1 = iatoms[stride * c + 1];
617 int a2 = iatoms[stride * c + 2];
619 // Each constraint will be represented as a tuple, containing index of the second
620 // constrained atom, index of the constraint and a sign that indicates the order of atoms in
621 // which they are listed. Sign is needed to compute the mass factors.
622 atomsAdjacencyList.at(a1).push_back(std::make_tuple(a2, c, +1));
623 atomsAdjacencyList.at(a2).push_back(std::make_tuple(a1, c, -1));
626 // Compute, how many coupled constraints are in front of each constraint.
627 // Needed to introduce splits in data so that all coupled constraints will be computed in a single GPU block.
628 // The position 'c' of the vector spaceNeeded should have the number of constraints that are coupled to a constraint
629 // 'c' and are after 'c' in the vector. Only first index of the connected group of the constraints is needed later in the
630 // code, hence the spaceNeeded vector is also used to keep track if the constrain was already counted.
631 std::vector<int> spaceNeeded;
632 spaceNeeded.resize(numConstraints, -1);
633 std::fill(spaceNeeded.begin(), spaceNeeded.end(), -1);
634 for (int c = 0; c < numConstraints; c++)
636 int a1 = iatoms[stride * c + 1];
637 int a2 = iatoms[stride * c + 2];
638 if (spaceNeeded.at(c) == -1)
640 spaceNeeded.at(c) = countCoupled(a1, &spaceNeeded, &atomsAdjacencyList)
641 + countCoupled(a2, &spaceNeeded, &atomsAdjacencyList);
645 // Map of splits in the constraints data. For each 'old' constraint index gives 'new' which
646 // takes into account the empty spaces which might be needed in the end of each thread block.
647 std::vector<int> splitMap;
648 splitMap.resize(numConstraints, -1);
649 int currentMapIndex = 0;
650 for (int c = 0; c < numConstraints; c++)
652 // Check if coupled constraints all fit in one block
654 spaceNeeded.at(c) < c_threadsPerBlock,
655 "Maximum number of coupled constraints exceedes the size of the CUDA thread block. "
656 "Most likely, you are trying to use GPU version of LINCS with constraints on "
658 "which is not supported. Try using H-bonds constraints only.");
659 if (currentMapIndex / c_threadsPerBlock != (currentMapIndex + spaceNeeded.at(c)) / c_threadsPerBlock)
661 currentMapIndex = ((currentMapIndex / c_threadsPerBlock) + 1) * c_threadsPerBlock;
663 splitMap.at(c) = currentMapIndex;
666 kernelParams_.numConstraintsThreads =
667 currentMapIndex + c_threadsPerBlock - currentMapIndex % c_threadsPerBlock;
668 GMX_RELEASE_ASSERT(kernelParams_.numConstraintsThreads % c_threadsPerBlock == 0,
669 "Number of threads should be a multiple of the block size");
671 // Initialize constraints and their target indexes taking into account the splits in the data arrays.
675 constraintsHost.resize(kernelParams_.numConstraintsThreads, pair);
676 std::fill(constraintsHost.begin(), constraintsHost.end(), pair);
677 constraintsTargetLengthsHost.resize(kernelParams_.numConstraintsThreads, 0.0);
678 std::fill(constraintsTargetLengthsHost.begin(), constraintsTargetLengthsHost.end(), 0.0);
679 for (int c = 0; c < numConstraints; c++)
681 int a1 = iatoms[stride * c + 1];
682 int a2 = iatoms[stride * c + 2];
683 int type = iatoms[stride * c];
688 constraintsHost.at(splitMap.at(c)) = pair;
689 constraintsTargetLengthsHost.at(splitMap.at(c)) = idef.iparams[type].constr.dA;
692 // The adjacency list of constraints (i.e. the list of coupled constraints for each constraint).
693 // We map a single thread to a single constraint, hence each thread 'c' will be using one
694 // element from coupledConstraintsCountsHost array, which is the number of constraints coupled
695 // to the constraint 'c'. The coupled constraints indexes are placed into the
696 // coupledConstraintsIndicesHost array. Latter is organized as a one-dimensional array to ensure
697 // good memory alignment. It is addressed as [c + i*numConstraintsThreads], where 'i' goes from
698 // zero to the number of constraints coupled to 'c'. 'numConstraintsThreads' is the width of the
699 // array --- a number, greater then total number of constraints, taking into account the splits
700 // in the constraints array due to the GPU block borders. This number can be adjusted to improve
701 // memory access pattern. Mass factors are saved in a similar data structure.
702 int maxCoupledConstraints = 0;
703 for (int c = 0; c < numConstraints; c++)
705 int a1 = iatoms[stride * c + 1];
706 int a2 = iatoms[stride * c + 2];
708 // Constraint 'c' is counted twice, but it should be excluded altogether. Hence '-2'.
709 int nCoupedConstraints = atomsAdjacencyList.at(a1).size() + atomsAdjacencyList.at(a2).size() - 2;
711 if (nCoupedConstraints > maxCoupledConstraints)
713 maxCoupledConstraints = nCoupedConstraints;
717 coupledConstraintsCountsHost.resize(kernelParams_.numConstraintsThreads, 0);
718 coupledConstraintsIndicesHost.resize(maxCoupledConstraints * kernelParams_.numConstraintsThreads, -1);
719 massFactorsHost.resize(maxCoupledConstraints * kernelParams_.numConstraintsThreads, -1);
721 for (int c1 = 0; c1 < numConstraints; c1++)
723 coupledConstraintsCountsHost.at(splitMap.at(c1)) = 0;
724 int c1a1 = iatoms[stride * c1 + 1];
725 int c1a2 = iatoms[stride * c1 + 2];
732 // Constraints, coupled trough the first atom.
734 for (unsigned j = 0; j < atomsAdjacencyList.at(c1a1).size(); j++)
737 std::tie(c2a2, c2, sign) = atomsAdjacencyList.at(c1a1).at(j);
741 int index = kernelParams_.numConstraintsThreads
742 * coupledConstraintsCountsHost.at(splitMap.at(c1))
745 coupledConstraintsIndicesHost.at(index) = splitMap.at(c2);
749 float sqrtmu1 = 1.0 / sqrt(md.invmass[c1a1] + md.invmass[c1a2]);
750 float sqrtmu2 = 1.0 / sqrt(md.invmass[c2a1] + md.invmass[c2a2]);
752 massFactorsHost.at(index) = -sign * md.invmass[center] * sqrtmu1 * sqrtmu2;
754 coupledConstraintsCountsHost.at(splitMap.at(c1))++;
758 // Constraints, coupled through the second atom.
760 for (unsigned j = 0; j < atomsAdjacencyList.at(c1a2).size(); j++)
763 std::tie(c2a2, c2, sign) = atomsAdjacencyList.at(c1a2).at(j);
767 int index = kernelParams_.numConstraintsThreads
768 * coupledConstraintsCountsHost.at(splitMap.at(c1))
771 coupledConstraintsIndicesHost.at(index) = splitMap.at(c2);
775 float sqrtmu1 = 1.0 / sqrt(md.invmass[c1a1] + md.invmass[c1a2]);
776 float sqrtmu2 = 1.0 / sqrt(md.invmass[c2a1] + md.invmass[c2a2]);
778 massFactorsHost.at(index) = sign * md.invmass[center] * sqrtmu1 * sqrtmu2;
780 coupledConstraintsCountsHost.at(splitMap.at(c1))++;
785 // (Re)allocate the memory, if the number of constraints has increased.
786 if (kernelParams_.numConstraintsThreads > numConstraintsThreadsAlloc_)
788 // Free memory if it was allocated before (i.e. if not the first time here).
789 if (numConstraintsThreadsAlloc_ > 0)
791 freeDeviceBuffer(&kernelParams_.d_constraints);
792 freeDeviceBuffer(&kernelParams_.d_constraintsTargetLengths);
794 freeDeviceBuffer(&kernelParams_.d_coupledConstraintsCounts);
795 freeDeviceBuffer(&kernelParams_.d_coupledConstraintsIndices);
796 freeDeviceBuffer(&kernelParams_.d_massFactors);
797 freeDeviceBuffer(&kernelParams_.d_matrixA);
800 numConstraintsThreadsAlloc_ = kernelParams_.numConstraintsThreads;
802 allocateDeviceBuffer(&kernelParams_.d_constraints, kernelParams_.numConstraintsThreads, nullptr);
803 allocateDeviceBuffer(&kernelParams_.d_constraintsTargetLengths,
804 kernelParams_.numConstraintsThreads, nullptr);
806 allocateDeviceBuffer(&kernelParams_.d_coupledConstraintsCounts,
807 kernelParams_.numConstraintsThreads, nullptr);
808 allocateDeviceBuffer(&kernelParams_.d_coupledConstraintsIndices,
809 maxCoupledConstraints * kernelParams_.numConstraintsThreads, nullptr);
810 allocateDeviceBuffer(&kernelParams_.d_massFactors,
811 maxCoupledConstraints * kernelParams_.numConstraintsThreads, nullptr);
812 allocateDeviceBuffer(&kernelParams_.d_matrixA,
813 maxCoupledConstraints * kernelParams_.numConstraintsThreads, nullptr);
816 // (Re)allocate the memory, if the number of atoms has increased.
817 if (numAtoms > numAtomsAlloc_)
819 if (numAtomsAlloc_ > 0)
821 freeDeviceBuffer(&kernelParams_.d_inverseMasses);
823 numAtomsAlloc_ = numAtoms;
824 allocateDeviceBuffer(&kernelParams_.d_inverseMasses, numAtoms, nullptr);
828 copyToDeviceBuffer(&kernelParams_.d_constraints, constraintsHost.data(), 0,
829 kernelParams_.numConstraintsThreads, commandStream_,
830 GpuApiCallBehavior::Sync, nullptr);
831 copyToDeviceBuffer(&kernelParams_.d_constraintsTargetLengths,
832 constraintsTargetLengthsHost.data(), 0, kernelParams_.numConstraintsThreads,
833 commandStream_, GpuApiCallBehavior::Sync, nullptr);
834 copyToDeviceBuffer(&kernelParams_.d_coupledConstraintsCounts,
835 coupledConstraintsCountsHost.data(), 0, kernelParams_.numConstraintsThreads,
836 commandStream_, GpuApiCallBehavior::Sync, nullptr);
837 copyToDeviceBuffer(&kernelParams_.d_coupledConstraintsIndices, coupledConstraintsIndicesHost.data(),
838 0, maxCoupledConstraints * kernelParams_.numConstraintsThreads,
839 commandStream_, GpuApiCallBehavior::Sync, nullptr);
840 copyToDeviceBuffer(&kernelParams_.d_massFactors, massFactorsHost.data(), 0,
841 maxCoupledConstraints * kernelParams_.numConstraintsThreads, commandStream_,
842 GpuApiCallBehavior::Sync, nullptr);
844 GMX_RELEASE_ASSERT(md.invmass != nullptr, "Masses of attoms should be specified.\n");
845 copyToDeviceBuffer(&kernelParams_.d_inverseMasses, md.invmass, 0, numAtoms, commandStream_,
846 GpuApiCallBehavior::Sync, nullptr);
849 void LincsCuda::setPbc(const t_pbc* pbc)
851 setPbcAiuc(pbc->ndim_ePBC, pbc->box, &kernelParams_.pbcAiuc);