OpenCLIntegrationUtilities.cpp 51.3 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
    lastStepSize = mm_double2(0.0, 0.0);
110
111
112
113
114
    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");
115
        stepSize->upload(&lastStepSize);
116
117
118
119
120
121
    }
    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");
122
123
        mm_float2 lastStepSizeFloat = mm_float2(0.0f, 0.0f);
        stepSize->upload(&lastStepSizeFloat);
124
    }
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
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
void OpenCLIntegrationUtilities::setNextStepSize(double size) {
    if (size != lastStepSize.x || size != lastStepSize.y) {
        lastStepSize = mm_double2(size, size);
        if (context.getUseDoublePrecision() || context.getUseMixedPrecision())
            stepSize->upload(&lastStepSize);
        else {
            mm_float2 lastStepSizeFloat = mm_float2((float) size, (float) size);
            stepSize->upload(&lastStepSizeFloat);
        }
    }
}

double OpenCLIntegrationUtilities::getLastStepSize() {
    if (context.getUseDoublePrecision() || context.getUseMixedPrecision())
        stepSize->download(&lastStepSize);
    else {
        mm_float2 lastStepSizeFloat;
        stepSize->download(&lastStepSizeFloat);
        lastStepSize = mm_double2(lastStepSizeFloat.x, lastStepSizeFloat.y);
    }
    return lastStepSize.y;
}

750
void OpenCLIntegrationUtilities::applyConstraints(double tol) {
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
    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;
    }
777
    if (settleAtoms != NULL) {
778
        if (!hasInitialized) {
779
780
            settleKernel.setArg<cl_int>(0, settleAtoms->getSize());
            settleKernel.setArg<cl::Buffer>(2, context.getPosq().getDeviceBuffer());
781
782
783
784
785
786
787
788
            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());
789
        }
790
791
792
793
        if (context.getUseDoublePrecision() || context.getUseMixedPrecision())
            settleKernel.setArg<cl_double>(1, (cl_double) tol);
        else
            settleKernel.setArg<cl_float>(1, (cl_float) tol);
794
795
        context.executeKernel(settleKernel, settleAtoms->getSize());
    }
796
    if (shakeAtoms != NULL) {
797
        if (!hasInitialized) {
798
799
            shakeKernel.setArg<cl_int>(0, shakeAtoms->getSize());
            shakeKernel.setArg<cl::Buffer>(2, context.getPosq().getDeviceBuffer());
800
801
802
803
804
805
806
            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());
807
        }
808
809
810
811
        if (context.getUseDoublePrecision() || context.getUseMixedPrecision())
            shakeKernel.setArg<cl_double>(1, (cl_double) tol);
        else
            shakeKernel.setArg<cl_float>(1, (cl_float) tol);
812
813
        context.executeKernel(shakeKernel, shakeAtoms->getSize());
    }
814
    if (ccmaAtoms != NULL) {
815
        if (!hasInitialized) {
816
817
818
            ccmaDirectionsKernel.setArg<cl::Buffer>(0, ccmaAtoms->getDeviceBuffer());
            ccmaDirectionsKernel.setArg<cl::Buffer>(1, ccmaDistance->getDeviceBuffer());
            ccmaDirectionsKernel.setArg<cl::Buffer>(2, context.getPosq().getDeviceBuffer());
819
820
821
822
            if (context.getUseMixedPrecision())
                ccmaDirectionsKernel.setArg<cl::Buffer>(3, context.getPosqCorrection().getDeviceBuffer());
            else
                ccmaDirectionsKernel.setArg<void*>(3, NULL);
823
            ccmaDirectionsKernel.setArg<cl::Buffer>(4, ccmaConverged->getDeviceBuffer());
824
825
            ccmaForceKernel.setArg<cl::Buffer>(0, ccmaAtoms->getDeviceBuffer());
            ccmaForceKernel.setArg<cl::Buffer>(1, ccmaDistance->getDeviceBuffer());
826
            ccmaForceKernel.setArg<cl::Buffer>(2, constrainVelocities ? context.getVelm().getDeviceBuffer() : posDelta->getDeviceBuffer());
827
828
829
            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
830
            ccmaForceKernel.setArg<cl::Buffer>(6, ccmaConvergedHostBuffer->getDeviceBuffer());
831
832
833
834
835
836
837
838
            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());
839
            ccmaUpdateKernel.setArg<cl::Buffer>(3, constrainVelocities ? context.getVelm().getDeviceBuffer() : posDelta->getDeviceBuffer());
840
841
842
843
844
            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());
        }
845
        if (context.getUseDoublePrecision() || context.getUseMixedPrecision())
Peter Eastman's avatar
Peter Eastman committed
846
            ccmaForceKernel.setArg<cl_double>(7, (cl_double) tol);
847
        else
Peter Eastman's avatar
Peter Eastman committed
848
            ccmaForceKernel.setArg<cl_float>(7, (cl_float) tol);
849
        context.executeKernel(ccmaDirectionsKernel, ccmaAtoms->getSize());
Peter Eastman's avatar
Peter Eastman committed
850
        const int checkInterval = 4;
851
        int* converged = (int*) context.getPinnedBuffer();
852
        int* ccmaConvergedHostMemory = (int*) context.getQueue().enqueueMapBuffer(ccmaConvergedHostBuffer->getDeviceBuffer(), CL_TRUE, CL_MAP_WRITE, 0, sizeof(cl_int));
