kLincs.cu 9.48 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
/* -------------------------------------------------------------------------- *
 *                                   OpenMM                                   *
 * -------------------------------------------------------------------------- *
 * This is part of the OpenMM molecular simulation toolkit originating from   *
 * Simbios, the NIH National Center for Physics-Based Simulation of           *
 * Biological Structures at Stanford, funded under the NIH Roadmap for        *
 * Medical Research, grant U54 GM072970. See https://simtk.org.               *
 *                                                                            *
 * Portions copyright (c) 2009 Stanford University and the Authors.           *
 * Authors: Scott Le Grand, Peter Eastman                                     *
 * Contributors:                                                              *
 *                                                                            *
13
14
15
16
 * This program is free software: you can redistribute it and/or modify       *
 * it under the terms of the GNU Lesser General Public License as published   *
 * by the Free Software Foundation, either version 3 of the License, or       *
 * (at your option) any later version.                                        *
17
 *                                                                            *
18
19
20
21
 * This program is distributed in the hope that it will be useful,            *
 * but WITHOUT ANY WARRANTY; without even the implied warranty of             *
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the              *
 * GNU Lesser General Public License for more details.                        *
22
 *                                                                            *
23
24
 * You should have received a copy of the GNU Lesser General Public License   *
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.      *
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
 * -------------------------------------------------------------------------- */

#include <cuda.h>
#include <vector_functions.h>
using namespace std;

#include "gputypes.h"


static __constant__ cudaGmxSimulation cSim;

void SetLincsSim(gpuContext gpu)
{
    cudaError_t status;
    status = cudaMemcpyToSymbol(cSim, &gpu->sim, sizeof(cudaGmxSimulation));
    RTERROR(status, "cudaMemcpyToSymbol: SetSim copy to cSim failed");
}

void GetLincsSim(gpuContext gpu)
{
    cudaError_t status;
    status = cudaMemcpyFromSymbol(&gpu->sim, cSim, sizeof(cudaGmxSimulation));
    RTERROR(status, "cudaMemcpyFromSymbol: SetSim copy from cSim failed");
}

50
51
52
53
/**
 * Synchronize all threads across all blocks.
 */
__device__ void kSyncAllThreads_kernel(short* syncCounter, short newCount)
54
{
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
//    short* syncCounter = &cSim.pSyncCounter[newCount%2 == 0 ? 0 : gridDim.x];
    __syncthreads();
    if (threadIdx.x == 0)
        syncCounter[blockIdx.x] = newCount;
    if (threadIdx.x < gridDim.x)
    {
        volatile short* counter = &syncCounter[threadIdx.x];
        do
        {
        } while (*counter != newCount);
    }
    __syncthreads();
}

__global__ void kSolveLincsMatrix_kernel(float4* atomPositions)
{
    for (unsigned int iteration = 0; iteration < cSim.lincsTerms; iteration++) {
        float* rhs1 = (iteration%2 == 0 ? cSim.pLincsRhs1 : cSim.pLincsRhs2);
        float* rhs2 = (iteration%2 == 0 ? cSim.pLincsRhs2 : cSim.pLincsRhs1);
        unsigned int pos = threadIdx.x + blockIdx.x * blockDim.x;
        while (pos < cSim.lincsConstraints)
        {
            float rhs = 0.0f;
            int num = cSim.pLincsNumConnections[pos];
            for (int i = 0; i < num; i++)
            {
                int index = pos+i*cSim.lincsConstraints;
                int otherConstraint = cSim.pLincsConnections[index];
                rhs += cSim.pLincsCoupling[index]*rhs1[otherConstraint];
            }
            rhs2[pos] = rhs;
            cSim.pLincsSolution[pos] += rhs;
            pos += blockDim.x * gridDim.x;
        }
        kSyncAllThreads_kernel(&cSim.pSyncCounter[iteration%2 == 0 ? 0 : gridDim.x], iteration);
    }

Peter Eastman's avatar
Peter Eastman committed
92
    // Update the atom positions based on the solution to the matrix equations.
93
94
95
96
97
98

    unsigned int pos = threadIdx.x + blockIdx.x * blockDim.x;
    while (pos < cSim.atoms)
    {
        float4 atomPos = atomPositions[pos];
        float invMass = cSim.pVelm4[pos].w;
99
100
        int num = cSim.pLincsNumAtomConstraints[pos];
        for (int i = 0; i < num; i++)
101
        {
102
103
            int index = pos+i*cSim.atoms;
            int constraint = cSim.pLincsAtomConstraints[index];
104
105
            bool forward = (constraint > 0);
            constraint = (forward ? constraint-1 : -constraint-1);
106
107
            float4 dir = cSim.pLincsDistance[constraint];
            float c = invMass*cSim.pLincsS[constraint]*cSim.pLincsSolution[constraint];
108
            c = (forward ? -c : c);
109
110
111
112
113
114
115
116
117
            atomPos.x += c*dir.x;
            atomPos.y += c*dir.y;
            atomPos.z += c*dir.z;
        }
        atomPositions[pos] = atomPos;
        pos += blockDim.x * gridDim.x;
    }
}

