kCalculateAmoebaCudaPME.cu 48.3 KB
Newer Older
Mark Friedrichs's avatar
Mark Friedrichs committed
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
/* -------------------------------------------------------------------------- *
 *                                   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/>.      *
 * -------------------------------------------------------------------------- */
26
27

#include "amoebaGpuTypes.h"
28
#include "cudaKernels.h"
29
#include "amoebaCudaKernels.h"
30
#include "bbsort.h"
31
32
33
34

static __constant__ cudaGmxSimulation cSim;
static __constant__ cudaAmoebaGmxSimulation cAmoebaSim;

35
/* Cuda compiler on Windows does not recognized "static const float" values */
36
#define LOCAL_HACK_PI 3.1415926535897932384626433832795f
37

38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
void SetCalculateAmoebaPMESim(amoebaGpuContext amoebaGpu)
{
    cudaError_t status;
    gpuContext gpu = amoebaGpu->gpuContext;
    status         = cudaMemcpyToSymbol(cSim, &gpu->sim, sizeof(cudaGmxSimulation));
    RTERROR(status, "SetCalculateAmoebaPMESim: cudaMemcpyToSymbol: SetSim copy to cSim failed");
    status         = cudaMemcpyToSymbol(cAmoebaSim, &amoebaGpu->amoebaSim, sizeof(cudaAmoebaGmxSimulation));
    RTERROR(status, "SetCalculateAmoebaPMESim: cudaMemcpyToSymbol: SetSim copy to cAmoebaSim failed");
}

#define ARRAY(x,y) array[(x)-1+((y)-1)*AMOEBA_PME_ORDER]

/**
 * This is called from computeBsplines().  It calculates the spline coefficients for a single atom along a single axis.
 */
__device__ void computeBSplinePoint(float4* thetai, float w, float* array)
{
    // initialization to get to 2nd order recursion

    ARRAY(2,2) = w;
    ARRAY(2,1) = 1.0f - w;

    // perform one pass to get to 3rd order recursion

    ARRAY(3,3) = 0.5f * w * ARRAY(2,2);
    ARRAY(3,2) = 0.5f * ((1.0f+w)*ARRAY(2,1)+(2.0f-w)*ARRAY(2,2));
    ARRAY(3,1) = 0.5f * (1.0f-w) * ARRAY(2,1);

    // compute standard B-spline recursion to desired order

    for (int i = 4; i <= AMOEBA_PME_ORDER; i++)
    {
        int k = i - 1;
        float denom = 1.0f / k;
        ARRAY(i,i) = denom * w * ARRAY(k,k);
        for (int j = 1; j <= i-2; j++)
            ARRAY(i,i-j) = denom * ((w+j)*ARRAY(k,i-j-1)+(i-j-w)*ARRAY(k,i-j));
        ARRAY(i,1) = denom * (1.0f-w) * ARRAY(k,1);
    }

    // get coefficients for the B-spline first derivative

    int k = AMOEBA_PME_ORDER - 1;
    ARRAY(k,AMOEBA_PME_ORDER) = ARRAY(k,AMOEBA_PME_ORDER-1);
    for (int i = AMOEBA_PME_ORDER-1; i >= 2; i--)
        ARRAY(k,i) = ARRAY(k,i-1) - ARRAY(k,i);
    ARRAY(k,1) = -ARRAY(k,1);

    // get coefficients for the B-spline second derivative

    k = AMOEBA_PME_ORDER - 2;
    ARRAY(k,AMOEBA_PME_ORDER-1) = ARRAY(k,AMOEBA_PME_ORDER-2);
    for (int i = AMOEBA_PME_ORDER-2; i >= 2; i--)
        ARRAY(k,i) = ARRAY(k,i-1) - ARRAY(k,i);
    ARRAY(k,1) = -ARRAY(k,1);
    ARRAY(k,AMOEBA_PME_ORDER) = ARRAY(k,AMOEBA_PME_ORDER-1);
    for (int i = AMOEBA_PME_ORDER-1; i >= 2; i--)
        ARRAY(k,i) = ARRAY(k,i-1) - ARRAY(k,i);
    ARRAY(k,1) = -ARRAY(k,1);

    // get coefficients for the B-spline third derivative

    k = AMOEBA_PME_ORDER - 3;
    ARRAY(k,AMOEBA_PME_ORDER-2) = ARRAY(k,AMOEBA_PME_ORDER-3);
    for (int i = AMOEBA_PME_ORDER-3; i >= 2; i--)
        ARRAY(k,i) = ARRAY(k,i-1) - ARRAY(k,i);
    ARRAY(k,1) = -ARRAY(k,1);
    ARRAY(k,AMOEBA_PME_ORDER-1) = ARRAY(k,AMOEBA_PME_ORDER-2);
    for (int i = AMOEBA_PME_ORDER-2; i >= 2; i--)
        ARRAY(k,i) = ARRAY(k,i-1) - ARRAY(k,i);
    ARRAY(k,1) = -ARRAY(k,1);
    ARRAY(k,AMOEBA_PME_ORDER) = ARRAY(k,AMOEBA_PME_ORDER-1);
    for (int i = AMOEBA_PME_ORDER-1; i >= 2; i--)
        ARRAY(k,i) = ARRAY(k,i-1) - ARRAY(k,i);
    ARRAY(k,1) = -ARRAY(k,1);

    // copy coefficients from temporary to permanent storage

    for (int i = 1; i <= AMOEBA_PME_ORDER; i++)
117
        thetai[i-1] = make_float4(ARRAY(AMOEBA_PME_ORDER,i), ARRAY(AMOEBA_PME_ORDER-1,i), ARRAY(AMOEBA_PME_ORDER-2,i), ARRAY(AMOEBA_PME_ORDER-3,i));
118
119
120
121
122
123
124
}

/**
 * Compute bspline coefficients.
 */
__global__
#if (__CUDA_ARCH__ >= 200)
Peter Eastman's avatar
Peter Eastman committed
125
__launch_bounds__(448, 1)
126
#elif (__CUDA_ARCH__ >= 120)
Peter Eastman's avatar
Peter Eastman committed
127
__launch_bounds__(160, 1)
128
#else
Peter Eastman's avatar
Peter Eastman committed
129
__launch_bounds__(160, 1)
130
#endif
131
void kComputeAmoebaBsplines_kernel()
132
133
134
135
136
137
138
139
{
    extern __shared__ float bsplines_cache[]; // size = block_size*pme_order*pme_order
    float* array = &bsplines_cache[threadIdx.x*AMOEBA_PME_ORDER*AMOEBA_PME_ORDER];

    //  get the B-spline coefficients for each multipole site

    for (int i = blockIdx.x*blockDim.x+threadIdx.x; i < cSim.atoms; i += blockDim.x*gridDim.x) {
        float4 posq = cSim.pPosq[i];
140
141
142
        posq.x -= floorf(posq.x*cSim.invPeriodicBoxSizeX)*cSim.periodicBoxSizeX;
        posq.y -= floorf(posq.y*cSim.invPeriodicBoxSizeY)*cSim.periodicBoxSizeY;
        posq.z -= floorf(posq.z*cSim.invPeriodicBoxSizeZ)*cSim.periodicBoxSizeZ;
143
144
145
146
147
148
149

        // First axis.

        float w = posq.x*cSim.invPeriodicBoxSizeX;
        float fr = cSim.pmeGridSize.x*(w-(int)(w+0.5f)+0.5f);
        int ifr = (int) fr;
        w = fr - ifr;
150
        int igrid1 = ifr-AMOEBA_PME_ORDER+1;
151
152
153
154
155
156
157
158
        computeBSplinePoint(&cAmoebaSim.pThetai1[i*AMOEBA_PME_ORDER], w, array);

        // Second axis.

        w = posq.y*cSim.invPeriodicBoxSizeY;
        fr = cSim.pmeGridSize.y*(w-(int)(w+0.5f)+0.5f);
        ifr = (int) fr;
        w = fr - ifr;
159
        int igrid2 = ifr-AMOEBA_PME_ORDER+1;
160
161
162
163
164
165
166
167
        computeBSplinePoint(&cAmoebaSim.pThetai2[i*AMOEBA_PME_ORDER], w, array);

        // Third axis.

        w = posq.z*cSim.invPeriodicBoxSizeZ;
        fr = cSim.pmeGridSize.z*(w-(int)(w+0.5f)+0.5f);
        ifr = (int) fr;
        w = fr - ifr;
168
        int igrid3 = ifr-AMOEBA_PME_ORDER+1;
169
170
        computeBSplinePoint(&cAmoebaSim.pThetai3[i*AMOEBA_PME_ORDER], w, array);

171
172
        // Record the grid point.

173
174
175
        igrid1 += (igrid1 < 0 ? cSim.pmeGridSize.x : 0);
        igrid2 += (igrid2 < 0 ? cSim.pmeGridSize.y : 0);
        igrid3 += (igrid3 < 0 ? cSim.pmeGridSize.z : 0);
176
        cAmoebaSim.pIgrid[i] = make_int4(igrid1, igrid2, igrid3, 0);
177
        cSim.pPmeAtomGridIndex[i] = make_int2(i, igrid1*cSim.pmeGridSize.y*cSim.pmeGridSize.z+igrid2*cSim.pmeGridSize.z+igrid3);
178
179
180
    }
}

