OpenCLNonbondedUtilities.cpp 39.9 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-2023 Stanford University and the Authors.      *
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
 * 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/>.      *
 * -------------------------------------------------------------------------- */

27
#include "openmm/OpenMMException.h"
28
29
#include "OpenCLNonbondedUtilities.h"
#include "OpenCLArray.h"
30
#include "OpenCLContext.h"
31
#include "OpenCLKernelSources.h"
32
#include "OpenCLExpressionUtilities.h"
33
34
#include "OpenCLSort.h"
#include <algorithm>
35
#include <map>
36
37
#include <set>
#include <utility>
38
39
40
41

using namespace OpenMM;
using namespace std;

42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
class OpenCLNonbondedUtilities::BlockSortTrait : public OpenCLSort::SortTrait {
public:
    BlockSortTrait(bool useDouble) : useDouble(useDouble) {
    }
    int getDataSize() const {return useDouble ? sizeof(mm_double2) : sizeof(mm_float2);}
    int getKeySize() const {return useDouble ? sizeof(cl_double) : sizeof(cl_float);}
    const char* getDataType() const {return "real2";}
    const char* getKeyType() const {return "real";}
    const char* getMinKey() const {return "-MAXFLOAT";}
    const char* getMaxKey() const {return "MAXFLOAT";}
    const char* getMaxValue() const {return "(real2) (MAXFLOAT, MAXFLOAT)";}
    const char* getSortKey() const {return "value.x";}
private:
    bool useDouble;
};

58
OpenCLNonbondedUtilities::OpenCLNonbondedUtilities(OpenCLContext& context) : context(context), useCutoff(false), usePeriodic(false), useNeighborList(false), anyExclusions(false), usePadding(true),
59
        blockSorter(NULL), pinnedCountBuffer(NULL), pinnedCountMemory(NULL), forceRebuildNeighborList(true), lastCutoff(0.0), groupFlags(0), tilesAfterReorder(0) {
60
    // Decide how many thread blocks and force buffers to use.
61

62
    deviceIsCpu = (context.getDevice().getInfo<CL_DEVICE_TYPE>() == CL_DEVICE_TYPE_CPU);
63
64
65
66
67
    if (deviceIsCpu) {
        numForceThreadBlocks = context.getNumThreadBlocks();
        forceThreadBlockSize = 1;
    }
    else if (context.getSIMDWidth() == 32) {
68
69
70
71
72
73
74
75
        int blocksPerComputeUnit = 4;
        std::string vendor = context.getDevice().getInfo<CL_DEVICE_VENDOR>();
        if (vendor.size() >= 5 && vendor.substr(0, 5) == "Apple") {
            // 1536 threads per GPU core.
            blocksPerComputeUnit = 6;
        }
        numForceThreadBlocks = blocksPerComputeUnit*context.getDevice().getInfo<CL_DEVICE_MAX_COMPUTE_UNITS>();
        forceThreadBlockSize = 256;
76
    }
77
    else {
78
        numForceThreadBlocks = context.getNumThreadBlocks();
79
        forceThreadBlockSize = (context.getSIMDWidth() >= 32 ? OpenCLContext::ThreadBlockSize : 32);
80
    }
81
82
    pinnedCountBuffer = new cl::Buffer(context.getContext(), CL_MEM_ALLOC_HOST_PTR, sizeof(unsigned int));
    pinnedCountMemory = (unsigned int*) context.getQueue().enqueueMapBuffer(*pinnedCountBuffer, CL_TRUE, CL_MAP_READ, 0, sizeof(int));
83
84
85
86
87
88
89
    
    // When building the neighbor list, we can optionally use large blocks (1024 atoms) to
    // accelerate the process.  This makes building the neighbor list faster, but it prevents
    // us from sorting atom blocks by size, which leads to a slightly less efficient neighbor
    // list.  We guess based on system size which will be faster.

    useLargeBlocks = (context.getNumAtoms() > 100000);
90
    setKernelSource(deviceIsCpu ? OpenCLKernelSources::nonbonded_cpu : OpenCLKernelSources::nonbonded);
91
92
93
}

OpenCLNonbondedUtilities::~OpenCLNonbondedUtilities() {
94
95
    if (blockSorter != NULL)
        delete blockSorter;
96
97
    if (pinnedCountBuffer != NULL)
        delete pinnedCountBuffer;
98
99
}

100
void OpenCLNonbondedUtilities::addInteraction(bool usesCutoff, bool usesPeriodic, bool usesExclusions, double cutoffDistance, const vector<vector<int> >& exclusionList, const string& kernel, int forceGroup, bool useNeighborList) {
101
    if (groupCutoff.size() > 0) {
102
103
104
105
        if (usesCutoff != useCutoff)
            throw OpenMMException("All Forces must agree on whether to use a cutoff");
        if (usesPeriodic != usePeriodic)
            throw OpenMMException("All Forces must agree on whether to use periodic boundary conditions");
106
107
        if (usesCutoff && groupCutoff.find(forceGroup) != groupCutoff.end() && groupCutoff[forceGroup] != cutoffDistance)
            throw OpenMMException("All Forces in a single force group must use the same cutoff distance");
108
    }
109
110
111
112
    if (usesExclusions)
        requestExclusions(exclusionList);
    useCutoff = usesCutoff;
    usePeriodic = usesPeriodic;
113
    this->useNeighborList |= ((useNeighborList || deviceIsCpu) && useCutoff);
114
115
116
117
118
119
120
121
122
123
    groupCutoff[forceGroup] = cutoffDistance;
    groupFlags |= 1<<forceGroup;
    if (kernel.size() > 0) {
        if (groupKernelSource.find(forceGroup) == groupKernelSource.end())
            groupKernelSource[forceGroup] = "";
        map<string, string> replacements;
        replacements["CUTOFF"] = "CUTOFF_"+context.intToString(forceGroup);
        replacements["CUTOFF_SQUARED"] = "CUTOFF_"+context.intToString(forceGroup)+"_SQUARED";
        groupKernelSource[forceGroup] += context.replaceStrings(kernel, replacements)+"\n";
    }
124
125
}

