OpenCLIntegrationUtilities.cpp 50.5 KB
Newer Older
1
2
3
4
5
6
7
8
/* -------------------------------------------------------------------------- *
 *                                   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.               *
 *                                                                            *
9
 * Portions copyright (c) 2009-2015 Stanford University and the Authors.      *
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
 * Authors: Peter Eastman                                                     *
 * Contributors:                                                              *
 *                                                                            *
 * This program is free software: you can redistribute it and/or modify       *
 * it under the terms of the GNU Lesser General Public License as published   *
 * by the Free Software Foundation, either version 3 of the License, or       *
 * (at your option) any later version.                                        *
 *                                                                            *
 * This program is distributed in the hope that it will be useful,            *
 * but WITHOUT ANY WARRANTY; without even the implied warranty of             *
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the              *
 * GNU Lesser General Public License for more details.                        *
 *                                                                            *
 * You should have received a copy of the GNU Lesser General Public License   *
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.      *
 * -------------------------------------------------------------------------- */

#include "OpenCLIntegrationUtilities.h"
#include "OpenCLArray.h"
29
#include "OpenCLKernelSources.h"
30
#include "openmm/internal/OSRngSeed.h"
31
#include "openmm/HarmonicAngleForce.h"
32
#include "openmm/VirtualSite.h"
33
34
#include "quern.h"
#include "OpenCLExpressionUtilities.h"
35
#include "ReferenceCCMAAlgorithm.h"
36
#include <algorithm>
37
38
#include <cmath>
#include <cstdlib>
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
#include <map>

using namespace OpenMM;
using namespace std;

struct OpenCLIntegrationUtilities::ShakeCluster {
    int centralID;
    int peripheralID[3];
    int size;
    bool valid;
    double distance;
    double centralInvMass, peripheralInvMass;
    ShakeCluster() : valid(true) {
    }
    ShakeCluster(int centralID, double invMass) : centralID(centralID), centralInvMass(invMass), size(0), valid(true) {
    }
    void addAtom(int id, double dist, double invMass) {
56
        if (size == 3 || (size > 0 && abs(dist-distance)/distance > 1e-8) || (size > 0 && abs(invMass-peripheralInvMass)/peripheralInvMass > 1e-8))
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
            valid = false;
        else {
            peripheralID[size++] = id;
            distance = dist;
            peripheralInvMass = invMass;
        }
    }
    void markInvalid(map<int, ShakeCluster>& allClusters, vector<bool>& invalidForShake)
    {
        valid = false;
        invalidForShake[centralID] = true;
        for (int i = 0; i < size; i++) {
            invalidForShake[peripheralID[i]] = true;
            map<int, ShakeCluster>::iterator otherCluster = allClusters.find(peripheralID[i]);
            if (otherCluster != allClusters.end() && otherCluster->second.valid)
                otherCluster->second.markInvalid(allClusters, invalidForShake);
        }
    }
};

77
78
79
struct OpenCLIntegrationUtilities::ConstraintOrderer : public binary_function<int, int, bool> {
    const vector<int>& atom1;
    const vector<int>& atom2;
80
81
    const vector<int>& constraints;
    ConstraintOrderer(const vector<int>& atom1, const vector<int>& atom2, const vector<int>& constraints) : atom1(atom1), atom2(atom2), constraints(constraints) {
82
83
    }
    bool operator()(int x, int y) {
84
85
86
87
88
        int ix = constraints[x];
        int iy = constraints[y];
        if (atom1[ix] != atom1[iy])
            return atom1[ix] < atom1[iy];
        return atom2[ix] < atom2[iy];
89
90
91
    }
};

92
93
94
95
96
97
98
static void setPosqCorrectionArg(OpenCLContext& cl, cl::Kernel& kernel, int index) {
    if (cl.getUseMixedPrecision())
        kernel.setArg<cl::Buffer>(index, cl.getPosqCorrection().getDeviceBuffer());
    else
        kernel.setArg<void*>(index, NULL);
}

99
OpenCLIntegrationUtilities::OpenCLIntegrationUtilities(OpenCLContext& context, const System& system) : context(context),
100
        posDelta(NULL), settleAtoms(NULL), settleParams(NULL), shakeAtoms(NULL), shakeParams(NULL),
101
        random(NULL), randomSeed(NULL), randomPos(0), stepSize(NULL), ccmaAtoms(NULL), ccmaDistance(NULL),
102
        ccmaReducedMass(NULL), ccmaAtomConstraints(NULL), ccmaNumAtomConstraints(NULL), ccmaConstraintMatrixColumn(NULL),
Peter Eastman's avatar
Peter Eastman committed
103
        ccmaConstraintMatrixValue(NULL), ccmaDelta1(NULL), ccmaDelta2(NULL), ccmaConverged(NULL), ccmaConvergedHostBuffer(NULL),
104
        vsite2AvgAtoms(NULL), vsite2AvgWeights(NULL), vsite3AvgAtoms(NULL), vsite3AvgWeights(NULL),
105
        vsiteOutOfPlaneAtoms(NULL), vsiteOutOfPlaneWeights(NULL), vsiteLocalCoordsAtoms(NULL), vsiteLocalCoordsParams(NULL),