181
182
183
184
185
186
/**
 * For each grid point, find the range of sorted atoms associated with that point.
 */
__global__
#if (__CUDA_ARCH__ >= 200)
__launch_bounds__(1024, 1)
187
#elif (__CUDA_ARCH__ >= 120)
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
__launch_bounds__(512, 1)
#else
__launch_bounds__(256, 1)
#endif
void kFindAmoebaAtomRangeForGrid_kernel()
{
    int thread = blockIdx.x*blockDim.x+threadIdx.x;
    int start = (cSim.atoms*thread)/(blockDim.x*gridDim.x);
    int end = (cSim.atoms*(thread+1))/(blockDim.x*gridDim.x);
    int last = (start == 0 ? -1 : cSim.pPmeAtomGridIndex[start-1].y);
    for (int i = start; i < end; ++i)
    {
        int2 atomData = cSim.pPmeAtomGridIndex[i];
        int gridIndex = atomData.y;
        if (gridIndex != last)
        {
            for (int j = last+1; j <= gridIndex; ++j)
                cSim.pPmeAtomRange[j] = i;
            last = gridIndex;
        }

        // The grid index won't be needed again.  Reuse that component to hold the z index, thus saving
        // some work in the charge spreading kernel.

        float posz = cSim.pPosq[atomData.x].z;
213
        posz -= floorf(posz*cSim.invPeriodicBoxSizeZ)*cSim.periodicBoxSizeZ;
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
        float w = posz*cSim.invPeriodicBoxSizeZ;
        float fr = cSim.pmeGridSize.z*(w-(int)(w+0.5f)+0.5f);
        int z = ((int) fr)-AMOEBA_PME_ORDER+1;
        cSim.pPmeAtomGridIndex[i].y = z;
    }

    // Fill in values beyond the last atom.

    if (thread == blockDim.x*gridDim.x-1)
    {
        int gridSize = cSim.pmeGridSize.x*cSim.pmeGridSize.y*cSim.pmeGridSize.z;
        for (int j = last+1; j <= gridSize; ++j)
            cSim.pPmeAtomRange[j] = cSim.atoms;
    }
}
229
__global__
Peter Eastman's avatar
Peter Eastman committed
230
__launch_bounds__(64, 10)
231
void kGridSpreadFixedMultipoles_kernel()
232
{
233
234
235
    const float xscale = cSim.pmeGridSize.x*cSim.invPeriodicBoxSizeX;
    const float yscale = cSim.pmeGridSize.y*cSim.invPeriodicBoxSizeY;
    const float zscale = cSim.pmeGridSize.z*cSim.invPeriodicBoxSizeZ;
236
237
238
239
240
241
242
243
244
245
    unsigned int numGridPoints = cSim.pmeGridSize.x*cSim.pmeGridSize.y*cSim.pmeGridSize.z;
    unsigned int numThreads = gridDim.x*blockDim.x;
    for (int gridIndex = blockIdx.x*blockDim.x+threadIdx.x; gridIndex < numGridPoints; gridIndex += numThreads)
    {
        int3 gridPoint;
        gridPoint.x = gridIndex/(cSim.pmeGridSize.y*cSim.pmeGridSize.z);
        int remainder = gridIndex-gridPoint.x*cSim.pmeGridSize.y*cSim.pmeGridSize.z;
        gridPoint.y = remainder/cSim.pmeGridSize.z;
        gridPoint.z = remainder-gridPoint.y*cSim.pmeGridSize.z;
        float result = 0.0f;
246
        for (int ix = 0; ix < AMOEBA_PME_ORDER; ++ix)
247
248
        {
            int x = gridPoint.x-ix+(gridPoint.x >= ix ? 0 : cSim.pmeGridSize.x);
249
            for (int iy = 0; iy < AMOEBA_PME_ORDER; ++iy)
250
251
            {
                int y = gridPoint.y-iy+(gridPoint.y >= iy ? 0 : cSim.pmeGridSize.y);
252
                int z1 = gridPoint.z-AMOEBA_PME_ORDER+1;
253
254
255
256
257
258
259
260
261
262
263
264
                z1 += (z1 >= 0 ? 0 : cSim.pmeGridSize.z);
                int z2 = (z1 < gridPoint.z ? gridPoint.z : cSim.pmeGridSize.z-1);
                int gridIndex1 = x*cSim.pmeGridSize.y*cSim.pmeGridSize.z+y*cSim.pmeGridSize.z+z1;
                int gridIndex2 = x*cSim.pmeGridSize.y*cSim.pmeGridSize.z+y*cSim.pmeGridSize.z+z2;
                int firstAtom = cSim.pPmeAtomRange[gridIndex1];
                int lastAtom = cSim.pPmeAtomRange[gridIndex2+1];
                for (int i = firstAtom; i < lastAtom; ++i)
                {
                    int2 atomData = cSim.pPmeAtomGridIndex[i];
                    int atomIndex = atomData.x;
                    int z = atomData.y;
                    int iz = gridPoint.z-z+(gridPoint.z >= z ? 0 : cSim.pmeGridSize.z);
Mark Friedrichs's avatar
Mark Friedrichs committed
265
266
267
                    if( iz >= cSim.pmeGridSize.z ){
                        iz -= cSim.pmeGridSize.z;
                    }
268
269
270
271
272
273
274
275
276
277
                    float atomCharge = cSim.pPosq[atomIndex].w;
                    float atomDipoleX = xscale*cAmoebaSim.pLabFrameDipole[atomIndex*3];
                    float atomDipoleY = yscale*cAmoebaSim.pLabFrameDipole[atomIndex*3+1];
                    float atomDipoleZ = zscale*cAmoebaSim.pLabFrameDipole[atomIndex*3+2];
                    float atomQuadrupoleXX = xscale*xscale*cAmoebaSim.pLabFrameQuadrupole[atomIndex*9];
                    float atomQuadrupoleXY = 2*xscale*yscale*cAmoebaSim.pLabFrameQuadrupole[atomIndex*9+1];
                    float atomQuadrupoleXZ = 2*xscale*zscale*cAmoebaSim.pLabFrameQuadrupole[atomIndex*9+2];
                    float atomQuadrupoleYY = yscale*yscale*cAmoebaSim.pLabFrameQuadrupole[atomIndex*9+4];
                    float atomQuadrupoleYZ = 2*yscale*zscale*cAmoebaSim.pLabFrameQuadrupole[atomIndex*9+5];
                    float atomQuadrupoleZZ = zscale*zscale*cAmoebaSim.pLabFrameQuadrupole[atomIndex*9+8];
278
279
280
                    float4 t = cAmoebaSim.pThetai1[atomIndex*AMOEBA_PME_ORDER+ix];
                    float4 u = cAmoebaSim.pThetai2[atomIndex*AMOEBA_PME_ORDER+iy];
                    float4 v = cAmoebaSim.pThetai3[atomIndex*AMOEBA_PME_ORDER+iz];
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
                    float term0 = atomCharge*u.x*v.x + atomDipoleY*u.y*v.x + atomDipoleZ*u.x*v.y + atomQuadrupoleYY*u.z*v.x + atomQuadrupoleZZ*u.x*v.z + atomQuadrupoleYZ*u.y*v.y;
                    float term1 = atomDipoleX*u.x*v.x + atomQuadrupoleXY*u.y*v.x + atomQuadrupoleXZ*u.x*v.y;
                    float term2 = atomQuadrupoleXX * u.x * v.x;
                    result += term0*t.x + term1*t.y + term2*t.z;
                }
                if (z1 > gridPoint.z)
                {
                    gridIndex1 = x*cSim.pmeGridSize.y*cSim.pmeGridSize.z+y*cSim.pmeGridSize.z;
                    gridIndex2 = x*cSim.pmeGridSize.y*cSim.pmeGridSize.z+y*cSim.pmeGridSize.z+gridPoint.z;
                    firstAtom = cSim.pPmeAtomRange[gridIndex1];
                    lastAtom = cSim.pPmeAtomRange[gridIndex2+1];
                    for (int i = firstAtom; i < lastAtom; ++i)
                    {
                        int2 atomData = cSim.pPmeAtomGridIndex[i];
                        int atomIndex = atomData.x;
                        int z = atomData.y;
                        int iz = gridPoint.z-z+(gridPoint.z >= z ? 0 : cSim.pmeGridSize.z);
Mark Friedrichs's avatar
Mark Friedrichs committed
298
299
300
                        if( iz >= cSim.pmeGridSize.z ){
                            iz -= cSim.pmeGridSize.z;
                        }
301
302
303
304
305
306
307
308
309
310
                        float atomCharge = cSim.pPosq[atomIndex].w;
                        float atomDipoleX = xscale*cAmoebaSim.pLabFrameDipole[atomIndex*3];
                        float atomDipoleY = yscale*cAmoebaSim.pLabFrameDipole[atomIndex*3+1];
                        float atomDipoleZ = zscale*cAmoebaSim.pLabFrameDipole[atomIndex*3+2];
                        float atomQuadrupoleXX = xscale*xscale*cAmoebaSim.pLabFrameQuadrupole[atomIndex*9];
                        float atomQuadrupoleXY = 2*xscale*yscale*cAmoebaSim.pLabFrameQuadrupole[atomIndex*9+1];
                        float atomQuadrupoleXZ = 2*xscale*zscale*cAmoebaSim.pLabFrameQuadrupole[atomIndex*9+2];
                        float atomQuadrupoleYY = yscale*yscale*cAmoebaSim.pLabFrameQuadrupole[atomIndex*9+4];
                        float atomQuadrupoleYZ = 2*yscale*zscale*cAmoebaSim.pLabFrameQuadrupole[atomIndex*9+5];
                        float atomQuadrupoleZZ = zscale*zscale*cAmoebaSim.pLabFrameQuadrupole[atomIndex*9+8];
311
312
313
                        float4 t = cAmoebaSim.pThetai1[atomIndex*AMOEBA_PME_ORDER+ix];
                        float4 u = cAmoebaSim.pThetai2[atomIndex*AMOEBA_PME_ORDER+iy];
                        float4 v = cAmoebaSim.pThetai3[atomIndex*AMOEBA_PME_ORDER+iz];
314
315
316
317
318
319
320
321
                        float term0 = atomCharge*u.x*v.x + atomDipoleY*u.y*v.x + atomDipoleZ*u.x*v.y + atomQuadrupoleYY*u.z*v.x + atomQuadrupoleZZ*u.x*v.z + atomQuadrupoleYZ*u.y*v.y;
                        float term1 = atomDipoleX*u.x*v.x + atomQuadrupoleXY*u.y*v.x + atomQuadrupoleXZ*u.x*v.y;
                        float term2 = atomQuadrupoleXX * u.x * v.x;
                        result += term0*t.x + term1*t.y + term2*t.z;
                    }
                }
            }
        }
322
323
324
325
        cSim.pPmeGrid[gridIndex] = make_cuComplex(result, 0.0f);
    }
}