126
127
void OpenCLNonbondedUtilities::addParameter(ComputeParameterInfo parameter) {
    parameters.push_back(ParameterInfo(parameter.getName(), parameter.getComponentType(), parameter.getNumComponents(),
128
            parameter.getSize(), context.unwrap(parameter.getArray()).getDeviceBuffer(), parameter.isConstant()));
129
130
}

131
132
133
134
void OpenCLNonbondedUtilities::addParameter(const ParameterInfo& parameter) {
    parameters.push_back(parameter);
}

135
136
void OpenCLNonbondedUtilities::addArgument(ComputeParameterInfo parameter) {
    arguments.push_back(ParameterInfo(parameter.getName(), parameter.getComponentType(), parameter.getNumComponents(),
137
            parameter.getSize(), context.unwrap(parameter.getArray()).getDeviceBuffer(), parameter.isConstant()));
138
139
}

140
141
142
143
void OpenCLNonbondedUtilities::addArgument(const ParameterInfo& parameter) {
    arguments.push_back(parameter);
}

144
145
146
147
148
149
150
151
152
153
154
155
156
string OpenCLNonbondedUtilities::addEnergyParameterDerivative(const string& param) {
    // See if the parameter has already been added.
    
    int index;
    for (index = 0; index < energyParameterDerivatives.size(); index++)
        if (param == energyParameterDerivatives[index])
            break;
    if (index == energyParameterDerivatives.size())
        energyParameterDerivatives.push_back(param);
    context.addEnergyParameterDerivative(param);
    return string("energyParamDeriv")+context.intToString(index);
}

157
158
void OpenCLNonbondedUtilities::requestExclusions(const vector<vector<int> >& exclusionList) {
    if (anyExclusions) {
159
        bool sameExclusions = (exclusionList.size() == atomExclusions.size());
160
        for (int i = 0; i < (int) exclusionList.size() && sameExclusions; i++) {
161
162
            if (exclusionList[i].size() != atomExclusions[i].size())
                sameExclusions = false;
163
164
            set<int> expectedExclusions;
            expectedExclusions.insert(atomExclusions[i].begin(), atomExclusions[i].end());
165
            for (int j = 0; j < (int) exclusionList[i].size(); j++)
166
                if (expectedExclusions.find(exclusionList[i][j]) == expectedExclusions.end())
167
168
169
170
171
                    sameExclusions = false;
        }
        if (!sameExclusions)
            throw OpenMMException("All Forces must have identical exceptions");
    }
172
    else {
173
        atomExclusions = exclusionList;
174
175
        anyExclusions = true;
    }
176
177
}

178
static bool compareInt2(mm_int2 a, mm_int2 b) {
peastman's avatar
peastman committed
179
180
181
182
183
    // This version is used on devices with SIMD width of 32 or less.  It sorts tiles to improve cache efficiency.

    return ((a.y < b.y) || (a.y == b.y && a.x < b.x));
}

184
static bool compareInt2LargeSIMD(mm_int2 a, mm_int2 b) {
peastman's avatar
peastman committed
185
186
187
188
189
190
191
192
193
194
    // This version is used on devices with SIMD width greater than 32.  It puts diagonal tiles before off-diagonal
    // ones to reduce thread divergence.
    
    if (a.x == a.y) {
        if (b.x == b.y)
            return (a.x < b.x);
        return true;
    }
    if (b.x == b.y)
        return false;
195
196
197
    return ((a.y < b.y) || (a.y == b.y && a.x < b.x));
}