Peter Eastman's avatar
Peter Eastman committed
853
        ccmaConvergedHostMemory[0] = 0;
854
        context.getQueue().enqueueUnmapMemObject(ccmaConvergedHostBuffer->getDeviceBuffer(), ccmaConvergedHostMemory);
855
        for (int i = 0; i < 150; i++) {
Peter Eastman's avatar
Peter Eastman committed
856
            ccmaForceKernel.setArg<cl_int>(8, i);
857
            context.executeKernel(ccmaForceKernel, ccmaAtoms->getSize());
858
859
860
            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
861
            ccmaMultiplyKernel.setArg<cl_int>(5, i);
862
863
864
            context.executeKernel(ccmaMultiplyKernel, ccmaAtoms->getSize());
            ccmaUpdateKernel.setArg<cl_int>(8, i);
            context.executeKernel(ccmaUpdateKernel, context.getNumAtoms());
Peter Eastman's avatar
Peter Eastman committed
865
            if ((i+1)%checkInterval == 0) {
Peter Eastman's avatar
Peter Eastman committed
866
                if (ccmaUseDirectBuffer) {
867
868
869
870
871
872
                    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
873
                }
874
875
876
877
                else
                    event.wait();
                if (converged[i%2])
                    break;
Peter Eastman's avatar
Peter Eastman committed
878
            }
879
880
        }
    }
881
882
}

883
884
885
886
887
888
void OpenCLIntegrationUtilities::computeVirtualSites() {
    if (numVsites > 0)
        context.executeKernel(vsitePositionKernel, numVsites);
}

void OpenCLIntegrationUtilities::distributeForcesFromVirtualSites() {
Peter Eastman's avatar
Peter Eastman committed
889
    if (numVsites > 0) {
890
891
        // Set arguments that didn't exist yet in the constructor.
        
892
        vsiteForceKernel.setArg<cl::Buffer>(1, context.getForce().getDeviceBuffer());
893
894
895
896
897
898
899
900
        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());
            }
        }
901
        context.executeKernel(vsiteForceKernel, numVsites);
902
903
904
905
906
907
908
        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
909
    }
910
911
}

912
913
914
915
916
917
918
919
920
921
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;
922
    random = OpenCLArray::create<mm_float4>(context, 4*context.getPaddedNumAtoms(), "random");
923
    randomSeed = OpenCLArray::create<mm_int4>(context, context.getNumThreadBlocks()*OpenCLContext::ThreadBlockSize, "randomSeed");
924
925
    randomPos = random->getSize();

926
    // Use a quick and dirty RNG to pick seeds for the real random number generator.
927
928

    vector<mm_int4> seed(randomSeed->getSize());
929
    unsigned int r = randomNumberSeed;
930
    // A seed of 0 means use a unique one
931
    if (r == 0) r = (unsigned int) osrngseed();
932
933
934
935
936
937
    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;
    }
938
939
940
941
    randomSeed->upload(seed);

    // Create the kernel.

942
    cl::Program randomProgram = context.createProgram(OpenCLKernelSources::random);
943
944
945
946
947
948
949
950
951
    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
952
953
    if (numValues > random->getSize()) {
        delete random;
954
        random = OpenCLArray::create<mm_float4>(context, numValues, "random");
Peter Eastman's avatar
Peter Eastman committed
955
    }
956
957
958
959
960
961
962
    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
963
964

void OpenCLIntegrationUtilities::createCheckpoint(ostream& stream) {
Yutong Zhao's avatar
Yutong Zhao committed
965
966
    if(random == NULL) 
        return;
Peter Eastman's avatar
Peter Eastman committed
967
968
969
970
971
972
973
974
975
976
    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
977
    if(random == NULL) 
978
        return; 
Peter Eastman's avatar
Peter Eastman committed
979
980
981
982
983
984
985
986
    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);
}
987
988
989

double OpenCLIntegrationUtilities::computeKineticEnergy(double timeShift) {
    int numParticles = context.getNumAtoms();
990
991
992
993
994
995
996
997
998
    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());
999
        if (context.getUseDoublePrecision())
1000
1001
1002
1003
1004
            timeShiftKernel.setArg<cl_double>(2, timeShift);
        else
            timeShiftKernel.setArg<cl_float>(2, (cl_float) timeShift);
        context.executeKernel(timeShiftKernel, numParticles);
        applyConstraints(true, 1e-4);
1005
    }
1006
1007
1008
1009
1010
    
    // Compute the kinetic energy.
    
    double energy = 0.0;
    if (context.getUseDoublePrecision() || context.getUseMixedPrecision()) {
1011
1012
        vector<mm_double4> velm;
        context.getVelm().download(velm);
1013
1014
        for (int i = 0; i < numParticles; i++) {
            mm_double4 v = velm[i];
1015
            if (v.w != 0)
1016
1017
1018
1019
1020
1021
1022
                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++) {
1023
1024
            mm_float4 v = velm[i];
            if (v.w != 0)
1025
1026
1027
                energy += (v.x*v.x+v.y*v.y+v.z*v.z)/v.w;
        }
    }
1028
1029
1030
1031
1032
    
    // Restore the velocities.
    
    if (timeShift != 0)
        posDelta->copyTo(context.getVelm());
1033
1034
    return 0.5*energy;
}