106
        hasInitializedPosConstraintKernels(false), hasInitializedVelConstraintKernels(false), hasOverlappingVsites(false) {
107
108
    // Create workspace arrays.

109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
    if (context.getUseDoublePrecision() || context.getUseMixedPrecision()) {
        posDelta = OpenCLArray::create<mm_double4>(context, context.getPaddedNumAtoms(), "posDelta");
        vector<mm_double4> deltas(posDelta->getSize(), mm_double4(0.0, 0.0, 0.0, 0.0));
        posDelta->upload(deltas);
        stepSize = OpenCLArray::create<mm_double2>(context, 1, "stepSize");
        vector<mm_double2> step(1, mm_double2(0.0, 0.0));
        stepSize->upload(step);
    }
    else {
        posDelta = OpenCLArray::create<mm_float4>(context, context.getPaddedNumAtoms(), "posDelta");
        vector<mm_float4> deltas(posDelta->getSize(), mm_float4(0.0f, 0.0f, 0.0f, 0.0f));
        posDelta->upload(deltas);
        stepSize = OpenCLArray::create<mm_float2>(context, 1, "stepSize");
        vector<mm_float2> step(1, mm_float2(0.0f, 0.0f));
        stepSize->upload(step);
    }
125
126
127
128
129
130
131
    
    // Create the time shift kernel for calculating kinetic energy.
    
    map<string, string> timeShiftDefines;
    timeShiftDefines["NUM_ATOMS"] = context.intToString(system.getNumParticles());
    cl::Program utilitiesProgram = context.createProgram(OpenCLKernelSources::integrationUtilities, timeShiftDefines);
    timeShiftKernel = cl::Kernel(utilitiesProgram, "timeShiftVelocities");
132
133
134

    // Create kernels for enforcing constraints.

135
136
    map<string, string> velocityDefines;
    velocityDefines["CONSTRAIN_VELOCITIES"] = "1";
137
    cl::Program settleProgram = context.createProgram(OpenCLKernelSources::settle);
138
139
    settlePosKernel = cl::Kernel(settleProgram, "applySettle");
    settleVelKernel = cl::Kernel(settleProgram, "constrainVelocities");
140
    cl::Program shakeProgram = context.createProgram(OpenCLKernelSources::shakeHydrogens);
141
142
143
    shakePosKernel = cl::Kernel(shakeProgram, "applyShakeToHydrogens");
    shakeProgram = context.createProgram(OpenCLKernelSources::shakeHydrogens, velocityDefines);
    shakeVelKernel = cl::Kernel(shakeProgram, "applyShakeToHydrogens");
144
145
146

    // Record the set of constraints and how many constraints each atom is involved in.

147
148
149
    vector<int> atom1;
    vector<int> atom2;
    vector<double> distance;
150
    vector<int> constraintCount(context.getNumAtoms(), 0);
151
152
153
154
155
156
157
158
159
160
161
    for (int i = 0; i < system.getNumConstraints(); i++) {
        int p1, p2;
        double d;
        system.getConstraintParameters(i, p1, p2, d);
        if (system.getParticleMass(p1) != 0 || system.getParticleMass(p2) != 0) {
            atom1.push_back(p1);
            atom2.push_back(p2);
            distance.push_back(d);
            constraintCount[p1]++;
            constraintCount[p2]++;
        }
162
163
164
165
166
167
    }

    // 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.

168
169
    int numAtoms = system.getNumParticles();
    vector<map<int, float> > settleConstraints(numAtoms);
170
171
    for (int i = 0; i < (int)atom1.size(); i++) {
        if (constraintCount[atom1[i]] == 2 && constraintCount[atom2[i]] == 2) {
172
173
            settleConstraints[atom1[i]][atom2[i]] = (float) distance[i];
            settleConstraints[atom2[i]][atom1[i]] = (float) distance[i];
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
        }
    }

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

    vector<int> settleClusters;
    for (int i = 0; i < (int)settleConstraints.size(); i++) {
        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 SETTLE clusters.

196
    vector<bool> isShakeAtom(numAtoms, false);
197
    if (settleClusters.size() > 0) {
198
199
        vector<mm_int4> atoms;
        vector<mm_float2> params;
200
201
202
203
204
205
206
207
208
        for (int i = 0; i < (int) settleClusters.size(); i++) {
            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
209
210
                atoms.push_back(mm_int4(atom1, atom2, atom3, 0));
                params.push_back(mm_float2(dist12, dist23));
211
212
213
            }
            else if (dist12 == dist23) {
                // atom2 is the central atom
214
215
                atoms.push_back(mm_int4(atom2, atom1, atom3, 0));
                params.push_back(mm_float2(dist12, dist13));
216
217
218
            }
            else if (dist13 == dist23) {
                // atom3 is the central atom
219
220
                atoms.push_back(mm_int4(atom3, atom1, atom2, 0));
                params.push_back(mm_float2(dist13, dist12));
221
222
            }
            else
223
                continue; // We can't handle this with SETTLE
224
225
226
227
            isShakeAtom[atom1] = true;
            isShakeAtom[atom2] = true;
            isShakeAtom[atom3] = true;
        }
228
229
230
231
232
233
        if (atoms.size() > 0) {
            settleAtoms = OpenCLArray::create<mm_int4>(context, atoms.size(), "settleAtoms");
            settleParams = OpenCLArray::create<mm_float2>(context, params.size(), "settleParams");
            settleAtoms->upload(atoms);
            settleParams->upload(params);
        }
234
235
236
237
238
    }

    // Find clusters consisting of a central atom with up to three peripheral atoms.

    map<int, ShakeCluster> clusters;
239
    vector<bool> invalidForShake(numAtoms, false);
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
    for (int i = 0; i < (int) atom1.size(); i++) {
        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;
        if (firstIsCentral) {
            centralID = atom1[i];
            peripheralID = atom2[i];
        }
        else {
            centralID = atom2[i];
            peripheralID = atom1[i];
        }

        // Add it to the cluster.

        if (clusters.find(centralID) == clusters.end()) {
            clusters[centralID] = ShakeCluster(centralID, 1.0/system.getParticleMass(centralID));
        }
        ShakeCluster& cluster = clusters[centralID];
        cluster.addAtom(peripheralID, distance[i], 1.0/system.getParticleMass(peripheralID));
        if (constraintCount[peripheralID] != 1 || invalidForShake[atom1[i]] || invalidForShake[atom2[i]]) {
            cluster.markInvalid(clusters, invalidForShake);
            map<int, ShakeCluster>::iterator otherCluster = clusters.find(peripheralID);
            if (otherCluster != clusters.end() && otherCluster->second.valid)
                otherCluster->second.markInvalid(clusters, invalidForShake);
        }
    }
    int validShakeClusters = 0;
    for (map<int, ShakeCluster>::iterator iter = clusters.begin(); iter != clusters.end(); ++iter) {
        ShakeCluster& cluster = iter->second;
        if (cluster.valid) {
283
            cluster.valid = !invalidForShake[cluster.centralID] && cluster.size == constraintCount[cluster.centralID];
284
285
286
287
288
289
290
291
292
293
294
            for (int i = 0; i < cluster.size; i++)
                if (invalidForShake[cluster.peripheralID[i]])
                    cluster.valid = false;
            if (cluster.valid)
                ++validShakeClusters;
        }
    }

    // Record the SHAKE clusters.

    if (validShakeClusters > 0) {
295
296
        vector<mm_int4> atoms;
        vector<mm_float4> params;
297
298
299
300
301
        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;
302
            atoms.push_back(mm_int4(cluster.centralID, cluster.peripheralID[0], (cluster.size > 1 ? cluster.peripheralID[1] : -1), (cluster.size > 2 ? cluster.peripheralID[2] : -1)));
303
            params.push_back(mm_float4((cl_float) cluster.centralInvMass, (cl_float) (0.5/(cluster.centralInvMass+cluster.peripheralInvMass)), (cl_float) (cluster.distance*cluster.distance), (cl_float) cluster.peripheralInvMass));
304
305
306
307
308
309
310
311
            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;
        }
312
313
        shakeAtoms = OpenCLArray::create<mm_int4>(context, atoms.size(), "shakeAtoms");
        shakeParams = OpenCLArray::create<mm_float4>(context, params.size(), "shakeParams");
314
315
316
        shakeAtoms->upload(atoms);
        shakeParams->upload(params);
    }
