kCalculateAmoebaCudaPmeFixedEField.h 19.2 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
28
29
30
/* -------------------------------------------------------------------------- *
 *                                   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/>.      *
 * -------------------------------------------------------------------------- */

#include "amoebaScaleFactors.h"

__global__ 
#if (__CUDA_ARCH__ >= 200)
31
32
33
__launch_bounds__(384, 1)
#elif (__CUDA_ARCH__ >= 120)
__launch_bounds__(192, 1)
34
#else
35
__launch_bounds__(64, 1)
36
37
38
39
40
41
42
43
44
45
46
#endif
void METHOD_NAME(kCalculateAmoebaPmeDirectFixedE_Field, _kernel)(
                            unsigned int* workUnit,
                            float* outputEField,
                            float* outputEFieldPolar
#ifdef AMOEBA_DEBUG
                           , float4* debugArray, unsigned int targetAtom
#endif
){

#ifdef AMOEBA_DEBUG
Mark Friedrichs's avatar
Mark Friedrichs committed
47
    int maxPullIndex = 1;
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
    float4 pullBack[12];
#endif

    extern __shared__ FixedFieldParticle sA[];

    unsigned int totalWarps      = gridDim.x*blockDim.x/GRID;
    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;
    unsigned int lasty           = 0xFFFFFFFF;

    while (pos < end)
    {

        unsigned int x;
        unsigned int y;
        bool bExclusionFlag;
Mark Friedrichs's avatar
Mark Friedrichs committed
66
67
68
69
        float dScaleValue;
        float pScaleValue;
        int  dScaleMask;
        int2 pScaleMask;
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101

        // extract cell coordinates

        decodeCell( workUnit[pos], &x, &y, &bExclusionFlag );

        unsigned int tgx           = threadIdx.x & (GRID - 1);
        unsigned int tbx           = threadIdx.x - tgx;
        unsigned int tj            = tgx;

        FixedFieldParticle* psA    = &sA[tbx];
        unsigned int atomI         = x + tgx;
        FixedFieldParticle localParticle;
        loadFixedFieldShared( &localParticle, atomI );

        float fieldSum[3];
        float fieldPolarSum[3];

        fieldSum[0]                = 0.0f;
        fieldSum[1]                = 0.0f;
        fieldSum[2]                = 0.0f;

        fieldPolarSum[0]           = 0.0f;
        fieldPolarSum[1]           = 0.0f;
        fieldPolarSum[2]           = 0.0f;

        if (x == y)
        {

            // load coordinates, charge, ...

            loadFixedFieldShared( &(sA[threadIdx.x]), atomI );

Mark Friedrichs's avatar
Mark Friedrichs committed
102
            if( bExclusionFlag ){
103
                unsigned int xi       = x >> GRIDBITS;
104
                unsigned int cell     = xi + xi*cSim.paddedNumberOfAtoms/GRID-xi*(xi+1)/2;
Mark Friedrichs's avatar
Mark Friedrichs committed
105
106
107
108
                dScaleMask            = cAmoebaSim.pD_ScaleIndices[cAmoebaSim.pScaleIndicesIndex[cell]+tgx];
                pScaleMask            = cAmoebaSim.pP_ScaleIndices[cAmoebaSim.pScaleIndicesIndex[cell]+tgx];
            } else {
                dScaleValue = pScaleValue = 1.0f;
109

Mark Friedrichs's avatar
Mark Friedrichs committed
110
            }
111

Mark Friedrichs's avatar
Mark Friedrichs committed
112
113
            for (unsigned int j = 0; j < GRID; j++)
            {
114

Mark Friedrichs's avatar
Mark Friedrichs committed
115
                if( bExclusionFlag ){
116
117
                    getMaskedDScaleFactor( j, dScaleMask, &dScaleValue );
                    getMaskedPScaleFactor( j, pScaleMask, &pScaleValue );
Mark Friedrichs's avatar
Mark Friedrichs committed
118
                }
119

120
                float4 ijField[3];
Mark Friedrichs's avatar
Mark Friedrichs committed
121
                calculateFixedFieldRealSpacePairIxn_kernel( localParticle, psA[j], dScaleValue, pScaleValue, ijField
122
#ifdef AMOEBA_DEBUG
Mark Friedrichs's avatar
Mark Friedrichs committed
123
                                            , pullBack
124
#endif
Mark Friedrichs's avatar
Mark Friedrichs committed
125
                );
126

Mark Friedrichs's avatar
Mark Friedrichs committed
127
128
                // nan*0.0 = nan not 0.0, so explicitly exclude (atomI == atomJ) contribution
                // by setting match flag
129

130
                unsigned int match      = ( (atomI == (y + j)) || (atomI >= cSim.atoms) || ((y+j) >= cSim.atoms) ) ? 1 : 0;
131

Mark Friedrichs's avatar
Mark Friedrichs committed
132
                // add to field at atomI the field due atomJ's charge/dipole/quadrupole
133

134
135
136
                fieldSum[0]            += match ? 0.0f : ijField[0].x;
                fieldSum[1]            += match ? 0.0f : ijField[1].x;
                fieldSum[2]            += match ? 0.0f : ijField[2].x;
137

138
139
140
                fieldPolarSum[0]       += match ? 0.0f : ijField[0].z;
                fieldPolarSum[1]       += match ? 0.0f : ijField[1].z;
                fieldPolarSum[2]       += match ? 0.0f : ijField[2].z;
141
142

#ifdef AMOEBA_DEBUG
Mark Friedrichs's avatar
Mark Friedrichs committed
143
if( atomI == targetAtom || targetAtom == (y+j) ){
Mark Friedrichs's avatar
Mark Friedrichs committed
144
145
146
    unsigned int index                 = atomI == targetAtom ? (y + j) : atomI;
    unsigned int indexI                = 0;
    unsigned int indexJ                = indexI ? 0 : 2;
Mark Friedrichs's avatar
Mark Friedrichs committed
147
    float flag                         = 7.0f;
Mark Friedrichs's avatar
Mark Friedrichs committed
148
149
150
151
152

    debugArray[index].x                = (float) atomI;
    debugArray[index].y                = (float) (y + j);
    debugArray[index].z                = dScaleValue;
    debugArray[index].w                = pScaleValue;
Mark Friedrichs's avatar
Mark Friedrichs committed
153
/*
154
    index                             += cSim.paddedNumberOfAtoms;
Mark Friedrichs's avatar
Mark Friedrichs committed
155
156
157
158
159
    debugArray[index].x                = (float) bExclusionFlag;
    debugArray[index].y                = (float) (tgx);
    debugArray[index].z                = (float) j;
    debugArray[index].w                = flag;

160
    index                             += cSim.paddedNumberOfAtoms;
Mark Friedrichs's avatar
Mark Friedrichs committed
161
162
163
164
    debugArray[index].x                = (float) dScaleMask;
    debugArray[index].y                = (float) pScaleMask.x;
    debugArray[index].z                = (float) pScaleMask.y;
    debugArray[index].w                = flag;
Mark Friedrichs's avatar
Mark Friedrichs committed
165
*/
166
    index                             += cSim.paddedNumberOfAtoms;
Mark Friedrichs's avatar
Mark Friedrichs committed
167
168
169
170
171
    debugArray[index].x                = match ? 0.0f : ijField[0].x;
    debugArray[index].y                = match ? 0.0f : ijField[1].x;
    debugArray[index].z                = match ? 0.0f : ijField[2].x;
    debugArray[index].w                = flag + 1.0f;

172
    index                             += cSim.paddedNumberOfAtoms;
Mark Friedrichs's avatar
Mark Friedrichs committed
173
174
175
176
177
    debugArray[index].x                = match ? 0.0f : ijField[0].z;
    debugArray[index].y                = match ? 0.0f : ijField[1].z;
    debugArray[index].z                = match ? 0.0f : ijField[2].z;
    debugArray[index].w                = flag + 2.0f;

178
    index                             += cSim.paddedNumberOfAtoms;
Mark Friedrichs's avatar
Mark Friedrichs committed
179
180
181
182
183
    debugArray[index].x                = match ? 0.0f : ijField[0].y;
    debugArray[index].y                = match ? 0.0f : ijField[1].y;
    debugArray[index].z                = match ? 0.0f : ijField[2].y;
    debugArray[index].w                = flag + 3.0f;

184
    index                             += cSim.paddedNumberOfAtoms;
Mark Friedrichs's avatar
Mark Friedrichs committed
185
186
187
188
    debugArray[index].x                = match ? 0.0f : ijField[0].w;
    debugArray[index].y                = match ? 0.0f : ijField[1].w;
    debugArray[index].z                = match ? 0.0f : ijField[2].w;
    debugArray[index].w                = flag + 4.0f;
Mark Friedrichs's avatar
Mark Friedrichs committed
189
190


Mark Friedrichs's avatar
Mark Friedrichs committed
191
    for( int pullIndex = 0; pullIndex < maxPullIndex; pullIndex++ ){
192
        index                             += cSim.paddedNumberOfAtoms;
Mark Friedrichs's avatar
Mark Friedrichs committed
193
194
195
196
197
        debugArray[index].x                = pullBack[pullIndex].x;
        debugArray[index].y                = pullBack[pullIndex].y;
        debugArray[index].z                = pullBack[pullIndex].z;
        debugArray[index].w                = pullBack[pullIndex].w;
    }   
198
199
200
201
202
203
204
205
206


}
#endif
            }

            // Write results

#ifdef USE_OUTPUT_BUFFER_PER_WARP
207
            unsigned int offset                 = 3*(x + tgx + warp*cSim.paddedNumberOfAtoms);
208
209
210
            load3dArrayBufferPerWarp( offset, fieldSum,       outputEField );
            load3dArrayBufferPerWarp( offset, fieldPolarSum,  outputEFieldPolar );
#else
211
            unsigned int offset                 = 3*(x + tgx + (x >> GRIDBITS) * cSim.paddedNumberOfAtoms);
212
213
214
215
            load3dArray( offset, fieldSum,       outputEField );
            load3dArray( offset, fieldPolarSum,  outputEFieldPolar );
#endif

Mark Friedrichs's avatar
Mark Friedrichs committed
216
        } else {
217

Mark Friedrichs's avatar
Mark Friedrichs committed
218
219
220
221
222
223
224
225
            if (lasty != y ) {
    
                // load coordinates, charge, ...
    
                loadFixedFieldShared( &(sA[threadIdx.x]), (y+tgx) );
    
            }

Mark Friedrichs's avatar
Mark Friedrichs committed
226
227
228
229
            unsigned int flags = cSim.pInteractionFlag[pos];
            if (flags == 0) {
                // No interactions in this block.
            } else {
230

Mark Friedrichs's avatar
Mark Friedrichs committed
231
                // zero shared fields
232

Mark Friedrichs's avatar
Mark Friedrichs committed
233
                zeroFixedFieldParticleSharedField( &(sA[threadIdx.x]) );
234

Mark Friedrichs's avatar
Mark Friedrichs committed
235
236
237
                if( bExclusionFlag ) {
                    unsigned int xi   = x >> GRIDBITS;
                    unsigned int yi   = y >> GRIDBITS;
238
                    unsigned int cell = xi+yi*cSim.paddedNumberOfAtoms/GRID-yi*(yi+1)/2;
Mark Friedrichs's avatar
Mark Friedrichs committed
239
240
241
242
                    dScaleMask        = cAmoebaSim.pD_ScaleIndices[cAmoebaSim.pScaleIndicesIndex[cell]+tgx];
                    pScaleMask        = cAmoebaSim.pP_ScaleIndices[cAmoebaSim.pScaleIndicesIndex[cell]+tgx];
                } else {
                    dScaleValue = pScaleValue  = 1.0f;
243
244
                }

Mark Friedrichs's avatar
Mark Friedrichs committed
245
                for (unsigned int j = 0; j < GRID; j++){
246

Peter Eastman's avatar
Peter Eastman committed
247
248
249
250
251
252
                    if ((flags&(1<<j)) != 0) {
                        unsigned int jIdx = (flags == 0xFFFFFFFF) ? tj : j;
                        if( bExclusionFlag ){
                            getMaskedDScaleFactor( jIdx, dScaleMask, &dScaleValue );
                            getMaskedPScaleFactor( jIdx, pScaleMask, &pScaleValue );
                        }
253

Peter Eastman's avatar
Peter Eastman committed
254
255
256
257
258
259
260
                        float4 ijField[3];
                        calculateFixedFieldRealSpacePairIxn_kernel( localParticle, psA[jIdx], dScaleValue, pScaleValue, ijField
    #ifdef AMOEBA_DEBUG
                                                    , pullBack
    #endif
                        );

261
                        unsigned int outOfBounds     = ( (atomI >= cSim.atoms) || ((y+jIdx) >= cSim.atoms) ) ? 1 : 0;
Peter Eastman's avatar
Peter Eastman committed
262
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

                        // add to field at atomI the field due atomJ's charge/dipole/quadrupole

                        fieldSum[0]                 += outOfBounds ? 0.0f : ijField[0].x;
                        fieldSum[1]                 += outOfBounds ? 0.0f : ijField[1].x;
                        fieldSum[2]                 += outOfBounds ? 0.0f : ijField[2].x;

                        fieldPolarSum[0]            += outOfBounds ? 0.0f : ijField[0].z;
                        fieldPolarSum[1]            += outOfBounds ? 0.0f : ijField[1].z;
                        fieldPolarSum[2]            += outOfBounds ? 0.0f : ijField[2].z;

                        if( flags == 0xFFFFFFFF ){

                            // add to field at atomJ the field due atomI's charge/dipole/quadrupole

                            psA[jIdx].eField[0]        += outOfBounds ? 0.0f : ijField[0].y;
                            psA[jIdx].eField[1]        += outOfBounds ? 0.0f : ijField[1].y;
                            psA[jIdx].eField[2]        += outOfBounds ? 0.0f : ijField[2].y;

                            psA[jIdx].eFieldP[0]       += outOfBounds ? 0.0f : ijField[0].w;
                            psA[jIdx].eFieldP[1]       += outOfBounds ? 0.0f : ijField[1].w;
                            psA[jIdx].eFieldP[2]       += outOfBounds ? 0.0f : ijField[2].w;

                        } else {

                            sA[threadIdx.x].tempBuffer[0]  = outOfBounds ? 0.0f : ijField[0].y;
                            sA[threadIdx.x].tempBuffer[1]  = outOfBounds ? 0.0f : ijField[1].y;
                            sA[threadIdx.x].tempBuffer[2]  = outOfBounds ? 0.0f : ijField[2].y;

                            sA[threadIdx.x].tempBufferP[0] = outOfBounds ? 0.0f : ijField[0].w;
                            sA[threadIdx.x].tempBufferP[1] = outOfBounds ? 0.0f : ijField[1].w;
                            sA[threadIdx.x].tempBufferP[2] = outOfBounds ? 0.0f : ijField[2].w;

                            if( tgx % 2 == 0 ){
                                sumTempBuffer( sA[threadIdx.x], sA[threadIdx.x+1] );
                            }
                            if( tgx % 4 == 0 ){
                                sumTempBuffer( sA[threadIdx.x], sA[threadIdx.x+2] );
                            }
                            if( tgx % 8 == 0 ){
                                sumTempBuffer( sA[threadIdx.x], sA[threadIdx.x+4] );
                            }
                            if( tgx % 16 == 0 ){
                                sumTempBuffer( sA[threadIdx.x], sA[threadIdx.x+8] );
                            }

                            if (tgx == 0)
                            {
                                psA[jIdx].eField[0]  += sA[threadIdx.x].tempBuffer[0]  + sA[threadIdx.x+16].tempBuffer[0];
                                psA[jIdx].eField[1]  += sA[threadIdx.x].tempBuffer[1]  + sA[threadIdx.x+16].tempBuffer[1];
                                psA[jIdx].eField[2]  += sA[threadIdx.x].tempBuffer[2]  + sA[threadIdx.x+16].tempBuffer[2];

                                psA[jIdx].eFieldP[0] += sA[threadIdx.x].tempBufferP[0] + sA[threadIdx.x+16].tempBufferP[0];
                                psA[jIdx].eFieldP[1] += sA[threadIdx.x].tempBufferP[1] + sA[threadIdx.x+16].tempBufferP[1];
                                psA[jIdx].eFieldP[2] += sA[threadIdx.x].tempBufferP[2] + sA[threadIdx.x+16].tempBufferP[2];
                            }
Mark Friedrichs's avatar
Mark Friedrichs committed
318
                        }
319
320

#ifdef AMOEBA_DEBUG
Mark Friedrichs's avatar
Mark Friedrichs committed
321
if( (atomI == targetAtom || (y + jIdx) == targetAtom) ){
322

Mark Friedrichs's avatar
Mark Friedrichs committed
323
            unsigned int index                 = (atomI == targetAtom) ? (y + jIdx) : atomI;
324
325
326
327
            unsigned int indexI                = (atomI == targetAtom) ? 0 : 2;
            unsigned int indexJ                = (atomI == targetAtom) ? 2 : 0;

            debugArray[index].x                = (float) atomI;
Mark Friedrichs's avatar
Mark Friedrichs committed
328
329
330
            debugArray[index].y                = (float) (y + jIdx);
            debugArray[index].z                = dScaleValue;
            debugArray[index].w                = pScaleValue;
331
332

            float flag                         = 9.0f;
Mark Friedrichs's avatar
Mark Friedrichs committed
333
/*
334
            index                             += cSim.paddedNumberOfAtoms;
Mark Friedrichs's avatar
Mark Friedrichs committed
335
336
337
338
339
            debugArray[index].x                = (float) bExclusionFlag;
            debugArray[index].y                = (float) (tgx);
            debugArray[index].z                = (float) j;
            debugArray[index].w                = jIdx;
        
340
            index                             += cSim.paddedNumberOfAtoms;
Mark Friedrichs's avatar
Mark Friedrichs committed
341
342
343
344
            debugArray[index].x                = (float) dScaleMask;
            debugArray[index].y                = (float) pScaleMask.x;
            debugArray[index].z                = (float) pScaleMask.y;
            debugArray[index].w                = (float) flags;
Mark Friedrichs's avatar
Mark Friedrichs committed
345
 */       
346
            index                             += cSim.paddedNumberOfAtoms;
Mark Friedrichs's avatar
Mark Friedrichs committed
347
348
349
350
            debugArray[index].x                =  outOfBounds ? 0.0f : ijField[0].x;
            debugArray[index].y                =  outOfBounds ? 0.0f : ijField[1].x;
            debugArray[index].z                =  outOfBounds ? 0.0f : ijField[2].x;
            debugArray[index].w                =  flag + 1.0f;
351

352
            index                             += cSim.paddedNumberOfAtoms;
Mark Friedrichs's avatar
Mark Friedrichs committed
353
354
355
356
            debugArray[index].x                =  outOfBounds ? 0.0f : ijField[0].y;
            debugArray[index].y                =  outOfBounds ? 0.0f : ijField[1].y;
            debugArray[index].z                =  outOfBounds ? 0.0f : ijField[2].y;
            debugArray[index].w                = flag + 2.0f;
357

358
            index                             += cSim.paddedNumberOfAtoms;
Mark Friedrichs's avatar
Mark Friedrichs committed
359
360
361
362
            debugArray[index].x                =  outOfBounds ? 0.0f : ijField[0].z;
            debugArray[index].y                =  outOfBounds ? 0.0f : ijField[1].z;
            debugArray[index].z                =  outOfBounds ? 0.0f : ijField[2].z;
            debugArray[index].w                = flag + 3.0f;
363

364
            index                             += cSim.paddedNumberOfAtoms;
Mark Friedrichs's avatar
Mark Friedrichs committed
365
366
367
368
            debugArray[index].x                =  outOfBounds ? 0.0f : ijField[0].w;
            debugArray[index].y                =  outOfBounds ? 0.0f : ijField[1].w;
            debugArray[index].z                =  outOfBounds ? 0.0f : ijField[2].w;
            debugArray[index].w                = flag + 4.0f;
369

Mark Friedrichs's avatar
Mark Friedrichs committed
370
            for( int pullIndex = 0; pullIndex < maxPullIndex; pullIndex++ ){
371
                index                             += cSim.paddedNumberOfAtoms;
Mark Friedrichs's avatar
Mark Friedrichs committed
372
373
374
375
376
377
                debugArray[index].x                = pullBack[pullIndex].x;
                debugArray[index].y                = pullBack[pullIndex].y;
                debugArray[index].z                = pullBack[pullIndex].z;
                debugArray[index].w                = pullBack[pullIndex].w;
            }
}        
378
#endif
Peter Eastman's avatar
Peter Eastman committed
379
                    }
380
381
                    tj                  = (tj + 1) & (GRID - 1);

Mark Friedrichs's avatar
Mark Friedrichs committed
382
383
384
385
                } // j-loop block
    
                // Write results
    
386
#ifdef USE_OUTPUT_BUFFER_PER_WARP
387
                unsigned int offset                 = 3*(x + tgx + warp*cSim.paddedNumberOfAtoms);
Mark Friedrichs's avatar
Mark Friedrichs committed
388
389
390
                load3dArrayBufferPerWarp( offset, fieldSum,       outputEField );
                load3dArrayBufferPerWarp( offset, fieldPolarSum,  outputEFieldPolar );
    
391
                offset                              = 3*(y + tgx + warp*cSim.paddedNumberOfAtoms);
Mark Friedrichs's avatar
Mark Friedrichs committed
392
393
394
                load3dArrayBufferPerWarp( offset, sA[threadIdx.x].eField,  outputEField );
                load3dArrayBufferPerWarp( offset, sA[threadIdx.x].eFieldP, outputEFieldPolar );
    
395
#else
396
                unsigned int offset                 = 3*(x + tgx + (y >> GRIDBITS) * cSim.paddedNumberOfAtoms);
Mark Friedrichs's avatar
Mark Friedrichs committed
397
398
399
                load3dArray( offset, fieldSum,       outputEField );
                load3dArray( offset, fieldPolarSum,  outputEFieldPolar );
    
400
                offset                              = 3*(y + tgx + (x >> GRIDBITS) * cSim.paddedNumberOfAtoms);
Mark Friedrichs's avatar
Mark Friedrichs committed
401
402
403
                load3dArray( offset, sA[threadIdx.x].eField,  outputEField );
                load3dArray( offset, sA[threadIdx.x].eFieldP, outputEFieldPolar );
     
404
#endif
Mark Friedrichs's avatar
Mark Friedrichs committed
405
            } // end of pInteractionFlag block 
406
            lasty = y;
Mark Friedrichs's avatar
Mark Friedrichs committed
407
        } // x == y block
408
409
410
411

        pos++;
    }
}