Peter Eastman's avatar
Peter Eastman committed
118
__global__ void kApplyLincsPart1_kernel(float4* atomPositions, bool addOldPosition)
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
{
   // Calculate the direction of each constraint, along with the initial RHS and solution vectors.

    unsigned int pos = threadIdx.x + blockIdx.x * blockDim.x;
    while (pos < cSim.lincsConstraints)
    {
        int2 atoms = cSim.pLincsAtoms[pos];
        float4 delta1 = atomPositions[atoms.x];
        float4 delta2 = atomPositions[atoms.y];
        float4 dir = cSim.pLincsDistance[pos];
        if (addOldPosition)
        {
            float4 oldPos1 = cSim.pOldPosq[atoms.x];
            float4 oldPos2 = cSim.pOldPosq[atoms.y];
            dir.x = (oldPos1.x-oldPos2.x)+(delta1.x-delta2.x);
            dir.y = (oldPos1.y-oldPos2.y)+(delta1.y-delta2.y);
            dir.z = (oldPos1.z-oldPos2.z)+(delta1.z-delta2.z);
        }
        else
        {
            dir.x = delta1.x-delta2.x;
            dir.y = delta1.y-delta2.y;
            dir.z = delta1.z-delta2.z;
        }
        float invLength = 1.0f/sqrt(dir.x*dir.x + dir.y*dir.y + dir.z*dir.z);
        dir.x *= invLength;
        dir.y *= invLength;
        dir.z *= invLength;
        cSim.pLincsDistance[pos] = dir;
        float diff = cSim.pLincsS[pos]*(1.0f/invLength-dir.w);
        cSim.pLincsRhs1[pos] = diff;
        cSim.pLincsSolution[pos] = diff;
        pos += blockDim.x * gridDim.x;
    }
153
    kSyncAllThreads_kernel(cSim.pSyncCounter, cSim.lincsTerms+1);
154
155
156

    // Build the coupling matrix.

157
    pos = threadIdx.x + blockIdx.x * blockDim.x;
158
159
160
161
    while (pos < cSim.lincsConstraints)
    {
        float4 dir1 = cSim.pLincsDistance[pos];
        int2 atoms1 = cSim.pLincsAtoms[pos];
162
        int num = cSim.pLincsNumConnections[pos];
163
164
        float s = cSim.pLincsS[pos];
        float invMass = cSim.pVelm4[atoms1.x].w;
165
        for (int i = 0; i < num; i++)
166
        {
167
168
            int index = pos+i*cSim.lincsConstraints;
            int otherConstraint = cSim.pLincsConnections[index];
169
170
            float4 dir2 = cSim.pLincsDistance[otherConstraint];
            int2 atoms2 = cSim.pLincsAtoms[otherConstraint];
171
            float signedMass = (atoms1.x == atoms2.x || atoms1.y == atoms2.y ? -invMass : cSim.pVelm4[atoms1.y].w);
172
            cSim.pLincsCoupling[index] = signedMass*s*(dir1.x*dir2.x+dir1.y*dir2.y+dir1.z*dir2.z)*cSim.pLincsS[otherConstraint];
173
174
175
        }
        pos += blockDim.x * gridDim.x;
    }
Peter Eastman's avatar
Peter Eastman committed
176
}
177

