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

#include <stdio.h>
28
#include <string.h>
Peter Eastman's avatar
Peter Eastman committed
29
30
31
32
33
34
35
36
37
#include <cuda.h>
#include <vector_functions.h>
#include <cstdlib>
#include <string>
#include <iostream>
#include <fstream>
#include <sstream>
#include <cmath>
#include <map>
38
#include <set>
39
#include <algorithm>
Peter Eastman's avatar
Peter Eastman committed
40
41
42
43
44
45
46
47
48
#ifdef WIN32
  #include <windows.h>
#else
  #include <stdint.h>
#endif
using namespace std;

#include "gputypes.h"
#include "cudaKernels.h"
49
#include "hilbert.h"
50
#include "openmm/OpenMMException.h"
51
#include "jama_svd.h"
52
#include "quern.h"
Peter Eastman's avatar
Peter Eastman committed
53
54

using OpenMM::OpenMMException;
55
56
using TNT::Array2D;
using JAMA::SVD;
Peter Eastman's avatar
Peter Eastman committed
57
58
59
60
61

struct ShakeCluster {
    int centralID;
    int peripheralID[3];
    int size;
62
    bool valid;
Peter Eastman's avatar
Peter Eastman committed
63
64
    float distance;
    float centralInvMass, peripheralInvMass;
65
    ShakeCluster() : valid(true) {
Peter Eastman's avatar
Peter Eastman committed
66
    }
67
    ShakeCluster(int centralID, float invMass) : centralID(centralID), centralInvMass(invMass), size(0), valid(true) {
Peter Eastman's avatar
Peter Eastman committed
68
69
    }
    void addAtom(int id, float dist, float invMass) {
70
71
72
73
74
75
76
        if (size == 3 || (size > 0 && dist != distance) || (size > 0 && invMass != peripheralInvMass))
            valid = false;
        else {
            peripheralID[size++] = id;
            distance = dist;
            peripheralInvMass = invMass;
        }
Peter Eastman's avatar
Peter Eastman committed
77
78
79
    }
};

80
81
82
83
84
85
86
87
struct Constraint
{
    Constraint(int atom1, int atom2, float distance2) : atom1(atom1), atom2(atom2), distance2(distance2) {
    }
    int atom1, atom2;
    float distance2;
};

88
89
90
91
92
93
94
95
96
97
98
99
struct ConstraintOrderer : public binary_function<int, int, bool> {
    const vector<int>& atom1;
    const vector<int>& atom2;
    ConstraintOrderer(const vector<int>& atom1, const vector<int>& atom2) : atom1(atom1), atom2(atom2) {
    }
    bool operator()(int x, int y) {
        if (atom1[x] != atom1[y])
            return atom1[x] < atom1[y];
        return atom2[x] < atom2[y];
    }
};

100
101
102
103
104
105
106
107
108
struct Molecule {
    vector<int> atoms;
    vector<int> bonds;
    vector<int> angles;
    vector<int> periodicTorsions;
    vector<int> rbTorsions;
    vector<int> constraints;
};

Peter Eastman's avatar
Peter Eastman committed
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
static const float dielectricOffset         =    0.009f;
static const float PI                       =    3.1415926535f;
static const float probeRadius              =    0.14f;
static const float forceConversionFactor    =    0.4184f;

//static const float surfaceAreaFactor        =   -6.0f * 0.06786f * forceConversionFactor * 1000.0f;  // PI * 4.0f * 0.0049f * 1000.0f;
//static const float surfaceAreaFactor        =   -6.0f * PI * 4.0f * 0.0049f * 1000.0f;
static const float surfaceAreaFactor        = -6.0f*PI*0.0216f*1000.0f*0.4184f;
//static const float surfaceAreaFactor        = -1.7035573959e+001;
//static const float surfaceAreaFactor        = -166.02691f;
//static const float surfaceAreaFactor        = 1.0f;

static const float alphaOBC                 =    1.0f;
static const float betaOBC                  =    0.8f;
static const float gammaOBC                 =    4.85f;
static const float kcalMolTokJNM            =   -0.4184f;
static const float electricConstant         = -166.02691f;
static const float defaultInnerDielectric   =    1.0f;
static const float defaultSolventDielectric =   78.3f;
static const float KILO                     =    1e3;                      // Thousand
static const float BOLTZMANN                =    1.380658e-23f;            // (J/K)    
static const float AVOGADRO                 =    6.0221367e23f;            // ()        
static const float RGAS                     =    BOLTZMANN * AVOGADRO;     // (J/(mol K))
static const float BOLTZ                    =    (RGAS / KILO);            // (kJ/(mol K)) 

#define DUMP_PARAMETERS 0

extern "C"
void gpuSetBondParameters(gpuContext gpu, const vector<int>& atom1, const vector<int>& atom2, const vector<float>& length, const vector<float>& k)
{
    int bonds = atom1.size();
    gpu->sim.bonds                              = bonds;
141
    CUDAStream<int4>* psBondID                  = new CUDAStream<int4>(bonds, 1, "BondID");
Peter Eastman's avatar
Peter Eastman committed
142
143
    gpu->psBondID                               = psBondID;
    gpu->sim.pBondID                            = psBondID->_pDevStream[0];
144
    CUDAStream<float2>* psBondParameter         = new CUDAStream<float2>(bonds, 1, "BondParameter");
Peter Eastman's avatar
Peter Eastman committed
145
146
147
148
    gpu->psBondParameter                        = psBondParameter;
    gpu->sim.pBondParameter                     = psBondParameter->_pDevStream[0];
    for (int i = 0; i < bonds; i++)
    {
149
150
151
152
153
154
        (*psBondID)[i].x = atom1[i];
        (*psBondID)[i].y = atom2[i];
        (*psBondParameter)[i].x = length[i];
        (*psBondParameter)[i].y = k[i];
        psBondID->_pSysData[i].z = gpu->pOutputBufferCounter[psBondID->_pSysData[i].x]++;
        psBondID->_pSysData[i].w = gpu->pOutputBufferCounter[psBondID->_pSysData[i].y]++;
Peter Eastman's avatar
Peter Eastman committed
155
156
157
#if (DUMP_PARAMETERS == 1)                
        cout << 
            i << " " << 
158
159
160
161
162
163
            (*psBondID)[i].x << " " <<
            (*psBondID)[i].y << " " <<
            (*psBondID)[i].z << " " <<
            (*psBondID)[i].w << " " <<
            (*psBondParameter)[i].x << " " <<
            (*psBondParameter)[i].y <<
Peter Eastman's avatar
Peter Eastman committed
164
165
166
167
168
169
170
171
172
173
174
175
176
            endl;
#endif
    }
    psBondID->Upload();
    psBondParameter->Upload();
}

extern "C"
void gpuSetBondAngleParameters(gpuContext gpu, const vector<int>& atom1, const vector<int>& atom2, const vector<int>& atom3,
        const vector<float>& angle, const vector<float>& k)
{
    int bond_angles = atom1.size();
    gpu->sim.bond_angles                        = bond_angles;
177
    CUDAStream<int4>* psBondAngleID1            = new CUDAStream<int4>(bond_angles, 1, "BondAngleID1");
Peter Eastman's avatar
Peter Eastman committed
178
179
    gpu->psBondAngleID1                         = psBondAngleID1;
    gpu->sim.pBondAngleID1                      = psBondAngleID1->_pDevStream[0];
180
    CUDAStream<int2>* psBondAngleID2            = new CUDAStream<int2>(bond_angles, 1, "BondAngleID2");
Peter Eastman's avatar
Peter Eastman committed
181
182
    gpu->psBondAngleID2                         = psBondAngleID2;
    gpu->sim.pBondAngleID2                      = psBondAngleID2->_pDevStream[0];
183
    CUDAStream<float2>* psBondAngleParameter    = new CUDAStream<float2>(bond_angles, 1, "BondAngleParameter");
Peter Eastman's avatar
Peter Eastman committed
184
185
186
187
188
    gpu->psBondAngleParameter                   = psBondAngleParameter;
    gpu->sim.pBondAngleParameter                = psBondAngleParameter->_pDevStream[0];        

    for (int i = 0; i < bond_angles; i++)
    {
189
190
191
192
193
194
195
196
        (*psBondAngleID1)[i].x = atom1[i];
        (*psBondAngleID1)[i].y = atom2[i];
        (*psBondAngleID1)[i].z = atom3[i];
        (*psBondAngleParameter)[i].x = angle[i];
        (*psBondAngleParameter)[i].y = k[i];
        psBondAngleID1->_pSysData[i].w = gpu->pOutputBufferCounter[psBondAngleID1->_pSysData[i].x]++;
        psBondAngleID2->_pSysData[i].x = gpu->pOutputBufferCounter[psBondAngleID1->_pSysData[i].y]++;
        psBondAngleID2->_pSysData[i].y = gpu->pOutputBufferCounter[psBondAngleID1->_pSysData[i].z]++;
Peter Eastman's avatar
Peter Eastman committed
197
198
199
#if (DUMP_PARAMETERS == 1)
         cout << 
            i << " " << 
200
201
202
203
204
205
206
207
            (*psBondAngleID1)[i].x << " " <<
            (*psBondAngleID1)[i].y << " " <<
            (*psBondAngleID1)[i].z << " " <<
            (*psBondAngleID1)[i].w << " " <<
            (*psBondAngleID2)[i].x << " " <<
            (*psBondAngleID2)[i].y << " " <<
            (*psBondAngleParameter)[i].x << " " <<
            (*psBondAngleParameter)[i].y <<
Peter Eastman's avatar
Peter Eastman committed
208
209
210
211
212
213
214
215
216
217
218
219
220
221
            endl;
#endif
    }
    psBondAngleID1->Upload();
    psBondAngleID2->Upload();
    psBondAngleParameter->Upload();
}

extern "C"
void gpuSetDihedralParameters(gpuContext gpu, const vector<int>& atom1, const vector<int>& atom2, const vector<int>& atom3, const vector<int>& atom4,
        const vector<float>& k, const vector<float>& phase, const vector<int>& periodicity)
{
        int dihedrals = atom1.size();
        gpu->sim.dihedrals = dihedrals;
222
        CUDAStream<int4>* psDihedralID1             = new CUDAStream<int4>(dihedrals, 1, "DihedralID1");
Peter Eastman's avatar
Peter Eastman committed
223
224
        gpu->psDihedralID1                          = psDihedralID1;
        gpu->sim.pDihedralID1                       = psDihedralID1->_pDevStream[0];
225
        CUDAStream<int4>* psDihedralID2             = new CUDAStream<int4>(dihedrals, 1, "DihedralID2");
Peter Eastman's avatar
Peter Eastman committed
226
227
        gpu->psDihedralID2                          = psDihedralID2;
        gpu->sim.pDihedralID2                       = psDihedralID2->_pDevStream[0];
228
        CUDAStream<float4>* psDihedralParameter     = new CUDAStream<float4>(dihedrals, 1, "DihedralParameter");
Peter Eastman's avatar
Peter Eastman committed
229
230
231
232
        gpu->psDihedralParameter                    = psDihedralParameter;
        gpu->sim.pDihedralParameter                 = psDihedralParameter->_pDevStream[0];
        for (int i = 0; i < dihedrals; i++)
        {
233
234
235
236
237
238
239
240
241
242
243
            (*psDihedralID1)[i].x = atom1[i];
            (*psDihedralID1)[i].y = atom2[i];
            (*psDihedralID1)[i].z = atom3[i];
            (*psDihedralID1)[i].w = atom4[i];
            (*psDihedralParameter)[i].x = k[i];
            (*psDihedralParameter)[i].y = phase[i];
            (*psDihedralParameter)[i].z = (float) periodicity[i];
            psDihedralID2->_pSysData[i].x = gpu->pOutputBufferCounter[psDihedralID1->_pSysData[i].x]++;
            psDihedralID2->_pSysData[i].y = gpu->pOutputBufferCounter[psDihedralID1->_pSysData[i].y]++;
            psDihedralID2->_pSysData[i].z = gpu->pOutputBufferCounter[psDihedralID1->_pSysData[i].z]++;
            psDihedralID2->_pSysData[i].w = gpu->pOutputBufferCounter[psDihedralID1->_pSysData[i].w]++;
Peter Eastman's avatar
Peter Eastman committed
244
245
246
#if (DUMP_PARAMETERS == 1)
            cout << 
                i << " " << 
247
248
249
250
251
252
253
254
255
256
257
                (*psDihedralID1)[i].x << " " <<
                (*psDihedralID1)[i].y << " " <<
                (*psDihedralID1)[i].z << " " <<
                (*psDihedralID1)[i].w << " " <<
                (*psDihedralID2)[i].x << " " <<
                (*psDihedralID2)[i].y << " " <<
                (*psDihedralID2)[i].z << " " <<
                (*psDihedralID2)[i].w << " " <<
                (*psDihedralParameter)[i].x << " " <<
                (*psDihedralParameter)[i].y << " " <<
                (*psDihedralParameter)[i].z << endl;
Peter Eastman's avatar
Peter Eastman committed
258
259
260
261
262
263
264
265
266
267
268
269
270
#endif
        }
        psDihedralID1->Upload();
        psDihedralID2->Upload();
        psDihedralParameter->Upload();
}

extern "C"
void gpuSetRbDihedralParameters(gpuContext gpu, const vector<int>& atom1, const vector<int>& atom2, const vector<int>& atom3, const vector<int>& atom4,
        const vector<float>& c0, const vector<float>& c1, const vector<float>& c2, const vector<float>& c3, const vector<float>& c4, const vector<float>& c5)
{
    int rb_dihedrals = atom1.size();
    gpu->sim.rb_dihedrals = rb_dihedrals;
271
    CUDAStream<int4>* psRbDihedralID1           = new CUDAStream<int4>(rb_dihedrals, 1, "RbDihedralID1");
Peter Eastman's avatar
Peter Eastman committed
272
273
    gpu->psRbDihedralID1                        = psRbDihedralID1;
    gpu->sim.pRbDihedralID1                     = psRbDihedralID1->_pDevStream[0];
274
    CUDAStream<int4>* psRbDihedralID2           = new CUDAStream<int4>(rb_dihedrals, 1, "RbDihedralID2");
Peter Eastman's avatar
Peter Eastman committed
275
276
    gpu->psRbDihedralID2                        = psRbDihedralID2;
    gpu->sim.pRbDihedralID2                     = psRbDihedralID2->_pDevStream[0];
277
    CUDAStream<float4>* psRbDihedralParameter1  = new CUDAStream<float4>(rb_dihedrals, 1, "RbDihedralParameter1");
Peter Eastman's avatar
Peter Eastman committed
278
279
    gpu->psRbDihedralParameter1                 = psRbDihedralParameter1;
    gpu->sim.pRbDihedralParameter1              = psRbDihedralParameter1->_pDevStream[0];
280
    CUDAStream<float2>* psRbDihedralParameter2  = new CUDAStream<float2>(rb_dihedrals, 1, "RbDihedralParameter2");
Peter Eastman's avatar
Peter Eastman committed
281
282
283
284
285
    gpu->psRbDihedralParameter2                 = psRbDihedralParameter2;
    gpu->sim.pRbDihedralParameter2              = psRbDihedralParameter2->_pDevStream[0];

    for (int i = 0; i < rb_dihedrals; i++)
    {
286
287
288
289
290
291
292
293
294
295
296
297
298
299
        (*psRbDihedralID1)[i].x = atom1[i];
        (*psRbDihedralID1)[i].y = atom2[i];
        (*psRbDihedralID1)[i].z = atom3[i];
        (*psRbDihedralID1)[i].w = atom4[i];
        (*psRbDihedralParameter1)[i].x = c0[i];
        (*psRbDihedralParameter1)[i].y = c1[i];
        (*psRbDihedralParameter1)[i].z = c2[i];
        (*psRbDihedralParameter1)[i].w = c3[i];
        (*psRbDihedralParameter2)[i].x = c4[i];
        (*psRbDihedralParameter2)[i].y = c5[i];
        psRbDihedralID2->_pSysData[i].x = gpu->pOutputBufferCounter[psRbDihedralID1->_pSysData[i].x]++;
        psRbDihedralID2->_pSysData[i].y = gpu->pOutputBufferCounter[psRbDihedralID1->_pSysData[i].y]++;
        psRbDihedralID2->_pSysData[i].z = gpu->pOutputBufferCounter[psRbDihedralID1->_pSysData[i].z]++;
        psRbDihedralID2->_pSysData[i].w = gpu->pOutputBufferCounter[psRbDihedralID1->_pSysData[i].w]++;
Peter Eastman's avatar
Peter Eastman committed
300
301
302
#if (DUMP_PARAMETERS == 1)
        cout << 
            i << " " << 
303
304
305
306
307
308
309
310
311
312
313
314
315
316
            (*psRbDihedralID1)[i].x << " " <<
            (*psRbDihedralID1)[i].y << " " <<
            (*psRbDihedralID1)[i].z << " " <<
            (*psRbDihedralID1)[i].w <<" " <<
            (*psRbDihedralID2)[i].x << " " <<
            (*psRbDihedralID2)[i].y << " " <<
            (*psRbDihedralID2)[i].z << " " <<
            (*psRbDihedralID2)[i].w <<" " <<
            (*psRbDihedralParameter1)[i].x << " " <<
            (*psRbDihedralParameter1)[i].y << " " <<
            (*psRbDihedralParameter1)[i].z << " " <<
            (*psRbDihedralParameter1)[i].w << " " <<
            (*psRbDihedralParameter2)[i].x << " " <<
            (*psRbDihedralParameter2)[i].y <<
Peter Eastman's avatar
Peter Eastman committed
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
            endl;
#endif
    }
    psRbDihedralID1->Upload();
    psRbDihedralID2->Upload();
    psRbDihedralParameter1->Upload();
    psRbDihedralParameter2->Upload();
}