326
__global__
Peter Eastman's avatar
Peter Eastman committed
327
__launch_bounds__(64, 10)
328
329
void kGridSpreadInducedDipoles_kernel()
{
330
331
332
    const float xscale = cSim.pmeGridSize.x*cSim.invPeriodicBoxSizeX;
    const float yscale = cSim.pmeGridSize.y*cSim.invPeriodicBoxSizeY;
    const float zscale = cSim.pmeGridSize.z*cSim.invPeriodicBoxSizeZ;
333
334
335
336
337
338
339
340
341
342
    unsigned int numGridPoints = cSim.pmeGridSize.x*cSim.pmeGridSize.y*cSim.pmeGridSize.z;
    unsigned int numThreads = gridDim.x*blockDim.x;
    for (int gridIndex = blockIdx.x*blockDim.x+threadIdx.x; gridIndex < numGridPoints; gridIndex += numThreads)
    {
        int3 gridPoint;
        gridPoint.x = gridIndex/(cSim.pmeGridSize.y*cSim.pmeGridSize.z);
        int remainder = gridIndex-gridPoint.x*cSim.pmeGridSize.y*cSim.pmeGridSize.z;
        gridPoint.y = remainder/cSim.pmeGridSize.z;
        gridPoint.z = remainder-gridPoint.y*cSim.pmeGridSize.z;
        cufftComplex result = make_cuComplex(0.0f, 0.0f);
343
        for (int ix = 0; ix < AMOEBA_PME_ORDER; ++ix)
344
345
        {
            int x = gridPoint.x-ix+(gridPoint.x >= ix ? 0 : cSim.pmeGridSize.x);
346
            for (int iy = 0; iy < AMOEBA_PME_ORDER; ++iy)
347
348
            {
                int y = gridPoint.y-iy+(gridPoint.y >= iy ? 0 : cSim.pmeGridSize.y);
349
                int z1 = gridPoint.z-AMOEBA_PME_ORDER+1;
350
351
352
353
354
355
356
357
358
359
360
361
                z1 += (z1 >= 0 ? 0 : cSim.pmeGridSize.z);
                int z2 = (z1 < gridPoint.z ? gridPoint.z : cSim.pmeGridSize.z-1);
                int gridIndex1 = x*cSim.pmeGridSize.y*cSim.pmeGridSize.z+y*cSim.pmeGridSize.z+z1;
                int gridIndex2 = x*cSim.pmeGridSize.y*cSim.pmeGridSize.z+y*cSim.pmeGridSize.z+z2;
                int firstAtom = cSim.pPmeAtomRange[gridIndex1];
                int lastAtom = cSim.pPmeAtomRange[gridIndex2+1];
                for (int i = firstAtom; i < lastAtom; ++i)
                {
                    int2 atomData = cSim.pPmeAtomGridIndex[i];
                    int atomIndex = atomData.x;
                    int z = atomData.y;
                    int iz = gridPoint.z-z+(gridPoint.z >= z ? 0 : cSim.pmeGridSize.z);
Mark Friedrichs's avatar
Mark Friedrichs committed
362
363
364
                    if( iz >= cSim.pmeGridSize.z ){
                        iz -= cSim.pmeGridSize.z;
                    }
365
366
367
368
369
370
                    float inducedDipoleX = xscale*cAmoebaSim.pInducedDipole[atomIndex*3];
                    float inducedDipoleY = yscale*cAmoebaSim.pInducedDipole[atomIndex*3+1];
                    float inducedDipoleZ = zscale*cAmoebaSim.pInducedDipole[atomIndex*3+2];
                    float inducedDipolePolarX = xscale*cAmoebaSim.pInducedDipolePolar[atomIndex*3];
                    float inducedDipolePolarY = yscale*cAmoebaSim.pInducedDipolePolar[atomIndex*3+1];
                    float inducedDipolePolarZ = zscale*cAmoebaSim.pInducedDipolePolar[atomIndex*3+2];
371
372
373
                    float4 t = cAmoebaSim.pThetai1[atomIndex*AMOEBA_PME_ORDER+ix];
                    float4 u = cAmoebaSim.pThetai2[atomIndex*AMOEBA_PME_ORDER+iy];
                    float4 v = cAmoebaSim.pThetai3[atomIndex*AMOEBA_PME_ORDER+iz];
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
                    float term01 = inducedDipoleY*u.y*v.x + inducedDipoleZ*u.x*v.y;
                    float term11 = inducedDipoleX*u.x*v.x;
                    float term02 = inducedDipolePolarY*u.y*v.x + inducedDipolePolarZ*u.x*v.y;
                    float term12 = inducedDipolePolarX*u.x*v.x;
                    result.x += term01*t.x + term11*t.y;
                    result.y += term02*t.x + term12*t.y;
                }
                if (z1 > gridPoint.z)
                {
                    gridIndex1 = x*cSim.pmeGridSize.y*cSim.pmeGridSize.z+y*cSim.pmeGridSize.z;
                    gridIndex2 = x*cSim.pmeGridSize.y*cSim.pmeGridSize.z+y*cSim.pmeGridSize.z+gridPoint.z;
                    firstAtom = cSim.pPmeAtomRange[gridIndex1];
                    lastAtom = cSim.pPmeAtomRange[gridIndex2+1];
                    for (int i = firstAtom; i < lastAtom; ++i)
                    {
                        int2 atomData = cSim.pPmeAtomGridIndex[i];
                        int atomIndex = atomData.x;
                        int z = atomData.y;
                        int iz = gridPoint.z-z+(gridPoint.z >= z ? 0 : cSim.pmeGridSize.z);
Mark Friedrichs's avatar
Mark Friedrichs committed
393
394
395
                        if( iz >= cSim.pmeGridSize.z ){
                            iz -= cSim.pmeGridSize.z;
                        }
396
397
398
399
400
401
                        float inducedDipoleX = xscale*cAmoebaSim.pInducedDipole[atomIndex*3];
                        float inducedDipoleY = yscale*cAmoebaSim.pInducedDipole[atomIndex*3+1];
                        float inducedDipoleZ = zscale*cAmoebaSim.pInducedDipole[atomIndex*3+2];
                        float inducedDipolePolarX = xscale*cAmoebaSim.pInducedDipolePolar[atomIndex*3];
                        float inducedDipolePolarY = yscale*cAmoebaSim.pInducedDipolePolar[atomIndex*3+1];
                        float inducedDipolePolarZ = zscale*cAmoebaSim.pInducedDipolePolar[atomIndex*3+2];
402
403
404
                        float4 t = cAmoebaSim.pThetai1[atomIndex*AMOEBA_PME_ORDER+ix];
                        float4 u = cAmoebaSim.pThetai2[atomIndex*AMOEBA_PME_ORDER+iy];
                        float4 v = cAmoebaSim.pThetai3[atomIndex*AMOEBA_PME_ORDER+iz];
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
                        float term01 = inducedDipoleY*u.y*v.x + inducedDipoleZ*u.x*v.y;
                        float term11 = inducedDipoleX*u.x*v.x;
                        float term02 = inducedDipolePolarY*u.y*v.x + inducedDipolePolarZ*u.x*v.y;
                        float term12 = inducedDipolePolarX*u.x*v.x;
                        result.x += term01*t.x + term11*t.y;
                        result.y += term02*t.x + term12*t.y;
                    }
                }
            }
        }
        cSim.pPmeGrid[gridIndex] = result;
    }
}