198
void OpenCLNonbondedUtilities::initialize(const System& system) {
199
200
    if (atomExclusions.size() == 0) {
        // No exclusions were specifically requested, so just mark every atom as not interacting with itself.
201

202
        atomExclusions.resize(context.getNumAtoms());
203
        for (int i = 0; i < (int) atomExclusions.size(); i++)
204
205
206
            atomExclusions[i].push_back(i);
    }

207
208
209
    // Create the list of tiles.

    int numAtomBlocks = context.getNumAtomBlocks();
210
    int numContexts = context.getPlatformData().contexts.size();
211
    setAtomBlockRange(context.getContextIndex()/(double) numContexts, (context.getContextIndex()+1)/(double) numContexts);
212

213
    // Build a list of tiles that contain exclusions.
214

215
216
217
218
219
220
221
222
223
    set<pair<int, int> > tilesWithExclusions;
    for (int atom1 = 0; atom1 < (int) atomExclusions.size(); ++atom1) {
        int x = atom1/OpenCLContext::TileSize;
        for (int j = 0; j < (int) atomExclusions[atom1].size(); ++j) {
            int atom2 = atomExclusions[atom1][j];
            int y = atom2/OpenCLContext::TileSize;
            tilesWithExclusions.insert(make_pair(max(x, y), min(x, y)));
        }
    }
224
    vector<mm_int2> exclusionTilesVec;
225
    for (set<pair<int, int> >::const_iterator iter = tilesWithExclusions.begin(); iter != tilesWithExclusions.end(); ++iter)
226
        exclusionTilesVec.push_back(mm_int2(iter->first, iter->second));
227
    sort(exclusionTilesVec.begin(), exclusionTilesVec.end(), context.getSIMDWidth() <= 32 || !useCutoff ? compareInt2 : compareInt2LargeSIMD);
228
    exclusionTiles.initialize<mm_int2>(context, exclusionTilesVec.size(), "exclusionTiles");
peastman's avatar
peastman committed
229
    exclusionTiles.upload(exclusionTilesVec);
230
231
    map<pair<int, int>, int> exclusionTileMap;
    for (int i = 0; i < (int) exclusionTilesVec.size(); i++) {
232
        mm_int2 tile = exclusionTilesVec[i];
233
234
235
236
237
238
239
        exclusionTileMap[make_pair(tile.x, tile.y)] = i;
    }
    vector<vector<int> > exclusionBlocksForBlock(numAtomBlocks);
    for (set<pair<int, int> >::const_iterator iter = tilesWithExclusions.begin(); iter != tilesWithExclusions.end(); ++iter) {
        exclusionBlocksForBlock[iter->first].push_back(iter->second);
        if (iter->first != iter->second)
            exclusionBlocksForBlock[iter->second].push_back(iter->first);
240
241
242
    }
    vector<cl_uint> exclusionRowIndicesVec(numAtomBlocks+1, 0);
    vector<cl_uint> exclusionIndicesVec;
243
244
245
    for (int i = 0; i < numAtomBlocks; i++) {
        exclusionIndicesVec.insert(exclusionIndicesVec.end(), exclusionBlocksForBlock[i].begin(), exclusionBlocksForBlock[i].end());
        exclusionRowIndicesVec[i+1] = exclusionIndicesVec.size();
246
    }
247
248
249
    maxExclusions = 0;
    for (int i = 0; i < (int) exclusionBlocksForBlock.size(); i++)
        maxExclusions = (maxExclusions > exclusionBlocksForBlock[i].size() ? maxExclusions : exclusionBlocksForBlock[i].size());
peastman's avatar
peastman committed
250
251
252
253
    exclusionIndices.initialize<cl_uint>(context, exclusionIndicesVec.size(), "exclusionIndices");
    exclusionRowIndices.initialize<cl_uint>(context, exclusionRowIndicesVec.size(), "exclusionRowIndices");
    exclusionIndices.upload(exclusionIndicesVec);
    exclusionRowIndices.upload(exclusionRowIndicesVec);
254
255
256

    // Record the exclusion data.

peastman's avatar
peastman committed
257
    exclusions.initialize<cl_uint>(context, tilesWithExclusions.size()*OpenCLContext::TileSize, "exclusions");
258
    cl_uint allFlags = (cl_uint) -1;
peastman's avatar
peastman committed
259
260
    vector<cl_uint> exclusionVec(exclusions.getSize(), allFlags);
    for (int i = 0; i < exclusions.getSize(); ++i)
261
262
263
264
265
266
267
268
269
        exclusionVec[i] = 0xFFFFFFFF;
    for (int atom1 = 0; atom1 < (int) atomExclusions.size(); ++atom1) {
        int x = atom1/OpenCLContext::TileSize;
        int offset1 = atom1-x*OpenCLContext::TileSize;
        for (int j = 0; j < (int) atomExclusions[atom1].size(); ++j) {
            int atom2 = atomExclusions[atom1][j];
            int y = atom2/OpenCLContext::TileSize;
            int offset2 = atom2-y*OpenCLContext::TileSize;
            if (x > y) {
270
271
                int index = exclusionTileMap[make_pair(x, y)]*OpenCLContext::TileSize;
                exclusionVec[index+offset1] &= allFlags-(1<<offset2);
272
273
            }
            else {
274
275
                int index = exclusionTileMap[make_pair(y, x)]*OpenCLContext::TileSize;
                exclusionVec[index+offset2] &= allFlags-(1<<offset1);
276
277
278
279
            }
        }
    }
    atomExclusions.clear(); // We won't use this again, so free the memory it used
peastman's avatar
peastman committed
280
    exclusions.upload(exclusionVec);
281
282
283
284

    // Create data structures for the neighbor list.

    if (useCutoff) {
285
286
287
288
289
290
291
292
293
        // Select a size for the arrays that hold the neighbor list.  We have to make a fairly
        // arbitrary guess, but if this turns out to be too small we'll increase it later.

        int maxTiles = 20*numAtomBlocks;
        if (maxTiles > numTiles)
            maxTiles = numTiles;
        if (maxTiles < 1)
            maxTiles = 1;
        int numAtoms = context.getNumAtoms();
peastman's avatar
peastman committed
294
295
296
        interactingTiles.initialize<cl_int>(context, maxTiles, "interactingTiles");
        interactingAtoms.initialize<cl_int>(context, OpenCLContext::TileSize*maxTiles, "interactingAtoms");
        interactionCount.initialize<cl_uint>(context, 1, "interactionCount");
297
        int elementSize = (context.getUseDoublePrecision() ? sizeof(cl_double) : sizeof(cl_float));
peastman's avatar
peastman committed
298
299
300
301
302
        blockCenter.initialize(context, numAtomBlocks, 4*elementSize, "blockCenter");
        blockBoundingBox.initialize(context, numAtomBlocks, 4*elementSize, "blockBoundingBox");
        sortedBlocks.initialize(context, numAtomBlocks, 2*elementSize, "sortedBlocks");
        sortedBlockCenter.initialize(context, numAtomBlocks+1, 4*elementSize, "sortedBlockCenter");
        sortedBlockBoundingBox.initialize(context, numAtomBlocks+1, 4*elementSize, "sortedBlockBoundingBox");
303
304
        largeBlockCenter.initialize(context, numAtomBlocks, 4*elementSize, "largeBlockCenter");
        largeBlockBoundingBox.initialize(context, numAtomBlocks, 4*elementSize, "largeBlockBoundingBox");
peastman's avatar
peastman committed
305
306
        oldPositions.initialize(context, numAtoms, 4*elementSize, "oldPositions");
        rebuildNeighborList.initialize<int>(context, 1, "rebuildNeighborList");
307
        blockSorter = new OpenCLSort(context, new BlockSortTrait(context.getUseDoublePrecision()), numAtomBlocks, false);
308
        vector<cl_uint> count(1, 0);
peastman's avatar
peastman committed
309
        interactionCount.upload(count);
peastman's avatar
peastman committed
310
        rebuildNeighborList.upload(count);
311
    }
312
313
}

314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
static void setPeriodicBoxArgs(OpenCLContext& cl, cl::Kernel& kernel, int index) {
    if (cl.getUseDoublePrecision()) {
        kernel.setArg<mm_double4>(index++, cl.getPeriodicBoxSizeDouble());
        kernel.setArg<mm_double4>(index++, cl.getInvPeriodicBoxSizeDouble());
        kernel.setArg<mm_double4>(index++, cl.getPeriodicBoxVecXDouble());
        kernel.setArg<mm_double4>(index++, cl.getPeriodicBoxVecYDouble());
        kernel.setArg<mm_double4>(index, cl.getPeriodicBoxVecZDouble());
    }
    else {
        kernel.setArg<mm_float4>(index++, cl.getPeriodicBoxSize());
        kernel.setArg<mm_float4>(index++, cl.getInvPeriodicBoxSize());
        kernel.setArg<mm_float4>(index++, cl.getPeriodicBoxVecX());
        kernel.setArg<mm_float4>(index++, cl.getPeriodicBoxVecY());
        kernel.setArg<mm_float4>(index, cl.getPeriodicBoxVecZ());
    }
329
330
}