178
__global__ void kApplyLincsPart2_kernel(float4* atomPositions, bool addOldPosition)
Peter Eastman's avatar
Peter Eastman committed
179
{
180
181
    // Correct for rotational lengthening.

Peter Eastman's avatar
Peter Eastman committed
182
    unsigned int pos = threadIdx.x + blockIdx.x * blockDim.x;
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
    while (pos < cSim.lincsConstraints)
    {
        int2 atoms = cSim.pLincsAtoms[pos];
        float4 delta1 = atomPositions[atoms.x];
        float4 delta2 = atomPositions[atoms.y];
        float3 delta;
        if (addOldPosition)
        {
            float4 oldPos1 = cSim.pOldPosq[atoms.x];
            float4 oldPos2 = cSim.pOldPosq[atoms.y];
            delta = make_float3((oldPos1.x-oldPos2.x)+(delta1.x-delta2.x),
                                (oldPos1.y-oldPos2.y)+(delta1.y-delta2.y),
                                (oldPos1.z-oldPos2.z)+(delta1.z-delta2.z));
        }
        else
        {
            delta = make_float3(delta1.x-delta2.x, delta1.y-delta2.y, delta1.z-delta2.z);
        }
        float distance = cSim.pLincsDistance[pos].w;
        float p2 = 2.0f*distance*distance-(delta.x*delta.x+delta.y*delta.y+delta.z*delta.z);
        p2 = (p2 < 0.0f ? 0.0f : p2);
        float diff = cSim.pLincsS[pos]*(distance-sqrt(p2));
        cSim.pLincsRhs1[pos] = diff;
        cSim.pLincsSolution[pos] = diff;
        pos += blockDim.x * gridDim.x;
    }
Peter Eastman's avatar
Peter Eastman committed
209
}
210

Peter Eastman's avatar
Peter Eastman committed
211
212
213
214
static void kApplyLincs(gpuContext gpu, float4* atomPositions, bool addOldPosition)
{
    kApplyLincsPart1_kernel<<<gpu->sim.blocks, gpu->sim.lincs_threads_per_block>>>(atomPositions, addOldPosition);
    LAUNCHERROR("kApplyLincsPart1");
215
216
217
    kSolveLincsMatrix_kernel<<<gpu->sim.blocks, gpu->sim.lincs_threads_per_block>>>(atomPositions);
    LAUNCHERROR("kSolveLincsMatrix_kernel");
    kApplyLincsPart2_kernel<<<gpu->sim.blocks, gpu->sim.lincs_threads_per_block>>>(atomPositions, addOldPosition);
Peter Eastman's avatar
Peter Eastman committed
218
    LAUNCHERROR("kApplyLincsPart2");
219
220
    kSolveLincsMatrix_kernel<<<gpu->sim.blocks, gpu->sim.lincs_threads_per_block>>>(atomPositions);
    LAUNCHERROR("kSolveLincsMatrix_kernel");
221
222
223
224
225
226
}

void kApplyFirstLincs(gpuContext gpu)
{
//    printf("kApplyFirstLincs\n");
    if (gpu->sim.lincsConstraints > 0)
Peter Eastman's avatar
Peter Eastman committed
227
        kApplyLincs(gpu, gpu->sim.pPosqP, true);
228
229
230
231
232
233
}

void kApplySecondLincs(gpuContext gpu)
{
//    printf("kApplySecondLincs\n");
    if (gpu->sim.lincsConstraints > 0)
Peter Eastman's avatar
Peter Eastman committed
234
        kApplyLincs(gpu, gpu->sim.pPosq, false);
235
}