__global__
#if (__CUDA_ARCH__ >= 200)
421
422
423
__launch_bounds__(768, 1)
#elif (__CUDA_ARCH__ >= 120)
__launch_bounds__(384, 1)
424
#else
425
__launch_bounds__(192, 1)
426
427
428
429
430
431
432
433
434
435
436
437
#endif
void kAmoebaReciprocalConvolution_kernel()
{
    const unsigned int gridSize = cSim.pmeGridSize.x*cSim.pmeGridSize.y*cSim.pmeGridSize.z;
    float expFactor = LOCAL_HACK_PI*LOCAL_HACK_PI/(cSim.alphaEwald*cSim.alphaEwald);
    float scaleFactor = 1.0/(LOCAL_HACK_PI*cSim.periodicBoxSizeX*cSim.periodicBoxSizeY*cSim.periodicBoxSizeZ);
    for (int index = blockIdx.x*blockDim.x+threadIdx.x; index < gridSize; index += blockDim.x*gridDim.x)
    {
        int kx = index/(cSim.pmeGridSize.y*cSim.pmeGridSize.z);
        int remainder = index-kx*cSim.pmeGridSize.y*cSim.pmeGridSize.z;
        int ky = remainder/cSim.pmeGridSize.z;
        int kz = remainder-ky*cSim.pmeGridSize.z;
438
439
        if (kx == 0 && ky == 0 && kz == 0){
            cSim.pPmeGrid[index] = make_cuComplex(0.0f, 0.0f);
440
            continue;
441
        }
442
443
444
445
446
447
448
449
450
451
452
453
        int mx = (kx < (cSim.pmeGridSize.x+1)/2) ? kx : (kx-cSim.pmeGridSize.x);
        int my = (ky < (cSim.pmeGridSize.y+1)/2) ? ky : (ky-cSim.pmeGridSize.y);
        int mz = (kz < (cSim.pmeGridSize.z+1)/2) ? kz : (kz-cSim.pmeGridSize.z);
        float mhx = mx*cSim.invPeriodicBoxSizeX;
        float mhy = my*cSim.invPeriodicBoxSizeY;
        float mhz = mz*cSim.invPeriodicBoxSizeZ;
        float bx = cSim.pPmeBsplineModuli[0][kx];
        float by = cSim.pPmeBsplineModuli[1][ky];
        float bz = cSim.pPmeBsplineModuli[2][kz];
        cuComplex grid = cSim.pPmeGrid[index];
        float m2 = mhx*mhx+mhy*mhy+mhz*mhz;
        float denom = m2*bx*by*bz;
454
        float eterm = scaleFactor*expf(-expFactor*m2)/denom;
455
456
457
458
        cSim.pPmeGrid[index] = make_cuComplex(grid.x*eterm, grid.y*eterm);
    }
}

459
460
__global__
#if (__CUDA_ARCH__ >= 200)
461
__launch_bounds__(384, 1)
Peter Eastman's avatar
Peter Eastman committed
462
#elif (__CUDA_ARCH__ >= 120)
463
__launch_bounds__(192, 1)
Peter Eastman's avatar
Peter Eastman committed
464
465
#else
__launch_bounds__(96, 1)
466
467
468
469
470
471
#endif
void kComputeFixedPotentialFromGrid_kernel()
{
    // extract the permanent multipole field at each site

    for (int m = blockIdx.x*blockDim.x+threadIdx.x; m < cSim.atoms; m += blockDim.x*gridDim.x) {
472
        int4 gridPoint = cAmoebaSim.pIgrid[m];
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
        float tuv000 = 0.0f;
        float tuv001 = 0.0f;
        float tuv010 = 0.0f;
        float tuv100 = 0.0f;
        float tuv200 = 0.0f;
        float tuv020 = 0.0f;
        float tuv002 = 0.0f;
        float tuv110 = 0.0f;
        float tuv101 = 0.0f;
        float tuv011 = 0.0f;
        float tuv300 = 0.0f;
        float tuv030 = 0.0f;
        float tuv003 = 0.0f;
        float tuv210 = 0.0f;
        float tuv201 = 0.0f;
        float tuv120 = 0.0f;
        float tuv021 = 0.0f;
        float tuv102 = 0.0f;
        float tuv012 = 0.0f;
        float tuv111 = 0.0f;
493
494
495
        for (int iz = 0; iz < AMOEBA_PME_ORDER; iz++) {
            int k = gridPoint.z+iz-(gridPoint.z+iz >= cSim.pmeGridSize.z ? cSim.pmeGridSize.z : 0);
            float4 v = cAmoebaSim.pThetai3[m*AMOEBA_PME_ORDER+iz];
496
497
498
499
500
501
502
503
504
505
            float tu00 = 0.0f;
            float tu10 = 0.0f;
            float tu01 = 0.0f;
            float tu20 = 0.0f;
            float tu11 = 0.0f;
            float tu02 = 0.0f;
            float tu30 = 0.0f;
            float tu21 = 0.0f;
            float tu12 = 0.0f;
            float tu03 = 0.0f;
506
507
508
            for (int iy = 0; iy < AMOEBA_PME_ORDER; iy++) {
                int j = gridPoint.y+iy-(gridPoint.y+iy >= cSim.pmeGridSize.y ? cSim.pmeGridSize.y : 0);
                float4 u = cAmoebaSim.pThetai2[m*AMOEBA_PME_ORDER+iy];
509
                float4 t = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
510
511
                for (int ix = 0; ix < AMOEBA_PME_ORDER; ix++) {
                    int i = gridPoint.x+ix-(gridPoint.x+ix >= cSim.pmeGridSize.x ? cSim.pmeGridSize.x : 0);
512
513
                    int gridIndex = i*cSim.pmeGridSize.y*cSim.pmeGridSize.z + j*cSim.pmeGridSize.z + k;
                    float tq = cSim.pPmeGrid[gridIndex].x;
514
                    float4 tadd = cAmoebaSim.pThetai1[m*AMOEBA_PME_ORDER+ix];
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
                    t.x += tq*tadd.x;
                    t.y += tq*tadd.y;
                    t.z += tq*tadd.z;
                    t.w += tq*tadd.w;
                }
                tu00 += t.x*u.x;
                tu10 += t.y*u.x;
                tu01 += t.x*u.y;
                tu20 += t.z*u.x;
                tu11 += t.y*u.y;
                tu02 += t.x*u.z;
                tu30 += t.w*u.x;
                tu21 += t.z*u.y;
                tu12 += t.y*u.z;
                tu03 += t.x*u.w;
            }
            tuv000 += tu00*v.x;
            tuv100 += tu10*v.x;
            tuv010 += tu01*v.x;
            tuv001 += tu00*v.y;
            tuv200 += tu20*v.x;
            tuv020 += tu02*v.x;
            tuv002 += tu00*v.z;
            tuv110 += tu11*v.x;
            tuv101 += tu10*v.y;
            tuv011 += tu01*v.y;
            tuv300 += tu30*v.x;
            tuv030 += tu03*v.x;
            tuv003 += tu00*v.w;
            tuv210 += tu21*v.x;
            tuv201 += tu20*v.y;
            tuv120 += tu12*v.x;
            tuv021 += tu02*v.y;
            tuv102 += tu10*v.z;
            tuv012 += tu01*v.z;
            tuv111 += tu11*v.y;
        }
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
        cAmoebaSim.pPhi[20*m] = tuv000;
        cAmoebaSim.pPhi[20*m+1] = tuv100;
        cAmoebaSim.pPhi[20*m+2] = tuv010;
        cAmoebaSim.pPhi[20*m+3] = tuv001;
        cAmoebaSim.pPhi[20*m+4] = tuv200;
        cAmoebaSim.pPhi[20*m+5] = tuv020;
        cAmoebaSim.pPhi[20*m+6] = tuv002;
        cAmoebaSim.pPhi[20*m+7] = tuv110;
        cAmoebaSim.pPhi[20*m+8] = tuv101;
        cAmoebaSim.pPhi[20*m+9] = tuv011;
        cAmoebaSim.pPhi[20*m+10] = tuv300;
        cAmoebaSim.pPhi[20*m+11] = tuv030;
        cAmoebaSim.pPhi[20*m+12] = tuv003;
        cAmoebaSim.pPhi[20*m+13] = tuv210;
        cAmoebaSim.pPhi[20*m+14] = tuv201;
        cAmoebaSim.pPhi[20*m+15] = tuv120;
        cAmoebaSim.pPhi[20*m+16] = tuv021;
        cAmoebaSim.pPhi[20*m+17] = tuv102;
        cAmoebaSim.pPhi[20*m+18] = tuv012;
        cAmoebaSim.pPhi[20*m+19] = tuv111;
572
573
574
575
576
    }
}

