kCalculateCustomNonbondedForces.h 18.3 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
/* -------------------------------------------------------------------------- *
 *                                   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:                                                              *
 *                                                                            *
 * 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.                                        *
 *                                                                            *
 * 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.                        *
 *                                                                            *
 * 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/>.      *
 * -------------------------------------------------------------------------- */

/**
28
 * This file contains the kernels for evalauating custom nonbonded forces.  It is included
29
30
31
32
33
34
35
 * several times in kCalculateCustomNonbondedForces.cu with different #defines to generate
 * different versions of the kernels.
 */

__global__ void METHOD_NAME(kCalculateCustomNonbonded, Forces_kernel)(unsigned int* workUnit)
{
    extern __shared__ float stack[];
36
    volatile Atom* sA = (volatile Atom*) &stack[cSim.customExpressionStackSize*blockDim.x];
37
    float* variables = (float*) &sA[blockDim.x];
38
    unsigned int totalWarps = gridDim.x*blockDim.x/GRID;
39
40
41
42
    unsigned int warp = (blockIdx.x*blockDim.x+threadIdx.x)/GRID;
    unsigned int numWorkUnits = cSim.pInteractionCount[0];
    unsigned int pos = warp*numWorkUnits/totalWarps;
    unsigned int end = (warp+1)*numWorkUnits/totalWarps;
43
    float totalEnergy = 0.0f;
44
#ifdef USE_CUTOFF
45
    volatile float3* tempBuffer = (volatile float3*) &variables[9*blockDim.x];
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
#endif

    unsigned int lasty = 0xFFFFFFFF;
    while (pos < end)
    {
        // Extract cell coordinates from appropriate work unit
        unsigned int x = workUnit[pos];
        unsigned int y = ((x >> 2) & 0x7fff) << GRIDBITS;
        bool bExclusionFlag = (x & 0x1);
        x = (x >> 17) << GRIDBITS;
        float4      apos;   // Local atom x, y, z, q
        float3      af;     // Local atom fx, fy, fz
        unsigned int tgx = threadIdx.x & (GRID - 1);
        unsigned int tbx = threadIdx.x - tgx;
        unsigned int tj = tgx;
61
        volatile Atom* psA = &sA[tbx];
62
63
        unsigned int i      = x + tgx;
        apos                = cSim.pPosq[i];
64
        float4 params       = cSim.pCustomParams[i];
65
66
67
68
69
70
71
72
73
        af.x                = 0.0f;
        af.y                = 0.0f;
        af.z                = 0.0f;
        if (x == y) // Handle diagonals uniquely at 50% efficiency
        {
            // Read fixed atom data into registers and GRF
            sA[threadIdx.x].x   = apos.x;
            sA[threadIdx.x].y   = apos.y;
            sA[threadIdx.x].z   = apos.z;
74
75
76
77
            sA[threadIdx.x].params.x = params.x;
            sA[threadIdx.x].params.y = params.y;
            sA[threadIdx.x].params.z = params.z;
            sA[threadIdx.x].params.w = params.w;
78
79
80
81
            unsigned int xi   = x>>GRIDBITS;
            unsigned int cell = xi+xi*cSim.paddedNumberOfAtoms/GRID-xi*(xi+1)/2;
            unsigned int excl = cSim.pExclusion[cSim.pExclusionIndex[cell]+tgx];
            for (unsigned int j = 0; j < GRID; j++)
82
            {
83
                // Record the parameters.
84

85
86
87
88
                VARIABLE(0) = params.x;
                VARIABLE(1) = params.y;
                VARIABLE(2) = params.z;
                VARIABLE(3) = params.w;
89
90
91
92
                VARIABLE(4) = psA[j].params.x;
                VARIABLE(5) = psA[j].params.y;
                VARIABLE(6) = psA[j].params.z;
                VARIABLE(7) = psA[j].params.w;
93
94
95
96
97
98

                // Compute the force.

                float dx        = psA[j].x - apos.x;
                float dy        = psA[j].y - apos.y;
                float dz        = psA[j].z - apos.z;
99
#ifdef USE_PERIODIC
100
101
102
                dx -= floorf(dx*cSim.invPeriodicBoxSizeX+0.5f)*cSim.periodicBoxSizeX;
                dy -= floorf(dy*cSim.invPeriodicBoxSizeY+0.5f)*cSim.periodicBoxSizeY;
                dz -= floorf(dz*cSim.invPeriodicBoxSizeZ+0.5f)*cSim.periodicBoxSizeZ;
103
#endif
104
                float r         = sqrtf(dx*dx + dy*dy + dz*dz);
105
                float invR      = 1.0f/r;
106
                VARIABLE(8)     = r;
107
108
                float dEdR      = -kEvaluateExpression_kernel(&forceExp, stack, variables)*invR;
                float energy    = kEvaluateExpression_kernel(&energyExp, stack, variables);
109
#ifdef USE_CUTOFF
110
                if (!(excl & 0x1) || r > cSim.nonbondedCutoff)
111
#else
112
                if (!(excl & 0x1))
113
#endif
114
115
                {
                    dEdR = 0.0f;
116
                    energy = 0.0f;
117
                }
118
                totalEnergy    += 0.5f*energy;
119
120
121
122
123
124
125
                dx             *= dEdR;
                dy             *= dEdR;
                dz             *= dEdR;
                af.x           -= dx;
                af.y           -= dy;
                af.z           -= dz;
                excl          >>= 1;
126
127
128
129
            }

            // Write results
#ifdef USE_OUTPUT_BUFFER_PER_WARP
130
131
132
133
134
            unsigned int offset                 = x + tgx + warp*cSim.stride;
#else
            unsigned int offset                 = x + tgx + (x >> GRIDBITS) * cSim.stride;
#endif
            float4 of                           = cSim.pForce4[offset];
135
136
137
            of.x                               += af.x;
            of.y                               += af.y;
            of.z                               += af.z;
138
            cSim.pForce4[offset]               = of;
139

140
141
142
143
144
145
146
147
148
149
150
        }
        else        // 100% utilization
        {
            // Read fixed atom data into registers and GRF
            if (lasty != y)
            {
                unsigned int j                   = y + tgx;
                float4 temp             = cSim.pPosq[j];
                sA[threadIdx.x].x       = temp.x;
                sA[threadIdx.x].y       = temp.y;
                sA[threadIdx.x].z       = temp.z;
151
152
153
154
                sA[threadIdx.x].params.x = cSim.pCustomParams[j].x;
                sA[threadIdx.x].params.y = cSim.pCustomParams[j].y;
                sA[threadIdx.x].params.z = cSim.pCustomParams[j].z;
                sA[threadIdx.x].params.w = cSim.pCustomParams[j].w;
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
            }
            sA[threadIdx.x].fx      = 0.0f;
            sA[threadIdx.x].fy      = 0.0f;
            sA[threadIdx.x].fz      = 0.0f;
            if (!bExclusionFlag)
            {
#ifdef USE_CUTOFF
                unsigned int flags = cSim.pInteractionFlag[pos];
                if (flags == 0)
                {
                    // No interactions in this block.
                }
                else if (flags == 0xFFFFFFFF)
#endif
                {
                    // Compute all interactions within this block.

                    for (unsigned int j = 0; j < GRID; j++)
                    {
174
                        // Record the parameters.
175

176
177
178
179
                        VARIABLE(0) = params.x;
                        VARIABLE(1) = params.y;
                        VARIABLE(2) = params.z;
                        VARIABLE(3) = params.w;
180
181
182
183
                        VARIABLE(4) = psA[tj].params.x;
                        VARIABLE(5) = psA[tj].params.y;
                        VARIABLE(6) = psA[tj].params.z;
                        VARIABLE(7) = psA[tj].params.w;
184
185
186

                        // Compute the force.

187
188
189
190
                        float dx        = psA[tj].x - apos.x;
                        float dy        = psA[tj].y - apos.y;
                        float dz        = psA[tj].z - apos.z;
#ifdef USE_PERIODIC
191
192
193
                        dx -= floorf(dx*cSim.invPeriodicBoxSizeX+0.5f)*cSim.periodicBoxSizeX;
                        dy -= floorf(dy*cSim.invPeriodicBoxSizeY+0.5f)*cSim.periodicBoxSizeY;
                        dz -= floorf(dz*cSim.invPeriodicBoxSizeZ+0.5f)*cSim.periodicBoxSizeZ;
194
#endif
195
                        float r         = sqrtf(dx*dx + dy*dy + dz*dz);
196
                        float invR      = 1.0f/r;
197
                        VARIABLE(8)     = r;
198
199
                        float dEdR      = -kEvaluateExpression_kernel(&forceExp, stack, variables)*invR;
                        float energy    = kEvaluateExpression_kernel(&energyExp, stack, variables);
200
201
202
203
#ifdef USE_CUTOFF
                        if (r > cSim.nonbondedCutoff)
                        {
                            dEdR = 0.0f;
204
                            energy = 0.0f;
205
206
                        }
#endif
207
                        totalEnergy    += energy;
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
                        dx             *= dEdR;
                        dy             *= dEdR;
                        dz             *= dEdR;
                        af.x           -= dx;
                        af.y           -= dy;
                        af.z           -= dz;
                        psA[tj].fx     += dx;
                        psA[tj].fy     += dy;
                        psA[tj].fz     += dz;
                        tj              = (tj + 1) & (GRID - 1);
                    }
                }
#ifdef USE_CUTOFF
                else
                {
                    // Compute only a subset of the interactions in this block.

                    for (unsigned int j = 0; j < GRID; j++)
                    {
                        if ((flags&(1<<j)) != 0)
                        {
229
                            // Record the parameters.
230

231
232
233
234
                            VARIABLE(0) = params.x;
                            VARIABLE(1) = params.y;
                            VARIABLE(2) = params.z;
                            VARIABLE(3) = params.w;
235
236
237
238
                            VARIABLE(4) = psA[j].params.x;
                            VARIABLE(5) = psA[j].params.y;
                            VARIABLE(6) = psA[j].params.z;
                            VARIABLE(7) = psA[j].params.w;
239
240
241

                            // Compute the force.

242
243
244
245
                            float dx        = psA[j].x - apos.x;
                            float dy        = psA[j].y - apos.y;
                            float dz        = psA[j].z - apos.z;
#ifdef USE_PERIODIC
246
247
248
                            dx -= floorf(dx*cSim.invPeriodicBoxSizeX+0.5f)*cSim.periodicBoxSizeX;
                            dy -= floorf(dy*cSim.invPeriodicBoxSizeY+0.5f)*cSim.periodicBoxSizeY;
                            dz -= floorf(dz*cSim.invPeriodicBoxSizeZ+0.5f)*cSim.periodicBoxSizeZ;
249
#endif
250
                            float r         = sqrtf(dx*dx + dy*dy + dz*dz);
251
                            float invR      = 1.0f/r;
252
                            VARIABLE(8)     = r;
253
254
                            float dEdR      = -kEvaluateExpression_kernel(&forceExp, stack, variables)*invR;
                            float energy    = kEvaluateExpression_kernel(&energyExp, stack, variables);
255
256
257
258
#ifdef USE_CUTOFF
                            if (r > cSim.nonbondedCutoff)
                            {
                                dEdR = 0.0f;
259
                                energy = 0.0f;
260
261
                            }
#endif
262
                            totalEnergy    += energy;
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
                            dx             *= dEdR;
                            dy             *= dEdR;
                            dz             *= dEdR;
                            af.x           -= dx;
                            af.y           -= dy;
                            af.z           -= dz;
                            tempBuffer[threadIdx.x].x = dx;
                            tempBuffer[threadIdx.x].y = dy;
                            tempBuffer[threadIdx.x].z = dz;

                            // Sum the forces on atom j.

                            if (tgx % 2 == 0)
                            {
                                tempBuffer[threadIdx.x].x += tempBuffer[threadIdx.x+1].x;
                                tempBuffer[threadIdx.x].y += tempBuffer[threadIdx.x+1].y;
                                tempBuffer[threadIdx.x].z += tempBuffer[threadIdx.x+1].z;
                            }
                            if (tgx % 4 == 0)
                            {
                                tempBuffer[threadIdx.x].x += tempBuffer[threadIdx.x+2].x;
                                tempBuffer[threadIdx.x].y += tempBuffer[threadIdx.x+2].y;
                                tempBuffer[threadIdx.x].z += tempBuffer[threadIdx.x+2].z;
                            }
                            if (tgx % 8 == 0)
                            {
                                tempBuffer[threadIdx.x].x += tempBuffer[threadIdx.x+4].x;
                                tempBuffer[threadIdx.x].y += tempBuffer[threadIdx.x+4].y;
                                tempBuffer[threadIdx.x].z += tempBuffer[threadIdx.x+4].z;
                            }
                            if (tgx % 16 == 0)
                            {
                                tempBuffer[threadIdx.x].x += tempBuffer[threadIdx.x+8].x;
                                tempBuffer[threadIdx.x].y += tempBuffer[threadIdx.x+8].y;
                                tempBuffer[threadIdx.x].z += tempBuffer[threadIdx.x+8].z;
                            }
                            if (tgx == 0)
                            {
                                psA[j].fx += tempBuffer[threadIdx.x].x + tempBuffer[threadIdx.x+16].x;
                                psA[j].fy += tempBuffer[threadIdx.x].y + tempBuffer[threadIdx.x+16].y;
                                psA[j].fz += tempBuffer[threadIdx.x].z + tempBuffer[threadIdx.x+16].z;
                            }
                        }
                    }
                }
#endif
            }
            else  // bExclusion
            {
                // Read fixed atom data into registers and GRF
                unsigned int xi   = x>>GRIDBITS;
                unsigned int yi   = y>>GRIDBITS;
                unsigned int cell          = xi+yi*cSim.paddedNumberOfAtoms/GRID-yi*(yi+1)/2;
                unsigned int excl = cSim.pExclusion[cSim.pExclusionIndex[cell]+tgx];
                excl              = (excl >> tgx) | (excl << (GRID - tgx));
                for (unsigned int j = 0; j < GRID; j++)
                {
320
                    // Record the parameters.
321

322
323
324
325
                    VARIABLE(0) = params.x;
                    VARIABLE(1) = params.y;
                    VARIABLE(2) = params.z;
                    VARIABLE(3) = params.w;
326
327
328
329
                    VARIABLE(4) = psA[tj].params.x;
                    VARIABLE(5) = psA[tj].params.y;
                    VARIABLE(6) = psA[tj].params.z;
                    VARIABLE(7) = psA[tj].params.w;
330
331
332

                    // Compute the force.

333
334
335
336
                    float dx        = psA[tj].x - apos.x;
                    float dy        = psA[tj].y - apos.y;
                    float dz        = psA[tj].z - apos.z;
#ifdef USE_PERIODIC
337
338
339
                    dx -= floorf(dx*cSim.invPeriodicBoxSizeX+0.5f)*cSim.periodicBoxSizeX;
                    dy -= floorf(dy*cSim.invPeriodicBoxSizeY+0.5f)*cSim.periodicBoxSizeY;
                    dz -= floorf(dz*cSim.invPeriodicBoxSizeZ+0.5f)*cSim.periodicBoxSizeZ;
340
#endif
341
                    float r         = sqrtf(dx*dx + dy*dy + dz*dz);
342
                    float invR      = 1.0f/r;
343
                    VARIABLE(8)     = r;
344
345
                    float dEdR      = -kEvaluateExpression_kernel(&forceExp, stack, variables)*invR;
                    float energy    = kEvaluateExpression_kernel(&energyExp, stack, variables);
346
347
348
349
350
351
352
#ifdef USE_CUTOFF
                    if (!(excl & 0x1) || r > cSim.nonbondedCutoff)
#else
                    if (!(excl & 0x1))
#endif
                    {
                        dEdR = 0.0f;
353
                        energy = 0.0f;
354
                    }
355
                    totalEnergy    += energy;
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
                    dx             *= dEdR;
                    dy             *= dEdR;
                    dz             *= dEdR;
                    af.x           -= dx;
                    af.y           -= dy;
                    af.z           -= dz;
                    psA[tj].fx     += dx;
                    psA[tj].fy     += dy;
                    psA[tj].fz     += dz;
                    excl          >>= 1;
                    tj              = (tj + 1) & (GRID - 1);
                }
            }

            // Write results
            float4 of;
#ifdef USE_OUTPUT_BUFFER_PER_WARP
            unsigned int offset                 = x + tgx + warp*cSim.stride;
374
375
376
#else
            unsigned int offset                 = x + tgx + (y >> GRIDBITS) * cSim.stride;
#endif
377
            of                                  = cSim.pForce4[offset];
378
379
380
            of.x                               += af.x;
            of.y                               += af.y;
            of.z                               += af.z;
381
            cSim.pForce4[offset]               = of;
382
#ifdef USE_OUTPUT_BUFFER_PER_WARP
383
            offset                              = y + tgx + warp*cSim.stride;
384
385
386
#else
            offset                              = y + tgx + (x >> GRIDBITS) * cSim.stride;
#endif
387
            of                                  = cSim.pForce4[offset];
388
389
390
            of.x                               += sA[threadIdx.x].fx;
            of.y                               += sA[threadIdx.x].fy;
            of.z                               += sA[threadIdx.x].fz;
391
            cSim.pForce4[offset]               = of;
392
393
394
395
396
            lasty = y;
        }

        pos++;
    }
397
    cSim.pEnergy[blockIdx.x*blockDim.x+threadIdx.x] += totalEnergy;
398
}