gbsaObc_nvidia.cl 33 KB
Newer Older
1
#pragma OPENCL EXTENSION cl_khr_global_int32_base_atomics : enable
2
3
4
#ifdef SUPPORTS_64_BIT_ATOMICS
#pragma OPENCL EXTENSION cl_khr_int64_base_atomics : enable
#endif
5
#define TILE_SIZE 32
6
#define WARPS_PER_GROUP (FORCE_WORK_GROUP_SIZE/TILE_SIZE)
7

8
9
10
11
12
typedef struct {
    float x, y, z;
    float q;
    float radius, scaledRadius;
    float bornSum;
Peter Eastman's avatar
Peter Eastman committed
13
} AtomData1;
14

15
16
17
/**
 * Compute the Born sum.
 */
18
19
__kernel void computeBornSum(
#ifdef SUPPORTS_64_BIT_ATOMICS
20
        __global long* restrict global_bornSum,
21
#else
22
        __global float* restrict global_bornSum,
23
#endif
24
        __global const float4* restrict posq, __global const float2* restrict global_params,
25
        __local AtomData1* restrict localData,
26
#ifdef USE_CUTOFF
27
        __global const ushort2* restrict tiles, __global const unsigned int* restrict interactionCount, float4 periodicBoxSize, float4 invPeriodicBoxSize, unsigned int maxTiles, __global const unsigned int* restrict interactionFlags,
28
#else
Peter Eastman's avatar
Peter Eastman committed
29
        unsigned int numTiles,
30
#endif
Peter Eastman's avatar
Peter Eastman committed
31
        __global unsigned int* exclusionIndices, __global unsigned int* exclusionRowIndices) {
32
33
    unsigned int totalWarps = get_global_size(0)/TILE_SIZE;
    unsigned int warp = get_global_id(0)/TILE_SIZE;
34
35
#ifdef USE_CUTOFF
    unsigned int numTiles = interactionCount[0];
36
37
    unsigned int pos = warp*(numTiles > maxTiles ? NUM_BLOCKS*(NUM_BLOCKS+1)/2 : numTiles)/totalWarps;
    unsigned int end = (warp+1)*(numTiles > maxTiles ? NUM_BLOCKS*(NUM_BLOCKS+1)/2 : numTiles)/totalWarps;
38
#else
39
40
    unsigned int pos = warp*numTiles/totalWarps;
    unsigned int end = (warp+1)*numTiles/totalWarps;
41
#endif
42
    unsigned int lasty = 0xFFFFFFFF;
43
    __local float tempBuffer[FORCE_WORK_GROUP_SIZE];
44
    __local int2 reservedBlocks[WARPS_PER_GROUP];
Peter Eastman's avatar
Peter Eastman committed
45
46
    __local unsigned int* exclusionRange = (__local unsigned int*) reservedBlocks;
    __local int exclusionIndex[WARPS_PER_GROUP];
47
48
    
    do {
49
        // Extract the coordinates of this tile
50
51
52
        const unsigned int tgx = get_local_id(0) & (TILE_SIZE-1);
        const unsigned int tbx = get_local_id(0) - tgx;
        const unsigned int localGroupIndex = get_local_id(0)/TILE_SIZE;
53
        unsigned int x, y;
54
55
        float bornSum = 0.0f;
        if (pos < end) {
56
#ifdef USE_CUTOFF
57
58
59
60
61
62
            if (numTiles <= maxTiles) {
                ushort2 tileIndices = tiles[pos];
                x = tileIndices.x;
                y = tileIndices.y;
            }
            else
63
#endif
64
65
            {
                y = (unsigned int) floor(NUM_BLOCKS+0.5f-sqrt((NUM_BLOCKS+0.5f)*(NUM_BLOCKS+0.5f)-2*pos));
66
                x = (pos-y*NUM_BLOCKS+y*(y+1)/2);
67
68
69
70
                if (x < y || x >= NUM_BLOCKS) { // Occasionally happens due to roundoff error.
                    y += (x < y ? -1 : 1);
                    x = (pos-y*NUM_BLOCKS+y*(y+1)/2);
                }
71
            }
72
73
74
75
76
77
78
            unsigned int atom1 = x*TILE_SIZE + tgx;
            float4 posq1 = posq[atom1];
            float2 params1 = global_params[atom1];
            if (pos >= end)
                ; // This warp is done.
            else if (x == y) {
                // This tile is on the diagonal.
79

80
81
82
83
84
85
86
87
                localData[get_local_id(0)].x = posq1.x;
                localData[get_local_id(0)].y = posq1.y;
                localData[get_local_id(0)].z = posq1.z;
                localData[get_local_id(0)].q = posq1.w;
                localData[get_local_id(0)].radius = params1.x;
                localData[get_local_id(0)].scaledRadius = params1.y;
                for (unsigned int j = 0; j < TILE_SIZE; j++) {
                    float4 delta = (float4) (localData[tbx+j].x-posq1.x, localData[tbx+j].y-posq1.y, localData[tbx+j].z-posq1.z, 0.0f);
88
#ifdef USE_PERIODIC
89
90
91
                    delta.x -= floor(delta.x*invPeriodicBoxSize.x+0.5f)*periodicBoxSize.x;
                    delta.y -= floor(delta.y*invPeriodicBoxSize.y+0.5f)*periodicBoxSize.y;
                    delta.z -= floor(delta.z*invPeriodicBoxSize.z+0.5f)*periodicBoxSize.z;
92
#endif
93
                    float r2 = delta.x*delta.x + delta.y*delta.y + delta.z*delta.z;
94
#ifdef USE_CUTOFF
95
                    if (atom1 < NUM_ATOMS && y*TILE_SIZE+j < NUM_ATOMS && r2 < CUTOFF_SQUARED) {
96
#else
97
                    if (atom1 < NUM_ATOMS && y*TILE_SIZE+j < NUM_ATOMS) {
98
#endif
99
100
101
102
103
104
105
106
107
108
                        float invR = RSQRT(r2);
                        float r = RECIP(invR);
                        float2 params2 = (float2) (localData[tbx+j].radius, localData[tbx+j].scaledRadius);
                        float rScaledRadiusJ = r+params2.y;
                        if ((j != tgx) && (params1.x < rScaledRadiusJ)) {
                            float l_ij = RECIP(max(params1.x, fabs(r-params2.y)));
                            float u_ij = RECIP(rScaledRadiusJ);
                            float l_ij2 = l_ij*l_ij;
                            float u_ij2 = u_ij*u_ij;
                            float ratio = LOG(u_ij * RECIP(l_ij));
Peter Eastman's avatar
Peter Eastman committed
109
110
                            bornSum += l_ij - u_ij + (0.50f*invR*ratio) + 0.25f*(r*(u_ij2-l_ij2) +
                                             (params2.y*params2.y*invR)*(l_ij2-u_ij2));
111
112
113
                            if (params1.x < params2.x-r)
                                bornSum += 2.0f*(RECIP(params1.x)-l_ij);
                        }
114
115
116
                    }
                }
            }
117
118
            else {
                // This is an off-diagonal tile.
119

120
121
122
123
124
125
126
127
128
129
                if (lasty != y) {
                    unsigned int j = y*TILE_SIZE + tgx;
                    float4 tempPosq = posq[j];
                    localData[get_local_id(0)].x = tempPosq.x;
                    localData[get_local_id(0)].y = tempPosq.y;
                    localData[get_local_id(0)].z = tempPosq.z;
                    localData[get_local_id(0)].q = tempPosq.w;
                    float2 tempParams = global_params[j];
                    localData[get_local_id(0)].radius = tempParams.x;
                    localData[get_local_id(0)].scaledRadius = tempParams.y;
130
                }
131
132
133
                localData[get_local_id(0)].bornSum = 0.0f;
#ifdef USE_CUTOFF
                unsigned int flags = (numTiles <= maxTiles ? interactionFlags[pos] : 0xFFFFFFFF);
Peter Eastman's avatar
Peter Eastman committed
134
135
136
137
138
139
140
141
142
143
144
145
                bool computeSubset = false;
                if (flags != 0xFFFFFFFF) {
                    if (tgx < 2)
                        exclusionRange[2*localGroupIndex+tgx] = exclusionRowIndices[x+tgx];
                    if (tgx == 0)
                        exclusionIndex[localGroupIndex] = -1;
                    for (int i = exclusionRange[2*localGroupIndex]+tgx; i < exclusionRange[2*localGroupIndex+1]; i += TILE_SIZE)
                        if (exclusionIndices[i] == y)
                            exclusionIndex[localGroupIndex] = i*TILE_SIZE;
                    computeSubset = (exclusionIndex[localGroupIndex] == -1);
                }
                if (computeSubset) {
146
147
148
149
150
                    if (flags == 0) {
                        // No interactions in this tile.
                    }
                    else {
                        // Compute only a subset of the interactions in this tile.
151

152
153
154
                        for (unsigned int j = 0; j < TILE_SIZE; j++) {
                            if ((flags&(1<<j)) != 0) {
                                float4 delta = (float4) (localData[tbx+j].x-posq1.x, localData[tbx+j].y-posq1.y, localData[tbx+j].z-posq1.z, 0.0f);
155
#ifdef USE_PERIODIC
156
157
158
                                delta.x -= floor(delta.x*invPeriodicBoxSize.x+0.5f)*periodicBoxSize.x;
                                delta.y -= floor(delta.y*invPeriodicBoxSize.y+0.5f)*periodicBoxSize.y;
                                delta.z -= floor(delta.z*invPeriodicBoxSize.z+0.5f)*periodicBoxSize.z;
159
#endif
160
161
                                float r2 = delta.x*delta.x + delta.y*delta.y + delta.z*delta.z;
                                tempBuffer[get_local_id(0)] = 0.0f;
162
#ifdef USE_CUTOFF
163
                                if (atom1 < NUM_ATOMS && y*TILE_SIZE+j < NUM_ATOMS && r2 < CUTOFF_SQUARED) {
164
#else
165
                                if (atom1 < NUM_ATOMS && y*TILE_SIZE+j < NUM_ATOMS) {
166
#endif
167
168
169
170
171
172
173
174
175
176
                                    float invR = RSQRT(r2);
                                    float r = RECIP(invR);
                                    float2 params2 = (float2) (localData[tbx+j].radius, localData[tbx+j].scaledRadius);
                                    float rScaledRadiusJ = r+params2.y;
                                    if (params1.x < rScaledRadiusJ) {
                                        float l_ij = RECIP(max(params1.x, fabs(r-params2.y)));
                                        float u_ij = RECIP(rScaledRadiusJ);
                                        float l_ij2 = l_ij*l_ij;
                                        float u_ij2 = u_ij*u_ij;
                                        float ratio = LOG(u_ij * RECIP(l_ij));
Peter Eastman's avatar
Peter Eastman committed
177
178
                                        bornSum += l_ij - u_ij + (0.50f*invR*ratio) + 0.25f*(r*(u_ij2-l_ij2) +
                                                         (params2.y*params2.y*invR)*(l_ij2-u_ij2));
179
180
181
182
183
184
185
186
187
188
                                        if (params1.x < params2.x-r)
                                            bornSum += 2.0f*(RECIP(params1.x)-l_ij);
                                    }
                                    float rScaledRadiusI = r+params1.y;
                                    if (params2.x < rScaledRadiusI) {
                                        float l_ij = RECIP(max(params2.x, fabs(r-params1.y)));
                                        float u_ij = RECIP(rScaledRadiusI);
                                        float l_ij2 = l_ij*l_ij;
                                        float u_ij2 = u_ij*u_ij;
                                        float ratio = LOG(u_ij * RECIP(l_ij));
Peter Eastman's avatar
Peter Eastman committed
189
190
                                        float term = l_ij - u_ij + (0.50f*invR*ratio) + 0.25f*(r*(u_ij2-l_ij2) +
                                                         (params1.y*params1.y*invR)*(l_ij2-u_ij2));
191
192
193
194
                                        if (params2.x < params1.x-r)
                                            term += 2.0f*(RECIP(params2.x)-l_ij);
                                        tempBuffer[get_local_id(0)] = term;
                                    }
195
196
                                }

197
                                // Sum the forces on atom j.
198

199
                                if (tgx % 4 == 0)
200
                                    tempBuffer[get_local_id(0)] += tempBuffer[get_local_id(0)+1]+tempBuffer[get_local_id(0)+2]+tempBuffer[get_local_id(0)+3];
201
                                if (tgx == 0)
202
                                    localData[tbx+j].bornSum += tempBuffer[get_local_id(0)]+tempBuffer[get_local_id(0)+4]+tempBuffer[get_local_id(0)+8]+tempBuffer[get_local_id(0)+12]+tempBuffer[get_local_id(0)+16]+tempBuffer[get_local_id(0)+20]+tempBuffer[get_local_id(0)+24]+tempBuffer[get_local_id(0)+28];
203
                            }
204
205
206
                        }
                    }
                }
207
                else
208
#endif
209
210
                {
                    // Compute the full set of interactions in this tile.
211

212
213
214
                    unsigned int tj = tgx;
                    for (unsigned int j = 0; j < TILE_SIZE; j++) {
                        float4 delta = (float4) (localData[tbx+tj].x-posq1.x, localData[tbx+tj].y-posq1.y, localData[tbx+tj].z-posq1.z, 0.0f);
215
#ifdef USE_PERIODIC
216
217
218
                        delta.x -= floor(delta.x*invPeriodicBoxSize.x+0.5f)*periodicBoxSize.x;
                        delta.y -= floor(delta.y*invPeriodicBoxSize.y+0.5f)*periodicBoxSize.y;
                        delta.z -= floor(delta.z*invPeriodicBoxSize.z+0.5f)*periodicBoxSize.z;
219
#endif
220
                        float r2 = delta.x*delta.x + delta.y*delta.y + delta.z*delta.z;
221
#ifdef USE_CUTOFF
222
                        if (atom1 < NUM_ATOMS && y*TILE_SIZE+tj < NUM_ATOMS && r2 < CUTOFF_SQUARED) {
223
#else
224
                        if (atom1 < NUM_ATOMS && y*TILE_SIZE+tj < NUM_ATOMS) {
225
#endif
226
227
228
229
230
231
232
233
234
235
                            float invR = RSQRT(r2);
                            float r = RECIP(invR);
                            float2 params2 = (float2) (localData[tbx+tj].radius, localData[tbx+tj].scaledRadius);
                            float rScaledRadiusJ = r+params2.y;
                            if (params1.x < rScaledRadiusJ) {
                                float l_ij = RECIP(max(params1.x, fabs(r-params2.y)));
                                float u_ij = RECIP(rScaledRadiusJ);
                                float l_ij2 = l_ij*l_ij;
                                float u_ij2 = u_ij*u_ij;
                                float ratio = LOG(u_ij * RECIP(l_ij));
Peter Eastman's avatar
Peter Eastman committed
236
237
                                bornSum += l_ij - u_ij + (0.50f*invR*ratio) + 0.25f*(r*(u_ij2-l_ij2) +
                                                 (params2.y*params2.y*invR)*(l_ij2-u_ij2));
238
239
240
241
242
243
244
245
246
247
                                if (params1.x < params2.x-r)
                                    bornSum += 2.0f*(RECIP(params1.x)-l_ij);
                            }
                            float rScaledRadiusI = r+params1.y;
                            if (params2.x < rScaledRadiusI) {
                                float l_ij = RECIP(max(params2.x, fabs(r-params1.y)));
                                float u_ij = RECIP(rScaledRadiusI);
                                float l_ij2 = l_ij*l_ij;
                                float u_ij2 = u_ij*u_ij;
                                float ratio = LOG(u_ij * RECIP(l_ij));
Peter Eastman's avatar
Peter Eastman committed
248
249
                                float term = l_ij - u_ij + (0.50f*invR*ratio) + 0.25f*(r*(u_ij2-l_ij2) +
                                                 (params1.y*params1.y*invR)*(l_ij2-u_ij2));
250
251
252
253
                                if (params2.x < params1.x-r)
                                    term += 2.0f*(RECIP(params2.x)-l_ij);
                                localData[tbx+tj].bornSum += term;
                            }
254
                        }
255
                        tj = (tj + 1) & (TILE_SIZE - 1);
256
257
258
                    }
                }
            }
259
260
261
262
263
        }
        
        // Write results.  We need to coordinate between warps to make sure no two of them
        // ever try to write to the same piece of memory at the same time.
        
264
265
266
267
268
269
270
271
272
273
#ifdef SUPPORTS_64_BIT_ATOMICS
        if (pos < end) {
            const unsigned int offset = x*TILE_SIZE + tgx;
            atom_add(&global_bornSum[offset], (long) (bornSum*0xFFFFFFFF));
        }
        if (pos < end && x != y) {
            const unsigned int offset = y*TILE_SIZE + tgx;
            atom_add(&global_bornSum[offset], (long) (localData[get_local_id(0)].bornSum*0xFFFFFFFF));
        }
#else
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
        int writeX = (pos < end ? x : -1);
        int writeY = (pos < end && x != y ? y : -1);
        if (tgx == 0)
            reservedBlocks[localGroupIndex] = (int2)(writeX, writeY);
        bool done = false;
        int doneIndex = 0;
        int checkIndex = 0;
        while (true) {
            // See if any warp still needs to write its data.

            bool allDone = true;
            barrier(CLK_LOCAL_MEM_FENCE);
            while (doneIndex < WARPS_PER_GROUP && allDone) {
                if (reservedBlocks[doneIndex].x != -1)
                    allDone = false;
                else
                    doneIndex++;
            }
            if (allDone)
                break;
            if (!done) {
                // See whether this warp can write its data.  This requires that no previous warp
                // is trying to write to the same block of the buffer.
297

298
299
300
301
302
303
304
305
306
307
                bool canWrite = (writeX != -1);
                while (checkIndex < localGroupIndex && canWrite) {
                    if ((reservedBlocks[checkIndex].x == x || reservedBlocks[checkIndex].y == x) ||
                            (writeY != -1 && (reservedBlocks[checkIndex].x == y || reservedBlocks[checkIndex].y == y)))
                        canWrite = false;
                    else
                        checkIndex++;
                }
                if (canWrite) {
                    // Write the data to global memory, then mark this warp as done.
308

309
310
311
312
313
314
315
316
317
318
319
320
321
                    if (writeX > -1) {
                        const unsigned int offset = x*TILE_SIZE + tgx + get_group_id(0)*PADDED_NUM_ATOMS;
                        global_bornSum[offset] += bornSum;
                    }
                    if (writeY > -1) {
                        const unsigned int offset = y*TILE_SIZE + tgx + get_group_id(0)*PADDED_NUM_ATOMS;
                        global_bornSum[offset] += localData[get_local_id(0)].bornSum;
                    }
                    done = true;
                    if (tgx == 0)
                        reservedBlocks[localGroupIndex] = (int2)(-1, -1);
                }
            }
322
        }
323
#endif
324
        lasty = y;
325
        pos++;
326
    } while (pos < end);
327
328
}

Peter Eastman's avatar
Peter Eastman committed
329
330
331
332
333
334
335
typedef struct {
    float x, y, z;
    float q;
    float fx, fy, fz, fw;
    float bornRadius;
} AtomData2;

336
337
338
339
/**
 * First part of computing the GBSA interaction.
 */

340
341
__kernel void computeGBSAForce1(
#ifdef SUPPORTS_64_BIT_ATOMICS
342
        __global long* restrict forceBuffers, __global long* restrict global_bornForce,
343
#else
344
        __global float4* restrict forceBuffers, __global float* restrict global_bornForce,
345
#endif
346
        __global float* restrict energyBuffer, __global const float4* restrict posq, __global const float* restrict global_bornRadii,
347
        __local AtomData2* restrict localData,
348
#ifdef USE_CUTOFF
349
        __global const ushort2* restrict tiles, __global const unsigned int* restrict interactionCount, float4 periodicBoxSize, float4 invPeriodicBoxSize, unsigned int maxTiles, __global const unsigned int* restrict interactionFlags,
350
#else
Peter Eastman's avatar
Peter Eastman committed
351
        unsigned int numTiles,
352
#endif
Peter Eastman's avatar
Peter Eastman committed
353
        __global unsigned int* exclusionIndices, __global unsigned int* exclusionRowIndices) {
354
355
    unsigned int totalWarps = get_global_size(0)/TILE_SIZE;
    unsigned int warp = get_global_id(0)/TILE_SIZE;
356
357
#ifdef USE_CUTOFF
    unsigned int numTiles = interactionCount[0];
358
359
    unsigned int pos = warp*(numTiles > maxTiles ? NUM_BLOCKS*(NUM_BLOCKS+1)/2 : numTiles)/totalWarps;
    unsigned int end = (warp+1)*(numTiles > maxTiles ? NUM_BLOCKS*(NUM_BLOCKS+1)/2 : numTiles)/totalWarps;
360
#else
361
362
    unsigned int pos = warp*numTiles/totalWarps;
    unsigned int end = (warp+1)*numTiles/totalWarps;
363
#endif
364
365
    float energy = 0.0f;
    unsigned int lasty = 0xFFFFFFFF;
366
    __local float4 tempBuffer[FORCE_WORK_GROUP_SIZE];
367
    __local int2 reservedBlocks[WARPS_PER_GROUP];
Peter Eastman's avatar
Peter Eastman committed
368
369
    __local unsigned int* exclusionRange = (__local unsigned int*) reservedBlocks;
    __local int exclusionIndex[WARPS_PER_GROUP];
370
371
    
    do {
372
        // Extract the coordinates of this tile
373
374
375
        const unsigned int tgx = get_local_id(0) & (TILE_SIZE-1);
        const unsigned int tbx = get_local_id(0) - tgx;
        const unsigned int localGroupIndex = get_local_id(0)/TILE_SIZE;
376
        unsigned int x, y;
377
378
        float4 force = 0.0f;
        if (pos < end) {
379
#ifdef USE_CUTOFF
380
381
382
383
384
385
            if (numTiles <= maxTiles) {
                ushort2 tileIndices = tiles[pos];
                x = tileIndices.x;
                y = tileIndices.y;
            }
            else
386
#endif
387
388
            {
                y = (unsigned int) floor(NUM_BLOCKS+0.5f-sqrt((NUM_BLOCKS+0.5f)*(NUM_BLOCKS+0.5f)-2*pos));
389
                x = (pos-y*NUM_BLOCKS+y*(y+1)/2);
390
391
392
393
                if (x < y || x >= NUM_BLOCKS) { // Occasionally happens due to roundoff error.
                    y += (x < y ? -1 : 1);
                    x = (pos-y*NUM_BLOCKS+y*(y+1)/2);
                }
394
            }
395
396
397
398
399
            unsigned int atom1 = x*TILE_SIZE + tgx;
            float4 posq1 = posq[atom1];
            float bornRadius1 = global_bornRadii[atom1];
            if (x == y) {
                // This tile is on the diagonal.
400

401
402
403
404
405
406
407
408
409
                localData[get_local_id(0)].x = posq1.x;
                localData[get_local_id(0)].y = posq1.y;
                localData[get_local_id(0)].z = posq1.z;
                localData[get_local_id(0)].q = posq1.w;
                localData[get_local_id(0)].bornRadius = bornRadius1;
                for (unsigned int j = 0; j < TILE_SIZE; j++) {
                    if (atom1 < NUM_ATOMS && y*TILE_SIZE+j < NUM_ATOMS) {
                        float4 posq2 = (float4) (localData[tbx+j].x, localData[tbx+j].y, localData[tbx+j].z, localData[tbx+j].q);
                        float4 delta = (float4) (posq2.xyz - posq1.xyz, 0.0f);
410
#ifdef USE_PERIODIC
411
412
413
                        delta.x -= floor(delta.x*invPeriodicBoxSize.x+0.5f)*periodicBoxSize.x;
                        delta.y -= floor(delta.y*invPeriodicBoxSize.y+0.5f)*periodicBoxSize.y;
                        delta.z -= floor(delta.z*invPeriodicBoxSize.z+0.5f)*periodicBoxSize.z;
414
#endif
415
                        float r2 = delta.x*delta.x + delta.y*delta.y + delta.z*delta.z;
Peter Eastman's avatar
Peter Eastman committed
416
417
418
#ifdef USE_CUTOFF
                        if (r2 < CUTOFF_SQUARED) {
#endif
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
                        float invR = RSQRT(r2);
                        float r = RECIP(invR);
                        float bornRadius2 = localData[tbx+j].bornRadius;
                        float alpha2_ij = bornRadius1*bornRadius2;
                        float D_ij = r2*RECIP(4.0f*alpha2_ij);
                        float expTerm = EXP(-D_ij);
                        float denominator2 = r2 + alpha2_ij*expTerm;
                        float denominator = SQRT(denominator2);
                        float tempEnergy = (PREFACTOR*posq1.w*posq2.w)*RECIP(denominator);
                        float Gpol = tempEnergy*RECIP(denominator2);
                        float dGpol_dalpha2_ij = -0.5f*Gpol*expTerm*(1.0f+D_ij);
                        float dEdR = Gpol*(1.0f - 0.25f*expTerm);
                        force.w += dGpol_dalpha2_ij*bornRadius2;
                        energy += 0.5f*tempEnergy;
                        delta.xyz *= dEdR;
                        force.xyz -= delta.xyz;
Peter Eastman's avatar
Peter Eastman committed
435
436
437
#ifdef USE_CUTOFF
                        }
#endif
438
                    }
439
440
                }
            }
441
442
            else {
                // This is an off-diagonal tile.
443

444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
                if (lasty != y) {
                    unsigned int j = y*TILE_SIZE + tgx;
                    float4 tempPosq = posq[j];
                    localData[get_local_id(0)].x = tempPosq.x;
                    localData[get_local_id(0)].y = tempPosq.y;
                    localData[get_local_id(0)].z = tempPosq.z;
                    localData[get_local_id(0)].q = tempPosq.w;
                    localData[get_local_id(0)].bornRadius = global_bornRadii[j];
                }
                localData[get_local_id(0)].fx = 0.0f;
                localData[get_local_id(0)].fy = 0.0f;
                localData[get_local_id(0)].fz = 0.0f;
                localData[get_local_id(0)].fw = 0.0f;
#ifdef USE_CUTOFF
                unsigned int flags = (numTiles <= maxTiles ? interactionFlags[pos] : 0xFFFFFFFF);
Peter Eastman's avatar
Peter Eastman committed
459
460
461
462
463
464
465
466
467
468
469
470
                bool computeSubset = false;
                if (flags != 0xFFFFFFFF) {
                    if (tgx < 2)
                        exclusionRange[2*localGroupIndex+tgx] = exclusionRowIndices[x+tgx];
                    if (tgx == 0)
                        exclusionIndex[localGroupIndex] = -1;
                    for (int i = exclusionRange[2*localGroupIndex]+tgx; i < exclusionRange[2*localGroupIndex+1]; i += TILE_SIZE)
                        if (exclusionIndices[i] == y)
                            exclusionIndex[localGroupIndex] = i*TILE_SIZE;
                    computeSubset = (exclusionIndex[localGroupIndex] == -1);
                }
                if (computeSubset) {
471
472
473
474
475
                    if (flags == 0) {
                        // No interactions in this tile.
                    }
                    else {
                        // Compute only a subset of the interactions in this tile.
476

477
478
479
480
481
482
483
484
485
486
                        for (unsigned int j = 0; j < TILE_SIZE; j++) {
                            if ((flags&(1<<j)) != 0) {
                                float4 posq2 = (float4) (localData[tbx+j].x, localData[tbx+j].y, localData[tbx+j].z, localData[tbx+j].q);
                                float4 delta = (float4) (posq2.xyz - posq1.xyz, 0.0f);
#ifdef USE_PERIODIC
                                delta.x -= floor(delta.x*invPeriodicBoxSize.x+0.5f)*periodicBoxSize.x;
                                delta.y -= floor(delta.y*invPeriodicBoxSize.y+0.5f)*periodicBoxSize.y;
                                delta.z -= floor(delta.z*invPeriodicBoxSize.z+0.5f)*periodicBoxSize.z;
#endif
                                float r2 = delta.x*delta.x + delta.y*delta.y + delta.z*delta.z;
Peter Eastman's avatar
Peter Eastman committed
487
488
489
#ifdef USE_CUTOFF
                                if (r2 < CUTOFF_SQUARED) {
#endif
490
491
492
493
494
495
496
497
498
499
500
501
                                float invR = RSQRT(r2);
                                float r = RECIP(invR);
                                float bornRadius2 = localData[tbx+j].bornRadius;
                                float alpha2_ij = bornRadius1*bornRadius2;
                                float D_ij = r2*RECIP(4.0f*alpha2_ij);
                                float expTerm = EXP(-D_ij);
                                float denominator2 = r2 + alpha2_ij*expTerm;
                                float denominator = SQRT(denominator2);
                                float tempEnergy = (PREFACTOR*posq1.w*posq2.w)*RECIP(denominator);
                                float Gpol = tempEnergy*RECIP(denominator2);
                                float dGpol_dalpha2_ij = -0.5f*Gpol*expTerm*(1.0f+D_ij);
                                float dEdR = Gpol*(1.0f - 0.25f*expTerm);
502
#ifdef USE_CUTOFF
503
504
505
506
507
508
509
510
511
512
513
514
515
                                if (atom1 >= NUM_ATOMS || y*TILE_SIZE+j >= NUM_ATOMS || r2 > CUTOFF_SQUARED) {
#else
                                if (atom1 >= NUM_ATOMS || y*TILE_SIZE+j >= NUM_ATOMS) {
#endif
                                    dEdR = 0.0f;
                                    dGpol_dalpha2_ij = 0.0f;
                                    tempEnergy = 0.0f;
                                }
                                energy += tempEnergy;
                                force.w += dGpol_dalpha2_ij*bornRadius2;
                                delta.xyz *= dEdR;
                                force.xyz -= delta.xyz;
                                tempBuffer[get_local_id(0)] = (float4) (delta.xyz, dGpol_dalpha2_ij*bornRadius1);
Peter Eastman's avatar
Peter Eastman committed
516
517
518
519
520
#ifdef USE_CUTOFF
                                }
                                else
                                    tempBuffer[get_local_id(0)] = (float4) 0.0f;
#endif
521
522
523
524

                                // Sum the forces on atom j.

                                if (tgx % 4 == 0)
525
                                    tempBuffer[get_local_id(0)] += tempBuffer[get_local_id(0)+1]+tempBuffer[get_local_id(0)+2]+tempBuffer[get_local_id(0)+3];
526
                                if (tgx == 0) {
527
                                    float4 sum = tempBuffer[get_local_id(0)]+tempBuffer[get_local_id(0)+4]+tempBuffer[get_local_id(0)+8]+tempBuffer[get_local_id(0)+12]+tempBuffer[get_local_id(0)+16]+tempBuffer[get_local_id(0)+20]+tempBuffer[get_local_id(0)+24]+tempBuffer[get_local_id(0)+28];
528
529
530
531
532
533
534
535
                                    localData[tbx+j].fx += sum.x;
                                    localData[tbx+j].fy += sum.y;
                                    localData[tbx+j].fz += sum.z;
                                    localData[tbx+j].fw += sum.w;
                                }
                            }
                        }
                    }
536
                }
537
538
539
540
                else
#endif
                {
                    // Compute the full set of interactions in this tile.
541

542
                    unsigned int tj = tgx;
543
                    for (unsigned int j = 0; j < TILE_SIZE; j++) {
544
545
                        if (atom1 < NUM_ATOMS && y*TILE_SIZE+tj < NUM_ATOMS) {
                            float4 posq2 = (float4) (localData[tbx+tj].x, localData[tbx+tj].y, localData[tbx+tj].z, localData[tbx+tj].q);
546
                            float4 delta = (float4) (posq2.xyz - posq1.xyz, 0.0f);
547
#ifdef USE_PERIODIC
548
549
550
                            delta.x -= floor(delta.x*invPeriodicBoxSize.x+0.5f)*periodicBoxSize.x;
                            delta.y -= floor(delta.y*invPeriodicBoxSize.y+0.5f)*periodicBoxSize.y;
                            delta.z -= floor(delta.z*invPeriodicBoxSize.z+0.5f)*periodicBoxSize.z;
551
552
#endif
                            float r2 = delta.x*delta.x + delta.y*delta.y + delta.z*delta.z;
Peter Eastman's avatar
Peter Eastman committed
553
554
555
#ifdef USE_CUTOFF
                            if (r2 < CUTOFF_SQUARED) {
#endif
556
557
                            float invR = RSQRT(r2);
                            float r = RECIP(invR);
558
                            float bornRadius2 = localData[tbx+tj].bornRadius;
559
                            float alpha2_ij = bornRadius1*bornRadius2;
560
                            float D_ij = r2*RECIP(4.0f*alpha2_ij);
561
                            float expTerm = EXP(-D_ij);
562
                            float denominator2 = r2 + alpha2_ij*expTerm;
563
                            float denominator = SQRT(denominator2);
564
565
                            float tempEnergy = (PREFACTOR*posq1.w*posq2.w)*RECIP(denominator);
                            float Gpol = tempEnergy*RECIP(denominator2);
566
567
                            float dGpol_dalpha2_ij = -0.5f*Gpol*expTerm*(1.0f+D_ij);
                            float dEdR = Gpol*(1.0f - 0.25f*expTerm);
568
                            force.w += dGpol_dalpha2_ij*bornRadius2;
569
                            energy += tempEnergy;
570
571
                            delta.xyz *= dEdR;
                            force.xyz -= delta.xyz;
572
573
574
575
                            localData[tbx+tj].fx += delta.x;
                            localData[tbx+tj].fy += delta.y;
                            localData[tbx+tj].fz += delta.z;
                            localData[tbx+tj].fw += dGpol_dalpha2_ij*bornRadius1;
Peter Eastman's avatar
Peter Eastman committed
576
577
578
#ifdef USE_CUTOFF
                            }
#endif
579
                        }
580
                        tj = (tj + 1) & (TILE_SIZE - 1);
581
582
583
                    }
                }
            }
584
585
586
587
588
        }
        
        // Write results.  We need to coordinate between warps to make sure no two of them
        // ever try to write to the same piece of memory at the same time.
        
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
#ifdef SUPPORTS_64_BIT_ATOMICS
        if (pos < end) {
            const unsigned int offset = x*TILE_SIZE + tgx;
            atom_add(&forceBuffers[offset], (long) (force.x*0xFFFFFFFF));
            atom_add(&forceBuffers[offset+PADDED_NUM_ATOMS], (long) (force.y*0xFFFFFFFF));
            atom_add(&forceBuffers[offset+2*PADDED_NUM_ATOMS], (long) (force.z*0xFFFFFFFF));
            atom_add(&global_bornForce[offset], (long) (force.w*0xFFFFFFFF));
        }
        if (pos < end && x != y) {
            const unsigned int offset = y*TILE_SIZE + tgx;
            atom_add(&forceBuffers[offset], (long) (localData[get_local_id(0)].fx*0xFFFFFFFF));
            atom_add(&forceBuffers[offset+PADDED_NUM_ATOMS], (long) (localData[get_local_id(0)].fy*0xFFFFFFFF));
            atom_add(&forceBuffers[offset+2*PADDED_NUM_ATOMS], (long) (localData[get_local_id(0)].fz*0xFFFFFFFF));
            atom_add(&global_bornForce[offset], (long) (localData[get_local_id(0)].fw*0xFFFFFFFF));
        }
#else
605
606
607
608
609
610
611
612
613
        int writeX = (pos < end ? x : -1);
        int writeY = (pos < end && x != y ? y : -1);
        if (tgx == 0)
            reservedBlocks[localGroupIndex] = (int2)(writeX, writeY);
        bool done = false;
        int doneIndex = 0;
        int checkIndex = 0;
        while (true) {
            // See if any warp still needs to write its data.
614

615
616
617
618
619
620
621
            bool allDone = true;
            barrier(CLK_LOCAL_MEM_FENCE);
            while (doneIndex < WARPS_PER_GROUP && allDone) {
                if (reservedBlocks[doneIndex].x != -1)
                    allDone = false;
                else
                    doneIndex++;
622
            }
623
624
625
626
627
            if (allDone)
                break;
            if (!done) {
                // See whether this warp can write its data.  This requires that no previous warp
                // is trying to write to the same block of the buffer.
628

629
630
631
632
633
634
635
636
637
638
                bool canWrite = (writeX != -1);
                while (checkIndex < localGroupIndex && canWrite) {
                    if ((reservedBlocks[checkIndex].x == x || reservedBlocks[checkIndex].y == x) ||
                            (writeY != -1 && (reservedBlocks[checkIndex].x == y || reservedBlocks[checkIndex].y == y)))
                        canWrite = false;
                    else
                        checkIndex++;
                }
                if (canWrite) {
                    // Write the data to global memory, then mark this warp as done.
639

640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
                    if (writeX > -1) {
                        const unsigned int offset = x*TILE_SIZE + tgx + get_group_id(0)*PADDED_NUM_ATOMS;
                        forceBuffers[offset].xyz += force.xyz;
                        global_bornForce[offset] += force.w;
                    }
                    if (writeY > -1) {
                        const unsigned int offset = y*TILE_SIZE + tgx + get_group_id(0)*PADDED_NUM_ATOMS;
                        forceBuffers[offset] += (float4) (localData[get_local_id(0)].fx, localData[get_local_id(0)].fy, localData[get_local_id(0)].fz, 0.0f);
                        global_bornForce[offset] += localData[get_local_id(0)].fw;
                    }
                    done = true;
                    if (tgx == 0)
                        reservedBlocks[localGroupIndex] = (int2)(-1, -1);
                }
            }
655
        }
656
#endif
657
        lasty = y;
658
        pos++;
659
    } while (pos < end);
660
661
    energyBuffer[get_global_id(0)] += energy;
}