__global__
#if (__CUDA_ARCH__ >= 200)
Peter Eastman's avatar
Peter Eastman committed
577
__launch_bounds__(256, 1)
578
#elif (__CUDA_ARCH__ >= 120)
Peter Eastman's avatar
Peter Eastman committed
579
__launch_bounds__(128, 1)
580
#else
Peter Eastman's avatar
Peter Eastman committed
581
__launch_bounds__(64, 1)
582
583
584
585
586
587
#endif
void kComputeInducedPotentialFromGrid_kernel()
{
    // extract the induced dipole field at each site

    for (int m = blockIdx.x*blockDim.x+threadIdx.x; m < cSim.atoms; m += blockDim.x*gridDim.x) {
588
        int4 gridPoint = cAmoebaSim.pIgrid[m];
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
        float tuv100_1 = 0.0f;
        float tuv010_1 = 0.0f;
        float tuv001_1 = 0.0f;
        float tuv200_1 = 0.0f;
        float tuv020_1 = 0.0f;
        float tuv002_1 = 0.0f;
        float tuv110_1 = 0.0f;
        float tuv101_1 = 0.0f;
        float tuv011_1 = 0.0f;
        float tuv100_2 = 0.0f;
        float tuv010_2 = 0.0f;
        float tuv001_2 = 0.0f;
        float tuv200_2 = 0.0f;
        float tuv020_2 = 0.0f;
        float tuv002_2 = 0.0f;
        float tuv110_2 = 0.0f;
        float tuv101_2 = 0.0f;
        float tuv011_2 = 0.0f;
        float tuv000 = 0.0f;
        float tuv001 = 0.0f;
        float tuv010 = 0.0f;
        float tuv100 = 0.0f;
        float tuv200 = 0.0f;
        float tuv020 = 0.0f;
        float tuv002 = 0.0f;
        float tuv110 = 0.0f;
        float tuv101 = 0.0f;
        float tuv011 = 0.0f;
        float tuv300 = 0.0f;
        float tuv030 = 0.0f;
        float tuv003 = 0.0f;
        float tuv210 = 0.0f;
        float tuv201 = 0.0f;
        float tuv120 = 0.0f;
        float tuv021 = 0.0f;
        float tuv102 = 0.0f;
        float tuv012 = 0.0f;
        float tuv111 = 0.0f;
627
628
629
        for (int iz = 0; iz < AMOEBA_PME_ORDER; iz++) {
            int k = gridPoint.z+iz-(gridPoint.z+iz >= cSim.pmeGridSize.z ? cSim.pmeGridSize.z : 0);
            float4 v = cAmoebaSim.pThetai3[m*AMOEBA_PME_ORDER+iz];
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
            float tu00_1 = 0.0f;
            float tu01_1 = 0.0f;
            float tu10_1 = 0.0f;
            float tu20_1 = 0.0f;
            float tu11_1 = 0.0f;
            float tu02_1 = 0.0f;
            float tu00_2 = 0.0f;
            float tu01_2 = 0.0f;
            float tu10_2 = 0.0f;
            float tu20_2 = 0.0f;
            float tu11_2 = 0.0f;
            float tu02_2 = 0.0f;
            float tu00 = 0.0f;
            float tu10 = 0.0f;
            float tu01 = 0.0f;
            float tu20 = 0.0f;
            float tu11 = 0.0f;
            float tu02 = 0.0f;
            float tu30 = 0.0f;
            float tu21 = 0.0f;
            float tu12 = 0.0f;
            float tu03 = 0.0f;
652
653
654
            for (int iy = 0; iy < AMOEBA_PME_ORDER; iy++) {
                int j = gridPoint.y+iy-(gridPoint.y+iy >= cSim.pmeGridSize.y ? cSim.pmeGridSize.y : 0);
                float4 u = cAmoebaSim.pThetai2[m*AMOEBA_PME_ORDER+iy];
655
656
657
658
659
660
661
                float t0_1 = 0.0f;
                float t1_1 = 0.0f;
                float t2_1 = 0.0f;
                float t0_2 = 0.0f;
                float t1_2 = 0.0f;
                float t2_2 = 0.0f;
                float t3 = 0.0f;
662
663
                for (int ix = 0; ix < AMOEBA_PME_ORDER; ix++) {
                    int i = gridPoint.x+ix-(gridPoint.x+ix >= cSim.pmeGridSize.x ? cSim.pmeGridSize.x : 0);
664
665
                    int gridIndex = i*cSim.pmeGridSize.y*cSim.pmeGridSize.z + j*cSim.pmeGridSize.z + k;
                    cufftComplex tq = cSim.pPmeGrid[gridIndex];
666
                    float4 tadd = cAmoebaSim.pThetai1[m*AMOEBA_PME_ORDER+ix];
667
668
669
670
671
672
                    t0_1 += tq.x*tadd.x;
                    t1_1 += tq.x*tadd.y;
                    t2_1 += tq.x*tadd.z;
                    t0_2 += tq.y*tadd.x;
                    t1_2 += tq.y*tadd.y;
                    t2_2 += tq.y*tadd.z;
673
                    t3 += (tq.x+tq.y)*tadd.w;
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
                }
                tu00_1 += t0_1*u.x;
                tu10_1 += t1_1*u.x;
                tu01_1 += t0_1*u.y;
                tu20_1 += t2_1*u.x;
                tu11_1 += t1_1*u.y;
                tu02_1 += t0_1*u.z;
                tu00_2 += t0_2*u.x;
                tu10_2 += t1_2*u.x;
                tu01_2 += t0_2*u.y;
                tu20_2 += t2_2*u.x;
                tu11_2 += t1_2*u.y;
                tu02_2 += t0_2*u.z;
                float t0 = t0_1 + t0_2;
                float t1 = t1_1 + t1_2;
                float t2 = t2_1 + t2_2;
                tu00 += t0*u.x;
                tu10 += t1*u.x;
                tu01 += t0*u.y;
                tu20 += t2*u.x;
                tu11 += t1*u.y;
                tu02 += t0*u.z;
                tu30 += t3*u.x;
                tu21 += t2*u.y;
                tu12 += t1*u.z;
                tu03 += t0*u.w;
            }
            tuv100_1 += tu10_1*v.x;
            tuv010_1 += tu01_1*v.x;
            tuv001_1 += tu00_1*v.y;
            tuv200_1 += tu20_1*v.x;
            tuv020_1 += tu02_1*v.x;
            tuv002_1 += tu00_1*v.z;
            tuv110_1 += tu11_1*v.x;
            tuv101_1 += tu10_1*v.y;
            tuv011_1 += tu01_1*v.y;
            tuv100_2 += tu10_2*v.x;
            tuv010_2 += tu01_2*v.x;
            tuv001_2 += tu00_2*v.y;
            tuv200_2 += tu20_2*v.x;
            tuv020_2 += tu02_2*v.x;
            tuv002_2 += tu00_2*v.z;
            tuv110_2 += tu11_2*v.x;
            tuv101_2 += tu10_2*v.y;
            tuv011_2 += tu01_2*v.y;
            tuv000 += tu00*v.x;
            tuv100 += tu10*v.x;
            tuv010 += tu01*v.x;
            tuv001 += tu00*v.y;
            tuv200 += tu20*v.x;
            tuv020 += tu02*v.x;
            tuv002 += tu00*v.z;
            tuv110 += tu11*v.x;
            tuv101 += tu10*v.y;
            tuv011 += tu01*v.y;
            tuv300 += tu30*v.x;
            tuv030 += tu03*v.x;
            tuv003 += tu00*v.w;
            tuv210 += tu21*v.x;
            tuv201 += tu20*v.y;
            tuv120 += tu12*v.x;
            tuv021 += tu02*v.y;
            tuv102 += tu10*v.z;
            tuv012 += tu01*v.z;
            tuv111 += tu11*v.y;
        }
740
741
742
        cAmoebaSim.pPhid[10*m+1] = tuv100_1;
        cAmoebaSim.pPhid[10*m+2] = tuv010_1;
        cAmoebaSim.pPhid[10*m+3] = tuv001_1;
Mark Friedrichs's avatar
Mark Friedrichs committed
743
744
        cAmoebaSim.pPhid[10*m+4] = tuv200_1;
        cAmoebaSim.pPhid[10*m+5] = tuv020_1;
745
746
747
748
        cAmoebaSim.pPhid[10*m+6] = tuv002_1;
        cAmoebaSim.pPhid[10*m+7] = tuv110_1;
        cAmoebaSim.pPhid[10*m+8] = tuv101_1;
        cAmoebaSim.pPhid[10*m+9] = tuv011_1;
Mark Friedrichs's avatar
Mark Friedrichs committed
749

750
751
752
        cAmoebaSim.pPhip[10*m+1] = tuv100_2;
        cAmoebaSim.pPhip[10*m+2] = tuv010_2;
        cAmoebaSim.pPhip[10*m+3] = tuv001_2;
Mark Friedrichs's avatar
Mark Friedrichs committed
753
754
        cAmoebaSim.pPhip[10*m+4] = tuv200_2;
        cAmoebaSim.pPhip[10*m+5] = tuv020_2;
755
756
757
758
        cAmoebaSim.pPhip[10*m+6] = tuv002_2;
        cAmoebaSim.pPhip[10*m+7] = tuv110_2;
        cAmoebaSim.pPhip[10*m+8] = tuv101_2;
        cAmoebaSim.pPhip[10*m+9] = tuv011_2;
Mark Friedrichs's avatar
Mark Friedrichs committed
759

760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
        cAmoebaSim.pPhidp[20*m] = tuv000;
        cAmoebaSim.pPhidp[20*m+1] = tuv100;
        cAmoebaSim.pPhidp[20*m+2] = tuv010;
        cAmoebaSim.pPhidp[20*m+3] = tuv001;
        cAmoebaSim.pPhidp[20*m+4] = tuv200;
        cAmoebaSim.pPhidp[20*m+5] = tuv020;
        cAmoebaSim.pPhidp[20*m+6] = tuv002;
        cAmoebaSim.pPhidp[20*m+7] = tuv110;
        cAmoebaSim.pPhidp[20*m+8] = tuv101;
        cAmoebaSim.pPhidp[20*m+9] = tuv011;
        cAmoebaSim.pPhidp[20*m+10] = tuv300;
        cAmoebaSim.pPhidp[20*m+11] = tuv030;
        cAmoebaSim.pPhidp[20*m+12] = tuv003;
        cAmoebaSim.pPhidp[20*m+13] = tuv210;
        cAmoebaSim.pPhidp[20*m+14] = tuv201;
        cAmoebaSim.pPhidp[20*m+15] = tuv120;
        cAmoebaSim.pPhidp[20*m+16] = tuv021;
        cAmoebaSim.pPhidp[20*m+17] = tuv102;
        cAmoebaSim.pPhidp[20*m+18] = tuv012;
        cAmoebaSim.pPhidp[20*m+19] = tuv111;
780
781
    }
}
782
783
784