317
318
319
320
321
322
323
324
325
326
327
328

    // Find connected constraints for CCMA.

    vector<int> ccmaConstraints;
    for (unsigned i = 0; i < atom1.size(); i++)
        if (!isShakeAtom[atom1[i]])
            ccmaConstraints.push_back(i);

    // Record the connections between constraints.

    int numCCMA = (int) ccmaConstraints.size();
    if (numCCMA > 0) {
329
330
331
332
        // Record information needed by ReferenceCCMAAlgorithm.
        
        vector<pair<int, int> > refIndices(numCCMA);
        vector<RealOpenMM> refDistance(numCCMA);
333
        for (int i = 0; i < numCCMA; i++) {
334
335
336
            int index = ccmaConstraints[i];
            refIndices[i] = make_pair(atom1[index], atom2[index]);
            refDistance[i] = distance[index];
337
        }
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
        vector<RealOpenMM> refMasses(numAtoms);
        for (int i = 0; i < numAtoms; ++i)
            refMasses[i] = (RealOpenMM) system.getParticleMass(i);

        // Look up angles for CCMA.
        
        vector<ReferenceCCMAAlgorithm::AngleInfo> angles;
        for (int i = 0; i < system.getNumForces(); i++) {
            const HarmonicAngleForce* force = dynamic_cast<const HarmonicAngleForce*>(&system.getForce(i));
            if (force != NULL) {
                for (int j = 0; j < force->getNumAngles(); j++) {
                    int atom1, atom2, atom3;
                    double angle, k;
                    force->getAngleParameters(j, atom1, atom2, atom3, angle, k);
                    angles.push_back(ReferenceCCMAAlgorithm::AngleInfo(atom1, atom2, atom3, (RealOpenMM) angle));
353
354
355
                }
            }
        }
356
357
358
359
360
        
        // Create a ReferenceCCMAAlgorithm.  It will build and invert the constraint matrix for us.
        
        ReferenceCCMAAlgorithm ccma(numAtoms, numCCMA, refIndices, refDistance, refMasses, angles, 0.1);
        vector<vector<pair<int, double> > > matrix = ccma.getMatrix();
361
362
363
364
365
        int maxRowElements = 0;
        for (unsigned i = 0; i < matrix.size(); i++)
            maxRowElements = max(maxRowElements, (int) matrix[i].size());
        maxRowElements++;

366
367
368
369
370
371
372
373
374
375
376
        // Build the list of constraints for each atom.

        vector<vector<int> > atomConstraints(context.getNumAtoms());
        for (int i = 0; i < numCCMA; i++) {
            atomConstraints[atom1[ccmaConstraints[i]]].push_back(i);
            atomConstraints[atom2[ccmaConstraints[i]]].push_back(i);
        }
        int maxAtomConstraints = 0;
        for (unsigned i = 0; i < atomConstraints.size(); i++)
            maxAtomConstraints = max(maxAtomConstraints, (int) atomConstraints[i].size());

377
378
379
380
381
        // Sort the constraints.

        vector<int> constraintOrder(numCCMA);
        for (int i = 0; i < numCCMA; ++i)
            constraintOrder[i] = i;
382
        sort(constraintOrder.begin(), constraintOrder.end(), ConstraintOrderer(atom1, atom2, ccmaConstraints));
383
384
385
386
387
388
389
390
391
        vector<int> inverseOrder(numCCMA);
        for (int i = 0; i < numCCMA; ++i)
            inverseOrder[constraintOrder[i]] = i;
        for (int i = 0; i < (int)matrix.size(); ++i)
            for (int j = 0; j < (int)matrix[i].size(); ++j)
                matrix[i][j].first = inverseOrder[matrix[i][j].first];

        // Record the CCMA data structures.

392
393
394
        ccmaAtoms = OpenCLArray::create<mm_int2>(context, numCCMA, "CcmaAtoms");
        ccmaAtomConstraints = OpenCLArray::create<cl_int>(context, numAtoms*maxAtomConstraints, "CcmaAtomConstraints");
        ccmaNumAtomConstraints = OpenCLArray::create<cl_int>(context, numAtoms, "CcmaAtomConstraintsIndex");
395
        ccmaConstraintMatrixColumn = OpenCLArray::create<cl_int>(context, numCCMA*maxRowElements, "ConstraintMatrixColumn");
396
        ccmaConverged = OpenCLArray::create<cl_int>(context, 2, "CcmaConverged");
Peter Eastman's avatar
Peter Eastman committed
397
398
399
400
        ccmaConvergedHostBuffer = OpenCLArray::create<cl_int>(context, 1, "CcmaConvergedHostBuffer", CL_MEM_WRITE_ONLY | CL_MEM_ALLOC_HOST_PTR);
        // Different communication mechanisms give optimal performance on AMD and on NVIDIA.
        string vendor = context.getDevice().getInfo<CL_DEVICE_VENDOR>();
        ccmaUseDirectBuffer = (vendor.size() >= 28 && vendor.substr(0, 28) == "Advanced Micro Devices, Inc.");
401
402
403
404
        vector<mm_int2> atomsVec(ccmaAtoms->getSize());
        vector<cl_int> atomConstraintsVec(ccmaAtomConstraints->getSize());
        vector<cl_int> numAtomConstraintsVec(ccmaNumAtomConstraints->getSize());
        vector<cl_int> constraintMatrixColumnVec(ccmaConstraintMatrixColumn->getSize());
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
        if (context.getUseDoublePrecision() || context.getUseMixedPrecision()) {
            ccmaDistance = OpenCLArray::create<mm_double4>(context, numCCMA, "CcmaDistance");
            ccmaDelta1 = OpenCLArray::create<cl_double>(context, numCCMA, "CcmaDelta1");
            ccmaDelta2 = OpenCLArray::create<cl_double>(context, numCCMA, "CcmaDelta2");
            ccmaReducedMass = OpenCLArray::create<cl_double>(context, numCCMA, "CcmaReducedMass");
            ccmaConstraintMatrixValue = OpenCLArray::create<cl_double>(context, numCCMA*maxRowElements, "ConstraintMatrixValue");
            vector<mm_double4> distanceVec(ccmaDistance->getSize());
            vector<cl_double> reducedMassVec(ccmaReducedMass->getSize());
            vector<cl_double> constraintMatrixValueVec(ccmaConstraintMatrixValue->getSize());
            for (int i = 0; i < numCCMA; i++) {
                int index = constraintOrder[i];
                int c = ccmaConstraints[index];
                atomsVec[i].x = atom1[c];
                atomsVec[i].y = atom2[c];
                distanceVec[i].w = distance[c];
                reducedMassVec[i] = (0.5/(1.0/system.getParticleMass(atom1[c])+1.0/system.getParticleMass(atom2[c])));
                for (unsigned int j = 0; j < matrix[index].size(); j++) {
                    constraintMatrixColumnVec[i+j*numCCMA] = matrix[index][j].first;
                    constraintMatrixValueVec[i+j*numCCMA] = matrix[index][j].second;
                }
                constraintMatrixColumnVec[i+matrix[index].size()*numCCMA] = numCCMA;
            }
            for (unsigned int i = 0; i < atomConstraints.size(); i++) {
                numAtomConstraintsVec[i] = atomConstraints[i].size();
                for (unsigned int j = 0; j < atomConstraints[i].size(); j++) {
                    bool forward = (atom1[ccmaConstraints[atomConstraints[i][j]]] == i);
                    atomConstraintsVec[i+j*numAtoms] = (forward ? inverseOrder[atomConstraints[i][j]]+1 : -inverseOrder[atomConstraints[i][j]]-1);
                }
433
            }
434
435
436
            ccmaDistance->upload(distanceVec);
            ccmaReducedMass->upload(reducedMassVec);
            ccmaConstraintMatrixValue->upload(constraintMatrixValueVec);
437
        }
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
        else {
            ccmaDistance = OpenCLArray::create<mm_float4>(context, numCCMA, "CcmaDistance");
            ccmaDelta1 = OpenCLArray::create<cl_float>(context, numCCMA, "CcmaDelta1");
            ccmaDelta2 = OpenCLArray::create<cl_float>(context, numCCMA, "CcmaDelta2");
            ccmaReducedMass = OpenCLArray::create<cl_float>(context, numCCMA, "CcmaReducedMass");
            ccmaConstraintMatrixValue = OpenCLArray::create<cl_float>(context, numCCMA*maxRowElements, "ConstraintMatrixValue");
            vector<mm_float4> distanceVec(ccmaDistance->getSize());
            vector<cl_float> reducedMassVec(ccmaReducedMass->getSize());
            vector<cl_float> constraintMatrixValueVec(ccmaConstraintMatrixValue->getSize());
            for (int i = 0; i < numCCMA; i++) {
                int index = constraintOrder[i];
                int c = ccmaConstraints[index];
                atomsVec[i].x = atom1[c];
                atomsVec[i].y = atom2[c];
                distanceVec[i].w = (float) distance[c];
                reducedMassVec[i] = (float) (0.5/(1.0/system.getParticleMass(atom1[c])+1.0/system.getParticleMass(atom2[c])));
                for (unsigned int j = 0; j < matrix[index].size(); j++) {
                    constraintMatrixColumnVec[i+j*numCCMA] = matrix[index][j].first;
                    constraintMatrixValueVec[i+j*numCCMA] = (float) matrix[index][j].second;
                }
                constraintMatrixColumnVec[i+matrix[index].size()*numCCMA] = numCCMA;
459
            }
460
461
462
463
464
465
466
467
468
469
            for (unsigned int i = 0; i < atomConstraints.size(); i++) {
                numAtomConstraintsVec[i] = atomConstraints[i].size();
                for (unsigned int j = 0; j < atomConstraints[i].size(); j++) {
                    bool forward = (atom1[ccmaConstraints[atomConstraints[i][j]]] == i);
                    atomConstraintsVec[i+j*numAtoms] = (forward ? inverseOrder[atomConstraints[i][j]]+1 : -inverseOrder[atomConstraints[i][j]]-1);
                }
            }
            ccmaDistance->upload(distanceVec);
            ccmaReducedMass->upload(reducedMassVec);
            ccmaConstraintMatrixValue->upload(constraintMatrixValueVec);
470
471
472
473
474
475
476
477
478
        }
        ccmaAtoms->upload(atomsVec);
        ccmaAtomConstraints->upload(atomConstraintsVec);
        ccmaNumAtomConstraints->upload(numAtomConstraintsVec);
        ccmaConstraintMatrixColumn->upload(constraintMatrixColumnVec);

        // Create the CCMA kernels.

        map<string, string> defines;
479
480
        defines["NUM_CONSTRAINTS"] = context.intToString(numCCMA);
        defines["NUM_ATOMS"] = context.intToString(numAtoms);
481
482
        cl::Program ccmaProgram = context.createProgram(OpenCLKernelSources::ccma, defines);
        ccmaDirectionsKernel = cl::Kernel(ccmaProgram, "computeConstraintDirections");
483
        ccmaPosForceKernel = cl::Kernel(ccmaProgram, "computeConstraintForce");
484
        ccmaMultiplyKernel = cl::Kernel(ccmaProgram, "multiplyByConstraintMatrix");
485
486
487
488
489
        ccmaPosUpdateKernel = cl::Kernel(ccmaProgram, "updateAtomPositions");
        defines["CONSTRAIN_VELOCITIES"] = "1";
        ccmaProgram = context.createProgram(OpenCLKernelSources::ccma, defines);
        ccmaVelForceKernel = cl::Kernel(ccmaProgram, "computeConstraintForce");
        ccmaVelUpdateKernel = cl::Kernel(ccmaProgram, "updateAtomPositions");
490
    }