extern "C"
void gpuSetLJ14Parameters(gpuContext gpu, float epsfac, float fudge, const vector<int>& atom1, const vector<int>& atom2,
        const vector<float>& c6, const vector<float>& c12, const vector<float>& q1, const vector<float>& q2)
{
    int LJ14s = atom1.size();
    float scale = epsfac * fudge;

    gpu->sim.LJ14s                              = LJ14s;
334
    CUDAStream<int4>* psLJ14ID                  = new CUDAStream<int4>(LJ14s, 1, "LJ14ID");
Peter Eastman's avatar
Peter Eastman committed
335
336
    gpu->psLJ14ID                               = psLJ14ID;
    gpu->sim.pLJ14ID                            = psLJ14ID->_pDevStream[0];
337
    CUDAStream<float4>* psLJ14Parameter         = new CUDAStream<float4>(LJ14s, 1, "LJ14Parameter");
Peter Eastman's avatar
Peter Eastman committed
338
339
340
341
342
    gpu->psLJ14Parameter                        = psLJ14Parameter;
    gpu->sim.pLJ14Parameter                     = psLJ14Parameter->_pDevStream[0];

    for (int i = 0; i < LJ14s; i++)
    {
343
344
345
346
        (*psLJ14ID)[i].x = atom1[i];
        (*psLJ14ID)[i].y = atom2[i];
        psLJ14ID->_pSysData[i].z = gpu->pOutputBufferCounter[psLJ14ID->_pSysData[i].x]++;
        psLJ14ID->_pSysData[i].w = gpu->pOutputBufferCounter[psLJ14ID->_pSysData[i].y]++;
Peter Eastman's avatar
Peter Eastman committed
347
348
349
350
351
352
353
354
355
356
357
358
        float p0, p1, p2;
        if (c12[i] == 0.0f)
        {
            p0 = 0.0f;
            p1 = 1.0f;
        }
        else
        {
            p0 = c6[i] * c6[i] / c12[i];
            p1 = pow(c12[i] / c6[i], 1.0f / 6.0f);
        }
        p2 = scale * q1[i] * q2[i];
359
360
361
        (*psLJ14Parameter)[i].x = p0;
        (*psLJ14Parameter)[i].y = p1;
        (*psLJ14Parameter)[i].z = p2;
Peter Eastman's avatar
Peter Eastman committed
362
363
364
365
    }
#if (DUMP_PARAMETERS == 1)
        cout << 
            i << " " <<
366
367
368
369
370
371
372
            (*psLJ14ID)[i].x << " " <<
            (*psLJ14ID)[i].y << " " <<
            (*psLJ14ID)[i].z << " " <<
            (*psLJ14ID)[i].w << " " <<
            (*psLJ14Parameter)[i].x << " " <<
            (*psLJ14Parameter)[i].y << " " <<
            (*psLJ14Parameter)[i].z << " " <<
Peter Eastman's avatar
Peter Eastman committed
373
374
375
376
377
378
379
380
381
382
383
            p0 << " " << 
            p1 << " " << 
            p2 << " " << 
            endl;
#endif
    psLJ14ID->Upload();
    psLJ14Parameter->Upload();
}

extern "C"
void gpuSetCoulombParameters(gpuContext gpu, float epsfac, const vector<int>& atom, const vector<float>& c6, const vector<float>& c12, const vector<float>& q,
384
        const vector<char>& symbol, const vector<vector<int> >& exclusions, CudaNonbondedMethod method)
Peter Eastman's avatar
Peter Eastman committed
385
386
387
{
    unsigned int coulombs = atom.size();
    gpu->sim.epsfac = epsfac;
388
389
    gpu->sim.nonbondedMethod = method;
    gpu->exclusions = exclusions;
Peter Eastman's avatar
Peter Eastman committed
390
391
392
393
394
395
396
397
398
399
400
401

    for (unsigned int i = 0; i < coulombs; i++)
    {
            float p0 = q[i];
            float p1 = 0.5f, p2 = 0.0f;               
            if ((c6[i] > 0.0f) && (c12[i] > 0.0f))
            {
                p1 = 0.5f * pow(c12[i] / c6[i], 1.0f / 6.0f);
                p2 = c6[i] * sqrt(1.0f / c12[i]);
            }
            if (symbol.size() > 0)
                gpu->pAtomSymbol[i] = symbol[i];
402
403
404
            (*gpu->psPosq4)[i].w = p0;
            (*gpu->psSigEps2)[i].x = p1;
            (*gpu->psSigEps2)[i].y = p2;
Peter Eastman's avatar
Peter Eastman committed
405
406
407
408
409
    }

    // Dummy out extra atom data
    for (unsigned int i = coulombs; i < gpu->sim.paddedNumberOfAtoms; i++)
    {
410
411
412
413
414
415
        (*gpu->psPosq4)[i].x       = 100000.0f + i * 10.0f;
        (*gpu->psPosq4)[i].y       = 100000.0f + i * 10.0f;
        (*gpu->psPosq4)[i].z       = 100000.0f + i * 10.0f;
        (*gpu->psPosq4)[i].w       = 0.0f;
        (*gpu->psSigEps2)[i].x     = 0.0f;
        (*gpu->psSigEps2)[i].y     = 0.0f;
Peter Eastman's avatar
Peter Eastman committed
416
417
418
419
    }

    gpu->psPosq4->Upload();
    gpu->psSigEps2->Upload();
420
}
Peter Eastman's avatar
Peter Eastman committed
421

422
423
424
425
426
427
extern "C"
void gpuSetNonbondedCutoff(gpuContext gpu, float cutoffDistance, float solventDielectric)
{
    gpu->sim.nonbondedCutoffSqr = cutoffDistance*cutoffDistance;
    gpu->sim.reactionFieldK = pow(cutoffDistance, -3.0f)*(solventDielectric-1.0f)/(2.0f*solventDielectric+1.0f);
}
Peter Eastman's avatar
Peter Eastman committed
428

Rossen Apostolov's avatar
Rossen Apostolov committed
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
extern "C"
void gpuSetEwaldParameters(gpuContext gpu)//, float alphaEwald, int kmax )
{

    // hard coded alphaEwald and kmax, no interface yet
    float alpha            = 3.123413;
    float PI               = 3.14159265358979323846f;
    float TWO_PI           = 2.0 * PI;

    gpu->sim.recipBoxSizeX = TWO_PI / gpu->sim.periodicBoxSizeX ;
    gpu->sim.recipBoxSizeY = TWO_PI / gpu->sim.periodicBoxSizeY ;
    gpu->sim.recipBoxSizeZ = TWO_PI / gpu->sim.periodicBoxSizeZ ;

    gpu->sim.cellVolume        = gpu->sim.periodicBoxSizeX * gpu->sim.periodicBoxSizeY * gpu->sim.periodicBoxSizeZ;

    gpu->sim.alphaEwald        = alpha;
    gpu->sim.factorEwald       = -1 / (4*alpha*alpha);
    gpu->sim.kmax              = 20+1;
}

449
450
451
452
453
454
extern "C"
void gpuSetPeriodicBoxSize(gpuContext gpu, float xsize, float ysize, float zsize)
{
    gpu->sim.periodicBoxSizeX = xsize;
    gpu->sim.periodicBoxSizeY = ysize;
    gpu->sim.periodicBoxSizeZ = zsize;
Peter Eastman's avatar
Peter Eastman committed
455
456
457
}

extern "C"
458
void gpuSetObcParameters(gpuContext gpu, float innerDielectric, float solventDielectric, const vector<float>& radius, const vector<float>& scale, const vector<float>& charge)
Peter Eastman's avatar
Peter Eastman committed
459
{
460
    unsigned int atoms = radius.size();
461
462

    gpu->bIncludeGBSA = true;
Peter Eastman's avatar
Peter Eastman committed
463
464
    for (unsigned int i = 0; i < atoms; i++)
    {
465
466
            (*gpu->psObcData)[i].x = radius[i] - dielectricOffset;
            (*gpu->psObcData)[i].y = scale[i] * (*gpu->psObcData)[i].x;
467
            (*gpu->psPosq4)[i].w = charge[i];
Peter Eastman's avatar
Peter Eastman committed
468
469
470
471

#if (DUMP_PARAMETERS == 1)
        cout << 
            i << " " << 
472
473
            (*gpu->psObcData)[i].x << " " <<
            (*gpu->psObcData)[i].y;
Peter Eastman's avatar
Peter Eastman committed
474
475
476
477
478
479
#endif
    }

    // Dummy out extra atom data
    for (unsigned int i = atoms; i < gpu->sim.paddedNumberOfAtoms; i++)
    {
480
481
482
        (*gpu->psBornRadii)[i]     = 0.2f;
        (*gpu->psObcData)[i].x     = 0.01f;
        (*gpu->psObcData)[i].y     = 0.01f;
Peter Eastman's avatar
Peter Eastman committed
483
484
485
486
    }

    gpu->psBornRadii->Upload();
    gpu->psObcData->Upload();
487
    gpu->psPosq4->Upload();
Peter Eastman's avatar
Peter Eastman committed
488
489
490
    gpu->sim.preFactor = 2.0f*electricConstant*((1.0f/innerDielectric)-(1.0f/solventDielectric))*gpu->sim.forceConversionFactor;
}

491
static void markShakeClusterInvalid(ShakeCluster& cluster, map<int, ShakeCluster>& allClusters, vector<bool>& invalidForShake)
492
493
494
495
496
497
498
499
500
501
502
{
    cluster.valid = false;
    invalidForShake[cluster.centralID] = true;
    for (int i = 0; i < cluster.size; i++) {
        invalidForShake[cluster.peripheralID[i]] = true;
        map<int, ShakeCluster>::iterator otherCluster = allClusters.find(cluster.peripheralID[i]);
        if (otherCluster != allClusters.end() && otherCluster->second.valid)
            markShakeClusterInvalid(otherCluster->second, allClusters, invalidForShake);
    }
}

503
static void findRigidClusters(gpuContext gpu, const vector<int>& firstAtom, const vector<int>& secondAtom, const vector<float>& invMass1, const vector<float>& invMass2, const vector<float>& distance, vector<int>& constraintIndices)
504
505
{
    vector<map<int, int> > atomConstraints(firstAtom.size());
506
    for (int i = 0; i < (int)constraintIndices.size(); i++) {
507
508
509
510
511
512
        atomConstraints[firstAtom[i]][secondAtom[i]] = constraintIndices[i];
        atomConstraints[secondAtom[i]][firstAtom[i]] = constraintIndices[i];
    }
    vector<vector<int> > rigidClusters;
    int totalConstraints = 0;
    int totalMatrixElements = 0;
513
    for (int i = 0; i < (int)firstAtom.size(); i++) {
514
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
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
        if (atomConstraints[i].size() < 2)
            continue;

        // Begin by looking for a triangle this atom is part of.

        set<int> atoms;
        atoms.insert(i);
        for (map<int, int>::const_iterator atom1 = atomConstraints[i].begin(); atom1 != atomConstraints[i].end() && atoms.size() == 1; ++atom1) {
            for (map<int, int>::const_iterator atom2 = atomConstraints[atom1->first].begin(); atom2 != atomConstraints[atom1->first].end(); ++atom2) {
                if (atomConstraints[i].count(atom2->first) != 0) {
                    atoms.insert(atom1->first);
                    atoms.insert(atom2->first);
                    break;
                }
            }
        }
        if (atoms.size() == 1)
            continue;

        // We have three atoms that are part of a cluster, so look for other atoms we can add.

        bool done = false;
        while (!done) {
           done = true;
            for (set<int>::const_iterator atom1 = atoms.begin(); atom1 != atoms.end(); ++atom1) {
                for (map<int, int>::const_iterator atom2 = atomConstraints[*atom1].begin(); atom2 != atomConstraints[*atom1].end(); ++atom2) {
                    if (atoms.find(atom2->first) != atoms.end())
                        continue; // This atom is already in the cluster.

                    // See if this atom is linked to three other atoms in the cluster.

                    int linkCount = 0;
                    for (map<int, int>::const_iterator atom3 = atomConstraints[atom2->first].begin(); atom3 != atomConstraints[atom2->first].end(); ++atom3)
                        if (atoms.find(atom3->first) != atoms.end())
                            linkCount++;
                    if (linkCount > 2) {
                        atoms.insert(atom2->first);
                        done = false;
                    }
                }
            }
        }

        // Record the cluster.

        vector<int> constraints;
        for (set<int>::const_iterator atom1 = atoms.begin(); atom1 != atoms.end(); ++atom1) {
            for (map<int, int>::const_iterator atom2 = atomConstraints[*atom1].begin(); atom2 != atomConstraints[*atom1].end(); ++atom2) {
                if (*atom1 < atom2->first && atoms.find(atom2->first) != atoms.end())
                    constraints.push_back(atom2->second);
            }
        }
        rigidClusters.push_back(constraints);
        totalConstraints += constraints.size();
        totalMatrixElements += constraints.size()*constraints.size();
        for (set<int>::const_iterator atom1 = atoms.begin(); atom1 != atoms.end(); ++atom1) {
            for (map<int, int>::const_iterator atom2 = atomConstraints[*atom1].begin(); atom2 != atomConstraints[*atom1].end(); ++atom2)
                atomConstraints[atom2->first].erase(*atom1);
            atomConstraints[*atom1].clear();
        }
    }

576
577
    // Reorder the constraints so those in a cluster are sequential.

578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
//    vector<int> constraintOrder(constraintIndices.size());
//    vector<int> clusterStartIndex(rigidClusters.size());
//    set<int> inCluster;
//    int nextIndex = 0;
//    for (int i = 0; i < (int) rigidClusters.size(); ++i) {
//        clusterStartIndex[i] = nextIndex;
//        for (int j = 0; j < (int) rigidClusters[i].size(); ++j) {
//            int constraint = rigidClusters[i][j];
//            constraintOrder[nextIndex++] = constraint;
//            inCluster.insert(constraint);
//        }
//    }
//    for (int i = 0; i < (int) constraintIndices.size(); ++i)
//        if (inCluster.find(constraintIndices[i]) == inCluster.end())
//            constraintOrder[nextIndex++] = constraintIndices[i];
//    constraintIndices = constraintOrder;
594

595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
    // Build the CUDA streams.

    CUDAStream<unsigned int>* psRigidClusterConstraintIndex = new CUDAStream<unsigned int>((int) rigidClusters.size()+1, 1, "RigidClusterConstraintIndex");
    gpu->psRigidClusterConstraintIndex                      = psRigidClusterConstraintIndex;
    gpu->sim.pRigidClusterConstraintIndex                   = psRigidClusterConstraintIndex->_pDevData;
    CUDAStream<float>* psRigidClusterMatrix = new CUDAStream<float>(totalMatrixElements, 1, "RigidClusterMatrix");
    gpu->psRigidClusterMatrix       = psRigidClusterMatrix;
    gpu->sim.pRigidClusterMatrix    = psRigidClusterMatrix->_pDevData;
    CUDAStream<unsigned int>* psRigidClusterMatrixIndex = new CUDAStream<unsigned int>((int) rigidClusters.size()+1, 1, "RigidClusterMatrixIndex");
    gpu->psRigidClusterMatrixIndex                      = psRigidClusterMatrixIndex;
    gpu->sim.pRigidClusterMatrixIndex                   = psRigidClusterMatrixIndex->_pDevData;
    unsigned int constraintIndex = 0;
    unsigned int maxClusterSize = 0;
    for (unsigned int i = 0; i < rigidClusters.size(); i++) {
        vector<int>& cluster = rigidClusters[i];
        (*psRigidClusterConstraintIndex)[i] = constraintIndex;
611
        constraintIndex += cluster.size();
612
613
614
615
616
617
618
619
620
621
        if (cluster.size() > maxClusterSize)
            maxClusterSize = cluster.size();
    }
    (*psRigidClusterConstraintIndex)[rigidClusters.size()] = constraintIndex;
    gpu->sim.rigidClusters = rigidClusters.size();
    gpu->sim.maxRigidClusterSize = maxClusterSize;
    gpu->sim.clusterShakeBlockSize = 1;
    while (gpu->sim.clusterShakeBlockSize < maxClusterSize)
        gpu->sim.clusterShakeBlockSize *= 2;
    psRigidClusterConstraintIndex->Upload();
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
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

    // Build the inverse coupling matrix for each cluster.

    unsigned int elementIndex = 0;
    for (unsigned int i = 0; i < rigidClusters.size(); i++) {
        // Compute the constraint coupling matrix for this cluster.

        const vector<int>& cluster = rigidClusters[i];
        unsigned int size = cluster.size();
        Array2D<double> matrix(size, size);
        for (int j = 0; j < (int)size; j++) {
            for (int k = 0; k < (int)size; k++) {
                if (j == k) {
                    matrix[j][j] = 1.0;
                    continue;
                }
                double scale;
                int atomj0 = firstAtom[cluster[j]];
                int atomj1 = secondAtom[cluster[j]];
                int atomk0 = firstAtom[cluster[k]];
                int atomk1 = secondAtom[cluster[k]];
                int atoma, atomb;
                if (atomj0 == atomk0) {
                    atoma = atomj1;
                    atomb = atomk1;
                    scale = invMass1[cluster[j]]/(invMass1[cluster[j]]+invMass2[cluster[j]]);
                }
                else if (atomj1 == atomk1) {
                    atoma = atomj0;
                    atomb = atomk0;
                    scale = invMass2[cluster[j]]/(invMass1[cluster[j]]+invMass2[cluster[j]]);
                }
                else if (atomj0 == atomk1) {
                    atoma = atomj1;
                    atomb = atomk0;
                    scale = invMass1[cluster[j]]/(invMass1[cluster[j]]+invMass2[cluster[j]]);
                }
                else if (atomj1 == atomk0) {
                    atoma = atomj0;
                    atomb = atomk1;
                    scale = invMass2[cluster[j]]/(invMass1[cluster[j]]+invMass2[cluster[j]]);
                }
                else {
                    matrix[j][k] = 0.0;
                    continue; // These constraints are not connected.
                }

                // Find the third constraint forming a triangle with these two.

                for (int m = 0; m < size; m++) {
                    int other = cluster[m];
                    if ((firstAtom[other] == atoma && secondAtom[other] == atomb) || (firstAtom[other] == atomb && secondAtom[other] == atoma)) {
                        double d1 = distance[cluster[j]];
                        double d2 = distance[cluster[k]];
                        double d3 = distance[other];
                        matrix[j][k] = scale*(d1*d1+d2*d2-d3*d3)/(2.0*d1*d2);
                        break;
                    }
                }
            }
        }

        // Invert it using SVD.

        Array2D<double> u, v;
        Array1D<double> w;
        SVD<double> svd(matrix);
        svd.getU(u);
        svd.getV(v);
        svd.getSingularValues(w);
        double singularValueCutoff = 0.01*w[0];
        for (int j = 0; j < (int)size; j++)
            w[j] = (w[j] < singularValueCutoff ? 0.0 : 1.0/w[j]);
        for (int j = 0; j < (int)size; j++) {
            for (int k = 0; k < (int)size; k++) {
                matrix[j][k] = 0.0;
                for (int m = 0; m < (int)size; m++)
                    matrix[j][k] += v[j][m]*w[m]*u[k][m];
            }
        }

        // Record the inverted matrix.

        (*gpu->psRigidClusterMatrixIndex)[i] = elementIndex;
        for (int j = 0; j < (int)size; j++)
            for (int k = 0; k < (int)size; k++)
                (*gpu->psRigidClusterMatrix)[elementIndex++] = (float)(matrix[k][j]*distance[cluster[j]]/distance[cluster[k]]);
    }
    (*gpu->psRigidClusterMatrixIndex)[gpu->sim.rigidClusters] = elementIndex;
    gpu->psRigidClusterMatrix->Upload();
    gpu->psRigidClusterMatrixIndex->Upload();
713
714
}