__global__
#if (__CUDA_ARCH__ >= 200)
785
__launch_bounds__(768, 1)
786
#elif (__CUDA_ARCH__ >= 120)
787
__launch_bounds__(384, 1)
788
#else
789
__launch_bounds__(192, 1)
790
791
792
793
#endif
void kComputeFixedMultipoleForceAndEnergy_kernel()
{
    float multipole[10];
794
795
796
    const int deriv1[] = {1, 4, 7, 8, 10, 15, 17, 13, 14, 19};
    const int deriv2[] = {2, 7, 5, 9, 13, 11, 18, 15, 19, 16};
    const int deriv3[] = {3, 8, 9, 6, 14, 16, 12, 19, 17, 18};
797
798
799
    const float xscale = cSim.pmeGridSize.x*cSim.invPeriodicBoxSizeX;
    const float yscale = cSim.pmeGridSize.y*cSim.invPeriodicBoxSizeY;
    const float zscale = cSim.pmeGridSize.z*cSim.invPeriodicBoxSizeZ;
800
801
    float energy = 0.0f;
    for (int i = blockIdx.x*blockDim.x+threadIdx.x; i < cSim.atoms; i += blockDim.x*gridDim.x) {
802
        // Compute the torque.
803

804
        multipole[0] = cSim.pPosq[i].w;
805
806
807
808
        multipole[1] = cAmoebaSim.pLabFrameDipole[i*3];
        multipole[2] = cAmoebaSim.pLabFrameDipole[i*3+1];
        multipole[3] = cAmoebaSim.pLabFrameDipole[i*3+2];
        multipole[4] = cAmoebaSim.pLabFrameQuadrupole[i*9];
809
810
811
812
813
        multipole[5] = cAmoebaSim.pLabFrameQuadrupole[i*9+4];
        multipole[6] = cAmoebaSim.pLabFrameQuadrupole[i*9+8];
        multipole[7] = 2*cAmoebaSim.pLabFrameQuadrupole[i*9+1];
        multipole[8] = 2*cAmoebaSim.pLabFrameQuadrupole[i*9+2];
        multipole[9] = 2*cAmoebaSim.pLabFrameQuadrupole[i*9+5];
814
        float* phi = &cAmoebaSim.pPhi[20*i];
Mark Friedrichs's avatar
Mark Friedrichs committed
815
        cAmoebaSim.pTorque[3*i] = cAmoebaSim.electric*(multipole[3]*yscale*phi[2] - multipole[2]*zscale*phi[3]
816
817
818
                      + 2.0f*(multipole[6]-multipole[5])*zscale*zscale*phi[9]
                      + multipole[8]*yscale*yscale*phi[7] + multipole[9]*xscale*yscale*phi[5]
                      - multipole[7]*yscale*zscale*phi[8] - multipole[9]*xscale*zscale*phi[6]);
Mark Friedrichs's avatar
Mark Friedrichs committed
819
        cAmoebaSim.pTorque[3*i+1] = cAmoebaSim.electric*(multipole[1]*zscale*phi[3] - multipole[3]*xscale*phi[1]
820
821
822
                      + 2.0f*(multipole[4]-multipole[6])*zscale*zscale*phi[8]
                      + multipole[7]*zscale*zscale*phi[9] + multipole[8]*xscale*zscale*phi[6]
                      - multipole[8]*xscale*xscale*phi[4] - multipole[9]*yscale*yscale*phi[7]);
Mark Friedrichs's avatar
Mark Friedrichs committed
823
        cAmoebaSim.pTorque[3*i+2] = cAmoebaSim.electric*(multipole[2]*xscale*phi[1] - multipole[1]*yscale*phi[2]
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
                      + 2.0f*(multipole[5]-multipole[4])*yscale*yscale*phi[7]
                      + multipole[7]*xscale*xscale*phi[4] + multipole[9]*yscale*zscale*phi[8]
                      - multipole[7]*xscale*yscale*phi[5] - multipole[8]*zscale*zscale*phi[9]);

        // Compute the force and energy.

        multipole[1] *= xscale;
        multipole[2] *= yscale;
        multipole[3] *= zscale;
        multipole[4] *= xscale*xscale;
        multipole[5] *= xscale*yscale;
        multipole[6] *= xscale*zscale;
        multipole[7] *= yscale*yscale;
        multipole[8] *= yscale*zscale;
        multipole[9] *= zscale*zscale;
839
840
        float4 f = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
        for (int k = 0; k < 10; k++) {
841
842
843
844
            energy += multipole[k]*phi[k];
            f.x += multipole[k]*phi[deriv1[k]];
            f.y += multipole[k]*phi[deriv2[k]];
            f.z += multipole[k]*phi[deriv3[k]];
845
        }
846
847
848
        f.x *= cAmoebaSim.electric*cSim.pmeGridSize.x*cSim.invPeriodicBoxSizeX;
        f.y *= cAmoebaSim.electric*cSim.pmeGridSize.y*cSim.invPeriodicBoxSizeY;
        f.z *= cAmoebaSim.electric*cSim.pmeGridSize.z*cSim.invPeriodicBoxSizeZ;
849
        float4 force = cSim.pForce4[i];
Mark Friedrichs's avatar
Mark Friedrichs committed
850
851
852
        force.x -= f.x;
        force.y -= f.y;
        force.z -= f.z;
853
        cSim.pForce4[i] = force;
854
855


856
    }
857
    cSim.pEnergy[blockIdx.x*blockDim.x+threadIdx.x] += 0.5f*cAmoebaSim.electric*energy;
858
859
860
861
}