331
332
333
334
335
336
337
double OpenCLNonbondedUtilities::getMaxCutoffDistance() {
    double cutoff = 0.0;
    for (map<int, double>::const_iterator iter = groupCutoff.begin(); iter != groupCutoff.end(); ++iter)
        cutoff = max(cutoff, iter->second);
    return cutoff;
}

338
339
340
341
342
double OpenCLNonbondedUtilities::padCutoff(double cutoff) {
    double padding = (usePadding ? 0.1*cutoff : 0.0);
    return cutoff+padding;
}

343
344
345
346
347
348
void OpenCLNonbondedUtilities::prepareInteractions(int forceGroups) {
    if ((forceGroups&groupFlags) == 0)
        return;
    if (groupKernels.find(forceGroups) == groupKernels.end())
        createKernelsForGroups(forceGroups);
    KernelSet& kernels = groupKernels[forceGroups];
349
    if (useCutoff && usePeriodic) {
350
        mm_float4 box = context.getPeriodicBoxSize();
351
        double minAllowedSize = 1.999999*kernels.cutoffDistance;
352
353
354
        if (box.x < minAllowedSize || box.y < minAllowedSize || box.z < minAllowedSize)
            throw OpenMMException("The periodic box size has decreased to less than twice the nonbonded cutoff.");
    }
355
356
357
358
    if (!useNeighborList)
        return;
    if (numTiles == 0)
        return;
359
360
361

    // Compute the neighbor list.

362
363
    if (lastCutoff != kernels.cutoffDistance)
        forceRebuildNeighborList = true;
364
365
    setPeriodicBoxArgs(context, kernels.findBlockBoundsKernel, 1);
    context.executeKernel(kernels.findBlockBoundsKernel, context.getNumAtoms());
366
367
368
369
    if (useLargeBlocks)
        setPeriodicBoxArgs(context, kernels.sortBoxDataKernel, 12);
    else
        blockSorter->sort(sortedBlocks);
370
371
372
373
374
    kernels.sortBoxDataKernel.setArg<cl_int>(9, forceRebuildNeighborList);
    context.executeKernel(kernels.sortBoxDataKernel, context.getNumAtoms());
    setPeriodicBoxArgs(context, kernels.findInteractingBlocksKernel, 0);
    context.executeKernel(kernels.findInteractingBlocksKernel, context.getNumAtoms(), interactingBlocksThreadBlockSize);
    forceRebuildNeighborList = false;
375
    lastCutoff = kernels.cutoffDistance;
376
377
378
379
380
381
382
    context.getQueue().enqueueReadBuffer(interactionCount.getDeviceBuffer(), CL_FALSE, 0, sizeof(int), pinnedCountMemory, NULL, &downloadCountEvent);

    #if __APPLE__ && defined(__aarch64__)
    // Segment the command stream to avoid stalls later.
    if (groupKernels[forceGroups].hasForces)
        context.getQueue().flush();
    #endif
383
384
}

385
void OpenCLNonbondedUtilities::computeInteractions(int forceGroups, bool includeForces, bool includeEnergy) {
386
387
388
389
    if ((forceGroups&groupFlags) == 0)
        return;
    KernelSet& kernels = groupKernels[forceGroups];
    if (kernels.hasForces) {
390
391
392
        cl::Kernel& kernel = (includeForces ? (includeEnergy ? kernels.forceEnergyKernel : kernels.forceKernel) : kernels.energyKernel);
        if (*reinterpret_cast<cl_kernel*>(&kernel) == NULL)
            kernel = createInteractionKernel(kernels.source, parameters, arguments, true, true, forceGroups, includeForces, includeEnergy);
393
        if (useCutoff)
394
395
            setPeriodicBoxArgs(context, kernel, 9);
        context.executeKernel(kernel, numForceThreadBlocks*forceThreadBlockSize, forceThreadBlockSize);
396
    }
397
    if (useNeighborList && numTiles > 0) {
398
399
400
401
402
        #if __APPLE__ && defined(__aarch64__)
        // Ensure cached up work executes while you're waiting.
        if (kernels.hasForces)
            context.getQueue().flush();
        #endif
403
404
405
        downloadCountEvent.wait();
        updateNeighborListSize();
    }
406
407
}