Peter Eastman's avatar
Peter Eastman committed
715
extern "C"
716
717
void gpuSetConstraintParameters(gpuContext gpu, const vector<int>& atom1, const vector<int>& atom2, const vector<float>& distance,
        const vector<float>& invMass1, const vector<float>& invMass2, float shakeTolerance, unsigned int lincsTerms)
Peter Eastman's avatar
Peter Eastman committed
718
{
719
720
721
722
    // Create a vector for recording which atoms are handled by SHAKE (or SETTLE).

    vector<bool> isShakeAtom(gpu->natoms, false);

Peter Eastman's avatar
Peter Eastman committed
723
724
725
    // Find how many constraints each atom is involved in.
    
    vector<int> constraintCount(gpu->natoms, 0);
726
    for (int i = 0; i < (int)atom1.size(); i++) {
Peter Eastman's avatar
Peter Eastman committed
727
728
729
        constraintCount[atom1[i]]++;
        constraintCount[atom2[i]]++;
    }
730
731
732
733
734
735

    // Identify clusters of three atoms that can be treated with SETTLE.  First, for every
    // atom that might be part of such a cluster, make a list of the two other atoms it is
    // connected to.

    vector<map<int, float> > settleConstraints(gpu->natoms);
736
    for (int i = 0; i < (int)atom1.size(); i++) {
737
738
739
740
741
742
743
744
745
        if (constraintCount[atom1[i]] == 2 && constraintCount[atom2[i]] == 2) {
            settleConstraints[atom1[i]][atom2[i]] = distance[i];
            settleConstraints[atom2[i]][atom1[i]] = distance[i];
        }
    }

    // Now remove the ones that don't actually form closed loops of three atoms.

    vector<int> settleClusters;
746
    for (int i = 0; i < (int)settleConstraints.size(); i++) {
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
        if (settleConstraints[i].size() == 2) {
            int partner1 = settleConstraints[i].begin()->first;
            int partner2 = (++settleConstraints[i].begin())->first;
            if (settleConstraints[partner1].size() != 2 || settleConstraints[partner2].size() != 2 ||
                    settleConstraints[partner1].find(partner2) == settleConstraints[partner1].end())
                settleConstraints[i].clear();
            else if (i < partner1 && i < partner2)
                settleClusters.push_back(i);
        }
        else
            settleConstraints[i].clear();
    }

    // Record the actual SETTLE clusters.

762
    CUDAStream<int4>* psSettleID          = new CUDAStream<int4>((int) settleClusters.size(), 1, "SettleID");
763
764
    gpu->psSettleID                       = psSettleID;
    gpu->sim.pSettleID                    = psSettleID->_pDevStream[0];
765
    CUDAStream<float2>* psSettleParameter = new CUDAStream<float2>((int) settleClusters.size(), 1, "SettleParameter");
766
767
768
    gpu->psSettleParameter                = psSettleParameter;
    gpu->sim.pSettleParameter             = psSettleParameter->_pDevStream[0];
    gpu->sim.settleConstraints            = settleClusters.size();
769
      for (int i = 0; i < (int)settleClusters.size(); i++) {
770
771
772
773
774
775
776
        int atom1 = settleClusters[i];
        int atom2 = settleConstraints[atom1].begin()->first;
        int atom3 = (++settleConstraints[atom1].begin())->first;
        float dist12 = settleConstraints[atom1].find(atom2)->second;
        float dist13 = settleConstraints[atom1].find(atom3)->second;
        float dist23 = settleConstraints[atom2].find(atom3)->second;
        if (dist12 == dist13) { // atom1 is the central atom
777
778
779
780
781
            (*psSettleID)[i].x = atom1;
            (*psSettleID)[i].y = atom2;
            (*psSettleID)[i].z = atom3;
            (*psSettleParameter)[i].x = dist12;
            (*psSettleParameter)[i].y = dist23;
782
783
        }
        else if (dist12 == dist23) { // atom2 is the central atom
784
785
786
787
788
            (*psSettleID)[i].x = atom2;
            (*psSettleID)[i].y = atom1;
            (*psSettleID)[i].z = atom3;
            (*psSettleParameter)[i].x = dist12;
            (*psSettleParameter)[i].y = dist13;
789
790
        }
        else if (dist13 == dist23) { // atom3 is the central atom
791
792
793
794
795
            (*psSettleID)[i].x = atom3;
            (*psSettleID)[i].y = atom1;
            (*psSettleID)[i].z = atom2;
            (*psSettleParameter)[i].x = dist13;
            (*psSettleParameter)[i].y = dist12;
796
797
798
        }
        else
            throw OpenMMException("Two of the three distances constrained with SETTLE must be the same.");
799
800
801
        isShakeAtom[atom1] = true;
        isShakeAtom[atom2] = true;
        isShakeAtom[atom3] = true;
802
803
804
805
806
807
808
809
810
    }
    psSettleID->Upload();
    psSettleParameter->Upload();
    gpu->sim.settle_threads_per_block     = (gpu->sim.settleConstraints + gpu->sim.blocks - 1) / gpu->sim.blocks;
    if (gpu->sim.settle_threads_per_block > gpu->sim.max_shake_threads_per_block)
        gpu->sim.settle_threads_per_block = gpu->sim.max_shake_threads_per_block;
    if (gpu->sim.settle_threads_per_block < 1)
        gpu->sim.settle_threads_per_block = 1;

811
812
813
814
    // Find clusters consisting of a central atom with up to three peripheral atoms.

    map<int, ShakeCluster> clusters;
    vector<bool> invalidForShake(gpu->natoms, false);
815
    for (int i = 0; i < (int)atom1.size(); i++) {
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
        if (isShakeAtom[atom1[i]])
            continue; // This is being taken care of with SETTLE.

        // Determine which is the central atom.

        bool firstIsCentral;
        if (constraintCount[atom1[i]] > 1)
            firstIsCentral = true;
        else if (constraintCount[atom2[i]] > 1)
            firstIsCentral = false;
        else if (atom1[i] < atom2[i])
            firstIsCentral = true;
        else
            firstIsCentral = false;
        int centralID, peripheralID;
        float centralInvMass, peripheralInvMass;
        if (firstIsCentral) {
            centralID = atom1[i];
            peripheralID = atom2[i];
            centralInvMass = invMass1[i];
            peripheralInvMass = invMass2[i];
        }
        else {
            centralID = atom2[i];
            peripheralID = atom1[i];
            centralInvMass = invMass2[i];
            peripheralInvMass = invMass1[i];
        }

        // Add it to the cluster.

        if (clusters.find(centralID) == clusters.end()) {
            clusters[centralID] = ShakeCluster(centralID, centralInvMass);
        }
        ShakeCluster& cluster = clusters[centralID];
        cluster.addAtom(peripheralID, distance[i], peripheralInvMass);
        if (constraintCount[peripheralID] != 1 || invalidForShake[atom1[i]] || invalidForShake[atom2[i]]) {
            markShakeClusterInvalid(cluster, clusters, invalidForShake);
            map<int, ShakeCluster>::iterator otherCluster = clusters.find(peripheralID);
            if (otherCluster != clusters.end() && otherCluster->second.valid)
                markShakeClusterInvalid(otherCluster->second, clusters, invalidForShake);
        }
    }
    int validShakeClusters = 0;
    for (map<int, ShakeCluster>::iterator iter = clusters.begin(); iter != clusters.end(); ++iter) {
        ShakeCluster& cluster = iter->second;
        if (cluster.valid) {
            cluster.valid = !invalidForShake[cluster.centralID];
            for (int i = 0; i < cluster.size; i++)
                if (invalidForShake[cluster.peripheralID[i]])
                    cluster.valid = false;
            if (cluster.valid)
                ++validShakeClusters;
        }
    }

    // Fill in the Cuda streams.

874
    CUDAStream<int4>* psShakeID             = new CUDAStream<int4>(validShakeClusters, 1, "ShakeID");
875
876
    gpu->psShakeID                          = psShakeID;
    gpu->sim.pShakeID                       = psShakeID->_pDevStream[0];
877
    CUDAStream<float4>* psShakeParameter    = new CUDAStream<float4>(validShakeClusters, 1, "ShakeParameter");
878
879
880
881
882
883
884
885
    gpu->psShakeParameter                   = psShakeParameter;
    gpu->sim.pShakeParameter                = psShakeParameter->_pDevStream[0];
    gpu->sim.ShakeConstraints               = validShakeClusters;
    int index = 0;
    for (map<int, ShakeCluster>::const_iterator iter = clusters.begin(); iter != clusters.end(); ++iter) {
        const ShakeCluster& cluster = iter->second;
        if (!cluster.valid)
            continue;
886
887
888
889
890
891
892
893
        (*psShakeID)[index].x = cluster.centralID;
        (*psShakeID)[index].y = cluster.peripheralID[0];
        (*psShakeID)[index].z = cluster.size > 1 ? cluster.peripheralID[1] : -1;
        (*psShakeID)[index].w = cluster.size > 2 ? cluster.peripheralID[2] : -1;
        (*psShakeParameter)[index].x = cluster.centralInvMass;
        (*psShakeParameter)[index].y = 0.5f/(cluster.centralInvMass+cluster.peripheralInvMass);
        (*psShakeParameter)[index].z = cluster.distance*cluster.distance;
        (*psShakeParameter)[index].w = cluster.peripheralInvMass;
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
        isShakeAtom[cluster.centralID] = true;
        isShakeAtom[cluster.peripheralID[0]] = true;
        if (cluster.size > 1)
            isShakeAtom[cluster.peripheralID[1]] = true;
        if (cluster.size > 2)
            isShakeAtom[cluster.peripheralID[2]] = true;
        ++index;
    }
    psShakeID->Upload();
    psShakeParameter->Upload();
    gpu->sim.shakeTolerance = shakeTolerance;
    gpu->sim.shake_threads_per_block     = (gpu->sim.ShakeConstraints + gpu->sim.blocks - 1) / gpu->sim.blocks;
    if (gpu->sim.shake_threads_per_block > gpu->sim.max_shake_threads_per_block)
        gpu->sim.shake_threads_per_block = gpu->sim.max_shake_threads_per_block;
    if (gpu->sim.shake_threads_per_block < 1)
        gpu->sim.shake_threads_per_block = 1;

911
912
913
    // Find connected constraints for LINCS.

    vector<int> lincsConstraints;
914
    for (unsigned i = 0; i < atom1.size(); i++)
915
        if (!isShakeAtom[atom1[i]])
916
            lincsConstraints.push_back(i);
917
918
919

    // Identify rigid clusters of atoms.

920
    findRigidClusters(gpu, atom1, atom2, invMass1, invMass2, distance, lincsConstraints);
921
922
923

    // Record the connections between constraints.

924
    int numLincs = (int) lincsConstraints.size();
925
    vector<vector<int> > atomConstraints(gpu->natoms);
926
    for (int i = 0; i < numLincs; i++) {
927
928
929
        atomConstraints[atom1[lincsConstraints[i]]].push_back(i);
        atomConstraints[atom2[lincsConstraints[i]]].push_back(i);
    }
930
    vector<vector<int> > linkedConstraints(numLincs);
931
932
933
    for (unsigned atom = 0; atom < atomConstraints.size(); atom++) {
        for (unsigned i = 0; i < atomConstraints[atom].size(); i++)
            for (unsigned j = 0; j < i; j++) {
934
935
936
937
938
939
                int c1 = atomConstraints[atom][i];
                int c2 = atomConstraints[atom][j];
                linkedConstraints[c1].push_back(c2);
                linkedConstraints[c2].push_back(c1);
            }
    }
940
    int maxLinks = 0;
941
    for (unsigned i = 0; i < linkedConstraints.size(); i++)
942
943
        maxLinks = max(maxLinks, (int) linkedConstraints[i].size());
    int maxAtomConstraints = 0;
944
    for (unsigned i = 0; i < atomConstraints.size(); i++)
945
        maxAtomConstraints = max(maxAtomConstraints, (int) atomConstraints[i].size());
946

947
948
949
950
951
952
953
954
955
956
957
958
959
960
    // Compute the constraint coupling matrix

    vector<vector<int> > atomAngles(gpu->natoms);
    for (int i = 0; i < gpu->sim.bond_angles; i++)
        atomAngles[(*gpu->psBondAngleID1)[i].y].push_back(i);
    vector<vector<pair<int, double> > > matrix(numLincs);
    if (numLincs > 0) {
        for (int j = 0; j < numLincs; j++) {
            for (int k = 0; k < numLincs; k++) {
                if (j == k) {
                    matrix[j].push_back(pair<int, double>(j, 1.0));
                    continue;
                }
                double scale;
961
962
963
964
965
966
                int cj = lincsConstraints[j];
                int ck = lincsConstraints[k];
                int atomj0 = atom1[cj];
                int atomj1 = atom2[cj];
                int atomk0 = atom1[ck];
                int atomk1 = atom2[ck];
967
968
969
970
971
                int atoma, atomb, atomc;
                if (atomj0 == atomk0) {
                    atoma = atomj1;
                    atomb = atomj0;
                    atomc = atomk1;
972
                    scale = invMass1[cj]/(invMass1[cj]+invMass2[cj]);
973
974
975
976
977
                }
                else if (atomj1 == atomk1) {
                    atoma = atomj0;
                    atomb = atomj1;
                    atomc = atomk0;
978
                    scale = invMass2[cj]/(invMass1[cj]+invMass2[cj]);
979
980
981
982
983
                }
                else if (atomj0 == atomk1) {
                    atoma = atomj1;
                    atomb = atomj0;
                    atomc = atomk0;
984
                    scale = invMass1[cj]/(invMass1[cj]+invMass2[cj]);
985
986
987
988
989
                }
                else if (atomj1 == atomk0) {
                    atoma = atomj0;
                    atomb = atomj1;
                    atomc = atomk1;
990
                    scale = invMass2[cj]/(invMass1[cj]+invMass2[cj]);
991
992
993
994
995
996
997
998
999
                }
                else
                    continue; // These constraints are not connected.

                // Look for a third constraint forming a triangle with these two.

                bool foundConstraint = false;
                for (int other = 0; other < numLincs; other++) {
                    if ((atom1[other] == atoma && atom2[other] == atomc) || (atom1[other] == atomc && atom2[other] == atoma)) {
1000
1001
                        double d1 = distance[cj];
                        double d2 = distance[ck];
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
                        double d3 = distance[other];
                        matrix[j].push_back(pair<int, double>(k, scale*(d1*d1+d2*d2-d3*d3)/(2.0*d1*d2)));
                        foundConstraint = true;
                        break;
                    }
                }
                if (!foundConstraint) {
                    // We didn't find one, so look for an angle force field term.

                    const vector<int>& angleCandidates = atomAngles[atomb];
                    for (vector<int>::const_iterator iter = angleCandidates.begin(); iter != angleCandidates.end(); iter++) {
                        int4 atoms = (*gpu->psBondAngleID1)[*iter];
                        if ((atoms.x == atoma && atoms.z == atomc) || (atoms.z == atoma && atoms.x == atomc)) {
                            double angle = (*gpu->psBondAngleParameter)[*iter].x;
                            matrix[j].push_back(pair<int, double>(k, scale*cos(angle*PI/180.0)));
                            break;
                        }
                    }
                }
            }
        }

        // Invert it using QR.

        vector<int> matrixRowStart;
        vector<int> matrixColIndex;
        vector<double> matrixValue;
        for (int i = 0; i < numLincs; i++) {
            matrixRowStart.push_back(matrixValue.size());
            for (int j = 0; j < (int) matrix[i].size(); j++) {
                pair<int, double> element = matrix[i][j];
                matrixColIndex.push_back(element.first);
                matrixValue.push_back(element.second);
            }
        }
        matrixRowStart.push_back(matrixValue.size());
        int *qRowStart, *qColIndex, *rRowStart, *rColIndex;
        double *qValue, *rValue;
        int result = QUERN_compute_qr(numLincs, numLincs, &matrixRowStart[0], &matrixColIndex[0], &matrixValue[0], NULL,
                &qRowStart, &qColIndex, &qValue, &rRowStart, &rColIndex, &rValue);
        vector<double> rhs(numLincs);
        matrix.clear();
        matrix.resize(numLincs);
        for (int i = 0; i < numLincs; i++) {
            // Extract column i of the inverse matrix.

            for (int j = 0; j < numLincs; j++)
                rhs[j] = (i == j ? 1.0 : 0.0);
            result = QUERN_multiply_with_q_transpose(numLincs, qRowStart, qColIndex, qValue, &rhs[0]);
            result = QUERN_solve_with_r(numLincs, rRowStart, rColIndex, rValue, &rhs[0], &rhs[0]);
            for (int j = 0; j < numLincs; j++) {
1053
1054
                double value = rhs[j]*distance[lincsConstraints[i]]/distance[lincsConstraints[j]];
                if (abs(value) > 0.1)
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
                    matrix[j].push_back(pair<int, double>(i, value));
            }
        }
        QUERN_free_result(qRowStart, qColIndex, qValue);
        QUERN_free_result(rRowStart, rColIndex, rValue);
    }
    int maxRowElements = 0;
    for (unsigned i = 0; i < matrix.size(); i++)
        maxRowElements = max(maxRowElements, (int) matrix[i].size());
    maxRowElements++;

1066
    // Sort the constraints.
1067

1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
    vector<int> constraintOrder(numLincs);
    for (int i = 0; i < numLincs; ++i)
        constraintOrder[i] = i;
    sort(constraintOrder.begin(), constraintOrder.end(), ConstraintOrderer(atom1, atom2));
    vector<int> inverseOrder(numLincs);
    for (int i = 0; i < numLincs; ++i)
        inverseOrder[constraintOrder[i]] = i;
    for (int i = 0; i < matrix.size(); ++i)
        for (int j = 0; j < matrix[i].size(); ++j)
            matrix[i][j].first = inverseOrder[matrix[i][j].first];
1078

1079
1080
    // Fill in the CUDA streams.

1081
    CUDAStream<int2>* psLincsAtoms = new CUDAStream<int2>(numLincs, 1, "LincsAtoms");
1082
1083
    gpu->psLincsAtoms              = psLincsAtoms;
    gpu->sim.pLincsAtoms           = psLincsAtoms->_pDevData;
1084
    CUDAStream<float4>* psLincsDistance = new CUDAStream<float4>(numLincs, 1, "LincsDistance");
1085
1086
    gpu->psLincsDistance                = psLincsDistance;
    gpu->sim.pLincsDistance             = psLincsDistance->_pDevData;
1087
    CUDAStream<int>* psLincsConnections = new CUDAStream<int>(numLincs*maxLinks, 1, "LincsConnections");
1088
1089
    gpu->psLincsConnections             = psLincsConnections;
    gpu->sim.pLincsConnections          = psLincsConnections->_pDevData;
1090
1091
1092
1093
    CUDAStream<int>* psLincsNumConnections = new CUDAStream<int>(numLincs, 1, "LincsConnectionsIndex");
    gpu->psLincsNumConnections             = psLincsNumConnections;
    gpu->sim.pLincsNumConnections          = psLincsNumConnections->_pDevData;
    CUDAStream<int>* psLincsAtomConstraints = new CUDAStream<int>(gpu->natoms*maxAtomConstraints, 1, "LincsAtomConstraints");
1094
1095
    gpu->psLincsAtomConstraints             = psLincsAtomConstraints;
    gpu->sim.pLincsAtomConstraints          = psLincsAtomConstraints->_pDevData;
1096
1097
1098
1099
    CUDAStream<int>* psLincsNumAtomConstraints = new CUDAStream<int>(gpu->natoms, 1, "LincsAtomConstraintsIndex");
    gpu->psLincsNumAtomConstraints             = psLincsNumAtomConstraints;
    gpu->sim.pLincsNumAtomConstraints          = psLincsNumAtomConstraints->_pDevData;
    CUDAStream<float>* psLincsS = new CUDAStream<float>(numLincs, 1, "LincsS");
1100
1101
    gpu->psLincsS             = psLincsS;
    gpu->sim.pLincsS          = psLincsS->_pDevData;
1102
    CUDAStream<float>* psLincsCoupling = new CUDAStream<float>(numLincs*maxLinks, 1, "LincsCoupling");
1103
1104
    gpu->psLincsCoupling               = psLincsCoupling;
    gpu->sim.pLincsCoupling            = psLincsCoupling->_pDevData;
1105
    CUDAStream<float>* psLincsRhs1 = new CUDAStream<float>(numLincs, 1, "LincsRhs1");
1106
1107
    gpu->psLincsRhs1             = psLincsRhs1;
    gpu->sim.pLincsRhs1          = psLincsRhs1->_pDevData;
1108
    CUDAStream<float>* psLincsRhs2 = new CUDAStream<float>(numLincs, 1, "LincsRhs2");
1109
1110
    gpu->psLincsRhs2             = psLincsRhs2;
    gpu->sim.pLincsRhs2          = psLincsRhs2->_pDevData;
1111
    CUDAStream<float>* psLincsSolution = new CUDAStream<float>(numLincs, 1, "LincsSolution");
1112
1113
    gpu->psLincsSolution             = psLincsSolution;
    gpu->sim.pLincsSolution          = psLincsSolution->_pDevData;
1114
    CUDAStream<short>* psSyncCounter = new CUDAStream<short>(3*gpu->sim.blocks, 1, "SyncCounter");
1115
1116
    gpu->psSyncCounter               = psSyncCounter;
    gpu->sim.pSyncCounter            = psSyncCounter->_pDevData;
1117
1118
1119
1120
1121
1122
    CUDAStream<unsigned int>* psRequiredIterations = new CUDAStream<unsigned int>(1, 1, "RequiredIterations");
    gpu->psRequiredIterations               = psRequiredIterations;
    gpu->sim.pRequiredIterations            = psRequiredIterations->_pDevData;
    CUDAStream<float>* psShakeReducedMass = new CUDAStream<float>(numLincs, 1, "LincsSolution");
    gpu->psShakeReducedMass             = psShakeReducedMass;
    gpu->sim.pShakeReducedMass          = psShakeReducedMass->_pDevData;
1123
1124
1125
1126
1127
1128
    CUDAStream<unsigned int>* psConstraintMatrixColumn = new CUDAStream<unsigned int>(numLincs*maxRowElements, 1, "ConstraintMatrixColumn");
    gpu->psConstraintMatrixColumn               = psConstraintMatrixColumn;
    gpu->sim.pConstraintMatrixColumn            = psConstraintMatrixColumn->_pDevData;
    CUDAStream<float>* psConstraintMatrixValue = new CUDAStream<float>(numLincs*maxRowElements, 1, "ConstraintMatrixValue");
    gpu->psConstraintMatrixValue             = psConstraintMatrixValue;
    gpu->sim.pConstraintMatrixValue          = psConstraintMatrixValue->_pDevData;
1129
1130
    gpu->sim.lincsConstraints = numLincs;
    for (int i = 0; i < numLincs; i++) {
1131
1132
        int index = constraintOrder[i];
        int c = lincsConstraints[index];
1133
1134
1135
1136
        (*psLincsAtoms)[i].x = atom1[c];
        (*psLincsAtoms)[i].y = atom2[c];
        (*psLincsDistance)[i].w = distance[c];
        (*psLincsS)[i] = 1.0f/sqrt(invMass1[c]+invMass2[c]);
1137
        (*psShakeReducedMass)[i] = 0.5f/(invMass1[c]+invMass2[c]);
1138
1139
1140
1141
1142
1143
        (*psLincsNumConnections)[i] = linkedConstraints[index].size();
        for (unsigned int j = 0; j < linkedConstraints[index].size(); j++)
            (*psLincsConnections)[i+j*numLincs] = linkedConstraints[index][j];
        for (unsigned int j = 0; j < matrix[index].size(); j++) {
            (*psConstraintMatrixColumn)[i+j*numLincs] = matrix[index][j].first;
            (*psConstraintMatrixValue)[i+j*numLincs] = matrix[index][j].second;
1144
            }
1145
        (*psConstraintMatrixColumn)[i+matrix[index].size()*numLincs] = numLincs;
1146
            }
1147
    for (unsigned int i = 0; i < psSyncCounter->_length; i++)
1148
        (*psSyncCounter)[i] = -1;
1149
    for (unsigned int i = 0; i < atomConstraints.size(); i++) {
1150
        (*psLincsNumAtomConstraints)[i] = atomConstraints[i].size();
1151
1152
        for (unsigned int j = 0; j < atomConstraints[i].size(); j++) {
            bool forward = (atom1[lincsConstraints[atomConstraints[i][j]]] == i);
1153
            (*psLincsAtomConstraints)[i+j*gpu->natoms] = (forward ? inverseOrder[atomConstraints[i][j]]+1 : -inverseOrder[atomConstraints[i][j]]-1);
1154
        }
1155
1156
1157
1158
    }
    psLincsAtoms->Upload();
    psLincsDistance->Upload();
    psLincsS->Upload();
1159
    psShakeReducedMass->Upload();
1160
    psLincsConnections->Upload();
1161
    psLincsNumConnections->Upload();
1162
    psLincsAtomConstraints->Upload();
1163
    psLincsNumAtomConstraints->Upload();
1164
    psSyncCounter->Upload();
1165
1166
    psConstraintMatrixColumn->Upload();
    psConstraintMatrixValue->Upload();
1167
    gpu->sim.lincsTerms = lincsTerms;
1168
    gpu->sim.lincs_threads_per_block = (gpu->sim.lincsConstraints + gpu->sim.blocks - 1) / gpu->sim.blocks;
1169
1170
    if (gpu->sim.lincs_threads_per_block > gpu->sim.threads_per_block)
        gpu->sim.lincs_threads_per_block = gpu->sim.threads_per_block;
1171
1172
    if (gpu->sim.lincs_threads_per_block < gpu->sim.blocks)
        gpu->sim.lincs_threads_per_block = gpu->sim.blocks;
1173
1174
    if (gpu->sim.lincs_threads_per_block%gpu->sim.clusterShakeBlockSize != 0)
        gpu->sim.lincs_threads_per_block += gpu->sim.clusterShakeBlockSize - gpu->sim.lincs_threads_per_block%gpu->sim.clusterShakeBlockSize;
Peter Eastman's avatar
Peter Eastman committed
1175
1176
1177
1178
1179

    // count number of atoms w/o constraint

    int count = 0;
    for (int i = 0; i < gpu->natoms; i++)
1180
       if (!isShakeAtom[i])
Peter Eastman's avatar
Peter Eastman committed
1181
1182
1183
1184
1185
1186
1187
          count++;

    // Allocate NonShake parameters

    gpu->sim.NonShakeConstraints                  = count;
    if( count || true ){

1188
       CUDAStream<int>* psNonShakeID              = new CUDAStream<int>(count, 1, "NonShakeID");
Peter Eastman's avatar
Peter Eastman committed
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
       gpu->psNonShakeID                          = psNonShakeID;
       gpu->sim.pNonShakeID                       = psNonShakeID->_pDevStream[0];

       gpu->sim.nonshake_threads_per_block        = (count + gpu->sim.blocks - 1) / gpu->sim.blocks;

       if (gpu->sim.nonshake_threads_per_block > gpu->sim.max_shake_threads_per_block)
           gpu->sim.nonshake_threads_per_block = gpu->sim.max_shake_threads_per_block;

       if (gpu->sim.nonshake_threads_per_block < 1)
               gpu->sim.nonshake_threads_per_block = 1;

       // load indices

       count = 0;
       for (int i = 0; i < gpu->natoms; i++){
1204
          if (!isShakeAtom[i]){
1205
             (*psNonShakeID)[count++] = i;
1206
          }
Peter Eastman's avatar
Peter Eastman committed
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
       }
       psNonShakeID->Upload();

    } else {
       gpu->sim.nonshake_threads_per_block           = 0;
    }
}