491
492
493
494
    
    // Build the list of virtual sites.
    
    vector<mm_int4> vsite2AvgAtomVec;
495
    vector<mm_double2> vsite2AvgWeightVec;
496
    vector<mm_int4> vsite3AvgAtomVec;
497
    vector<mm_double4> vsite3AvgWeightVec;
498
    vector<mm_int4> vsiteOutOfPlaneAtomVec;
499
    vector<mm_double4> vsiteOutOfPlaneWeightVec;
500
501
    vector<mm_int4> vsiteLocalCoordsAtomVec;
    vector<cl_double> vsiteLocalCoordsParamVec;
502
503
504
505
506
507
508
    for (int i = 0; i < numAtoms; i++) {
        if (system.isVirtualSite(i)) {
            if (dynamic_cast<const TwoParticleAverageSite*>(&system.getVirtualSite(i)) != NULL) {
                // A two particle average.
                
                const TwoParticleAverageSite& site = dynamic_cast<const TwoParticleAverageSite&>(system.getVirtualSite(i));
                vsite2AvgAtomVec.push_back(mm_int4(i, site.getParticle(0), site.getParticle(1), 0));
509
                vsite2AvgWeightVec.push_back(mm_double2(site.getWeight(0), site.getWeight(1)));
510
511
512
513
514
515
            }
            else if (dynamic_cast<const ThreeParticleAverageSite*>(&system.getVirtualSite(i)) != NULL) {
                // A three particle average.
                
                const ThreeParticleAverageSite& site = dynamic_cast<const ThreeParticleAverageSite&>(system.getVirtualSite(i));
                vsite3AvgAtomVec.push_back(mm_int4(i, site.getParticle(0), site.getParticle(1), site.getParticle(2)));
516
                vsite3AvgWeightVec.push_back(mm_double4(site.getWeight(0), site.getWeight(1), site.getWeight(2), 0.0));
517
518
519
520
521
522
            }
            else if (dynamic_cast<const OutOfPlaneSite*>(&system.getVirtualSite(i)) != NULL) {
                // An out of plane site.
                
                const OutOfPlaneSite& site = dynamic_cast<const OutOfPlaneSite&>(system.getVirtualSite(i));
                vsiteOutOfPlaneAtomVec.push_back(mm_int4(i, site.getParticle(0), site.getParticle(1), site.getParticle(2)));
523
                vsiteOutOfPlaneWeightVec.push_back(mm_double4(site.getWeight12(), site.getWeight13(), site.getWeightCross(), 0.0));
524
            }
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
            else if (dynamic_cast<const LocalCoordinatesSite*>(&system.getVirtualSite(i)) != NULL) {
                // An out of plane site.
                
                const LocalCoordinatesSite& site = dynamic_cast<const LocalCoordinatesSite&>(system.getVirtualSite(i));
                vsiteLocalCoordsAtomVec.push_back(mm_int4(i, site.getParticle(0), site.getParticle(1), site.getParticle(2)));
                Vec3 origin = site.getOriginWeights();
                Vec3 x = site.getXWeights();
                Vec3 y = site.getYWeights();
                Vec3 pos = site.getLocalPosition();
                vsiteLocalCoordsParamVec.push_back(origin[0]);
                vsiteLocalCoordsParamVec.push_back(origin[1]);
                vsiteLocalCoordsParamVec.push_back(origin[2]);
                vsiteLocalCoordsParamVec.push_back(x[0]);
                vsiteLocalCoordsParamVec.push_back(x[1]);
                vsiteLocalCoordsParamVec.push_back(x[2]);
                vsiteLocalCoordsParamVec.push_back(y[0]);
                vsiteLocalCoordsParamVec.push_back(y[1]);
                vsiteLocalCoordsParamVec.push_back(y[2]);
                vsiteLocalCoordsParamVec.push_back(pos[0]);
                vsiteLocalCoordsParamVec.push_back(pos[1]);
                vsiteLocalCoordsParamVec.push_back(pos[2]);
            }
547
548
549
550
551
        }
    }
    int num2Avg = vsite2AvgAtomVec.size();
    int num3Avg = vsite3AvgAtomVec.size();
    int numOutOfPlane = vsiteOutOfPlaneAtomVec.size();
552
    int numLocalCoords = vsiteLocalCoordsAtomVec.size();
553
    numVsites = num2Avg+num3Avg+numOutOfPlane+numLocalCoords;
554
555
556
    vsite2AvgAtoms = OpenCLArray::create<mm_int4>(context, max(1, num2Avg), "vsite2AvgAtoms");
    vsite3AvgAtoms = OpenCLArray::create<mm_int4>(context, max(1, num3Avg), "vsite3AvgAtoms");
    vsiteOutOfPlaneAtoms = OpenCLArray::create<mm_int4>(context, max(1, numOutOfPlane), "vsiteOutOfPlaneAtoms");