408
bool OpenCLNonbondedUtilities::updateNeighborListSize() {
409
    if (!useCutoff)
410
        return false;
411
    if (context.getStepsSinceReorder() == 0 || tilesAfterReorder == 0)
412
413
414
        tilesAfterReorder = pinnedCountMemory[0];
    else if (context.getStepsSinceReorder() > 25 && pinnedCountMemory[0] > 1.1*tilesAfterReorder)
        context.forceReorder();
415
    if (pinnedCountMemory[0] <= interactingTiles.getSize())
416
        return false;
417
418
419
420

    // The most recent timestep had too many interactions to fit in the arrays.  Make the arrays bigger to prevent
    // this from happening in the future.

421
422
423
    unsigned int maxTiles = (unsigned int) (1.2*pinnedCountMemory[0]);
    unsigned int numBlocks = context.getNumAtomBlocks();
    int totalTiles = numBlocks*(numBlocks+1)/2;
424
425
    if (maxTiles > totalTiles)
        maxTiles = totalTiles;
peastman's avatar
peastman committed
426
    interactingTiles.resize(maxTiles);
427
    interactingAtoms.resize(OpenCLContext::TileSize*(size_t) maxTiles);
428
    for (map<int, KernelSet>::iterator iter = groupKernels.begin(); iter != groupKernels.end(); ++iter) {
429
430
        KernelSet& kernels = iter->second;
        if (*reinterpret_cast<cl_kernel*>(&kernels.forceKernel) != NULL) {
peastman's avatar
peastman committed
431
            kernels.forceKernel.setArg<cl::Buffer>(7, interactingTiles.getDeviceBuffer());
432
            kernels.forceKernel.setArg<cl_uint>(14, maxTiles);
peastman's avatar
peastman committed
433
            kernels.forceKernel.setArg<cl::Buffer>(17, interactingAtoms.getDeviceBuffer());
434
435
        }
        if (*reinterpret_cast<cl_kernel*>(&kernels.energyKernel) != NULL) {
peastman's avatar
peastman committed
436
            kernels.energyKernel.setArg<cl::Buffer>(7, interactingTiles.getDeviceBuffer());
437
            kernels.energyKernel.setArg<cl_uint>(14, maxTiles);
peastman's avatar
peastman committed
438
            kernels.energyKernel.setArg<cl::Buffer>(17, interactingAtoms.getDeviceBuffer());
439
440
        }
        if (*reinterpret_cast<cl_kernel*>(&kernels.forceEnergyKernel) != NULL) {
peastman's avatar
peastman committed
441
            kernels.forceEnergyKernel.setArg<cl::Buffer>(7, interactingTiles.getDeviceBuffer());
442
            kernels.forceEnergyKernel.setArg<cl_uint>(14, maxTiles);
peastman's avatar
peastman committed
443
            kernels.forceEnergyKernel.setArg<cl::Buffer>(17, interactingAtoms.getDeviceBuffer());
444
        }
peastman's avatar
peastman committed
445
446
        kernels.findInteractingBlocksKernel.setArg<cl::Buffer>(6, interactingTiles.getDeviceBuffer());
        kernels.findInteractingBlocksKernel.setArg<cl::Buffer>(7, interactingAtoms.getDeviceBuffer());
447
        kernels.findInteractingBlocksKernel.setArg<cl_uint>(9, maxTiles);
448
449
    }
    forceRebuildNeighborList = true;
450
    context.setForcesValid(false);
451
    return true;
452
453
}

454
455
456
457
458
459
460
461
void OpenCLNonbondedUtilities::setUsePadding(bool padding) {
    usePadding = padding;
}

void OpenCLNonbondedUtilities::setAtomBlockRange(double startFraction, double endFraction) {
    int numAtomBlocks = context.getNumAtomBlocks();
    startBlockIndex = (int) (startFraction*numAtomBlocks);
    numBlocks = (int) (endFraction*numAtomBlocks)-startBlockIndex;
462
    long long totalTiles = context.getNumAtomBlocks()*((long long)context.getNumAtomBlocks()+1)/2;
463
    startTileIndex = (int) (startFraction*totalTiles);;
464
    numTiles = (long long) (endFraction*totalTiles)-startTileIndex;
465
    if (useCutoff) {
466
        // We are using a cutoff, and the kernels have already been created.
467

468
        for (map<int, KernelSet>::iterator iter = groupKernels.begin(); iter != groupKernels.end(); ++iter) {
469
470
471
            KernelSet& kernels = iter->second;
            if (*reinterpret_cast<cl_kernel*>(&kernels.forceKernel) != NULL) {
                kernels.forceKernel.setArg<cl_uint>(5, startTileIndex);
472
                kernels.forceKernel.setArg<cl_ulong>(6, numTiles);
473
474
475
            }
            if (*reinterpret_cast<cl_kernel*>(&kernels.energyKernel) != NULL) {
                kernels.energyKernel.setArg<cl_uint>(5, startTileIndex);
476
                kernels.energyKernel.setArg<cl_ulong>(6, numTiles);
477
478
479
            }
            if (*reinterpret_cast<cl_kernel*>(&kernels.forceEnergyKernel) != NULL) {
                kernels.forceEnergyKernel.setArg<cl_uint>(5, startTileIndex);
480
                kernels.forceEnergyKernel.setArg<cl_ulong>(6, numTiles);
481
482
483
            }
            kernels.findInteractingBlocksKernel.setArg<cl_uint>(10, startBlockIndex);
            kernels.findInteractingBlocksKernel.setArg<cl_uint>(11, numBlocks);
484
485
        }
        forceRebuildNeighborList = true;
486
487
488
    }
}