extern "C"
int gpuAllocateInitialBuffers(gpuContext gpu)
{
    gpu->sim.atoms                      = gpu->natoms;
    gpu->sim.paddedNumberOfAtoms        = ((gpu->sim.atoms + GRID - 1) >> GRIDBITS) << GRIDBITS;
    gpu->sim.degreesOfFreedom           = 3 * gpu->sim.atoms - 6;
    gpu->gpAtomTable                    = NULL;
    gpu->gAtomTypes                     = 0;
1223
    gpu->psPosq4                        = new CUDAStream<float4>(gpu->sim.paddedNumberOfAtoms, 1, "Posq");
Peter Eastman's avatar
Peter Eastman committed
1224
1225
1226
1227
1228
1229
1230
1231
1232
    gpu->sim.stride                     = gpu->psPosq4->_stride;
    gpu->sim.stride2                    = gpu->sim.stride * 2;
    gpu->sim.stride3                    = gpu->sim.stride * 3;
    gpu->sim.stride4                    = gpu->sim.stride * 4;
    gpu->sim.pPosq                      = gpu->psPosq4->_pDevStream[0];
    gpu->sim.stride                     = gpu->psPosq4->_stride;
    gpu->sim.stride2                    = 2 * gpu->sim.stride;
    gpu->sim.stride3                    = 3 * gpu->sim.stride;
    gpu->sim.stride4                    = 4 * gpu->sim.stride;
1233
    gpu->psPosqP4                       = new CUDAStream<float4>(gpu->sim.paddedNumberOfAtoms, 1, "PosqP");
Peter Eastman's avatar
Peter Eastman committed
1234
    gpu->sim.pPosqP                     = gpu->psPosqP4->_pDevStream[0];
1235
    gpu->psOldPosq4                     = new CUDAStream<float4>(gpu->sim.paddedNumberOfAtoms, 1, "OldPosq");
Peter Eastman's avatar
Peter Eastman committed
1236
    gpu->sim.pOldPosq                   = gpu->psOldPosq4->_pDevStream[0];
1237
    gpu->psVelm4                        = new CUDAStream<float4>(gpu->sim.paddedNumberOfAtoms, 1, "Velm");
Peter Eastman's avatar
Peter Eastman committed
1238
    gpu->sim.pVelm4                     = gpu->psVelm4->_pDevStream[0];
1239
    gpu->psvVector4                     = new CUDAStream<float4>(gpu->sim.paddedNumberOfAtoms, 1, "vVector");
Peter Eastman's avatar
Peter Eastman committed
1240
    gpu->sim.pvVector4                  = gpu->psvVector4->_pDevStream[0];
1241
    gpu->psxVector4                     = new CUDAStream<float4>(gpu->sim.paddedNumberOfAtoms, 1, "xVector");
Peter Eastman's avatar
Peter Eastman committed
1242
    gpu->sim.pxVector4                  = gpu->psxVector4->_pDevStream[0];
1243
    gpu->psBornRadii                    = new CUDAStream<float>(gpu->sim.paddedNumberOfAtoms, 1, "BornRadii");
Peter Eastman's avatar
Peter Eastman committed
1244
    gpu->sim.pBornRadii                 = gpu->psBornRadii->_pDevStream[0];
1245
    gpu->psObcChain                     = new CUDAStream<float>(gpu->sim.paddedNumberOfAtoms, 1, "ObcChain");
Peter Eastman's avatar
Peter Eastman committed
1246
    gpu->sim.pObcChain                  = gpu->psObcChain->_pDevStream[0];
1247
    gpu->psSigEps2                      = new CUDAStream<float2>(gpu->sim.paddedNumberOfAtoms, 1, "SigEps2");
Peter Eastman's avatar
Peter Eastman committed
1248
    gpu->sim.pAttr                      = gpu->psSigEps2->_pDevStream[0];
1249
1250
1251
1252
1253
1254
    gpu->psEwaldEikr                    = new CUDAStream<float2>(gpu->sim.paddedNumberOfAtoms, 1, "EwaldEikr");
    gpu->sim.pEikr                      = gpu->psEwaldEikr->_pDevStream[0];
    gpu->psEwaldStructureFactor         = new CUDAStream<float2>(gpu->sim.paddedNumberOfAtoms, 1, "EwaldStructureFactor");
    gpu->sim.pStructureFactor           = gpu->psEwaldStructureFactor->_pDevStream[0];
    gpu->psEwaldCosSinSum               = new CUDAStream<float2>(gpu->sim.paddedNumberOfAtoms, 1, "EwaldCosSinSum");
    gpu->sim.pCosSinSum                 = gpu->psEwaldCosSinSum->_pDevStream[0];
1255
    gpu->psObcData                      = new CUDAStream<float2>(gpu->sim.paddedNumberOfAtoms, 1, "ObcData");
Peter Eastman's avatar
Peter Eastman committed
1256
1257
    gpu->sim.pObcData                   = gpu->psObcData->_pDevStream[0];
    gpu->pAtomSymbol                    = new unsigned char[gpu->natoms];
1258
    gpu->psAtomIndex                    = new CUDAStream<int>(gpu->sim.paddedNumberOfAtoms, 1, "AtomIndex");
1259
1260
    gpu->sim.pAtomIndex                 = gpu->psAtomIndex->_pDevStream[0];
    for (int i = 0; i < (int) gpu->sim.paddedNumberOfAtoms; i++)
1261
        (*gpu->psAtomIndex)[i] = i;
1262
    gpu->psAtomIndex->Upload();
Peter Eastman's avatar
Peter Eastman committed
1263
    // Determine randoms
1264
    gpu->seed                           = 1;
1265
    gpu->sim.randomFrames               = 20;
Peter Eastman's avatar
Peter Eastman committed
1266
    gpu->sim.randomIterations           = gpu->sim.randomFrames;
1267
    gpu->sim.randoms                    = gpu->sim.randomFrames * gpu->sim.paddedNumberOfAtoms;
Peter Eastman's avatar
Peter Eastman committed
1268
1269
    gpu->sim.totalRandoms               = gpu->sim.randoms + gpu->sim.paddedNumberOfAtoms;
    gpu->sim.totalRandomsTimesTwo       = gpu->sim.totalRandoms * 2;
1270
1271
1272
1273
    gpu->psRandom4                      = new CUDAStream<float4>(gpu->sim.totalRandomsTimesTwo, 1, "Random4");
    gpu->psRandom2                      = new CUDAStream<float2>(gpu->sim.totalRandomsTimesTwo, 1, "Random2");
    gpu->psRandomPosition               = new CUDAStream<int>(gpu->sim.blocks, 1, "RandomPosition");
    gpu->psRandomSeed                   = new CUDAStream<uint4>(gpu->sim.blocks * gpu->sim.random_threads_per_block, 1, "RandomSeed");
Peter Eastman's avatar
Peter Eastman committed
1274
1275
1276
1277
1278
1279
1280
1281
    gpu->sim.pRandom4a                  = gpu->psRandom4->_pDevStream[0];
    gpu->sim.pRandom2a                  = gpu->psRandom2->_pDevStream[0];
    gpu->sim.pRandom4b                  = gpu->psRandom4->_pDevStream[0] + gpu->sim.totalRandoms;
    gpu->sim.pRandom2b                  = gpu->psRandom2->_pDevStream[0] + gpu->sim.totalRandoms;
    gpu->sim.pRandomPosition            = gpu->psRandomPosition->_pDevStream[0];
    gpu->sim.pRandomSeed                = gpu->psRandomSeed->_pDevStream[0];

    // Allocate and clear linear momentum buffer
1282
    gpu->psLinearMomentum = new CUDAStream<float4>(gpu->sim.blocks, 1, "LinearMomentum");
Peter Eastman's avatar
Peter Eastman committed
1283
1284
1285
    gpu->sim.pLinearMomentum = gpu->psLinearMomentum->_pDevStream[0];
    for (int i = 0; i < (int) gpu->sim.blocks; i++)
    {
1286
1287
1288
1289
        (*gpu->psLinearMomentum)[i].x = 0.0f;
        (*gpu->psLinearMomentum)[i].y = 0.0f;
        (*gpu->psLinearMomentum)[i].z = 0.0f;
        (*gpu->psLinearMomentum)[i].w = 0.0f;
Peter Eastman's avatar
Peter Eastman committed
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
    }
    gpu->psLinearMomentum->Upload();

    return 1;
}