557
    vsiteLocalCoordsAtoms = OpenCLArray::create<mm_int4>(context, max(1, numLocalCoords), "vsiteLocalCoordinatesAtoms");
558
    if (num2Avg > 0)
559
        vsite2AvgAtoms->upload(vsite2AvgAtomVec);
560
    if (num3Avg > 0)
561
        vsite3AvgAtoms->upload(vsite3AvgAtomVec);
562
    if (numOutOfPlane > 0)
563
        vsiteOutOfPlaneAtoms->upload(vsiteOutOfPlaneAtomVec);
564
565
    if (numLocalCoords > 0)
        vsiteLocalCoordsAtoms->upload(vsiteLocalCoordsAtomVec);
566
567
568
569
    if (context.getUseDoublePrecision()) {
        vsite2AvgWeights = OpenCLArray::create<mm_double2>(context, max(1, num2Avg), "vsite2AvgWeights");
        vsite3AvgWeights = OpenCLArray::create<mm_double4>(context, max(1, num3Avg), "vsite3AvgWeights");
        vsiteOutOfPlaneWeights = OpenCLArray::create<mm_double4>(context, max(1, numOutOfPlane), "vsiteOutOfPlaneWeights");
570
        vsiteLocalCoordsParams = OpenCLArray::create<cl_double>(context, max(1, 12*numLocalCoords), "vsiteLocalCoordinatesParams");
571
572
573
574
575
576
        if (num2Avg > 0)
            vsite2AvgWeights->upload(vsite2AvgWeightVec);
        if (num3Avg > 0)
            vsite3AvgWeights->upload(vsite3AvgWeightVec);
        if (numOutOfPlane > 0)
            vsiteOutOfPlaneWeights->upload(vsiteOutOfPlaneWeightVec);
577
578
        if (numLocalCoords > 0)
            vsiteLocalCoordsParams->upload(vsiteLocalCoordsParamVec);
579
580
581
582
583
    }
    else {
        vsite2AvgWeights = OpenCLArray::create<mm_float2>(context, max(1, num2Avg), "vsite2AvgWeights");
        vsite3AvgWeights = OpenCLArray::create<mm_float4>(context, max(1, num3Avg), "vsite3AvgWeights");
        vsiteOutOfPlaneWeights = OpenCLArray::create<mm_float4>(context, max(1, numOutOfPlane), "vsiteOutOfPlaneWeights");
584
        vsiteLocalCoordsParams = OpenCLArray::create<float>(context, max(1, 12*numLocalCoords), "vsiteLocalCoordinatesParams");
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
        if (num2Avg > 0) {
            vector<mm_float2> floatWeights(num2Avg);
            for (int i = 0; i < num2Avg; i++)
                floatWeights[i] = mm_float2((float) vsite2AvgWeightVec[i].x, (float) vsite2AvgWeightVec[i].y);
            vsite2AvgWeights->upload(floatWeights);
        }
        if (num3Avg > 0) {
            vector<mm_float4> floatWeights(num3Avg);
            for (int i = 0; i < num3Avg; i++)
                floatWeights[i] = mm_float4((float) vsite3AvgWeightVec[i].x, (float) vsite3AvgWeightVec[i].y, (float) vsite3AvgWeightVec[i].z, 0.0f);
            vsite3AvgWeights->upload(floatWeights);
        }
        if (numOutOfPlane > 0) {
            vector<mm_float4> floatWeights(numOutOfPlane);
            for (int i = 0; i < numOutOfPlane; i++)
                floatWeights[i] = mm_float4((float) vsiteOutOfPlaneWeightVec[i].x, (float) vsiteOutOfPlaneWeightVec[i].y, (float) vsiteOutOfPlaneWeightVec[i].z, 0.0f);
            vsiteOutOfPlaneWeights->upload(floatWeights);
        }
603
604
605
606
607
608
        if (numLocalCoords > 0) {
            vector<cl_float> floatParams(vsiteLocalCoordsParamVec.size());
            for (int i = 0; i < (int) vsiteLocalCoordsParamVec.size(); i++)
                floatParams[i] = (cl_float) vsiteLocalCoordsParamVec[i];
            vsiteLocalCoordsParams->upload(floatParams);
        }
609
610
    }
    
611
612
613
614
615
616
617
618
619
620
621
622
623
624
    // If multiple virtual sites depend on the same particle, make sure the force distribution
    // can be done safely.
    
    vector<int> atomCounts(numAtoms, 0);
    for (int i = 0; i < numAtoms; i++)
        if (system.isVirtualSite(i))
            for (int j = 0; j < system.getVirtualSite(i).getNumParticles(); j++)
                atomCounts[system.getVirtualSite(i).getParticle(j)]++;
    for (int i = 0; i < numAtoms; i++)
        if (atomCounts[i] > 1)
            hasOverlappingVsites = true;
    if (hasOverlappingVsites && context.getUseDoublePrecision() && !context.getSupports64BitGlobalAtomics())
        throw OpenMMException("This device does not support 64 bit atomics.  Cannot use double precision when multiple virtual sites depend on the same atom.");
    
625
626
627
    // Create the kernels for virtual sites.

    map<string, string> defines;
628
629
630
    defines["NUM_2_AVERAGE"] = context.intToString(num2Avg);
    defines["NUM_3_AVERAGE"] = context.intToString(num3Avg);
    defines["NUM_OUT_OF_PLANE"] = context.intToString(numOutOfPlane);
631
    defines["NUM_LOCAL_COORDS"] = context.intToString(numLocalCoords);
632
633
634
635
    defines["NUM_ATOMS"] = context.intToString(numAtoms);
    defines["PADDED_NUM_ATOMS"] = context.intToString(context.getPaddedNumAtoms());
    if (hasOverlappingVsites)
        defines["HAS_OVERLAPPING_VSITES"] = "1";
Peter Eastman's avatar
Peter Eastman committed
636
637
    cl::Program vsiteProgram = context.createProgram(OpenCLKernelSources::virtualSites, defines);
    vsitePositionKernel = cl::Kernel(vsiteProgram, "computeVirtualSites");
638
639
640
641
642
643
644
645
646
647
    int index = 0;
    vsitePositionKernel.setArg<cl::Buffer>(index++, context.getPosq().getDeviceBuffer());
    if (context.getUseMixedPrecision())
        vsitePositionKernel.setArg<cl::Buffer>(index++, context.getPosqCorrection().getDeviceBuffer());
    vsitePositionKernel.setArg<cl::Buffer>(index++, vsite2AvgAtoms->getDeviceBuffer());
    vsitePositionKernel.setArg<cl::Buffer>(index++, vsite2AvgWeights->getDeviceBuffer());
    vsitePositionKernel.setArg<cl::Buffer>(index++, vsite3AvgAtoms->getDeviceBuffer());
    vsitePositionKernel.setArg<cl::Buffer>(index++, vsite3AvgWeights->getDeviceBuffer());
    vsitePositionKernel.setArg<cl::Buffer>(index++, vsiteOutOfPlaneAtoms->getDeviceBuffer());
    vsitePositionKernel.setArg<cl::Buffer>(index++, vsiteOutOfPlaneWeights->getDeviceBuffer());
648
649
    vsitePositionKernel.setArg<cl::Buffer>(index++, vsiteLocalCoordsAtoms->getDeviceBuffer());
    vsitePositionKernel.setArg<cl::Buffer>(index++, vsiteLocalCoordsParams->getDeviceBuffer());
Peter Eastman's avatar
Peter Eastman committed
650
    vsiteForceKernel = cl::Kernel(vsiteProgram, "distributeForces");