489
490
491
492
493
494
495
496
497
498
499
500
void OpenCLNonbondedUtilities::createKernelsForGroups(int groups) {
    KernelSet kernels;
    double cutoff = 0.0;
    string source;
    for (int i = 0; i < 32; i++) {
        if ((groups&(1<<i)) != 0) {
            cutoff = max(cutoff, groupCutoff[i]);
            source += groupKernelSource[i];
        }
    }
    kernels.hasForces = (source.size() > 0);
    kernels.cutoffDistance = cutoff;
501
    kernels.source = source;
502
    if (useCutoff) {
503
        double paddedCutoff = padCutoff(cutoff);
504
505
506
        map<string, string> defines;
        defines["TILE_SIZE"] = context.intToString(OpenCLContext::TileSize);
        defines["NUM_ATOMS"] = context.intToString(context.getNumAtoms());
507
        defines["PADDING"] = context.doubleToString(paddedCutoff-cutoff);
508
509
        defines["PADDED_CUTOFF"] = context.doubleToString(paddedCutoff);
        defines["PADDED_CUTOFF_SQUARED"] = context.doubleToString(paddedCutoff*paddedCutoff);
peastman's avatar
peastman committed
510
        defines["NUM_TILES_WITH_EXCLUSIONS"] = context.intToString(exclusionTiles.getSize());
511
512
513
514
        defines["NUM_BLOCKS"] = context.intToString(context.getNumAtomBlocks());
        defines["SIMD_WIDTH"] = context.intToString(context.getSIMDWidth());
        if (usePeriodic)
            defines["USE_PERIODIC"] = "1";
515
516
        if (context.getBoxIsTriclinic())
            defines["TRICLINIC"] = "1";
517
518
        if (useLargeBlocks)
            defines["USE_LARGE_BLOCKS"] = "1";
519
520
521
522
523
524
525
526
527
528
        defines["MAX_EXCLUSIONS"] = context.intToString(maxExclusions);
        defines["BUFFER_GROUPS"] = (deviceIsCpu ? "4" : "2");
        string file = (deviceIsCpu ? OpenCLKernelSources::findInteractingBlocks_cpu : OpenCLKernelSources::findInteractingBlocks);
        int groupSize = (deviceIsCpu || context.getSIMDWidth() < 32 ? 32 : 256);
        while (true) {
            defines["GROUP_SIZE"] = context.intToString(groupSize);
            cl::Program interactingBlocksProgram = context.createProgram(file, defines);
            kernels.findBlockBoundsKernel = cl::Kernel(interactingBlocksProgram, "findBlockBounds");
            kernels.findBlockBoundsKernel.setArg<cl_int>(0, context.getNumAtoms());
            kernels.findBlockBoundsKernel.setArg<cl::Buffer>(6, context.getPosq().getDeviceBuffer());
peastman's avatar
peastman committed
529
530
531
532
            kernels.findBlockBoundsKernel.setArg<cl::Buffer>(7, blockCenter.getDeviceBuffer());
            kernels.findBlockBoundsKernel.setArg<cl::Buffer>(8, blockBoundingBox.getDeviceBuffer());
            kernels.findBlockBoundsKernel.setArg<cl::Buffer>(9, rebuildNeighborList.getDeviceBuffer());
            kernels.findBlockBoundsKernel.setArg<cl::Buffer>(10, sortedBlocks.getDeviceBuffer());
533
            kernels.sortBoxDataKernel = cl::Kernel(interactingBlocksProgram, "sortBoxData");
peastman's avatar
peastman committed
534
535
536
537
538
            kernels.sortBoxDataKernel.setArg<cl::Buffer>(0, sortedBlocks.getDeviceBuffer());
            kernels.sortBoxDataKernel.setArg<cl::Buffer>(1, blockCenter.getDeviceBuffer());
            kernels.sortBoxDataKernel.setArg<cl::Buffer>(2, blockBoundingBox.getDeviceBuffer());
            kernels.sortBoxDataKernel.setArg<cl::Buffer>(3, sortedBlockCenter.getDeviceBuffer());
            kernels.sortBoxDataKernel.setArg<cl::Buffer>(4, sortedBlockBoundingBox.getDeviceBuffer());
539
            kernels.sortBoxDataKernel.setArg<cl::Buffer>(5, context.getPosq().getDeviceBuffer());
peastman's avatar
peastman committed
540
541
542
            kernels.sortBoxDataKernel.setArg<cl::Buffer>(6, oldPositions.getDeviceBuffer());
            kernels.sortBoxDataKernel.setArg<cl::Buffer>(7, interactionCount.getDeviceBuffer());
            kernels.sortBoxDataKernel.setArg<cl::Buffer>(8, rebuildNeighborList.getDeviceBuffer());
543
            kernels.sortBoxDataKernel.setArg<cl_int>(9, true);
544
545
546
547
            if (useLargeBlocks) {
                kernels.sortBoxDataKernel.setArg<cl::Buffer>(10, largeBlockCenter.getDeviceBuffer());
                kernels.sortBoxDataKernel.setArg<cl::Buffer>(11, largeBlockBoundingBox.getDeviceBuffer());
            }
548
            kernels.findInteractingBlocksKernel = cl::Kernel(interactingBlocksProgram, "findBlocksWithInteractions");
peastman's avatar
peastman committed
549
550
551
            kernels.findInteractingBlocksKernel.setArg<cl::Buffer>(5, interactionCount.getDeviceBuffer());
            kernels.findInteractingBlocksKernel.setArg<cl::Buffer>(6, interactingTiles.getDeviceBuffer());
            kernels.findInteractingBlocksKernel.setArg<cl::Buffer>(7, interactingAtoms.getDeviceBuffer());
552
            kernels.findInteractingBlocksKernel.setArg<cl::Buffer>(8, context.getPosq().getDeviceBuffer());
peastman's avatar
peastman committed
553
            kernels.findInteractingBlocksKernel.setArg<cl_uint>(9, interactingTiles.getSize());
554
555
            kernels.findInteractingBlocksKernel.setArg<cl_uint>(10, startBlockIndex);
            kernels.findInteractingBlocksKernel.setArg<cl_uint>(11, numBlocks);
peastman's avatar
peastman committed
556
557
558
559
560
561
562
            kernels.findInteractingBlocksKernel.setArg<cl::Buffer>(12, sortedBlocks.getDeviceBuffer());
            kernels.findInteractingBlocksKernel.setArg<cl::Buffer>(13, sortedBlockCenter.getDeviceBuffer());
            kernels.findInteractingBlocksKernel.setArg<cl::Buffer>(14, sortedBlockBoundingBox.getDeviceBuffer());
            kernels.findInteractingBlocksKernel.setArg<cl::Buffer>(15, exclusionIndices.getDeviceBuffer());
            kernels.findInteractingBlocksKernel.setArg<cl::Buffer>(16, exclusionRowIndices.getDeviceBuffer());
            kernels.findInteractingBlocksKernel.setArg<cl::Buffer>(17, oldPositions.getDeviceBuffer());
            kernels.findInteractingBlocksKernel.setArg<cl::Buffer>(18, rebuildNeighborList.getDeviceBuffer());
563
564
565
566
            if (useLargeBlocks) {
                kernels.findInteractingBlocksKernel.setArg<cl::Buffer>(19, largeBlockCenter.getDeviceBuffer());
                kernels.findInteractingBlocksKernel.setArg<cl::Buffer>(20, largeBlockBoundingBox.getDeviceBuffer());
            }
567
568
            if (kernels.findInteractingBlocksKernel.getWorkGroupInfo<CL_KERNEL_WORK_GROUP_SIZE>(context.getDevice()) < groupSize) {
                // The device can't handle this block size, so reduce it.
569

570
571
572
573
574
575
576
577
578
579
580
581
                groupSize -= 32;
                if (groupSize < 32)
                    throw OpenMMException("Failed to create findInteractingBlocks kernel");
                continue;
            }
            break;
        }
        interactingBlocksThreadBlockSize = (deviceIsCpu ? 1 : groupSize);
    }
    groupKernels[groups] = kernels;
}