extern "C"
void gpuSetPositions(gpuContext gpu, const vector<float>& x, const vector<float>& y, const vector<float>& z)
{
    for (int i = 0; i < gpu->natoms; i++)
    {
1301
1302
1303
        (*gpu->psPosq4)[i].x = x[i];
        (*gpu->psPosq4)[i].y = y[i];
        (*gpu->psPosq4)[i].z = z[i];
Peter Eastman's avatar
Peter Eastman committed
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
    }
    gpu->psPosq4->Upload();

	 // set flag to recalculate Born radii

	 gpu->bRecalculateBornRadii = true;
} 

extern "C"
void gpuSetVelocities(gpuContext gpu, const vector<float>& x, const vector<float>& y, const vector<float>& z)
{
    for (int i = 0; i < gpu->natoms; i++)
    {
1317
1318
1319
        (*gpu->psVelm4)[i].x = x[i];
        (*gpu->psVelm4)[i].y = y[i];
        (*gpu->psVelm4)[i].z = z[i];
Peter Eastman's avatar
Peter Eastman committed
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
    }
    gpu->psVelm4->Upload();
} 

extern "C"
void gpuSetMass(gpuContext gpu, const vector<float>& mass)
{
    float totalMass = 0.0f;
    for (int i = 0; i < gpu->natoms; i++)
    {
1330
        (*gpu->psVelm4)[i].w = 1.0f/mass[i];
Peter Eastman's avatar
Peter Eastman committed
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
        totalMass += mass[i];
    }
    gpu->sim.inverseTotalMass = 1.0f / totalMass;
    gpu->psVelm4->Upload();
} 

extern "C"
void gpuInitializeRandoms(gpuContext gpu)
{
    for (int i = 0; i < (int) gpu->sim.blocks; i++)
    {
1342
        (*gpu->psRandomPosition)[i] = 0;
Peter Eastman's avatar
Peter Eastman committed
1343
1344
1345
1346
1347
    }
    int seed = gpu->seed | ((gpu->seed ^ 0xffffffff) << 16);
    srand(seed);
    for (int i = 0; i < (int) (gpu->sim.blocks * gpu->sim.random_threads_per_block); i++)
    {
1348
1349
1350
1351
        (*gpu->psRandomSeed)[i].x = rand();
        (*gpu->psRandomSeed)[i].y = rand();
        (*gpu->psRandomSeed)[i].z = rand();
        (*gpu->psRandomSeed)[i].w = rand();
Peter Eastman's avatar
Peter Eastman committed
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
    }
    gpu->psRandomPosition->Upload();
    gpu->psRandomSeed->Upload();
    gpuSetConstants(gpu);
    kGenerateRandoms(gpu);
    return;
}

extern "C"
bool gpuIsAvailable()
{
    int deviceCount;
    cudaGetDeviceCount(&deviceCount);
    return (deviceCount > 0);
}

extern "C"
void* gpuInit(int numAtoms)
{
    gpuContext gpu = new _gpuContext;
    int LRFSize = 0;
    int SMCount = 0;
    int SMMajor = 0;
    int SMMinor = 0;

    // Get adapter
    unsigned int device = 0;
    char * pAdapter;
    pAdapter = getenv ("NV_FAH_DEVICE");
    if (pAdapter != NULL)
    {
        sscanf(pAdapter, "%d", &device);
    }
1385
1386
//    cudaError_t status = cudaSetDevice(device);
//    RTERROR(status, "Error setting CUDA device")
Peter Eastman's avatar
Peter Eastman committed
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453

    // Determine which core to run on
#if 0
    SYSTEM_INFO info;
    GetSystemInfo(&info);
    unsigned int cores = info.dwNumberOfProcessors;
    if (cores > 1)
    {
        HANDLE hproc = GetCurrentProcess();
        unsigned int core = (cores - 1) - (device % (cores - 1)); 
        unsigned int mask = 1 << core;
        SetProcessAffinityMask(hproc, mask);
    }
#endif

    // Determine kernel call configuration
    cudaDeviceProp deviceProp;
    cudaGetDeviceProperties(&deviceProp, 0);

    // Determine SM version
    if (deviceProp.major == 1)
    {
        switch (deviceProp.minor)
        {
        case 0:
        case 1:
            gpu->sm_version = SM_10;
            gpu->sim.workUnitsPerSM = G8X_NONBOND_WORKUNITS_PER_SM;
            break;

        default:
            gpu->sm_version = SM_12;
            gpu->sim.workUnitsPerSM = GT2XX_NONBOND_WORKUNITS_PER_SM;
            break;
        }
    }

    gpu->sim.nonbond_blocks = deviceProp.multiProcessorCount;
    gpu->sim.bornForce2_blocks = deviceProp.multiProcessorCount;
    gpu->sim.blocks = deviceProp.multiProcessorCount;
    if (deviceProp.regsPerBlock == 8192)
    {
        gpu->sim.nonbond_threads_per_block          = G8X_NONBOND_THREADS_PER_BLOCK;
        gpu->sim.bornForce2_threads_per_block       = G8X_BORNFORCE2_THREADS_PER_BLOCK;
        gpu->sim.max_shake_threads_per_block        = G8X_SHAKE_THREADS_PER_BLOCK;
        gpu->sim.max_update_threads_per_block       = G8X_UPDATE_THREADS_PER_BLOCK;
        gpu->sim.max_localForces_threads_per_block  = G8X_LOCALFORCES_THREADS_PER_BLOCK;
        gpu->sim.threads_per_block                  = G8X_THREADS_PER_BLOCK;
        gpu->sim.random_threads_per_block           = G8X_RANDOM_THREADS_PER_BLOCK;
    }
    else
    {
        gpu->sim.nonbond_threads_per_block          = GT2XX_NONBOND_THREADS_PER_BLOCK;
        gpu->sim.bornForce2_threads_per_block       = GT2XX_BORNFORCE2_THREADS_PER_BLOCK;
        gpu->sim.max_shake_threads_per_block        = GT2XX_SHAKE_THREADS_PER_BLOCK;
        gpu->sim.max_update_threads_per_block       = GT2XX_UPDATE_THREADS_PER_BLOCK;
        gpu->sim.max_localForces_threads_per_block  = GT2XX_LOCALFORCES_THREADS_PER_BLOCK;
        gpu->sim.threads_per_block                  = GT2XX_NONBOND_THREADS_PER_BLOCK;
        gpu->sim.random_threads_per_block           = GT2XX_RANDOM_THREADS_PER_BLOCK;
    }
    gpu->sim.shake_threads_per_block                = gpu->sim.max_shake_threads_per_block;
    gpu->sim.localForces_threads_per_block          = gpu->sim.max_localForces_threads_per_block;

    gpu->natoms = numAtoms;
    gpuAllocateInitialBuffers(gpu);
    for (int i = 0; i < gpu->natoms; i++)
    {
1454
1455
1456
1457
        (*gpu->psxVector4)[i].x = 0.0f;
        (*gpu->psxVector4)[i].y = 0.0f;
        (*gpu->psxVector4)[i].z = 0.0f;
        (*gpu->psxVector4)[i].w = 0.0f;
Peter Eastman's avatar
Peter Eastman committed
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
    }
    gpu->psxVector4->Upload();

    gpu->iterations = 0;
    gpu->sim.update_threads_per_block               = (gpu->natoms + gpu->sim.blocks - 1) / gpu->sim.blocks;
    if (gpu->sim.update_threads_per_block > gpu->sim.max_update_threads_per_block)
        gpu->sim.update_threads_per_block = gpu->sim.max_update_threads_per_block;
    if (gpu->sim.update_threads_per_block < 1)
            gpu->sim.update_threads_per_block = 1;
    gpu->sim.bf_reduce_threads_per_block = gpu->sim.update_threads_per_block;
    gpu->sim.bsf_reduce_threads_per_block = (gpu->sim.stride4 + gpu->natoms + gpu->sim.blocks - 1) / gpu->sim.blocks;
    gpu->sim.bsf_reduce_threads_per_block = ((gpu->sim.bsf_reduce_threads_per_block + (GRID - 1)) / GRID) * GRID;
    if (gpu->sim.bsf_reduce_threads_per_block > gpu->sim.threads_per_block)
        gpu->sim.bsf_reduce_threads_per_block = gpu->sim.threads_per_block;
    if (gpu->sim.bsf_reduce_threads_per_block < 1)
        gpu->sim.bsf_reduce_threads_per_block = 1;

    // Initialize constants to reasonable values
    gpu->sim.probeRadius            = probeRadius;
    gpu->sim.surfaceAreaFactor      = surfaceAreaFactor;
    gpu->sim.electricConstant       = electricConstant;
1479
1480
    gpu->sim.nonbondedMethod        = NO_CUTOFF;
    gpu->sim.nonbondedCutoffSqr     = 0.0f;
Peter Eastman's avatar
Peter Eastman committed
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496

    gpu->sim.bigFloat               = 99999999.0f;
    gpu->sim.forceConversionFactor  = forceConversionFactor;
    gpu->sim.preFactor              = 2.0f*electricConstant*((1.0f/defaultInnerDielectric)-(1.0f/defaultSolventDielectric))*gpu->sim.forceConversionFactor;
    gpu->sim.dielectricOffset       = dielectricOffset;
    gpu->sim.alphaOBC               = alphaOBC;
    gpu->sim.betaOBC                = betaOBC;
    gpu->sim.gammaOBC               = gammaOBC;
    gpuSetIntegrationParameters(gpu, 1.0f, 2.0e-3f, 300.0f);
    gpu->sim.maxShakeIterations     = 15;
    gpu->sim.shakeTolerance         = 1.0e-04f * 2.0f;
    gpu->sim.InvMassJ               = 9.920635e-001f;
    gpu->grid                       = GRID;
    gpu->bCalculateCM               = false;
    gpu->bRemoveCM                  = false;
    gpu->bRecalculateBornRadii      = true;
1497
    gpu->bIncludeGBSA               = false;
Peter Eastman's avatar
Peter Eastman committed
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
    gpuInitializeRandoms(gpu);

    // To be determined later
    gpu->psLJ14ID                   = NULL;
    gpu->psForce4                   = NULL;
    gpu->sim.pForce4                = NULL;
    gpu->sim.pForce4a               = NULL;
    gpu->sim.pForce4b               = NULL;
    gpu->psBornForce                = NULL;
    gpu->sim.pBornForce             = NULL;
    gpu->psBornSum                  = NULL;
    gpu->sim.pBornSum               = NULL;
    gpu->psBondID                   = NULL;
    gpu->psBondParameter            = NULL;
    gpu->psBondAngleID1             = NULL;
    gpu->psBondAngleID2             = NULL;
    gpu->psBondAngleParameter       = NULL;
    gpu->psDihedralID1              = NULL;
    gpu->psDihedralID2              = NULL;
    gpu->psDihedralParameter        = NULL;
    gpu->psRbDihedralID1            = NULL;
    gpu->psRbDihedralID2            = NULL;
    gpu->psRbDihedralParameter1     = NULL;
    gpu->psRbDihedralParameter2     = NULL;
    gpu->psLJ14ID                   = NULL;
    gpu->psLJ14Parameter            = NULL;
    gpu->psShakeID                  = NULL;
    gpu->psShakeParameter           = NULL;
1526
1527
    gpu->psSettleID                 = NULL;
    gpu->psSettleParameter          = NULL;
Peter Eastman's avatar
Peter Eastman committed
1528
    gpu->psExclusion                = NULL;
1529
    gpu->psExclusionIndex           = NULL;
Peter Eastman's avatar
Peter Eastman committed
1530
    gpu->psWorkUnit                 = NULL;
1531
1532
1533
1534
1535
    gpu->psInteractingWorkUnit      = NULL;
    gpu->psInteractionFlag          = NULL;
    gpu->psInteractionCount         = NULL;
    gpu->psGridBoundingBox          = NULL;
    gpu->psGridCenter               = NULL;
1536
1537
1538
    gpu->psLincsAtoms               = NULL;
    gpu->psLincsDistance            = NULL;
    gpu->psLincsConnections         = NULL;
1539
    gpu->psLincsNumConnections      = NULL;
1540
    gpu->psLincsAtomConstraints     = NULL;
1541
    gpu->psLincsNumAtomConstraints  = NULL;
1542
1543
1544
1545
1546
1547
    gpu->psLincsS                   = NULL;
    gpu->psLincsCoupling            = NULL;
    gpu->psLincsRhs1                = NULL;
    gpu->psLincsRhs2                = NULL;
    gpu->psLincsSolution            = NULL;
    gpu->psSyncCounter              = NULL;
1548
1549
    gpu->psRequiredIterations       = NULL;
    gpu->psShakeReducedMass         = NULL;
1550
1551
1552
    gpu->psRigidClusterConstraintIndex = NULL;
    gpu->psRigidClusterMatrix       = NULL;
    gpu->psRigidClusterMatrixIndex  = NULL;
1553
1554
    gpu->psConstraintMatrixColumn   = NULL;
    gpu->psConstraintMatrixValue    = NULL;
Peter Eastman's avatar
Peter Eastman committed
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655

    // Initialize output buffer before reading parameters
    gpu->pOutputBufferCounter       = new unsigned int[gpu->sim.paddedNumberOfAtoms];
    memset(gpu->pOutputBufferCounter, 0, gpu->sim.paddedNumberOfAtoms * sizeof(unsigned int));

    return (void*)gpu;
}