651
652
653
    index = 0;
    vsiteForceKernel.setArg<cl::Buffer>(index++, context.getPosq().getDeviceBuffer());
    index++; // Skip argument 1: the force array hasn't been created yet.
654
655
    if (context.getSupports64BitGlobalAtomics())
        index++; // Skip argument 2: the force array hasn't been created yet.
656
657
658
659
660
661
662
663
    if (context.getUseMixedPrecision())
        vsiteForceKernel.setArg<cl::Buffer>(index++, context.getPosqCorrection().getDeviceBuffer());
    vsiteForceKernel.setArg<cl::Buffer>(index++, vsite2AvgAtoms->getDeviceBuffer());
    vsiteForceKernel.setArg<cl::Buffer>(index++, vsite2AvgWeights->getDeviceBuffer());
    vsiteForceKernel.setArg<cl::Buffer>(index++, vsite3AvgAtoms->getDeviceBuffer());
    vsiteForceKernel.setArg<cl::Buffer>(index++, vsite3AvgWeights->getDeviceBuffer());
    vsiteForceKernel.setArg<cl::Buffer>(index++, vsiteOutOfPlaneAtoms->getDeviceBuffer());
    vsiteForceKernel.setArg<cl::Buffer>(index++, vsiteOutOfPlaneWeights->getDeviceBuffer());
664
665
    vsiteForceKernel.setArg<cl::Buffer>(index++, vsiteLocalCoordsAtoms->getDeviceBuffer());
    vsiteForceKernel.setArg<cl::Buffer>(index++, vsiteLocalCoordsParams->getDeviceBuffer());
666
667
    if (hasOverlappingVsites && context.getSupports64BitGlobalAtomics())
        vsiteAddForcesKernel = cl::Kernel(vsiteProgram, "addDistributedForces");
668
669
670
671
672
673
674
675
676
677
678
679
680
}

OpenCLIntegrationUtilities::~OpenCLIntegrationUtilities() {
    if (posDelta != NULL)
        delete posDelta;
    if (settleAtoms != NULL)
        delete settleAtoms;
    if (settleParams != NULL)
        delete settleParams;
    if (shakeAtoms != NULL)
        delete shakeAtoms;
    if (shakeParams != NULL)
        delete shakeParams;
681
682
683
684
    if (random != NULL)
        delete random;
    if (randomSeed != NULL)
        delete randomSeed;
685
686
    if (stepSize != NULL)
        delete stepSize;
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
    if (ccmaAtoms != NULL)
        delete ccmaAtoms;
    if (ccmaDistance != NULL)
        delete ccmaDistance;
    if (ccmaReducedMass != NULL)
        delete ccmaReducedMass;
    if (ccmaAtomConstraints != NULL)
        delete ccmaAtomConstraints;
    if (ccmaNumAtomConstraints != NULL)
        delete ccmaNumAtomConstraints;
    if (ccmaConstraintMatrixColumn != NULL)
        delete ccmaConstraintMatrixColumn;
    if (ccmaConstraintMatrixValue != NULL)
        delete ccmaConstraintMatrixValue;
    if (ccmaDelta1 != NULL)
        delete ccmaDelta1;
    if (ccmaDelta2 != NULL)
        delete ccmaDelta2;
    if (ccmaConverged != NULL)
        delete ccmaConverged;
Peter Eastman's avatar
Peter Eastman committed
707
708
    if (ccmaConvergedHostBuffer != NULL)
        delete ccmaConvergedHostBuffer;
709
710
711
712
713
714
715
716
717
718
719
720
    if (vsite2AvgAtoms != NULL)
        delete vsite2AvgAtoms;
    if (vsite2AvgWeights != NULL)
        delete vsite2AvgWeights;
    if (vsite3AvgAtoms != NULL)
        delete vsite3AvgAtoms;
    if (vsite3AvgWeights != NULL)
        delete vsite3AvgWeights;
    if (vsiteOutOfPlaneAtoms != NULL)
        delete vsiteOutOfPlaneAtoms;
    if (vsiteOutOfPlaneWeights != NULL)
        delete vsiteOutOfPlaneWeights;
721
722
723
724
    if (vsiteLocalCoordsAtoms != NULL)
        delete vsiteLocalCoordsAtoms;
    if (vsiteLocalCoordsParams != NULL)
        delete vsiteLocalCoordsParams;
725
726
}

727
void OpenCLIntegrationUtilities::applyConstraints(double tol) {
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
    applyConstraints(false, tol);
}

void OpenCLIntegrationUtilities::applyVelocityConstraints(double tol) {
    applyConstraints(true, tol);
}