__global__
#if (__CUDA_ARCH__ >= 200)
862
__launch_bounds__(768, 1)
863
#elif (__CUDA_ARCH__ >= 120)
864
__launch_bounds__(384, 1)
865
#else
866
__launch_bounds__(192, 1)
867
868
869
870
871
872
#endif
void kComputeInducedDipoleForceAndEnergy_kernel()
{
    float multipole[10];
    float inducedDipole[3];
    float inducedDipolePolar[3];
873
874
875
    const int deriv1[] = {1, 4, 7, 8, 10, 15, 17, 13, 14, 19};
    const int deriv2[] = {2, 7, 5, 9, 13, 11, 18, 15, 19, 16};
    const int deriv3[] = {3, 8, 9, 6, 14, 16, 12, 19, 17, 18};
876
877
878
    const float xscale = cSim.pmeGridSize.x*cSim.invPeriodicBoxSizeX;
    const float yscale = cSim.pmeGridSize.y*cSim.invPeriodicBoxSizeY;
    const float zscale = cSim.pmeGridSize.z*cSim.invPeriodicBoxSizeZ;
879
880
    float energy = 0.0f;
    for (int i = blockIdx.x*blockDim.x+threadIdx.x; i < cSim.atoms; i += blockDim.x*gridDim.x) {
881
        // Compute the torque.
882

883
        multipole[0] = cSim.pPosq[i].w;
884
885
886
887
        multipole[1] = cAmoebaSim.pLabFrameDipole[i*3];
        multipole[2] = cAmoebaSim.pLabFrameDipole[i*3+1];
        multipole[3] = cAmoebaSim.pLabFrameDipole[i*3+2];
        multipole[4] = cAmoebaSim.pLabFrameQuadrupole[i*9];
888
889
890
891
892
893
        multipole[5] = cAmoebaSim.pLabFrameQuadrupole[i*9+4];
        multipole[6] = cAmoebaSim.pLabFrameQuadrupole[i*9+8];
        multipole[7] = 2*cAmoebaSim.pLabFrameQuadrupole[i*9+1];
        multipole[8] = 2*cAmoebaSim.pLabFrameQuadrupole[i*9+2];
        multipole[9] = 2*cAmoebaSim.pLabFrameQuadrupole[i*9+5];
        float* phidp = &cAmoebaSim.pPhidp[20*i];
894
        cAmoebaSim.pTorque[3*i] += 0.5f*cAmoebaSim.electric*(multipole[3]*yscale*phidp[2] - multipole[2]*zscale*phidp[3]
895
896
897
                      + 2.0f*(multipole[6]-multipole[5])*zscale*zscale*phidp[9]
                      + multipole[8]*yscale*yscale*phidp[7] + multipole[9]*xscale*yscale*phidp[5]
                      - multipole[7]*yscale*zscale*phidp[8] - multipole[9]*xscale*zscale*phidp[6]);
898
        cAmoebaSim.pTorque[3*i+1] += 0.5f*cAmoebaSim.electric*(multipole[1]*zscale*phidp[3] - multipole[3]*xscale*phidp[1]
899
900
901
                      + 2.0f*(multipole[4]-multipole[6])*zscale*zscale*phidp[8]
                      + multipole[7]*zscale*zscale*phidp[9] + multipole[8]*xscale*zscale*phidp[6]
                      - multipole[8]*xscale*xscale*phidp[4] - multipole[9]*yscale*yscale*phidp[7]);
902
        cAmoebaSim.pTorque[3*i+2] += 0.5f*cAmoebaSim.electric*(multipole[2]*xscale*phidp[1] - multipole[1]*yscale*phidp[2]
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
                      + 2.0f*(multipole[5]-multipole[4])*yscale*yscale*phidp[7]
                      + multipole[7]*xscale*xscale*phidp[4] + multipole[9]*yscale*zscale*phidp[8]
                      - multipole[7]*xscale*yscale*phidp[5] - multipole[8]*zscale*zscale*phidp[9]);

        // Compute the force and energy.

        multipole[1] *= xscale;
        multipole[2] *= yscale;
        multipole[3] *= zscale;
        multipole[4] *= xscale*xscale;
        multipole[5] *= xscale*yscale;
        multipole[6] *= xscale*zscale;
        multipole[7] *= yscale*yscale;
        multipole[8] *= yscale*zscale;
        multipole[9] *= zscale*zscale;
918
919
920
921
922
923
        inducedDipole[0] = cAmoebaSim.pInducedDipole[i*3];
        inducedDipole[1] = cAmoebaSim.pInducedDipole[i*3+1];
        inducedDipole[2] = cAmoebaSim.pInducedDipole[i*3+2];
        inducedDipolePolar[0] = cAmoebaSim.pInducedDipolePolar[i*3];
        inducedDipolePolar[1] = cAmoebaSim.pInducedDipolePolar[i*3+1];
        inducedDipolePolar[2] = cAmoebaSim.pInducedDipolePolar[i*3+2];
924
        float* phi = &cAmoebaSim.pPhi[20*i];
925
926
        float* phip = &cAmoebaSim.pPhip[10*i];
        float* phid = &cAmoebaSim.pPhid[10*i];
927
        float4 f = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
Mark Friedrichs's avatar
Mark Friedrichs committed
928

Mark Friedrichs's avatar
Mark Friedrichs committed
929
930
931
932
        energy += cSim.pmeGridSize.x*cSim.invPeriodicBoxSizeX*inducedDipole[0]*phi[1];
        energy += cSim.pmeGridSize.y*cSim.invPeriodicBoxSizeY*inducedDipole[1]*phi[2];
        energy += cSim.pmeGridSize.z*cSim.invPeriodicBoxSizeZ*inducedDipole[2]*phi[3];

933
934
        for (int k = 0; k < 3; k++) {
            int j1 = deriv1[k+1];
935
936
            int j2 = deriv2[k+1];
            int j3 = deriv3[k+1];
Mark Friedrichs's avatar
Mark Friedrichs committed
937

938
939
940
941
942
943
944
945
946
947
            f.x += (inducedDipole[k]+inducedDipolePolar[k])*phi[j1];
            f.y += (inducedDipole[k]+inducedDipolePolar[k])*phi[j2];
            f.z += (inducedDipole[k]+inducedDipolePolar[k])*phi[j3];
 
            if( cAmoebaSim.polarizationType == 0 )
            {
                f.x += inducedDipole[k]*phip[j1] + inducedDipolePolar[k]*phid[j1];
                f.y += inducedDipole[k]*phip[j2] + inducedDipolePolar[k]*phid[j2];
                f.z += inducedDipole[k]*phip[j3] + inducedDipolePolar[k]*phid[j3];
            }
Mark Friedrichs's avatar
Mark Friedrichs committed
948

949
        }
Mark Friedrichs's avatar
Mark Friedrichs committed
950
951
952
953
954

        f.x *= cSim.pmeGridSize.x*cSim.invPeriodicBoxSizeX;
        f.y *= cSim.pmeGridSize.y*cSim.invPeriodicBoxSizeY;
        f.z *= cSim.pmeGridSize.z*cSim.invPeriodicBoxSizeZ;

955
        for (int k = 0; k < 10; k++) {
956
957
958
            f.x += multipole[k]*phidp[deriv1[k]];
            f.y += multipole[k]*phidp[deriv2[k]];
            f.z += multipole[k]*phidp[deriv3[k]];
959
        }
Mark Friedrichs's avatar
Mark Friedrichs committed
960

961
962
963
        f.x *= 0.5f*cAmoebaSim.electric*cSim.pmeGridSize.x*cSim.invPeriodicBoxSizeX;
        f.y *= 0.5f*cAmoebaSim.electric*cSim.pmeGridSize.y*cSim.invPeriodicBoxSizeY;
        f.z *= 0.5f*cAmoebaSim.electric*cSim.pmeGridSize.z*cSim.invPeriodicBoxSizeZ;
Mark Friedrichs's avatar
Mark Friedrichs committed
964

965
        float4 force = cSim.pForce4[i];
Mark Friedrichs's avatar
Mark Friedrichs committed
966
967
968
        force.x -= f.x;
        force.y -= f.y;
        force.z -= f.z;
969
970
        cSim.pForce4[i] = force;
    }
971
    cSim.pEnergy[blockIdx.x*blockDim.x+threadIdx.x] += 0.5f*cAmoebaSim.electric*energy;
972
}
973