extern "C"
void gpuSetIntegrationParameters(gpuContext gpu, float tau, float deltaT, float temperature) {
    gpu->sim.deltaT                 = deltaT;
    gpu->sim.oneOverDeltaT          = 1.0f/deltaT;
    gpu->sim.tau                    = tau;
    gpu->sim.GDT                    = gpu->sim.deltaT / gpu->sim.tau;
    gpu->sim.EPH                    = exp(0.5f * gpu->sim.GDT);
    gpu->sim.EMH                    = exp(-0.5f * gpu->sim.GDT);
    gpu->sim.EP                     = exp(gpu->sim.GDT);
    gpu->sim.EM                     = exp(-gpu->sim.GDT);
    gpu->sim.OneMinusEM             = 1.0f - gpu->sim.EM;
    gpu->sim.TauOneMinusEM          = gpu->sim.tau * gpu->sim.OneMinusEM;
    if (gpu->sim.GDT >= 0.1f)
    {
        float term1                 = gpu->sim.EPH - 1.0f;
        term1                      *= term1;
        gpu->sim.B                  = gpu->sim.GDT * (gpu->sim.EP - 1.0f) - 4.0f * term1;
        gpu->sim.C                  = gpu->sim.GDT - 3.0f + 4.0f * gpu->sim.EMH - gpu->sim.EM;
        gpu->sim.D                  = 2.0f - gpu->sim.EPH - gpu->sim.EMH;
    }
    else
    {
        float term1                 = 0.5f * gpu->sim.GDT;
        float term2                 = term1 * term1;
        float term4                 = term2 * term2;

        float third                 = 1.0f / 3.0f;
        float o7_9                  = 7.0f / 9.0f;
        float o1_12                 = 1.0f / 12.0f;
        float o17_90                = 17.0f / 90.0f;
        float o7_30                 = 7.0f / 30.0f;
        float o31_1260              = 31.0f / 1260.0f;
        float o_360                 = 1.0f / 360.0f;

        gpu->sim.B                  = term4 * (third + term1 * (third + term1 * (o17_90 + term1 * o7_9)));
        gpu->sim.C                  = term2 * term1 * (2.0f * third + term1 * (-0.5f + term1 * (o7_30 + term1 * (-o1_12 + term1 * o31_1260))));
        gpu->sim.D                  = term2 * (-1.0f + term2 * (-o1_12 - term2 * o_360));   
    }
    gpu->sim.TauDOverEMMinusOne     = gpu->sim.tau * gpu->sim.D / (gpu->sim.EM - 1.0f);
    gpu->sim.DOverTauC              = gpu->sim.D / (gpu->sim.tau * gpu->sim.C);
    gpu->sim.fix1                   = gpu->sim.tau * (gpu->sim.EPH - gpu->sim.EMH);
    gpu->sim.oneOverFix1            = 1.0f / (gpu->sim.tau * (gpu->sim.EPH - gpu->sim.EMH));
    gpu->sim.T                      = temperature;
    gpu->sim.kT                     = BOLTZ * gpu->sim.T;
    gpu->sim.V                      = sqrt(gpu->sim.kT * (1.0f - gpu->sim.EM));
    gpu->sim.X                      = gpu->sim.tau * sqrt(gpu->sim.kT * gpu->sim.C);
    gpu->sim.Yv                     = sqrt(gpu->sim.kT * gpu->sim.B / gpu->sim.C);
    gpu->sim.Yx                     = gpu->sim.tau * sqrt(gpu->sim.kT * gpu->sim.B / (1.0f - gpu->sim.EM));
}

extern "C"
void gpuSetVerletIntegrationParameters(gpuContext gpu, float deltaT) {
    gpu->sim.deltaT                 = deltaT;
    gpu->sim.oneOverDeltaT          = 1.0f/deltaT;
}

extern "C"
void gpuSetBrownianIntegrationParameters(gpuContext gpu, float tau, float deltaT, float temperature) {
    gpu->sim.deltaT                 = deltaT;
    gpu->sim.oneOverDeltaT          = 1.0f/deltaT;
    gpu->sim.tau                    = tau;
    gpu->sim.GDT                    = gpu->sim.deltaT * gpu->sim.tau;
    gpu->sim.T                      = temperature;
    gpu->sim.kT                     = BOLTZ * gpu->sim.T;
    gpu->sim.Yv = gpu->sim.Yx       = sqrt(2.0f*gpu->sim.kT*deltaT*tau);
}

extern "C"
void gpuSetAndersenThermostatParameters(gpuContext gpu, float temperature, float collisionProbability) {
    gpu->sim.T                      = temperature;
    gpu->sim.kT                     = BOLTZ * gpu->sim.T;
    gpu->sim.collisionProbability   = collisionProbability;
    gpu->sim.Yv = gpu->sim.Yx       = 1.0f;
    gpu->sim.V = gpu->sim.X         = 1.0f;
}

extern "C"
void gpuShutDown(gpuContext gpu)
{
    // Delete sysmem pointers
    delete[] gpu->pOutputBufferCounter;
    delete[] gpu->gpAtomTable;
    delete[] gpu->pAtomSymbol;

    // Delete device pointers
    delete gpu->psPosq4;
    delete gpu->psPosqP4;
    delete gpu->psOldPosq4;
    delete gpu->psVelm4;
    delete gpu->psForce4;
    delete gpu->psxVector4;
    delete gpu->psvVector4;
    delete gpu->psSigEps2; 
1656
1657
1658
    delete gpu->psEwaldEikr;
    delete gpu->psEwaldStructureFactor;
    delete gpu->psEwaldCosSinSum;
Peter Eastman's avatar
Peter Eastman committed
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
    delete gpu->psObcData; 
    delete gpu->psObcChain;
    delete gpu->psBornForce;
    delete gpu->psBornRadii;
    delete gpu->psBornSum;
    delete gpu->psBondID;
    delete gpu->psBondParameter;
    delete gpu->psBondAngleID1;
    delete gpu->psBondAngleID2;
    delete gpu->psBondAngleParameter;
    delete gpu->psDihedralID1;
    delete gpu->psDihedralID2;
    delete gpu->psDihedralParameter;
    delete gpu->psRbDihedralID1;
    delete gpu->psRbDihedralID2;
    delete gpu->psRbDihedralParameter1;
    delete gpu->psRbDihedralParameter2;
    delete gpu->psLJ14ID;
    delete gpu->psLJ14Parameter;
    delete gpu->psShakeID;
    delete gpu->psShakeParameter;
1680
1681
    delete gpu->psSettleID;
    delete gpu->psSettleParameter;
Peter Eastman's avatar
Peter Eastman committed
1682
    delete gpu->psExclusion;
1683
    delete gpu->psExclusionIndex;
Peter Eastman's avatar
Peter Eastman committed
1684
    delete gpu->psWorkUnit;
1685
1686
1687
    delete gpu->psInteractingWorkUnit;
    delete gpu->psInteractionFlag;
    delete gpu->psInteractionCount;
Peter Eastman's avatar
Peter Eastman committed
1688
1689
1690
1691
1692
    delete gpu->psRandom4;
    delete gpu->psRandom2;
    delete gpu->psRandomPosition;    
    delete gpu->psRandomSeed;
    delete gpu->psLinearMomentum;
1693
1694
1695
    delete gpu->psAtomIndex;
    delete gpu->psGridBoundingBox;
    delete gpu->psGridCenter;
1696
1697
1698
    delete gpu->psLincsAtoms;
    delete gpu->psLincsDistance;
    delete gpu->psLincsConnections;
1699
    delete gpu->psLincsNumConnections;
1700
    delete gpu->psLincsAtomConstraints;
1701
    delete gpu->psLincsNumAtomConstraints;
1702
1703
1704
1705
1706
1707
    delete gpu->psLincsS;
    delete gpu->psLincsCoupling;
    delete gpu->psLincsRhs1;
    delete gpu->psLincsRhs2;
    delete gpu->psLincsSolution;
    delete gpu->psSyncCounter;
1708
1709
    delete gpu->psRequiredIterations;
    delete gpu->psShakeReducedMass;
1710
1711
1712
    delete gpu->psRigidClusterConstraintIndex;
    delete gpu->psRigidClusterMatrix;
    delete gpu->psRigidClusterMatrixIndex;
1713
1714
    delete gpu->psConstraintMatrixColumn;
    delete gpu->psConstraintMatrixValue;
1715
1716
    if (gpu->cudpp != 0)
        cudppDestroyPlan(gpu->cudpp);
Peter Eastman's avatar
Peter Eastman committed
1717
1718
1719
1720
1721
1722
1723
1724
1725

    // Wrap up
    delete gpu;
    return;
}

extern "C"
int gpuBuildOutputBuffers(gpuContext gpu)
{
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
    // Select the number of output buffer to use.
    gpu->bOutputBufferPerWarp           = true;
    gpu->sim.nonbondOutputBuffers       = gpu->sim.nonbond_blocks * gpu->sim.nonbond_threads_per_block / GRID;
    if (gpu->sim.nonbondOutputBuffers >= gpu->sim.paddedNumberOfAtoms/GRID)
    {
        // For small systems, it is more efficient to have one output buffer per block of 32 atoms instead of one per warp.
        gpu->bOutputBufferPerWarp           = false;
        gpu->sim.nonbondOutputBuffers       = gpu->sim.paddedNumberOfAtoms / GRID;
    }
    gpu->sim.totalNonbondOutputBuffers  = (gpu->bIncludeGBSA ? 2 * gpu->sim.nonbondOutputBuffers : gpu->sim.nonbondOutputBuffers);
    gpu->sim.outputBuffers              = gpu->sim.totalNonbondOutputBuffers;


Peter Eastman's avatar
Peter Eastman committed
1739
1740
1741
1742
1743
1744
1745
1746
1747
    unsigned int outputBuffers = gpu->sim.totalNonbondOutputBuffers;
    for (unsigned int i = 0; i < gpu->sim.paddedNumberOfAtoms; i++)
    {
        if (outputBuffers < gpu->pOutputBufferCounter[i])
        {
            outputBuffers = gpu->pOutputBufferCounter[i];
        }
    }    
    gpu->sim.outputBuffers      = outputBuffers;
1748
1749
1750
    gpu->psForce4               = new CUDAStream<float4>(gpu->sim.paddedNumberOfAtoms, outputBuffers, "Force");
    gpu->psBornForce            = new CUDAStream<float>(gpu->sim.paddedNumberOfAtoms, gpu->sim.nonbondOutputBuffers, "BornForce");
    gpu->psBornSum              = new CUDAStream<float>(gpu->sim.paddedNumberOfAtoms, gpu->sim.nonbondOutputBuffers, "BornSum");
Peter Eastman's avatar
Peter Eastman committed
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
    gpu->sim.pForce4            = gpu->psForce4->_pDevStream[0];
    gpu->sim.pForce4a           = gpu->sim.pForce4;
    gpu->sim.pForce4b           = gpu->sim.pForce4 + 1 * gpu->sim.nonbondOutputBuffers * gpu->sim.stride;
    gpu->sim.pBornForce         = gpu->psBornForce->_pDevStream[0];
    gpu->sim.pBornSum           = gpu->psBornSum->_pDevStream[0];

    // Determine local energy paramter offsets for bonded interactions
    gpu->sim.bond_offset        =                                  gpu->psBondParameter->_stride;
    gpu->sim.bond_angle_offset  = gpu->sim.bond_offset           + gpu->psBondAngleParameter->_stride;
    gpu->sim.dihedral_offset    = gpu->sim.bond_angle_offset     + gpu->psDihedralParameter->_stride;
    gpu->sim.rb_dihedral_offset = gpu->sim.dihedral_offset       + gpu->psRbDihedralParameter1->_stride;
    gpu->sim.LJ14_offset        = gpu->sim.rb_dihedral_offset    + gpu->psLJ14Parameter->_stride;
    gpu->sim.localForces_threads_per_block  = (gpu->sim.LJ14_offset / gpu->sim.blocks + 15) & 0xfffffff0;
    if (gpu->sim.localForces_threads_per_block > gpu->sim.max_localForces_threads_per_block)
        gpu->sim.localForces_threads_per_block = gpu->sim.max_localForces_threads_per_block;
    if (gpu->sim.localForces_threads_per_block < 1)
        gpu->sim.localForces_threads_per_block = 1;

    // Flip local force output buffers
    int flip = outputBuffers - 1;
    for (int i = 0; i < (int) gpu->sim.bonds; i++)
    {
1773
1774
        (*gpu->psBondID)[i].z = flip - (*gpu->psBondID)[i].z;
        (*gpu->psBondID)[i].w = flip - (*gpu->psBondID)[i].w;
Peter Eastman's avatar
Peter Eastman committed
1775
1776
1777
    }
    for (int i = 0; i < (int) gpu->sim.bond_angles; i++)
    {
1778
1779
1780
        (*gpu->psBondAngleID1)[i].w = flip - (*gpu->psBondAngleID1)[i].w;
        (*gpu->psBondAngleID2)[i].x = flip - (*gpu->psBondAngleID2)[i].x;
        (*gpu->psBondAngleID2)[i].y = flip - (*gpu->psBondAngleID2)[i].y;
Peter Eastman's avatar
Peter Eastman committed
1781
1782
1783
    }
    for (int i = 0; i < (int) gpu->sim.dihedrals; i++)
    {
1784
1785
1786
1787
        (*gpu->psDihedralID2)[i].x = flip - (*gpu->psDihedralID2)[i].x;
        (*gpu->psDihedralID2)[i].y = flip - (*gpu->psDihedralID2)[i].y;
        (*gpu->psDihedralID2)[i].z = flip - (*gpu->psDihedralID2)[i].z;
        (*gpu->psDihedralID2)[i].w = flip - (*gpu->psDihedralID2)[i].w;
Peter Eastman's avatar
Peter Eastman committed
1788
1789
1790
    }
    for (int i = 0; i < (int) gpu->sim.rb_dihedrals; i++)
    {
1791
1792
1793
1794
        (*gpu->psRbDihedralID2)[i].x = flip - (*gpu->psRbDihedralID2)[i].x;
        (*gpu->psRbDihedralID2)[i].y = flip - (*gpu->psRbDihedralID2)[i].y;
        (*gpu->psRbDihedralID2)[i].z = flip - (*gpu->psRbDihedralID2)[i].z;
        (*gpu->psRbDihedralID2)[i].w = flip - (*gpu->psRbDihedralID2)[i].w;
Peter Eastman's avatar
Peter Eastman committed
1795
1796
1797
    }
    for (int i = 0; i < (int) gpu->sim.LJ14s; i++)
    {
1798
1799
        (*gpu->psLJ14ID)[i].z = flip - (*gpu->psLJ14ID)[i].z;
        (*gpu->psLJ14ID)[i].w = flip - (*gpu->psLJ14ID)[i].w;
Peter Eastman's avatar
Peter Eastman committed
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
    }
    gpu->psBondID->Upload();
    gpu->psBondAngleID1->Upload();
    gpu->psBondAngleID2->Upload();
    gpu->psDihedralID2->Upload();
    gpu->psRbDihedralID2->Upload();
    gpu->psLJ14ID->Upload();

    return 1;
}