void OpenCLIntegrationUtilities::applyConstraints(bool constrainVelocities, double tol) {
    bool hasInitialized;
    cl::Kernel settleKernel, shakeKernel, ccmaForceKernel, ccmaUpdateKernel;
    if (constrainVelocities) {
        hasInitialized = hasInitializedVelConstraintKernels;
        settleKernel = settleVelKernel;
        shakeKernel = shakeVelKernel;
        ccmaForceKernel = ccmaVelForceKernel;
        ccmaUpdateKernel = ccmaVelUpdateKernel;
        hasInitializedVelConstraintKernels = true;
    }
    else {
        hasInitialized = hasInitializedPosConstraintKernels;
        settleKernel = settlePosKernel;
        shakeKernel = shakePosKernel;
        ccmaForceKernel = ccmaPosForceKernel;
        ccmaUpdateKernel = ccmaPosUpdateKernel;
        hasInitializedPosConstraintKernels = true;
    }
754
    if (settleAtoms != NULL) {
755
        if (!hasInitialized) {
756
757
            settleKernel.setArg<cl_int>(0, settleAtoms->getSize());
            settleKernel.setArg<cl::Buffer>(2, context.getPosq().getDeviceBuffer());
758
759
760
761
762
763
764
765
            if (context.getUseMixedPrecision())
                settleKernel.setArg<cl::Buffer>(3, context.getPosqCorrection().getDeviceBuffer());
            else
                settleKernel.setArg<void*>(3, NULL);
            settleKernel.setArg<cl::Buffer>(4, posDelta->getDeviceBuffer());
            settleKernel.setArg<cl::Buffer>(5, context.getVelm().getDeviceBuffer());
            settleKernel.setArg<cl::Buffer>(6, settleAtoms->getDeviceBuffer());
            settleKernel.setArg<cl::Buffer>(7, settleParams->getDeviceBuffer());
766
        }
767
768
769
770
        if (context.getUseDoublePrecision() || context.getUseMixedPrecision())
            settleKernel.setArg<cl_double>(1, (cl_double) tol);
        else
            settleKernel.setArg<cl_float>(1, (cl_float) tol);
771
772
        context.executeKernel(settleKernel, settleAtoms->getSize());
    }
773
    if (shakeAtoms != NULL) {
774
        if (!hasInitialized) {
775
776
            shakeKernel.setArg<cl_int>(0, shakeAtoms->getSize());
            shakeKernel.setArg<cl::Buffer>(2, context.getPosq().getDeviceBuffer());
777
778
779
780
781
782
783
            if (context.getUseMixedPrecision())
                shakeKernel.setArg<cl::Buffer>(3, context.getPosqCorrection().getDeviceBuffer());
            else
                shakeKernel.setArg<void*>(3, NULL);
            shakeKernel.setArg<cl::Buffer>(4, constrainVelocities ? context.getVelm().getDeviceBuffer() : posDelta->getDeviceBuffer());
            shakeKernel.setArg<cl::Buffer>(5, shakeAtoms->getDeviceBuffer());
            shakeKernel.setArg<cl::Buffer>(6, shakeParams->getDeviceBuffer());
784
        }
785
786
787
788
        if (context.getUseDoublePrecision() || context.getUseMixedPrecision())
            shakeKernel.setArg<cl_double>(1, (cl_double) tol);
        else
            shakeKernel.setArg<cl_float>(1, (cl_float) tol);
789
790
        context.executeKernel(shakeKernel, shakeAtoms->getSize());
    }
791
    if (ccmaAtoms != NULL) {
792
        if (!hasInitialized) {
793
794
795
            ccmaDirectionsKernel.setArg<cl::Buffer>(0, ccmaAtoms->getDeviceBuffer());
            ccmaDirectionsKernel.setArg<cl::Buffer>(1, ccmaDistance->getDeviceBuffer());
            ccmaDirectionsKernel.setArg<cl::Buffer>(2, context.getPosq().getDeviceBuffer());
796
797
798
799
            if (context.getUseMixedPrecision())
                ccmaDirectionsKernel.setArg<cl::Buffer>(3, context.getPosqCorrection().getDeviceBuffer());
            else
                ccmaDirectionsKernel.setArg<void*>(3, NULL);
800
            ccmaDirectionsKernel.setArg<cl::Buffer>(4, ccmaConverged->getDeviceBuffer());
801
802
            ccmaForceKernel.setArg<cl::Buffer>(0, ccmaAtoms->getDeviceBuffer());
            ccmaForceKernel.setArg<cl::Buffer>(1, ccmaDistance->getDeviceBuffer());
803
            ccmaForceKernel.setArg<cl::Buffer>(2, constrainVelocities ? context.getVelm().getDeviceBuffer() : posDelta->getDeviceBuffer());
804
805
806
            ccmaForceKernel.setArg<cl::Buffer>(3, ccmaReducedMass->getDeviceBuffer());
            ccmaForceKernel.setArg<cl::Buffer>(4, ccmaDelta1->getDeviceBuffer());
            ccmaForceKernel.setArg<cl::Buffer>(5, ccmaConverged->getDeviceBuffer());
Peter Eastman's avatar
Peter Eastman committed
807
            ccmaForceKernel.setArg<cl::Buffer>(6, ccmaConvergedHostBuffer->getDeviceBuffer());
808
809
810
811
812
813
814
815
            ccmaMultiplyKernel.setArg<cl::Buffer>(0, ccmaDelta1->getDeviceBuffer());
            ccmaMultiplyKernel.setArg<cl::Buffer>(1, ccmaDelta2->getDeviceBuffer());
            ccmaMultiplyKernel.setArg<cl::Buffer>(2, ccmaConstraintMatrixColumn->getDeviceBuffer());
            ccmaMultiplyKernel.setArg<cl::Buffer>(3, ccmaConstraintMatrixValue->getDeviceBuffer());
            ccmaMultiplyKernel.setArg<cl::Buffer>(4, ccmaConverged->getDeviceBuffer());
            ccmaUpdateKernel.setArg<cl::Buffer>(0, ccmaNumAtomConstraints->getDeviceBuffer());
            ccmaUpdateKernel.setArg<cl::Buffer>(1, ccmaAtomConstraints->getDeviceBuffer());
            ccmaUpdateKernel.setArg<cl::Buffer>(2, ccmaDistance->getDeviceBuffer());
816
            ccmaUpdateKernel.setArg<cl::Buffer>(3, constrainVelocities ? context.getVelm().getDeviceBuffer() : posDelta->getDeviceBuffer());
817
818
819
820
821
            ccmaUpdateKernel.setArg<cl::Buffer>(4, context.getVelm().getDeviceBuffer());
            ccmaUpdateKernel.setArg<cl::Buffer>(5, ccmaDelta1->getDeviceBuffer());
            ccmaUpdateKernel.setArg<cl::Buffer>(6, ccmaDelta2->getDeviceBuffer());
            ccmaUpdateKernel.setArg<cl::Buffer>(7, ccmaConverged->getDeviceBuffer());
        }
822
        if (context.getUseDoublePrecision() || context.getUseMixedPrecision())
Peter Eastman's avatar
Peter Eastman committed
823
            ccmaForceKernel.setArg<cl_double>(7, (cl_double) tol);
824
        else
Peter Eastman's avatar
Peter Eastman committed
825
            ccmaForceKernel.setArg<cl_float>(7, (cl_float) tol);
826
        context.executeKernel(ccmaDirectionsKernel, ccmaAtoms->getSize());
Peter Eastman's avatar
Peter Eastman committed
827
        const int checkInterval = 4;
828
        int* converged = (int*) context.getPinnedBuffer();
829
        int* ccmaConvergedHostMemory = (int*) context.getQueue().enqueueMapBuffer(ccmaConvergedHostBuffer->getDeviceBuffer(), CL_TRUE, CL_MAP_WRITE, 0, sizeof(cl_int));
Peter Eastman's avatar
Peter Eastman committed
830
        ccmaConvergedHostMemory[0] = 0;
831
        context.getQueue().enqueueUnmapMemObject(ccmaConvergedHostBuffer->getDeviceBuffer(), ccmaConvergedHostMemory);
832
        for (int i = 0; i < 150; i++) {
Peter Eastman's avatar
Peter Eastman committed
833
            ccmaForceKernel.setArg<cl_int>(8, i);
834
            context.executeKernel(ccmaForceKernel, ccmaAtoms->getSize());
835
836
837
            cl::Event event;
            if ((i+1)%checkInterval == 0 && !ccmaUseDirectBuffer)
                context.getQueue().enqueueReadBuffer(ccmaConverged->getDeviceBuffer(), CL_FALSE, 0, 2*sizeof(cl_int), converged, NULL, &event);
Peter Eastman's avatar
Peter Eastman committed
838
            ccmaMultiplyKernel.setArg<cl_int>(5, i);
839
840
841
            context.executeKernel(ccmaMultiplyKernel, ccmaAtoms->getSize());
            ccmaUpdateKernel.setArg<cl_int>(8, i);
            context.executeKernel(ccmaUpdateKernel, context.getNumAtoms());
Peter Eastman's avatar
Peter Eastman committed
842
            if ((i+1)%checkInterval == 0) {
Peter Eastman's avatar
Peter Eastman committed
843
                if (ccmaUseDirectBuffer) {
844
845
846
847
848
849
                    ccmaConvergedHostMemory = (int*) context.getQueue().enqueueMapBuffer(ccmaConvergedHostBuffer->getDeviceBuffer(), CL_FALSE, CL_MAP_READ, 0, sizeof(cl_int), NULL, &event);
                    context.getQueue().flush();
                    while (event.getInfo<CL_EVENT_COMMAND_EXECUTION_STATUS>() != CL_COMPLETE)
                        ;
                    converged[i%2] = ccmaConvergedHostMemory[0];
                    context.getQueue().enqueueUnmapMemObject(ccmaConvergedHostBuffer->getDeviceBuffer(), ccmaConvergedHostMemory);
Peter Eastman's avatar
Peter Eastman committed
850
                }
851
852
853
854
                else
                    event.wait();
                if (converged[i%2])
                    break;
Peter Eastman's avatar
Peter Eastman committed
855
            }
856
857
        }
    }
858
859
}

860
861
862
863
864
865
void OpenCLIntegrationUtilities::computeVirtualSites() {
    if (numVsites > 0)
        context.executeKernel(vsitePositionKernel, numVsites);
}