974
975
976
977
978
979
980
981
982
983
__global__
#if (__CUDA_ARCH__ >= 200)
__launch_bounds__(768, 1)
#elif (__CUDA_ARCH__ >= 120)
__launch_bounds__(384, 1)
#else
__launch_bounds__(192, 1)
#endif
void kRecordFixedMultipoleField_kernel(float* output)
{
984
985
986
    const float xscale = cSim.pmeGridSize.x*cSim.invPeriodicBoxSizeX;
    const float yscale = cSim.pmeGridSize.y*cSim.invPeriodicBoxSizeY;
    const float zscale = cSim.pmeGridSize.z*cSim.invPeriodicBoxSizeZ;
987
    for (int i = blockIdx.x*blockDim.x+threadIdx.x; i < cSim.atoms; i += blockDim.x*gridDim.x) {
988
989
990
        output[3*i] = -xscale*cAmoebaSim.pPhi[20*i+1];
        output[3*i+1] = -yscale*cAmoebaSim.pPhi[20*i+2];
        output[3*i+2] = -zscale*cAmoebaSim.pPhi[20*i+3];
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
    }
}

__global__
#if (__CUDA_ARCH__ >= 200)
__launch_bounds__(768, 1)
#elif (__CUDA_ARCH__ >= 120)
__launch_bounds__(384, 1)
#else
__launch_bounds__(192, 1)
#endif
void kRecordInducedDipoleField_kernel(float* output, float* outputPolar)
{
1004
1005
1006
    const float xscale = cSim.pmeGridSize.x*cSim.invPeriodicBoxSizeX;
    const float yscale = cSim.pmeGridSize.y*cSim.invPeriodicBoxSizeY;
    const float zscale = cSim.pmeGridSize.z*cSim.invPeriodicBoxSizeZ;
1007
    for (int i = blockIdx.x*blockDim.x+threadIdx.x; i < cSim.atoms; i += blockDim.x*gridDim.x) {
1008
1009
1010
1011
1012
1013
        output[3*i] -= xscale*cAmoebaSim.pPhid[10*i+1];
        output[3*i+1] -= yscale*cAmoebaSim.pPhid[10*i+2];
        output[3*i+2] -= zscale*cAmoebaSim.pPhid[10*i+3];
        outputPolar[3*i] -= xscale*cAmoebaSim.pPhip[10*i+1];
        outputPolar[3*i+1] -= yscale*cAmoebaSim.pPhip[10*i+2];
        outputPolar[3*i+2] -= zscale*cAmoebaSim.pPhip[10*i+3];
1014
1015
1016
    }
}

1017
extern void cudaComputeAmoebaMapTorqueAndAddToForce(amoebaGpuContext gpu, CUDAStream<float>* psTorque);
1018

1019
/**
1020
 * Compute the potential and forces due to the reciprocal space PME calculation for fixed multipoles.
1021
 */
1022
void kCalculateAmoebaPMEFixedMultipoles(amoebaGpuContext amoebaGpu)
1023
{
1024
1025
1026
    // Compute B-spline coefficients and sort the atoms.

    gpuContext gpu = amoebaGpu->gpuContext;
Peter Eastman's avatar
Peter Eastman committed
1027
1028
    int bsplineThreads = (gpu->sm_version >= SM_20 ? 448 : (gpu->sm_version >= SM_12 ? 160 : 160));
    kComputeAmoebaBsplines_kernel<<<gpu->sim.blocks, bsplineThreads, bsplineThreads*AMOEBA_PME_ORDER*AMOEBA_PME_ORDER*sizeof(float)>>>();
1029
1030
    LAUNCHERROR("kComputeAmoebaBsplines");
    bbSort(gpu->psPmeAtomGridIndex->_pDevData, gpu->natoms);
1031
1032
    kFindAmoebaAtomRangeForGrid_kernel<<<gpu->sim.blocks, gpu->sim.update_threads_per_block>>>();
    LAUNCHERROR("kFindAmoebaAtomRangeForGrid");
1033
1034
1035

    // Perform PME for the fixed multipoles.

Peter Eastman's avatar
Peter Eastman committed
1036
    kGridSpreadFixedMultipoles_kernel<<<10*gpu->sim.blocks, 64>>>();
1037
1038
1039
1040
1041
    LAUNCHERROR("kGridSpreadFixedMultipoles");
    cufftExecC2C(gpu->fftplan, gpu->psPmeGrid->_pDevData, gpu->psPmeGrid->_pDevData, CUFFT_FORWARD);
    kAmoebaReciprocalConvolution_kernel<<<gpu->sim.blocks, gpu->sim.nonbond_threads_per_block>>>();
    LAUNCHERROR("kAmoebaReciprocalConvolution");
    cufftExecC2C(gpu->fftplan, gpu->psPmeGrid->_pDevData, gpu->psPmeGrid->_pDevData, CUFFT_INVERSE);
Peter Eastman's avatar
Peter Eastman committed
1042
1043
    int potentialThreads = (gpu->sm_version >= SM_20 ? 384 : (gpu->sm_version >= SM_12 ? 192 : 96));
    kComputeFixedPotentialFromGrid_kernel<<<gpu->sim.blocks, potentialThreads>>>();
1044
    LAUNCHERROR("kComputeFixedPotentialFromGrid");
1045
1046
    kRecordFixedMultipoleField_kernel<<<gpu->sim.blocks, gpu->sim.update_threads_per_block>>>(amoebaGpu->psE_Field->_pDevData);
    LAUNCHERROR("kRecordFixedMultipoleField");
1047
1048
    kComputeFixedMultipoleForceAndEnergy_kernel<<<gpu->sim.blocks, gpu->sim.update_threads_per_block>>>();
    LAUNCHERROR("kComputeFixedMultipoleForceAndEnergy");
1049

1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
}

/**
 * Compute the potential due to the reciprocal space PME calculation for induced dipoles.
 */
void kCalculateAmoebaPMEInducedDipoleField(amoebaGpuContext amoebaGpu)
{
    // Perform PME for the induced dipoles.

    gpuContext gpu = amoebaGpu->gpuContext;
Peter Eastman's avatar
Peter Eastman committed
1060
    kGridSpreadInducedDipoles_kernel<<<10*gpu->sim.blocks, 64>>>();
1061
1062
1063
1064
1065
    LAUNCHERROR("kGridSpreadInducedDipoles");
    cufftExecC2C(gpu->fftplan, gpu->psPmeGrid->_pDevData, gpu->psPmeGrid->_pDevData, CUFFT_FORWARD);
    kAmoebaReciprocalConvolution_kernel<<<gpu->sim.blocks, gpu->sim.nonbond_threads_per_block>>>();
    LAUNCHERROR("kAmoebaReciprocalConvolution");
    cufftExecC2C(gpu->fftplan, gpu->psPmeGrid->_pDevData, gpu->psPmeGrid->_pDevData, CUFFT_INVERSE);
Peter Eastman's avatar
Peter Eastman committed
1066
1067
    int potentialThreads = (gpu->sm_version >= SM_20 ? 256 : (gpu->sm_version >= SM_12 ? 128 : 64));
    kComputeInducedPotentialFromGrid_kernel<<<gpu->sim.blocks, potentialThreads>>>();
1068
    LAUNCHERROR("kComputeInducedPotentialFromGrid");
1069
    kRecordInducedDipoleField_kernel<<<gpu->sim.blocks, gpu->sim.update_threads_per_block>>>(amoebaGpu->psWorkVector[0]->_pDevData, amoebaGpu->psWorkVector[1]->_pDevData);
1070
    LAUNCHERROR("kRecordInducedDipoleField");
1071
1072
1073
}

/**
1074
 * Compute the forces due to the reciprocal space PME calculation for induced dipoles.
1075
 */
1076
void kCalculateAmoebaPMEInducedDipoleForces(amoebaGpuContext amoebaGpu)
1077
{
1078
    gpuContext gpu = amoebaGpu->gpuContext;
1079
1080
    kComputeInducedDipoleForceAndEnergy_kernel<<<gpu->sim.blocks, gpu->sim.update_threads_per_block>>>();
    LAUNCHERROR("kComputeInducedDipoleForceAndEnergy");
Mark Friedrichs's avatar
Mark Friedrichs committed
1081

1082
}