extern "C"
int gpuBuildThreadBlockWorkList(gpuContext gpu)
{
    const unsigned int atoms = gpu->sim.paddedNumberOfAtoms;
    const unsigned int grid = gpu->grid;
    const unsigned int dim = (atoms + (grid - 1)) / grid;
    const unsigned int cells = dim * (dim + 1) / 2;
1818
1819
    CUDAStream<unsigned int>* psWorkUnit = new CUDAStream<unsigned int>(cells, 1u, "WorkUnit");
    unsigned int* pWorkList = psWorkUnit->_pSysData;
Peter Eastman's avatar
Peter Eastman committed
1820
1821
    gpu->psWorkUnit = psWorkUnit;
    gpu->sim.pWorkUnit = psWorkUnit->_pDevStream[0];
1822
    CUDAStream<unsigned int>* psInteractingWorkUnit = new CUDAStream<unsigned int>(cells, 1u, "InteractingWorkUnit");
1823
1824
    gpu->psInteractingWorkUnit = psInteractingWorkUnit;
    gpu->sim.pInteractingWorkUnit = psInteractingWorkUnit->_pDevStream[0];
1825
    CUDAStream<unsigned int>* psInteractionFlag = new CUDAStream<unsigned int>(cells, 1u, "InteractionFlag");
1826
1827
    gpu->psInteractionFlag = psInteractionFlag;
    gpu->sim.pInteractionFlag = psInteractionFlag->_pDevStream[0];
1828
    CUDAStream<size_t>* psInteractionCount = new CUDAStream<size_t>(1, 1u, "InteractionCount");
1829
1830
    gpu->psInteractionCount = psInteractionCount;
    gpu->sim.pInteractionCount = psInteractionCount->_pDevStream[0];
1831
    CUDAStream<float4>* psGridBoundingBox = new CUDAStream<float4>(dim, 1u, "GridBoundingBox");
1832
1833
    gpu->psGridBoundingBox = psGridBoundingBox;
    gpu->sim.pGridBoundingBox = psGridBoundingBox->_pDevStream[0];
1834
    CUDAStream<float4>* psGridCenter = new CUDAStream<float4>(dim, 1u, "GridCenter");
1835
1836
    gpu->psGridCenter = psGridCenter;
    gpu->sim.pGridCenter = psGridCenter->_pDevStream[0];
Peter Eastman's avatar
Peter Eastman committed
1837
1838
1839
1840
    gpu->sim.nonbond_workBlock      = gpu->sim.nonbond_threads_per_block / GRID;
    gpu->sim.bornForce2_workBlock   = gpu->sim.bornForce2_threads_per_block / GRID;
    gpu->sim.workUnits = cells;

1841
1842
1843
1844
1845
1846
1847
1848
    // Initialize the CUDPP workspace.
    gpu->cudpp = 0;
    CUDPPConfiguration config;
    config.datatype = CUDPP_UINT;
    config.algorithm = CUDPP_COMPACT;
    config.options = CUDPP_OPTION_FORWARD;
    CUDPPResult result = cudppPlan(&gpu->cudpp, config, cells, 1, 0);
    if (CUDPP_SUCCESS != result)
Peter Eastman's avatar
Peter Eastman committed
1849
    {
1850
1851
        printf("Error initializing CUDPP: %d\n", result);
        exit(-1);
Peter Eastman's avatar
Peter Eastman committed
1852
    }
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864

    // Increase block count if necessary for extra large molecules that would
    // otherwise overflow the SM workunit buffers
//    int minimumBlocks = (cells + gpu->sim.workUnitsPerSM - 1) / gpu->sim.workUnitsPerSM;
//    if ((int) gpu->sim.nonbond_blocks < minimumBlocks)
//    {
//        gpu->sim.nonbond_blocks = gpu->sim.nonbond_blocks * ((minimumBlocks + gpu->sim.nonbond_blocks - 1) / gpu->sim.nonbond_blocks);
//    }
//    if ((int) gpu->sim.bornForce2_blocks < minimumBlocks)
//    {
//        gpu->sim.bornForce2_blocks = gpu->sim.bornForce2_blocks * ((minimumBlocks + gpu->sim.bornForce2_blocks - 1) / gpu->sim.bornForce2_blocks);
//    }
Peter Eastman's avatar
Peter Eastman committed
1865
1866
1867
1868
    gpu->sim.nbWorkUnitsPerBlock            = cells / gpu->sim.nonbond_blocks;
    gpu->sim.nbWorkUnitsPerBlockRemainder   = cells - gpu->sim.nonbond_blocks * gpu->sim.nbWorkUnitsPerBlock;
    gpu->sim.bf2WorkUnitsPerBlock           = cells / gpu->sim.bornForce2_blocks;
    gpu->sim.bf2WorkUnitsPerBlockRemainder  = cells - gpu->sim.bornForce2_blocks * gpu->sim.bf2WorkUnitsPerBlock;
1869
1870
    gpu->sim.interaction_threads_per_block = 64;
    gpu->sim.interaction_blocks = (gpu->sim.workUnits + gpu->sim.interaction_threads_per_block - 1) / gpu->sim.interaction_threads_per_block;
Peter Eastman's avatar
Peter Eastman committed
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897

    // Decrease thread count for extra small molecules to spread computation
    // across entire chip
    int activeWorkUnits = gpu->sim.nonbond_blocks * gpu->sim.nonbond_workBlock;
    if (activeWorkUnits > (int) cells)
    {
        int balancedWorkBlock                   = (cells + gpu->sim.nonbond_blocks - 1) / gpu->sim.nonbond_blocks;
        gpu->sim.nonbond_threads_per_block      = balancedWorkBlock * GRID;
        gpu->sim.nonbond_workBlock              = balancedWorkBlock;
    }
    activeWorkUnits = gpu->sim.bornForce2_blocks * gpu->sim.bornForce2_workBlock;
    if (activeWorkUnits > (int) cells)
    {
        int balancedWorkBlock                   = (cells + gpu->sim.bornForce2_blocks - 1) / gpu->sim.bornForce2_blocks;
        gpu->sim.bornForce2_threads_per_block   = balancedWorkBlock * GRID;
        gpu->sim.bornForce2_workBlock           = balancedWorkBlock;
    }

    unsigned int count = 0;
    for (unsigned int y = 0; y < dim; y++)
    {
        for (unsigned int x = y; x < dim; x++)
        {
            pWorkList[count] = (x << 17) | (y << 2);
            count++;
        }
    }
1898
    (*gpu->psInteractionCount)[0] = gpu->sim.workUnits;
Peter Eastman's avatar
Peter Eastman committed
1899

1900
    gpu->psInteractionCount->Upload();
Peter Eastman's avatar
Peter Eastman committed
1901
1902
1903
1904
1905
1906
    psWorkUnit->Upload();
    gpuSetConstants(gpu);
    return cells;
}

extern "C"
1907
void gpuBuildExclusionList(gpuContext gpu)
Peter Eastman's avatar
Peter Eastman committed
1908
{
1909
1910
    const unsigned int atoms = gpu->sim.paddedNumberOfAtoms;
    const unsigned int grid = gpu->grid;
1911
    const unsigned int dim = atoms/grid;
1912
    unsigned int* pWorkList = gpu->psWorkUnit->_pSysData;
1913

1914
    // Mark which work units have exclusions.
Peter Eastman's avatar
Peter Eastman committed
1915

1916
    for (int atom1 = 0; atom1 < (int)gpu->exclusions.size(); ++atom1)
Peter Eastman's avatar
Peter Eastman committed
1917
    {
1918
        int x = atom1/grid;
1919
        for (int j = 0; j < (int)gpu->exclusions[atom1].size(); ++j)
1920
1921
1922
1923
1924
1925
1926
        {
            int atom2 = gpu->exclusions[atom1][j];
            int y = atom2/grid;
            int cell = (x > y ? x+y*dim-y*(y+1)/2 : y+x*dim-x*(x+1)/2);
            pWorkList[cell] |= 1;
        }
    }
1927
    if ((int)gpu->sim.paddedNumberOfAtoms > gpu->natoms)
1928
1929
    {
        int lastBlock = gpu->natoms/grid;
1930
        for (int i = 0; i < (int)gpu->sim.workUnits; ++i)
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
        {
            int x = pWorkList[i]>>17;
            int y = (pWorkList[i]>>2)&0x7FFF;
            if (x == lastBlock || y == lastBlock)
                pWorkList[i] |= 1;
        }
    }

    // Build a list of indexes for the work units with exclusions.

1941
    CUDAStream<unsigned int>* psExclusionIndex = new CUDAStream<unsigned int>(gpu->sim.workUnits, 1u, "ExclusionIndex");
1942
1943
1944
1945
    gpu->psExclusionIndex = psExclusionIndex;
    unsigned int* pExclusionIndex = psExclusionIndex->_pSysData;
    gpu->sim.pExclusionIndex = psExclusionIndex->_pDevData;
    int numWithExclusions = 0;
1946
    for (int i = 0; i < (int)psExclusionIndex->_length; ++i)
1947
1948
1949
1950
1951
        if ((pWorkList[i]&1) == 1)
            pExclusionIndex[i] = (numWithExclusions++)*grid;

    // Record the exclusion data.

1952
    CUDAStream<unsigned int>* psExclusion = new CUDAStream<unsigned int>(numWithExclusions*grid, 1u, "Exclusion");
1953
1954
1955
    gpu->psExclusion = psExclusion;
    unsigned int* pExclusion = psExclusion->_pSysData;
    gpu->sim.pExclusion = psExclusion->_pDevData;
1956
    for (int i = 0; i < (int)psExclusion->_length; ++i)
1957
        pExclusion[i] = 0xFFFFFFFF;
1958
    for (int atom1 = 0; atom1 < (int)gpu->exclusions.size(); ++atom1)
1959
1960
1961
    {
        int x = atom1/grid;
        int offset1 = atom1-x*grid;
1962
        for (int j = 0; j < (int)gpu->exclusions[atom1].size(); ++j)
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
        {
            int atom2 = gpu->exclusions[atom1][j];
            int y = atom2/grid;
            int offset2 = atom2-y*grid;
            if (x > y)
            {
                int cell = x+y*dim-y*(y+1)/2;
                pExclusion[pExclusionIndex[cell]+offset1] &= 0xFFFFFFFF-(1<<offset2);
            }
            else
            {
                int cell = y+x*dim-x*(x+1)/2;
                pExclusion[pExclusionIndex[cell]+offset2] &= 0xFFFFFFFF-(1<<offset1);
            }
        }
    }
1979
1980
1981

    // Mark all interactions that involve a padding atom as being excluded.

1982
    for (int atom1 = gpu->natoms; atom1 < (int)atoms; ++atom1)
1983
1984
1985
    {
        int x = atom1/grid;
        int offset1 = atom1-x*grid;
1986
        for (int atom2 = 0; atom2 < (int)atoms; ++atom2)
1987
1988
1989
1990
        {
            int y = atom2/grid;
            int index = x*atoms+y*grid+offset1;
            int offset2 = atom2-y*grid;
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
            if (x >= y)
            {
                int cell = x+y*dim-y*(y+1)/2;
                pExclusion[pExclusionIndex[cell]+offset1] &= 0xFFFFFFFF-(1<<offset2);
            }
            if (y >= x)
            {
                int cell = y+x*dim-x*(x+1)/2;
                pExclusion[pExclusionIndex[cell]+offset2] &= 0xFFFFFFFF-(1<<offset1);
            }
Peter Eastman's avatar
Peter Eastman committed
2001
2002
2003
2004
        }
    }
    
    psExclusion->Upload();
2005
    psExclusionIndex->Upload();
2006
    gpu->psWorkUnit->Upload();
Peter Eastman's avatar
Peter Eastman committed
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
2022
    gpuSetConstants(gpu);
}

extern "C"
int gpuSetConstants(gpuContext gpu)
{
    SetCalculateCDLJForcesSim(gpu);
    SetCalculateCDLJObcGbsaForces1Sim(gpu);
    SetCalculateLocalForcesSim(gpu);
    SetCalculateObcGbsaBornSumSim(gpu);
    SetCalculateObcGbsaForces2Sim(gpu);
    SetCalculateAndersenThermostatSim(gpu);
    SetForcesSim(gpu);
    SetUpdateShakeHSim(gpu);
    SetVerletUpdateSim(gpu);
    SetBrownianUpdateSim(gpu);
2023
    SetSettleSim(gpu);
2024
    SetCShakeSim(gpu);
2025
    SetLincsSim(gpu);
Peter Eastman's avatar
Peter Eastman committed
2026
2027
2028
2029
    SetRandomSim(gpu);
    return 1;
}

2030
2031
2032
2033
2034
static void tagAtomsInMolecule(int atom, int molecule, vector<int>& atomMolecule, vector<vector<int> >& atomBonds)
{
    // Recursively tag atoms as belonging to a particular molecule.

    atomMolecule[atom] = molecule;
2035
    for (int i = 0; i < (int)atomBonds[atom].size(); i++)
2036
2037
2038
2039
2040
2041
2042
2043
2044
        if (atomMolecule[atomBonds[atom][i]] == -1)
            tagAtomsInMolecule(atomBonds[atom][i], molecule, atomMolecule, atomBonds);
}