void OpenCLIntegrationUtilities::distributeForcesFromVirtualSites() {
Peter Eastman's avatar
Peter Eastman committed
866
    if (numVsites > 0) {
867
868
        // Set arguments that didn't exist yet in the constructor.
        
869
        vsiteForceKernel.setArg<cl::Buffer>(1, context.getForce().getDeviceBuffer());
870
871
872
873
874
875
876
877
        if (context.getSupports64BitGlobalAtomics()) {
            vsiteForceKernel.setArg<cl::Buffer>(2, context.getLongForceBuffer().getDeviceBuffer());
            if (hasOverlappingVsites) {
                // We'll be using 64 bit atomics for the force redistribution, so clear the buffer.
                
                context.clearBuffer(context.getLongForceBuffer());
            }
        }
878
        context.executeKernel(vsiteForceKernel, numVsites);
879
880
881
882
883
884
885
        if (context.getSupports64BitGlobalAtomics() && hasOverlappingVsites) {
            // Add the redistributed forces from the virtual sites to the main force array.
            
            vsiteAddForcesKernel.setArg<cl::Buffer>(0, context.getLongForceBuffer().getDeviceBuffer());
            vsiteAddForcesKernel.setArg<cl::Buffer>(1, context.getForce().getDeviceBuffer());
            context.executeKernel(vsiteAddForcesKernel, context.getNumAtoms());
        }
Peter Eastman's avatar
Peter Eastman committed
886
    }
887
888
}

889
890
891
892
893
894
895
896
897
898
void OpenCLIntegrationUtilities::initRandomNumberGenerator(unsigned int randomNumberSeed) {
    if (random != NULL) {
        if (randomNumberSeed != lastSeed)
           throw OpenMMException("OpenCLIntegrationUtilities::initRandomNumberGenerator(): Requested two different values for the random number seed");
        return;
    }

    // Create the random number arrays.

    lastSeed = randomNumberSeed;
899
    random = OpenCLArray::create<mm_float4>(context, 4*context.getPaddedNumAtoms(), "random");
900
    randomSeed = OpenCLArray::create<mm_int4>(context, context.getNumThreadBlocks()*OpenCLContext::ThreadBlockSize, "randomSeed");
901
902
    randomPos = random->getSize();

903
    // Use a quick and dirty RNG to pick seeds for the real random number generator.
904
905

    vector<mm_int4> seed(randomSeed->getSize());
906
    unsigned int r = randomNumberSeed;
907
    // A seed of 0 means use a unique one
908
    if (r == 0) r = (unsigned int) osrngseed();
909
910
911
912
913
914
    for (int i = 0; i < randomSeed->getSize(); i++) {
        seed[i].x = r = (1664525*r + 1013904223) & 0xFFFFFFFF;
        seed[i].y = r = (1664525*r + 1013904223) & 0xFFFFFFFF;
        seed[i].z = r = (1664525*r + 1013904223) & 0xFFFFFFFF;
        seed[i].w = r = (1664525*r + 1013904223) & 0xFFFFFFFF;
    }
915
916
917
918
    randomSeed->upload(seed);

    // Create the kernel.

919
    cl::Program randomProgram = context.createProgram(OpenCLKernelSources::random);
920
921
922
923
924
925
926
927
928
    randomKernel = cl::Kernel(randomProgram, "generateRandomNumbers");
}

int OpenCLIntegrationUtilities::prepareRandomNumbers(int numValues) {
    if (randomPos+numValues <= random->getSize()) {
        int oldPos = randomPos;
        randomPos += numValues;
        return oldPos;
    }
Peter Eastman's avatar
Peter Eastman committed
929
930
    if (numValues > random->getSize()) {
        delete random;
931
        random = OpenCLArray::create<mm_float4>(context, numValues, "random");
Peter Eastman's avatar
Peter Eastman committed
932
    }
933
934
935
936
937
938
939
    randomKernel.setArg<cl_int>(0, random->getSize());
    randomKernel.setArg<cl::Buffer>(1, random->getDeviceBuffer());
    randomKernel.setArg<cl::Buffer>(2, randomSeed->getDeviceBuffer());
    context.executeKernel(randomKernel, random->getSize());
    randomPos = numValues;
    return 0;
}
Peter Eastman's avatar
Peter Eastman committed
940
941

void OpenCLIntegrationUtilities::createCheckpoint(ostream& stream) {
Yutong Zhao's avatar
Yutong Zhao committed
942
943
    if(random == NULL) 
        return;
Peter Eastman's avatar
Peter Eastman committed
944
945
946
947
948
949
950
951
952
953
    stream.write((char*) &randomPos, sizeof(int));
    vector<mm_float4> randomVec;
    random->download(randomVec);
    stream.write((char*) &randomVec[0], sizeof(mm_float4)*random->getSize());
    vector<mm_int4> randomSeedVec;
    randomSeed->download(randomSeedVec);
    stream.write((char*) &randomSeedVec[0], sizeof(mm_int4)*randomSeed->getSize());
}

void OpenCLIntegrationUtilities::loadCheckpoint(istream& stream) {
Yutong Zhao's avatar
Yutong Zhao committed
954
    if(random == NULL) 
955
        return; 
Peter Eastman's avatar
Peter Eastman committed
956
957
958
959
960
961
962
963
    stream.read((char*) &randomPos, sizeof(int));
    vector<mm_float4> randomVec(random->getSize());
    stream.read((char*) &randomVec[0], sizeof(mm_float4)*random->getSize());
    random->upload(randomVec);
    vector<mm_int4> randomSeedVec(randomSeed->getSize());
    stream.read((char*) &randomSeedVec[0], sizeof(mm_int4)*randomSeed->getSize());
    randomSeed->upload(randomSeedVec);
}
964
965
966

double OpenCLIntegrationUtilities::computeKineticEnergy(double timeShift) {
    int numParticles = context.getNumAtoms();
967
968
969
970
971
972
973
974
975
    if (timeShift != 0) {
        // Copy the velocities into the posDelta array while we temporarily modify them.

        context.getVelm().copyTo(*posDelta);

        // Apply the time shift.

        timeShiftKernel.setArg<cl::Buffer>(0, context.getVelm().getDeviceBuffer());
        timeShiftKernel.setArg<cl::Buffer>(1, context.getForce().getDeviceBuffer());
976
        if (context.getUseDoublePrecision())
977
978
979
980
981
            timeShiftKernel.setArg<cl_double>(2, timeShift);
        else
            timeShiftKernel.setArg<cl_float>(2, (cl_float) timeShift);
        context.executeKernel(timeShiftKernel, numParticles);
        applyConstraints(true, 1e-4);
982
    }
983
984
985
986
987
    
    // Compute the kinetic energy.
    
    double energy = 0.0;
    if (context.getUseDoublePrecision() || context.getUseMixedPrecision()) {
988
989
        vector<mm_double4> velm;
        context.getVelm().download(velm);
990
991
        for (int i = 0; i < numParticles; i++) {
            mm_double4 v = velm[i];
992
            if (v.w != 0)
993
994
995
996
997
998
999
                energy += (v.x*v.x+v.y*v.y+v.z*v.z)/v.w;
        }
    }
    else {
        vector<mm_float4> velm;
        context.getVelm().download(velm);
        for (int i = 0; i < numParticles; i++) {
1000
1001
            mm_float4 v = velm[i];
            if (v.w != 0)
1002
1003
1004
                energy += (v.x*v.x+v.y*v.y+v.z*v.z)/v.w;
        }
    }
1005
1006
1007
1008
1009
    
    // Restore the velocities.
    
    if (timeShift != 0)
        posDelta->copyTo(context.getVelm());
1010
1011
    return 0.5*energy;
}