582
cl::Kernel OpenCLNonbondedUtilities::createInteractionKernel(const string& source, const vector<ParameterInfo>& params, const vector<ParameterInfo>& arguments, bool useExclusions, bool isSymmetric, int groups, bool includeForces, bool includeEnergy) {
583
584
    map<string, string> replacements;
    replacements["COMPUTE_INTERACTION"] = source;
585
    const string suffixes[] = {"x", "y", "z", "w"};
586
    stringstream localData;
587
    int localDataSize = 0;
588
589
590
    for (const ParameterInfo& param : params) {
        if (param.getNumComponents() == 1)
            localData<<param.getType()<<" "<<param.getName()<<";\n";
591
        else {
592
593
            for (int j = 0; j < param.getNumComponents(); ++j)
                localData<<param.getComponentType()<<" "<<param.getName()<<"_"<<suffixes[j]<<";\n";
594
        }
595
        localDataSize += param.getSize();
596
597
    }
    replacements["ATOM_PARAMETER_DATA"] = localData.str();
598
    stringstream args;
599
600
601
602
603
604
605
606
    for (const ParameterInfo& param : params) {
        args << ", __global ";
        if (param.isConstant())
            args << "const ";
        if (param.getNumComponents() == 3)
            args << param.getComponentType();
        else
            args << param.getType();
607
        args << "* restrict global_";
608
        args << param.getName();
609
    }
610
611
    for (const ParameterInfo& arg : arguments) {
        if (arg.getMemory().getInfo<CL_MEM_TYPE>() == CL_MEM_OBJECT_IMAGE2D) {
612
            args << ", __read_only image2d_t ";
613
            args << arg.getName();
614
615
        }
        else {
616
617
618
619
620
            if ((arg.getMemory().getInfo<CL_MEM_FLAGS>() & CL_MEM_READ_ONLY) == 0) {
                args << ", __global ";
                if (arg.isConstant())
                    args << "const ";
            }
621
622
            else
                args << ", __constant ";
623
            args << arg.getType();
624
            args << "* restrict ";
625
            args << arg.getName();
626
        }
627
    }
628
    if (energyParameterDerivatives.size() > 0)
629
        args << ", __global mixed* restrict energyParamDerivs";
630
631
    replacements["PARAMETER_ARGUMENTS"] = args.str();
    stringstream loadLocal1;
632
633
634
    for (const ParameterInfo& param : params) {
        if (param.getNumComponents() == 1) {
            loadLocal1<<"localData[localAtomIndex]."<<param.getName()<<" = "<<param.getName()<<"1;\n";
635
636
        }
        else {
637
638
            for (int j = 0; j < param.getNumComponents(); ++j)
                loadLocal1<<"localData[localAtomIndex]."<<param.getName()<<"_"<<suffixes[j]<<" = "<<param.getName()<<"1."<<suffixes[j]<<";\n";
639
        }
640
641
    }
    replacements["LOAD_LOCAL_PARAMETERS_FROM_1"] = loadLocal1.str();
642
    replacements["DECLARE_LOCAL_PARAMETERS"] = "";
643
    stringstream loadLocal2;
644
645
646
    for (const ParameterInfo& param : params) {
        if (param.getNumComponents() == 1) {
            loadLocal2<<"localData[localAtomIndex]."<<param.getName()<<" = global_"<<param.getName()<<"[j];\n";
647
648
        }
        else {
649
650
651
652
653
654
            if (param.getNumComponents() == 3)
                loadLocal2<<param.getType()<<" temp_"<<param.getName()<<" = make_"<<param.getType()<<"(global_"<<param.getName()<<"[3*j], global_"<<param.getName()<<"[3*j+1], global_"<<param.getName()<<"[3*j+2]);\n";
            else
                loadLocal2<<param.getType()<<" temp_"<<param.getName()<<" = global_"<<param.getName()<<"[j];\n";
            for (int j = 0; j < param.getNumComponents(); ++j)
                loadLocal2<<"localData[localAtomIndex]."<<param.getName()<<"_"<<suffixes[j]<<" = temp_"<<param.getName()<<"."<<suffixes[j]<<";\n";
655
        }
656
657
658
    }
    replacements["LOAD_LOCAL_PARAMETERS_FROM_GLOBAL"] = loadLocal2.str();
    stringstream load1;
659
660
661
662
663
664
    for (const ParameterInfo& param : params) {
        load1<<param.getType()<<" "<<param.getName()<<"1 = ";
        if (param.getNumComponents() == 3)
            load1<<"make_"<<param.getType()<<"(global_"<<param.getName()<<"[3*atom1], global_"<<param.getName()<<"[3*atom1+1], global_"<<param.getName()<<"[3*atom1+2]);\n";
        else
            load1<<"global_"<<param.getName()<<"[atom1];\n";
665
666
667
    }
    replacements["LOAD_ATOM1_PARAMETERS"] = load1.str();
    stringstream load2j;
668
669
670
    for (const ParameterInfo& param : params) {
        if (param.getNumComponents() == 1) {
            load2j<<param.getType()<<" "<<param.getName()<<"2 = localData[atom2]."<<param.getName()<<";\n";
671
672
        }
        else {
673
674
            load2j<<param.getType()<<" "<<param.getName()<<"2 = ("<<param.getType()<<") (";
            for (int j = 0; j < param.getNumComponents(); ++j) {
675
676
                if (j > 0)
                    load2j<<", ";
677
                load2j<<"localData[atom2]."<<param.getName()<<"_"<<suffixes[j];
678
679
680
            }
            load2j<<");\n";
        }
681
    }
682
    replacements["LOAD_ATOM2_PARAMETERS"] = load2j.str();
683
    stringstream clearLocal;
684
685
686
    for (const ParameterInfo& param : params) {
        if (param.getNumComponents() == 1)
            clearLocal<<"localData[localAtomIndex]."<<param.getName()<<" = 0;\n";
687
        else
688
689
            for (int j = 0; j < param.getNumComponents(); ++j)
                clearLocal<<"localData[localAtomIndex]."<<param.getName()<<"_"<<suffixes[j]<<" = 0;\n";
690
691
    }
    replacements["CLEAR_LOCAL_PARAMETERS"] = clearLocal.str();
692
693
694
695
696
697
698
699
700
701
    stringstream initDerivs;
    for (int i = 0; i < energyParameterDerivatives.size(); i++)
        initDerivs<<"mixed energyParamDeriv"<<i<<" = 0;\n";
    replacements["INIT_DERIVATIVES"] = initDerivs.str();
    stringstream saveDerivs;
    const vector<string>& allParamDerivNames = context.getEnergyParamDerivNames();
    int numDerivs = allParamDerivNames.size();
    for (int i = 0; i < energyParameterDerivatives.size(); i++)
        for (int index = 0; index < numDerivs; index++)
            if (allParamDerivNames[index] == energyParameterDerivatives[i])
702
                saveDerivs<<"energyParamDerivs[GLOBAL_ID*"<<numDerivs<<"+"<<index<<"] += energyParamDeriv"<<i<<";\n";
703
    replacements["SAVE_DERIVATIVES"] = saveDerivs.str();
704
705
706
707
708
709
710
    map<string, string> defines;
    if (useCutoff)
        defines["USE_CUTOFF"] = "1";
    if (usePeriodic)
        defines["USE_PERIODIC"] = "1";
    if (useExclusions)
        defines["USE_EXCLUSIONS"] = "1";
711
712
    if (isSymmetric)
        defines["USE_SYMMETRIC"] = "1";
713
714
    if (useNeighborList)
        defines["USE_NEIGHBOR_LIST"] = "1";
715
716
    if (useCutoff && context.getSIMDWidth() < 32)
        defines["PRUNE_BY_CUTOFF"] = "1";
717
718
719
720
    if (includeForces)
        defines["INCLUDE_FORCES"] = "1";
    if (includeEnergy)
        defines["INCLUDE_ENERGY"] = "1";
721
    defines["THREAD_BLOCK_SIZE"] = context.intToString(forceThreadBlockSize);
722
    defines["FORCE_WORK_GROUP_SIZE"] = context.intToString(forceThreadBlockSize);
723
724
725
726
727
728
729
730
731
732
    double maxCutoff = 0.0;
    for (int i = 0; i < 32; i++) {
        if ((groups&(1<<i)) != 0) {
            double cutoff = groupCutoff[i];
            maxCutoff = max(maxCutoff, cutoff);
            defines["CUTOFF_"+context.intToString(i)+"_SQUARED"] = context.doubleToString(cutoff*cutoff);
            defines["CUTOFF_"+context.intToString(i)] = context.doubleToString(cutoff);
        }
    }
    defines["MAX_CUTOFF"] = context.doubleToString(maxCutoff);
733
734
735
    defines["NUM_ATOMS"] = context.intToString(context.getNumAtoms());
    defines["PADDED_NUM_ATOMS"] = context.intToString(context.getPaddedNumAtoms());
    defines["NUM_BLOCKS"] = context.intToString(context.getNumAtomBlocks());
736
    defines["TILE_SIZE"] = context.intToString(OpenCLContext::TileSize);
peastman's avatar
peastman committed
737
    int numExclusionTiles = exclusionTiles.getSize();
738
739
740
741
742
743
    defines["NUM_TILES_WITH_EXCLUSIONS"] = context.intToString(numExclusionTiles);
    int numContexts = context.getPlatformData().contexts.size();
    int startExclusionIndex = context.getContextIndex()*numExclusionTiles/numContexts;
    int endExclusionIndex = (context.getContextIndex()+1)*numExclusionTiles/numContexts;
    defines["FIRST_EXCLUSION_TILE"] = context.intToString(startExclusionIndex);
    defines["LAST_EXCLUSION_TILE"] = context.intToString(endExclusionIndex);
744
745
    if ((localDataSize/4)%2 == 0)
        defines["PARAMETER_SIZE_IS_EVEN"] = "1";
746
    cl::Program program = context.createProgram(context.replaceStrings(kernelSource, replacements), defines);
747
748
749
    cl::Kernel kernel(program, "computeNonbonded");

    // Set arguments to the Kernel.
750

751
    int index = 0;
752
    kernel.setArg<cl::Memory>(index++, context.getLongForceBuffer().getDeviceBuffer());
753
754
    kernel.setArg<cl::Buffer>(index++, context.getEnergyBuffer().getDeviceBuffer());
    kernel.setArg<cl::Buffer>(index++, context.getPosq().getDeviceBuffer());
peastman's avatar
peastman committed
755
756
    kernel.setArg<cl::Buffer>(index++, exclusions.getDeviceBuffer());
    kernel.setArg<cl::Buffer>(index++, exclusionTiles.getDeviceBuffer());
757
    kernel.setArg<cl_uint>(index++, startTileIndex);
758
    kernel.setArg<cl_ulong>(index++, numTiles);
759
    if (useCutoff) {
peastman's avatar
peastman committed
760
761
        kernel.setArg<cl::Buffer>(index++, interactingTiles.getDeviceBuffer());
        kernel.setArg<cl::Buffer>(index++, interactionCount.getDeviceBuffer());
762
        index += 5; // The periodic box size arguments are set when the kernel is executed.
peastman's avatar
peastman committed
763
764
765
766
        kernel.setArg<cl_uint>(index++, interactingTiles.getSize());
        kernel.setArg<cl::Buffer>(index++, blockCenter.getDeviceBuffer());
        kernel.setArg<cl::Buffer>(index++, blockBoundingBox.getDeviceBuffer());
        kernel.setArg<cl::Buffer>(index++, interactingAtoms.getDeviceBuffer());
767
    }
768
769
770
771
    for (const ParameterInfo& param : params)
        kernel.setArg<cl::Memory>(index++, param.getMemory());
    for (const ParameterInfo& arg : arguments)
        kernel.setArg<cl::Memory>(index++, arg.getMemory());
772
773
    if (energyParameterDerivatives.size() > 0)
        kernel.setArg<cl::Memory>(index++, context.getEnergyParamDerivBuffer().getDeviceBuffer());
774
    return kernel;
775
}
776
777
778
779

void OpenCLNonbondedUtilities::setKernelSource(const string& source) {
    kernelSource = source;
}