static void findMoleculeGroups(gpuContext gpu)
{
    // First make a list of constraints for future use.

    vector<Constraint> constraints;
2045
    for (int i = 0; i < (int)gpu->sim.ShakeConstraints; i++)
2046
    {
2047
2048
2049
2050
2051
        int atom1 = (*gpu->psShakeID)[i].x;
        int atom2 = (*gpu->psShakeID)[i].y;
        int atom3 = (*gpu->psShakeID)[i].z;
        int atom4 = (*gpu->psShakeID)[i].w;
        float distance2 = (*gpu->psShakeParameter)[i].z;
2052
2053
2054
2055
        constraints.push_back(Constraint(atom1, atom2, distance2));
        if (atom3 != -1)
            constraints.push_back(Constraint(atom1, atom3, distance2));
        if (atom4 != -1)
2056
            constraints.push_back(Constraint(atom1, atom4, distance2));
2057
    }
2058
    for (int i = 0; i < (int)gpu->sim.settleConstraints; i++)
2059
    {
2060
2061
2062
2063
2064
        int atom1 = (*gpu->psSettleID)[i].x;
        int atom2 = (*gpu->psSettleID)[i].y;
        int atom3 = (*gpu->psSettleID)[i].z;
        float distance12 = (*gpu->psSettleParameter)[i].x;
        float distance23 = (*gpu->psSettleParameter)[i].y;
2065
2066
2067
2068
        constraints.push_back(Constraint(atom1, atom2, distance12*distance12));
        constraints.push_back(Constraint(atom1, atom3, distance12*distance12));
        constraints.push_back(Constraint(atom2, atom3, distance23*distance23));
    }
2069
    for (int i = 0; i < (int)gpu->sim.lincsConstraints; i++)
Peter Eastman's avatar
Peter Eastman committed
2070
2071
2072
2073
2074
2075
    {
        int atom1 = (*gpu->psLincsAtoms)[i].x;
        int atom2 = (*gpu->psLincsAtoms)[i].y;
        float distance2 = (*gpu->psLincsDistance)[i].w;
        constraints.push_back(Constraint(atom1, atom2, distance2));
    }
2076
2077
2078
2079
2080

    // First make a list of every other atom to which each atom is connect by a bond or constraint.

    int numAtoms = gpu->natoms;
    vector<vector<int> > atomBonds(numAtoms);
2081
    for (int i = 0; i < (int)gpu->sim.bonds; i++)
2082
    {
2083
2084
        int atom1 = (*gpu->psBondID)[i].x;
        int atom2 = (*gpu->psBondID)[i].y;
2085
2086
2087
        atomBonds[atom1].push_back(atom2);
        atomBonds[atom2].push_back(atom1);
    }
2088
    for (int i = 0; i < (int)constraints.size(); i++)
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
    {
        int atom1 = constraints[i].atom1;
        int atom2 = constraints[i].atom2;
        atomBonds[atom1].push_back(atom2);
        atomBonds[atom2].push_back(atom1);
    }

    // Now tag atoms by which molecule they belong to.

    vector<int> atomMolecule(numAtoms, -1);
    int numMolecules = 0;
    for (int i = 0; i < numAtoms; i++)
        if (atomMolecule[i] == -1)
            tagAtomsInMolecule(i, numMolecules++, atomMolecule, atomBonds);
    vector<vector<int> > atomIndices(numMolecules);
    for (int i = 0; i < numAtoms; i++)
        atomIndices[atomMolecule[i]].push_back(i);

    // Construct a description of each molecule.

    vector<Molecule> molecules(numMolecules);
    for (int i = 0; i < numMolecules; i++)
        molecules[i].atoms = atomIndices[i];
2112
    for (int i = 0; i < (int)gpu->sim.bonds; i++)
2113
    {
2114
        int atom1 = (*gpu->psBondID)[i].x;
2115
2116
        molecules[atomMolecule[atom1]].bonds.push_back(i);
    }
2117
    for (int i = 0; i < (int)gpu->sim.bond_angles; i++)
2118
    {
2119
        int atom1 = (*gpu->psBondAngleID1)[i].x;
2120
2121
        molecules[atomMolecule[atom1]].angles.push_back(i);
    }
2122
    for (int i = 0; i < (int)gpu->sim.dihedrals; i++)
2123
    {
2124
        int atom1 = (*gpu->psDihedralID1)[i].x;
2125
2126
        molecules[atomMolecule[atom1]].periodicTorsions.push_back(i);
    }
2127
    for (int i = 0; i < (int)gpu->sim.rb_dihedrals; i++)
2128
    {
2129
        int atom1 = (*gpu->psRbDihedralID1)[i].x;
2130
2131
        molecules[atomMolecule[atom1]].rbTorsions.push_back(i);
    }
2132
    for (int i = 0; i < (int)constraints.size(); i++)
2133
2134
2135
2136
2137
2138
2139
2140
    {
        molecules[atomMolecule[constraints[i].atom1]].constraints.push_back(i);
    }

    // Sort them into groups of identical molecules.

    vector<Molecule> uniqueMolecules;
    vector<vector<int> > moleculeInstances;
2141
    for (int molIndex = 0; molIndex < (int)molecules.size(); molIndex++)
2142
2143
2144
2145
2146
2147
    {
        Molecule& mol = molecules[molIndex];

        // See if it is identical to another molecule.

        bool isNew = true;
2148
        for (int j = 0; j < (int)uniqueMolecules.size() && isNew; j++)
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
        {
            Molecule& mol2 = uniqueMolecules[j];
            bool identical = true;
            if (mol.atoms.size() != mol2.atoms.size() || mol.bonds.size() != mol2.bonds.size()
                    || mol.angles.size() != mol2.angles.size() || mol.periodicTorsions.size() != mol2.periodicTorsions.size()
                    || mol.rbTorsions.size() != mol2.rbTorsions.size() || mol.constraints.size() != mol2.constraints.size())
                identical = false;
            int atomOffset = mol2.atoms[0]-mol.atoms[0];
            float4* posq = gpu->psPosq4->_pSysData;
            float4* velm = gpu->psVelm4->_pSysData;
            float2* sigeps = gpu->psSigEps2->_pSysData;
2160
            for (int i = 0; i < (int)mol.atoms.size() && identical; i++)
2161
2162
2163
2164
2165
2166
                if (mol.atoms[i] != mol2.atoms[i]-atomOffset || posq[mol.atoms[i]].w != posq[mol2.atoms[i]].w ||
                        velm[mol.atoms[i]].w != velm[mol2.atoms[i]].w || sigeps[mol.atoms[i]].x != sigeps[mol2.atoms[i]].x ||
                        sigeps[mol.atoms[i]].y != sigeps[mol2.atoms[i]].y)
                    identical = false;
            int4* bondID = gpu->psBondID->_pSysData;
            float2* bondParam = gpu->psBondParameter->_pSysData;
2167
            for (int i = 0; i < (int)mol.bonds.size() && identical; i++)
2168
2169
2170
2171
2172
                if (bondID[mol.bonds[i]].x != bondID[mol2.bonds[i]].x-atomOffset || bondID[mol.bonds[i]].y != bondID[mol2.bonds[i]].y-atomOffset ||
                        bondParam[mol.bonds[i]].x != bondParam[mol2.bonds[i]].x || bondParam[mol.bonds[i]].y != bondParam[mol2.bonds[i]].y)
                    identical = false;
            int4* angleID = gpu->psBondAngleID1->_pSysData;
            float2* angleParam = gpu->psBondAngleParameter->_pSysData;
2173
            for (int i = 0; i < (int)mol.angles.size() && identical; i++)
2174
2175
2176
2177
2178
2179
2180
2181
                if (angleID[mol.angles[i]].x != angleID[mol2.angles[i]].x-atomOffset ||
                        angleID[mol.angles[i]].y != angleID[mol2.angles[i]].y-atomOffset ||
                        angleID[mol.angles[i]].z != angleID[mol2.angles[i]].z-atomOffset ||
                        angleParam[mol.angles[i]].x != angleParam[mol2.angles[i]].x ||
                        angleParam[mol.angles[i]].y != angleParam[mol2.angles[i]].y)
                    identical = false;
            int4* periodicID = gpu->psDihedralID1->_pSysData;
            float4* periodicParam = gpu->psDihedralParameter->_pSysData;
2182
            for (int i = 0; i < (int)mol.periodicTorsions.size() && identical; i++)
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
                if (periodicID[mol.periodicTorsions[i]].x != periodicID[mol2.periodicTorsions[i]].x-atomOffset ||
                        periodicID[mol.periodicTorsions[i]].y != periodicID[mol2.periodicTorsions[i]].y-atomOffset ||
                        periodicID[mol.periodicTorsions[i]].z != periodicID[mol2.periodicTorsions[i]].z-atomOffset ||
                        periodicID[mol.periodicTorsions[i]].w != periodicID[mol2.periodicTorsions[i]].w-atomOffset ||
                        periodicParam[mol.periodicTorsions[i]].x != periodicParam[mol2.periodicTorsions[i]].x ||
                        periodicParam[mol.periodicTorsions[i]].y != periodicParam[mol2.periodicTorsions[i]].y ||
                        periodicParam[mol.periodicTorsions[i]].z != periodicParam[mol2.periodicTorsions[i]].z)
                    identical = false;
            int4* rbID = gpu->psRbDihedralID1->_pSysData;
            float4* rbParam1 = gpu->psRbDihedralParameter1->_pSysData;
            float2* rbParam2 = gpu->psRbDihedralParameter2->_pSysData;
2194
            for (int i = 0; i < (int)mol.rbTorsions.size() && identical; i++)
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
                if (rbID[mol.rbTorsions[i]].x != rbID[mol2.rbTorsions[i]].x-atomOffset ||
                        rbID[mol.rbTorsions[i]].y != rbID[mol2.rbTorsions[i]].y-atomOffset ||
                        rbID[mol.rbTorsions[i]].z != rbID[mol2.rbTorsions[i]].z-atomOffset ||
                        rbID[mol.rbTorsions[i]].w != rbID[mol2.rbTorsions[i]].w-atomOffset ||
                        rbParam1[mol.rbTorsions[i]].x != rbParam1[mol2.rbTorsions[i]].x ||
                        rbParam1[mol.rbTorsions[i]].y != rbParam1[mol2.rbTorsions[i]].y ||
                        rbParam1[mol.rbTorsions[i]].z != rbParam1[mol2.rbTorsions[i]].z ||
                        rbParam1[mol.rbTorsions[i]].w != rbParam1[mol2.rbTorsions[i]].w ||
                        rbParam2[mol.rbTorsions[i]].x != rbParam2[mol2.rbTorsions[i]].x ||
                        rbParam2[mol.rbTorsions[i]].y != rbParam2[mol2.rbTorsions[i]].y)
                    identical = false;
2206
            for (int i = 0; i < (int)mol.constraints.size() && identical; i++)
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
                if (constraints[mol.constraints[i]].atom1 != constraints[mol2.constraints[i]].atom1-atomOffset ||
                        constraints[mol.constraints[i]].atom2 != constraints[mol2.constraints[i]].atom2-atomOffset ||
                        constraints[mol.constraints[i]].distance2 != constraints[mol2.constraints[i]].distance2)
                    identical = false;
            if (identical)
            {
                moleculeInstances[j].push_back(mol.atoms[0]);
                isNew = false;
            }
        }
        if (isNew)
        {
            uniqueMolecules.push_back(mol);
            moleculeInstances.push_back(vector<int>());
            moleculeInstances[moleculeInstances.size()-1].push_back(mol.atoms[0]);
        }
    }
    gpu->moleculeGroups.resize(moleculeInstances.size());
2225
    for (int i = 0; i < (int)moleculeInstances.size(); i++)
2226
2227
2228
2229
    {
        gpu->moleculeGroups[i].instances = moleculeInstances[i];
        vector<int>& atoms = uniqueMolecules[i].atoms;
        gpu->moleculeGroups[i].atoms.resize(atoms.size());
2230
        for (int j = 0; j < (int)atoms.size(); j++)
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
            gpu->moleculeGroups[i].atoms[j] = atoms[j]-atoms[0];
    }
}

extern "C"
void gpuReorderAtoms(gpuContext gpu)
{
    if (gpu->natoms == 0 || gpu->sim.nonbondedCutoffSqr == 0.0)
        return;
    if (gpu->moleculeGroups.size() == 0)
        findMoleculeGroups(gpu);

    // Find the range of positions and the number of bins along each axis.

    int numAtoms = gpu->natoms;
    gpu->psPosq4->Download();
    gpu->psVelm4->Download();
    float4* posq = gpu->psPosq4->_pSysData;
    float4* velm = gpu->psVelm4->_pSysData;
    float minx = posq[0].x, maxx = posq[0].x;
    float miny = posq[0].y, maxy = posq[0].y;
    float minz = posq[0].z, maxz = posq[0].z;
    if (gpu->sim.nonbondedMethod == PERIODIC)
    {
        minx = miny = minz = 0.0;
        maxx = gpu->sim.periodicBoxSizeX;
        maxy = gpu->sim.periodicBoxSizeY;
        maxz = gpu->sim.periodicBoxSizeZ;
    }
    else
    {
        for (int i = 1; i < numAtoms; i++)
        {
            minx = min(minx, posq[i].x);
            maxx = max(maxx, posq[i].x);
            miny = min(miny, posq[i].y);
            maxy = max(maxy, posq[i].y);
            minz = min(minz, posq[i].z);
            maxz = max(maxz, posq[i].z);
        }
    }

    // Loop over each group of identical molecules and reorder them.

    vector<int> originalIndex(numAtoms);
    vector<float4> newPosq(numAtoms);
    vector<float4> newVelm(numAtoms);
2278
    for (int group = 0; group < (int)gpu->moleculeGroups.size(); group++)
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
    {
        // Find the center of each molecule.

        gpuMoleculeGroup& mol = gpu->moleculeGroups[group];
        int numMolecules = mol.instances.size();
        vector<int>& atoms = mol.atoms;
        vector<float3> molPos(numMolecules);
        for (int i = 0; i < numMolecules; i++)
        {
            molPos[i].x = 0.0f;
            molPos[i].y = 0.0f;
            molPos[i].z = 0.0f;
2291
            for (int j = 0; j < (int)atoms.size(); j++)
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
            {
                int atom = atoms[j]+mol.instances[i];
                molPos[i].x += posq[atom].x;
                molPos[i].y += posq[atom].y;
                molPos[i].z += posq[atom].z;
            }
            molPos[i].x /= atoms.size();
            molPos[i].y /= atoms.size();
            molPos[i].z /= atoms.size();
        }
        if (gpu->sim.nonbondedMethod == PERIODIC)
        {
            // Move each molecule position into the same box.

            for (int i = 0; i < numMolecules; i++)
            {
2308
2309
2310
2311
2312
2313
2314
2315
                float dx = floor(molPos[i].x/gpu->sim.periodicBoxSizeX)*gpu->sim.periodicBoxSizeX;
                float dy = floor(molPos[i].y/gpu->sim.periodicBoxSizeY)*gpu->sim.periodicBoxSizeY;
                float dz = floor(molPos[i].z/gpu->sim.periodicBoxSizeZ)*gpu->sim.periodicBoxSizeZ;
                if (dx != 0.0f || dy != 0.0f || dz != 0.0f)
                {
                    molPos[i].x -= dx;
                    molPos[i].y -= dy;
                    molPos[i].z -= dz;
2316
                    for (int j = 0; j < (int)atoms.size(); j++)
2317
2318
2319
2320
2321
2322
2323
                    {
                        int atom = atoms[j]+mol.instances[i];
                        posq[atom].x -= dx;
                        posq[atom].y -= dy;
                        posq[atom].z -= dz;
                    }
                }
2324
2325
2326
2327
2328
            }
        }

        // Select a bin for each molecule, then sort them by bin.

2329
        bool useHilbert = (numMolecules > 5000 || atoms.size() > 8); // For small systems, a simple zigzag curve works better than a Hilbert curve.
2330
2331
        float binWidth;
        if (useHilbert)
2332
            binWidth = (float)(max(max(maxx-minx, maxy-miny), maxz-minz)/255.0);
2333
        else
2334
            binWidth = (float)(0.2*sqrt(gpu->sim.nonbondedCutoffSqr));
2335
2336
        int xbins = 1 + (int) ((maxx-minx)/binWidth);
        int ybins = 1 + (int) ((maxy-miny)/binWidth);
2337
        vector<pair<int, int> > molBins(numMolecules);
2338
        bitmask_t coords[3];
2339
2340
2341
2342
2343
        for (int i = 0; i < numMolecules; i++)
        {
            int x = (int) ((molPos[i].x-minx)/binWidth);
            int y = (int) ((molPos[i].y-miny)/binWidth);
            int z = (int) ((molPos[i].z-minz)/binWidth);
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
            int bin;
            if (useHilbert)
            {
                coords[0] = x;
                coords[1] = y;
                coords[2] = z;
                bin = (int) hilbert_c2i(3, 8, coords);
            }
            else
            {
                int yodd = y&1;
                int zodd = z&1;
                bin = z*xbins*ybins;
                bin += (zodd ? ybins-y : y)*xbins;
                bin += (yodd ? xbins-x : x);
            }
2360
2361
2362
2363
2364
2365
2366
2367
            molBins[i] = pair<int, int>(bin, i);
        }
        sort(molBins.begin(), molBins.end());

        // Reorder the atoms.

        for (int i = 0; i < numMolecules; i++)
        {
2368
            for (int j = 0; j < (int)atoms.size(); j++)
2369
2370
2371
            {
                int oldIndex = mol.instances[molBins[i].second]+atoms[j];
                int newIndex = mol.instances[i]+atoms[j];
2372
                originalIndex[newIndex] = (*gpu->psAtomIndex)[oldIndex];
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
                newPosq[newIndex] = posq[oldIndex];
                newVelm[newIndex] = velm[oldIndex];
            }
        }
    }

    // Update the streams.

    for (int i = 0; i < numAtoms; i++)
        posq[i] = newPosq[i];
    gpu->psPosq4->Upload();
    for (int i = 0; i < numAtoms; i++)
        velm[i] = newVelm[i];
    gpu->psVelm4->Upload();
    for (int i = 0; i < numAtoms; i++)
2388
        (*gpu->psAtomIndex)[i] = originalIndex[i];
2389
2390
    gpu->psAtomIndex->Upload();
}