CudaContext.cpp 36.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-2025 Stanford University and the Authors.      *
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
 * 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/>.      *
 * -------------------------------------------------------------------------- */

#ifdef WIN32
  #define _USE_MATH_DEFINES // Needed to get M_PI
#endif
#include <cmath>
#include "CudaContext.h"
#include "CudaArray.h"
33
#include "CudaBondedUtilities.h"
34
#include "CudaEvent.h"
35
#include "CudaFFT3D.h"
36
#include "CudaIntegrationUtilities.h"
37
#include "CudaKernels.h"
38
#include "CudaKernelSources.h"
39
#include "CudaNonbondedUtilities.h"
40
#include "CudaProgram.h"
41
#include "CudaSort.h"
42
#include "openmm/common/ComputeArray.h"
43
#include "openmm/common/ContextSelector.h"
44
#include "SHA1.h"
45
#include "openmm/MonteCarloFlexibleBarostat.h"
46
47
48
#include "openmm/Platform.h"
#include "openmm/System.h"
#include "openmm/VirtualSite.h"
49
#include "CudaExpressionUtilities.h"
50
#include "openmm/internal/ContextImpl.h"
51
52
53
#include <algorithm>
#include <cstdlib>
#include <fstream>
54
#include <iomanip>
55
#include <iostream>
56
#include <set>
57
58
#include <sstream>
#include <typeinfo>
59
#include <sys/stat.h>
60
#include <cudaProfiler.h>
61
#include <nvrtc.h>
62
63
64
#ifndef WIN32
  #include <unistd.h>
#endif
peastman's avatar
peastman committed
65

66
67
68
69
70

#define CHECK_RESULT(result) CHECK_RESULT2(result, errorMessage);
#define CHECK_RESULT2(result, prefix) \
    if (result != CUDA_SUCCESS) { \
        std::stringstream m; \
71
        m<<prefix<<": "<<getErrorString(result)<<" ("<<result<<")"<<" at "<<__FILE__<<":"<<__LINE__; \
72
73
        throw OpenMMException(m.str());\
    }
74
75
76
77
78
79
#define CHECK_NVRTC_RESULT(result, prefix) \
    if (result != NVRTC_SUCCESS) { \
        stringstream m; \
        m<<prefix<<": "<<nvrtcGetErrorString(result)<<" ("<<result<<")"<<" at "<<__FILE__<<":"<<__LINE__; \
        throw OpenMMException(m.str());\
    }
80
81
82
83
84

using namespace OpenMM;
using namespace std;

const int CudaContext::ThreadBlockSize = 64;
85
const int CudaContext::TileSize = sizeof(tileflags)*8;
86
87
bool CudaContext::hasInitializedCuda = false;

88
CudaContext::CudaContext(const System& system, int deviceIndex, bool useBlockingSync, const string& precision, const string& tempDir, CudaPlatform::PlatformData& platformData,
89
        CudaContext* originalContext) : ComputeContext(system), platformData(platformData), contextIsValid(false), hasAssignedPosqCharges(false),
90
        pinnedBuffer(NULL), integration(NULL), expression(NULL), bonded(NULL), nonbonded(NULL), useBlockingSync(useBlockingSync) {
91
92
    int cudaDriverVersion;
    cuDriverGetVersion(&cudaDriverVersion);
93
94
95
96
97
98
    if (!hasInitializedCuda) {
        CHECK_RESULT2(cuInit(0), "Error initializing CUDA");
        hasInitializedCuda = true;
    }
    if (precision == "single") {
        useDoublePrecision = false;
99
        useMixedPrecision = false;
100
101
102
    }
    else if (precision == "mixed") {
        useDoublePrecision = false;
103
        useMixedPrecision = true;
104
105
106
    }
    else if (precision == "double") {
        useDoublePrecision = true;
107
        useMixedPrecision = false;
108
109
    }
    else
110
        throw OpenMMException("Illegal value for Precision: "+precision);
111
112
    char* cacheVariable = getenv("OPENMM_CACHE_DIR");
    cacheDir = (cacheVariable == NULL ? tempDir : string(cacheVariable));
113
#ifdef WIN32
114
    this->tempDir = tempDir+"\\";
115
    cacheDir = cacheDir+"\\";
116
117
#else
    this->tempDir = tempDir+"/";
118
    cacheDir = cacheDir+"/";
119
120
121
#endif
    contextIndex = platformData.contexts.size();
    string errorMessage = "Error initializing Context";
122
123
124
125
126
127
    if (originalContext == NULL) {
        isLinkedContext = false;
        int numDevices;
        CHECK_RESULT(cuDeviceGetCount(&numDevices));
        if (deviceIndex < -1 || deviceIndex >= numDevices)
            throw OpenMMException("Illegal value for DeviceIndex: "+intToString(deviceIndex));
128

129
130
131
132
133
134
        vector<int> devicePrecedence;
        if (deviceIndex == -1) {
            devicePrecedence = getDevicePrecedence();
        } else {
            devicePrecedence.push_back(deviceIndex);
        }
135

136
137
138
139
140
141
142
143
144
145
        this->deviceIndex = -1;
        for (int i = 0; i < static_cast<int>(devicePrecedence.size()); i++) {
            int trialDeviceIndex = devicePrecedence[i];
            CHECK_RESULT(cuDeviceGet(&device, trialDeviceIndex));
            defaultOptimizationOptions = "--use_fast_math";
            unsigned int flags = CU_CTX_MAP_HOST;
            if (useBlockingSync)
                flags += CU_CTX_SCHED_BLOCKING_SYNC;
            else
                flags += CU_CTX_SCHED_SPIN;
146

147
148
            if (cuCtxCreate(&context, flags, device) == CUDA_SUCCESS) {
                this->deviceIndex = trialDeviceIndex;
149
150
                CUcontext popped;
                cuCtxPopCurrent(&popped);
151
152
                break;
            }
153
        }
154
        if (this->deviceIndex == -1) {
155
156
157
158
            if (deviceIndex != -1)
                throw OpenMMException("The requested CUDA device could not be loaded");
            else
                throw OpenMMException("No compatible CUDA device is available");
159
        }
160
161
162
163
164
165
    }
    else {
        isLinkedContext = true;
        context = originalContext->context;
        this->deviceIndex = originalContext->deviceIndex;
        this->device = originalContext->device;
166
    }
167

168
    int major, minor;
169
170
    CHECK_RESULT(cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device));
    CHECK_RESULT(cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, device));
peastman's avatar
peastman committed
171
    int numThreadBlocksPerComputeUnit = (major == 6 ? 4 : 6);
172
    if (cudaDriverVersion < 7000) {
173
174
175
176
177
        // This is a workaround to support GTX 980 with CUDA 6.5.  It reports
        // its compute capability as 5.2, but the compiler doesn't support
        // anything beyond 5.0.
        if (major == 5)
            minor = 0;
178
179
    }
    if (cudaDriverVersion < 8000) {
180
181
182
183
184
185
186
        // This is a workaround to support Pascal with CUDA 7.5.  It reports
        // its compute capability as 6.x, but the compiler doesn't support
        // anything beyond 5.3.
        if (major == 6) {
            major = 5;
            minor = 3;
        }
187
    }
188
    gpuArchitecture = 10*major+minor;
189
    computeCapability = major+0.1*minor;
190

191
    contextIsValid = true;
192
    ContextSelector selector(*this);
193
    CHECK_RESULT(cuCtxSetCacheConfig(CU_FUNC_CACHE_PREFER_SHARED));
194
    if (contextIndex > 0 && originalContext == NULL) {
root's avatar
root committed
195
196
197
        int canAccess;
        cuDeviceCanAccessPeer(&canAccess, getDevice(), platformData.contexts[0]->getDevice());
        if (canAccess) {
198
199
200
201
            {
                ContextSelector selector2(*platformData.contexts[0]);
                CHECK_RESULT(cuCtxEnablePeerAccess(getContext(), 0));
            }
root's avatar
root committed
202
203
204
            CHECK_RESULT(cuCtxEnablePeerAccess(platformData.contexts[0]->getContext(), 0));
        }
    }
205
206
    defaultQueue = shared_ptr<ComputeQueueImpl>(new CudaQueue(0));
    currentQueue = defaultQueue;
207
208
209
210
211
212
    numAtoms = system.getNumParticles();
    paddedNumAtoms = TileSize*((numAtoms+TileSize-1)/TileSize);
    numAtomBlocks = (paddedNumAtoms+(TileSize-1))/TileSize;
    int multiprocessors;
    CHECK_RESULT(cuDeviceGetAttribute(&multiprocessors, CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, device));
    numThreadBlocks = numThreadBlocksPerComputeUnit*multiprocessors;
213
    if (cudaDriverVersion >= 9000) {
Peter Eastman's avatar
Peter Eastman committed
214
215
        compilationDefines["SYNC_WARPS"] = "__syncwarp();";
        compilationDefines["SHFL(var, srcLane)"] = "__shfl_sync(0xffffffff, var, srcLane);";
216
        compilationDefines["BALLOT(var)"] = "__ballot_sync(0xffffffff, var);";
Peter Eastman's avatar
Peter Eastman committed
217
218
219
220
    }
    else {
        compilationDefines["SYNC_WARPS"] = "";
        compilationDefines["SHFL(var, srcLane)"] = "__shfl(var, srcLane);";
221
        compilationDefines["BALLOT(var)"] = "__ballot(var);";
Peter Eastman's avatar
Peter Eastman committed
222
    }
223
    if (useDoublePrecision) {
Peter Eastman's avatar
Peter Eastman committed
224
225
        posq.initialize<double4>(*this, paddedNumAtoms, "posq");
        velm.initialize<double4>(*this, paddedNumAtoms, "velm");
226
        compilationDefines["USE_DOUBLE_PRECISION"] = "1";
227
228
229
        compilationDefines["make_real2"] = "make_double2";
        compilationDefines["make_real3"] = "make_double3";
        compilationDefines["make_real4"] = "make_double4";
230
231
232
        compilationDefines["make_mixed2"] = "make_double2";
        compilationDefines["make_mixed3"] = "make_double3";
        compilationDefines["make_mixed4"] = "make_double4";
233
    }
234
    else if (useMixedPrecision) {
Peter Eastman's avatar
Peter Eastman committed
235
236
237
        posq.initialize<float4>(*this, paddedNumAtoms, "posq");
        posqCorrection.initialize<float4>(*this, paddedNumAtoms, "posqCorrection");
        velm.initialize<double4>(*this, paddedNumAtoms, "velm");
238
239
240
241
242
243
244
245
        compilationDefines["USE_MIXED_PRECISION"] = "1";
        compilationDefines["make_real2"] = "make_float2";
        compilationDefines["make_real3"] = "make_float3";
        compilationDefines["make_real4"] = "make_float4";
        compilationDefines["make_mixed2"] = "make_double2";
        compilationDefines["make_mixed3"] = "make_double3";
        compilationDefines["make_mixed4"] = "make_double4";
    }
246
    else {
Peter Eastman's avatar
Peter Eastman committed
247
248
        posq.initialize<float4>(*this, paddedNumAtoms, "posq");
        velm.initialize<float4>(*this, paddedNumAtoms, "velm");
249
250
251
        compilationDefines["make_real2"] = "make_float2";
        compilationDefines["make_real3"] = "make_float3";
        compilationDefines["make_real4"] = "make_float4";
252
253
254
        compilationDefines["make_mixed2"] = "make_float2";
        compilationDefines["make_mixed3"] = "make_float3";
        compilationDefines["make_mixed4"] = "make_float4";
255
    }
256
257
    force.initialize<long long>(*this, paddedNumAtoms*3, "force");
    posCellOffsets.resize(paddedNumAtoms, mm_int4(0, 0, 0, 0));
258
259
260
261
262
    atomIndexDevice.initialize<int>(*this, paddedNumAtoms, "atomIndex");
    atomIndex.resize(paddedNumAtoms);
    for (int i = 0; i < paddedNumAtoms; ++i)
        atomIndex[i] = i;
    atomIndexDevice.upload(atomIndex);
263
264
265
266

    // Create utility kernels that are used in multiple places.

    CUmodule utilities = createModule(CudaKernelSources::vectorOps+CudaKernelSources::utilities);
267
268
269
270
271
272
    clearBufferKernel = getKernel(utilities, "clearBuffer");
    clearTwoBuffersKernel = getKernel(utilities, "clearTwoBuffers");
    clearThreeBuffersKernel = getKernel(utilities, "clearThreeBuffers");
    clearFourBuffersKernel = getKernel(utilities, "clearFourBuffers");
    clearFiveBuffersKernel = getKernel(utilities, "clearFiveBuffers");
    clearSixBuffersKernel = getKernel(utilities, "clearSixBuffers");
Peter Eastman's avatar
Peter Eastman committed
273
    reduceEnergyKernel = getKernel(utilities, "reduceEnergy");
274
    setChargesKernel = getKernel(utilities, "setCharges");
275
276
277
278
279
280
281
282

    // Set defines based on the requested precision.

    compilationDefines["SQRT"] = useDoublePrecision ? "sqrt" : "sqrtf";
    compilationDefines["RSQRT"] = useDoublePrecision ? "rsqrt" : "rsqrtf";
    compilationDefines["RECIP"] = useDoublePrecision ? "1.0/" : "1.0f/";
    compilationDefines["EXP"] = useDoublePrecision ? "exp" : "expf";
    compilationDefines["LOG"] = useDoublePrecision ? "log" : "logf";
283
    compilationDefines["POW"] = useDoublePrecision ? "pow" : "powf";
284
285
286
287
288
289
    compilationDefines["COS"] = useDoublePrecision ? "cos" : "cosf";
    compilationDefines["SIN"] = useDoublePrecision ? "sin" : "sinf";
    compilationDefines["TAN"] = useDoublePrecision ? "tan" : "tanf";
    compilationDefines["ACOS"] = useDoublePrecision ? "acos" : "acosf";
    compilationDefines["ASIN"] = useDoublePrecision ? "asin" : "asinf";
    compilationDefines["ATAN"] = useDoublePrecision ? "atan" : "atanf";
290
291
    compilationDefines["ERF"] = useDoublePrecision ? "erf" : "erff";
    compilationDefines["ERFC"] = useDoublePrecision ? "erfc" : "erfcf";
292
    compilationDefines["FMA"] = useDoublePrecision ? "fma" : "fmaf";
293

294
    // Set defines for applying periodic boundary conditions.
295

296
297
298
299
300
    Vec3 boxVectors[3];
    system.getDefaultPeriodicBoxVectors(boxVectors[0], boxVectors[1], boxVectors[2]);
    boxIsTriclinic = (boxVectors[0][1] != 0.0 || boxVectors[0][2] != 0.0 ||
                      boxVectors[1][0] != 0.0 || boxVectors[1][2] != 0.0 ||
                      boxVectors[2][0] != 0.0 || boxVectors[2][1] != 0.0);
301
302
303
    for (int i = 0; i < system.getNumForces(); i++)
        if (dynamic_cast<const MonteCarloFlexibleBarostat*>(&system.getForce(i)) != NULL)
            boxIsTriclinic = true;
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
    if (boxIsTriclinic) {
        compilationDefines["APPLY_PERIODIC_TO_DELTA(delta)"] =
            "{"
            "real scale3 = floor(delta.z*invPeriodicBoxSize.z+0.5f); \\\n"
            "delta.x -= scale3*periodicBoxVecZ.x; \\\n"
            "delta.y -= scale3*periodicBoxVecZ.y; \\\n"
            "delta.z -= scale3*periodicBoxVecZ.z; \\\n"
            "real scale2 = floor(delta.y*invPeriodicBoxSize.y+0.5f); \\\n"
            "delta.x -= scale2*periodicBoxVecY.x; \\\n"
            "delta.y -= scale2*periodicBoxVecY.y; \\\n"
            "real scale1 = floor(delta.x*invPeriodicBoxSize.x+0.5f); \\\n"
            "delta.x -= scale1*periodicBoxVecX.x;}";
        compilationDefines["APPLY_PERIODIC_TO_POS(pos)"] =
            "{"
            "real scale3 = floor(pos.z*invPeriodicBoxSize.z); \\\n"
            "pos.x -= scale3*periodicBoxVecZ.x; \\\n"
            "pos.y -= scale3*periodicBoxVecZ.y; \\\n"
            "pos.z -= scale3*periodicBoxVecZ.z; \\\n"
            "real scale2 = floor(pos.y*invPeriodicBoxSize.y); \\\n"
            "pos.x -= scale2*periodicBoxVecY.x; \\\n"
            "pos.y -= scale2*periodicBoxVecY.y; \\\n"
            "real scale1 = floor(pos.x*invPeriodicBoxSize.x); \\\n"
            "pos.x -= scale1*periodicBoxVecX.x;}";
        compilationDefines["APPLY_PERIODIC_TO_POS_WITH_CENTER(pos, center)"] =
            "{"
            "real scale3 = floor((pos.z-center.z)*invPeriodicBoxSize.z+0.5f); \\\n"
            "pos.x -= scale3*periodicBoxVecZ.x; \\\n"
            "pos.y -= scale3*periodicBoxVecZ.y; \\\n"
            "pos.z -= scale3*periodicBoxVecZ.z; \\\n"
            "real scale2 = floor((pos.y-center.y)*invPeriodicBoxSize.y+0.5f); \\\n"
            "pos.x -= scale2*periodicBoxVecY.x; \\\n"
            "pos.y -= scale2*periodicBoxVecY.y; \\\n"
            "real scale1 = floor((pos.x-center.x)*invPeriodicBoxSize.x+0.5f); \\\n"
            "pos.x -= scale1*periodicBoxVecX.x;}";
    }
    else {
        compilationDefines["APPLY_PERIODIC_TO_DELTA(delta)"] =
            "{"
            "delta.x -= floor(delta.x*invPeriodicBoxSize.x+0.5f)*periodicBoxSize.x; \\\n"
            "delta.y -= floor(delta.y*invPeriodicBoxSize.y+0.5f)*periodicBoxSize.y; \\\n"
            "delta.z -= floor(delta.z*invPeriodicBoxSize.z+0.5f)*periodicBoxSize.z;}";
        compilationDefines["APPLY_PERIODIC_TO_POS(pos)"] =
            "{"
            "pos.x -= floor(pos.x*invPeriodicBoxSize.x)*periodicBoxSize.x; \\\n"
            "pos.y -= floor(pos.y*invPeriodicBoxSize.y)*periodicBoxSize.y; \\\n"
            "pos.z -= floor(pos.z*invPeriodicBoxSize.z)*periodicBoxSize.z;}";
        compilationDefines["APPLY_PERIODIC_TO_POS_WITH_CENTER(pos, center)"] =
            "{"
            "pos.x -= floor((pos.x-center.x)*invPeriodicBoxSize.x+0.5f)*periodicBoxSize.x; \\\n"
            "pos.y -= floor((pos.y-center.y)*invPeriodicBoxSize.y+0.5f)*periodicBoxSize.y; \\\n"
            "pos.z -= floor((pos.z-center.z)*invPeriodicBoxSize.z+0.5f)*periodicBoxSize.z;}";
    }

357
    // Create utilities objects.
358

359
360
    bonded = new CudaBondedUtilities(*this);
    nonbonded = new CudaNonbondedUtilities(*this);
361
362
    integration = new CudaIntegrationUtilities(*this, system);
    expression = new CudaExpressionUtilities(*this);
363
    clearBuffer(posq);
364
365
366
}

CudaContext::~CudaContext() {
367
    pushAsCurrent();
peastman's avatar
peastman committed
368
369
370
371
372
373
374
375
    for (auto force : forces)
        delete force;
    for (auto listener : reorderListeners)
        delete listener;
    for (auto computation : preComputations)
        delete computation;
    for (auto computation : postComputations)
        delete computation;
376
377
    if (pinnedBuffer != NULL)
        cuMemFreeHost(pinnedBuffer);
378
379
380
381
    if (integration != NULL)
        delete integration;
    if (expression != NULL)
        delete expression;
382
383
    if (bonded != NULL)
        delete bonded;
384
385
    if (nonbonded != NULL)
        delete nonbonded;
386
387
    if (contextIsValid && !isLinkedContext)
        cuProfilerStop();
388
    popAsCurrent();
389
    string errorMessage = "Error deleting Context";
390
    if (contextIsValid && !isLinkedContext)
391
        cuCtxDestroy(context);
392
    contextIsValid = false;
393
394
}

395
void CudaContext::initialize() {
396
    ContextSelector selector(*this);
397
    string errorMessage = "Error initializing Context";
398
    int numEnergyBuffers = max(numThreadBlocks*ThreadBlockSize, nonbonded->getNumEnergyBuffers());
399
400
    int multiprocessors;
    CHECK_RESULT2(cuDeviceGetAttribute(&multiprocessors, CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, device), "Error checking GPU properties");
401
    if (useDoublePrecision) {
Peter Eastman's avatar
Peter Eastman committed
402
        energyBuffer.initialize<double>(*this, numEnergyBuffers, "energyBuffer");
403
        energySum.initialize<double>(*this, multiprocessors, "energySum");
404
405
406
407
        int pinnedBufferSize = max(paddedNumAtoms*4, numEnergyBuffers);
        CHECK_RESULT(cuMemHostAlloc(&pinnedBuffer, pinnedBufferSize*sizeof(double), 0));
    }
    else if (useMixedPrecision) {
Peter Eastman's avatar
Peter Eastman committed
408
        energyBuffer.initialize<double>(*this, numEnergyBuffers, "energyBuffer");
409
        energySum.initialize<double>(*this, multiprocessors, "energySum");
410
411
412
413
        int pinnedBufferSize = max(paddedNumAtoms*4, numEnergyBuffers);
        CHECK_RESULT(cuMemHostAlloc(&pinnedBuffer, pinnedBufferSize*sizeof(double), 0));
    }
    else {
Peter Eastman's avatar
Peter Eastman committed
414
        energyBuffer.initialize<float>(*this, numEnergyBuffers, "energyBuffer");
415
        energySum.initialize<float>(*this, multiprocessors, "energySum");
416
417
418
        int pinnedBufferSize = max(paddedNumAtoms*6, numEnergyBuffers);
        CHECK_RESULT(cuMemHostAlloc(&pinnedBuffer, pinnedBufferSize*sizeof(float), 0));
    }
419
420
    for (int i = 0; i < numAtoms; i++) {
        double mass = system.getParticleMass(i);
421
        if (useDoublePrecision || useMixedPrecision)
422
423
424
425
            ((double4*) pinnedBuffer)[i] = make_double4(0.0, 0.0, 0.0, mass == 0.0 ? 0.0 : 1.0/mass);
        else
            ((float4*) pinnedBuffer)[i] = make_float4(0.0f, 0.0f, 0.0f, mass == 0.0 ? 0.0f : (float) (1.0/mass));
    }
Peter Eastman's avatar
Peter Eastman committed
426
    velm.upload(pinnedBuffer);
427
    bonded->initialize(system);
Peter Eastman's avatar
Peter Eastman committed
428
429
    addAutoclearBuffer(force.getDevicePointer(), force.getSize()*force.getElementSize());
    addAutoclearBuffer(energyBuffer.getDevicePointer(), energyBuffer.getSize()*energyBuffer.getElementSize());
430
431
432
    int numEnergyParamDerivs = energyParamDerivNames.size();
    if (numEnergyParamDerivs > 0) {
        if (useDoublePrecision || useMixedPrecision)
Peter Eastman's avatar
Peter Eastman committed
433
            energyParamDerivBuffer.initialize<double>(*this, numEnergyParamDerivs*numEnergyBuffers, "energyParamDerivBuffer");
434
        else
Peter Eastman's avatar
Peter Eastman committed
435
436
            energyParamDerivBuffer.initialize<float>(*this, numEnergyParamDerivs*numEnergyBuffers, "energyParamDerivBuffer");
        addAutoclearBuffer(energyParamDerivBuffer);
437
    }
438
    findMoleculeGroups();
439
    nonbonded->initialize(system);
440
}
441

442
443
void CudaContext::initializeContexts() {
    getPlatformData().initializeContexts(system);
444
445
}

446
447
FFT3D CudaContext::createFFT(int xsize, int ysize, int zsize, bool realToComplex) {
    return FFT3D(new CudaFFT3D(*this, xsize, ysize, zsize, realToComplex));
448
449
}

450
451
452
453
454
void CudaContext::setAsCurrent() {
    if (contextIsValid)
        cuCtxSetCurrent(context);
}

455
456
457
458
459
460
461
462
463
464
465
void CudaContext::pushAsCurrent() {
    if (contextIsValid)
        cuCtxPushCurrent(context);
}

void CudaContext::popAsCurrent() {
    CUcontext popped;
    if (contextIsValid)
        cuCtxPopCurrent(&popped);
}

466
467
468
469
470
CUmodule CudaContext::createModule(const string source, const char* optimizationFlags) {
    return createModule(source, map<string, string>(), optimizationFlags);
}

CUmodule CudaContext::createModule(const string source, const map<string, string>& defines, const char* optimizationFlags) {
471
    string bits = intToString(8*sizeof(void*));
472
473
474
475
    string options = (optimizationFlags == NULL ? defaultOptimizationOptions : string(optimizationFlags));
    stringstream src;
    if (!options.empty())
        src << "// Compilation Options: " << options << endl << endl;
peastman's avatar
peastman committed
476
    for (auto& pair : compilationDefines) {
477
478
479
480
481
482
483
        // Query defines to avoid duplicate variables
        if (defines.find(pair.first) == defines.end()) {
            src << "#define " << pair.first;
            if (!pair.second.empty())
                src << " " << pair.second;
            src << endl;
        }
484
485
486
    }
    if (!compilationDefines.empty())
        src << endl;
487
488
489
490
491
492
493
494
495
496
497
498
    if (useDoublePrecision) {
        src << "typedef double real;\n";
        src << "typedef double2 real2;\n";
        src << "typedef double3 real3;\n";
        src << "typedef double4 real4;\n";
    }
    else {
        src << "typedef float real;\n";
        src << "typedef float2 real2;\n";
        src << "typedef float3 real3;\n";
        src << "typedef float4 real4;\n";
    }
499
500
501
502
503
504
505
506
507
508
509
510
    if (useDoublePrecision || useMixedPrecision) {
        src << "typedef double mixed;\n";
        src << "typedef double2 mixed2;\n";
        src << "typedef double3 mixed3;\n";
        src << "typedef double4 mixed4;\n";
    }
    else {
        src << "typedef float mixed;\n";
        src << "typedef float2 mixed2;\n";
        src << "typedef float3 mixed3;\n";
        src << "typedef float4 mixed4;\n";
    }
511
    src << "typedef unsigned int tileflags;\n";
512
    src << CudaKernelSources::common << endl;
peastman's avatar
peastman committed
513
514
515
516
    for (auto& pair : defines) {
        src << "#define " << pair.first;
        if (!pair.second.empty())
            src << " " << pair.second;
517
518
519
520
521
        src << endl;
    }
    if (!defines.empty())
        src << endl;
    src << source << endl;
522
523
    
    // Determine what architecture to compile for.
524
525
526
527

    int maxCompilerArchitecture;
#if CUDA_VERSION < 11020
    // CUDA versions before 11.2 can't query the compiler to see what it supports.
528
    
529
530
531
532
533
534
535
536
537
    maxCompilerArchitecture = 75;
#else
    int numArchs;
    CHECK_NVRTC_RESULT(nvrtcGetNumSupportedArchs(&numArchs), "Error querying supported architectures");
    vector<int> archs(numArchs);
    CHECK_NVRTC_RESULT(nvrtcGetSupportedArchs(archs.data()), "Error querying supported architectures");
    maxCompilerArchitecture = archs.back();
#endif
    string compileArchitecture = intToString(min(gpuArchitecture, maxCompilerArchitecture));
538

539
    // See whether we already have PTX for this kernel cached.
540

541
542
543
544
545
546
    CSHA1 sha1;
    sha1.Update((const UINT_8*) src.str().c_str(), src.str().size());
    sha1.Final();
    UINT_8 hash[20];
    sha1.GetHash(hash);
    stringstream cacheFile;
547
    cacheFile << cacheDir;
548
549
550
    cacheFile.flags(ios::hex);
    for (int i = 0; i < 20; i++)
        cacheFile << setw(2) << setfill('0') << (int) hash[i];
551
    cacheFile << '_' << compileArchitecture << '_' << bits;
552
553
554
    CUmodule module;
    if (cuModuleLoad(&module, cacheFile.str().c_str()) == CUDA_SUCCESS)
        return module;
555

556
    // Select a name for the output file.
557

558
559
    stringstream tempFileName;
    tempFileName << "openmmTempKernel" << this; // Include a pointer to this context as part of the filename to avoid collisions.
560
    tempFileName << "_" << this_thread::get_id();
561
    string outputFile = (tempDir+tempFileName.str()+".ptx");
562

563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
    // Split the command line flags into an array of options.
    
    string flags = "-arch=compute_"+compileArchitecture+" "+options;
    stringstream flagsStream(flags);
    string flag;
    vector<string> splitFlags;
    while (flagsStream >> flag)
        splitFlags.push_back(flag);
    int numOptions = splitFlags.size();
    vector<const char*> optionsVec(numOptions);
    for (int i = 0; i < numOptions; i++)
        optionsVec[i] = &splitFlags[i][0];
    
    // Compile the program to PTX.
    
    nvrtcProgram program;
    CHECK_NVRTC_RESULT(nvrtcCreateProgram(&program, src.str().c_str(), NULL, 0, NULL, NULL), "Error creating program");
    try {
        nvrtcResult result = nvrtcCompileProgram(program, optionsVec.size(), &optionsVec[0]);
        if (result != NVRTC_SUCCESS) {
            size_t logSize;
            nvrtcGetProgramLogSize(program, &logSize);
            vector<char> log(logSize);
            nvrtcGetProgramLog(program, &log[0]);
            throw OpenMMException("Error compiling program: "+string(&log[0]));
        }
        size_t ptxSize;
        nvrtcGetPTXSize(program, &ptxSize);
        vector<char> ptx(ptxSize);
        nvrtcGetPTX(program, &ptx[0]);
        nvrtcDestroyProgram(&program);
594

595
        // If possible, write the PTX out to a temporary file so we can cache it for later use.
596

597
        bool wroteCache = false;
598
599
        try {
            ofstream out(outputFile.c_str());
600
            out << string(&ptx[0]);
601
            out.close();
602
603
            if (!out.fail())
                wroteCache = true;
604
605
        }
        catch (...) {
606
607
608
609
            // Ignore.
        }
        if (!wroteCache) {
            // An error occurred.  Possibly we don't have permission to write to the temp directory.  Just try to load the module directly.
610

611
612
613
614
            CHECK_RESULT2(cuModuleLoadDataEx(&module, &ptx[0], 0, NULL, NULL), "Error loading CUDA module");
            return module;
        }
    }
615
616
617
    catch (...) {
        nvrtcDestroyProgram(&program);
        throw;
618
    }
619
620
621
622
    try {
        CUresult result = cuModuleLoad(&module, outputFile.c_str());
        if (result != CUDA_SUCCESS) {
            std::stringstream m;
623
            m<<"Error loading CUDA module: "<<getErrorString(result)<<" ("<<result<<")";
624
625
            throw OpenMMException(m.str());
        }
626
627
        if (rename(outputFile.c_str(), cacheFile.str().c_str()) != 0)
            remove(outputFile.c_str());
628
629
630
631
632
633
634
        return module;
    }
    catch (...) {
        remove(outputFile.c_str());
        throw;
    }
}
635
636
637
638
639
640
641
642
643
644
645
646

CUfunction CudaContext::getKernel(CUmodule& module, const string& name) {
    CUfunction function;
    CUresult result = cuModuleGetFunction(&function, module, name.c_str());
    if (result != CUDA_SUCCESS) {
        std::stringstream m;
        m<<"Error creating kernel "<<name<<": "<<getErrorString(result)<<" ("<<result<<")";
        throw OpenMMException(m.str());
    }
    return function;
}

647
648
649
650
651
652
653
vector<ComputeContext*> CudaContext::getAllContexts() {
    vector<ComputeContext*> result;
    for (CudaContext* c : platformData.contexts)
        result.push_back(c);
    return result;
}

Peter Eastman's avatar
Peter Eastman committed
654
655
656
657
double& CudaContext::getEnergyWorkspace() {
    return platformData.contextEnergy[contextIndex];
}

658
659
ComputeQueue CudaContext::createQueue() {
    return shared_ptr<ComputeQueueImpl>(new CudaQueue());
660
661
}

662
663
CUstream CudaContext::getCurrentStream() {
    return dynamic_cast<CudaQueue*>(currentQueue.get())->getStream();
664
665
}

666
667
CudaArray* CudaContext::createArray() {
    return new CudaArray();
668
669
}

670
671
672
673
ComputeEvent CudaContext::createEvent() {
    return shared_ptr<ComputeEventImpl>(new CudaEvent(*this));
}

674
675
676
677
ComputeSort CudaContext::createSort(ComputeSortImpl::SortTrait* trait, unsigned int length, bool uniform) {
    return shared_ptr<ComputeSortImpl>(new CudaSort(*this, trait, length, uniform));
}

678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
ComputeProgram CudaContext::compileProgram(const std::string source, const std::map<std::string, std::string>& defines) {
    CUmodule module = createModule(CudaKernelSources::vectorOps+source, defines);
    return shared_ptr<ComputeProgramImpl>(new CudaProgram(*this, module));
}

CudaArray& CudaContext::unwrap(ArrayInterface& array) const {
    CudaArray* cuarray;
    ComputeArray* wrapper = dynamic_cast<ComputeArray*>(&array);
    if (wrapper != NULL)
        cuarray = dynamic_cast<CudaArray*>(&wrapper->getArray());
    else
        cuarray = dynamic_cast<CudaArray*>(&array);
    if (cuarray == NULL)
        throw OpenMMException("Array argument is not an CudaArray");
    return *cuarray;
693
694
695
}

std::string CudaContext::getErrorString(CUresult result) {
Peter Eastman's avatar
Peter Eastman committed
696
697
698
    const char* message;
    if (cuGetErrorName(result, &message) == CUDA_SUCCESS)
        return string(message);
peastman's avatar
peastman committed
699
    return "CUDA error";
700
701
702
703
704
705
}

void CudaContext::executeKernel(CUfunction kernel, void** arguments, int threads, int blockSize, unsigned int sharedSize) {
    if (blockSize == -1)
        blockSize = ThreadBlockSize;
    int gridSize = std::min((threads+blockSize-1)/blockSize, numThreadBlocks);
706
    CUresult result = cuLaunchKernel(kernel, gridSize, 1, 1, blockSize, 1, 1, sharedSize, getCurrentStream(), arguments, NULL);
707
708
709
710
711
712
713
    if (result != CUDA_SUCCESS) {
        stringstream str;
        str<<"Error invoking kernel: "<<getErrorString(result)<<" ("<<result<<")";
        throw OpenMMException(str.str());
    }
}

714
715
716
int CudaContext::computeThreadBlockSize(double memory) const {
    int maxShared;
    CHECK_RESULT2(cuDeviceGetAttribute(&maxShared, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK, device), "Error querying device property");
717
718
719
720
721
722
723
724
725
    int max = (int) (maxShared/memory);
    if (max < 64)
        return 32;
    int threads = 64;
    while (threads+64 < max)
        threads += 64;
    return threads;
}

726
727
void CudaContext::clearBuffer(ArrayInterface& array) {
    clearBuffer(unwrap(array).getDevicePointer(), array.getSize()*array.getElementSize());
728
729
730
}

void CudaContext::clearBuffer(CUdeviceptr memory, int size) {
731
732
    int words = size/4;
    void* args[] = {&memory, &words};
Peter Eastman's avatar
Peter Eastman committed
733
    executeKernel(clearBufferKernel, args, words, 128);
734
735
}

736
737
void CudaContext::addAutoclearBuffer(ArrayInterface& array) {
    addAutoclearBuffer(unwrap(array).getDevicePointer(), array.getSize()*array.getElementSize());
738
739
}

740
741
void CudaContext::addAutoclearBuffer(CUdeviceptr memory, int size) {
    autoclearBuffers.push_back(memory);
742
    autoclearBufferSizes.push_back(size/4);
743
744
}

745
746
747
748
749
750
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
777
778
779
780
781
782
783
784
785
786
787
void CudaContext::clearAutoclearBuffers() {
    int base = 0;
    int total = autoclearBufferSizes.size();
    while (total-base >= 6) {
        void* args[] = {&autoclearBuffers[base], &autoclearBufferSizes[base],
                        &autoclearBuffers[base+1], &autoclearBufferSizes[base+1],
                        &autoclearBuffers[base+2], &autoclearBufferSizes[base+2],
                        &autoclearBuffers[base+3], &autoclearBufferSizes[base+3],
                        &autoclearBuffers[base+4], &autoclearBufferSizes[base+4],
                        &autoclearBuffers[base+5], &autoclearBufferSizes[base+5]};
        executeKernel(clearSixBuffersKernel, args, max(max(max(max(max(autoclearBufferSizes[base], autoclearBufferSizes[base+1]), autoclearBufferSizes[base+2]), autoclearBufferSizes[base+3]), autoclearBufferSizes[base+4]), autoclearBufferSizes[base+5]), 128);
        base += 6;
    }
    if (total-base == 5) {
        void* args[] = {&autoclearBuffers[base], &autoclearBufferSizes[base],
                        &autoclearBuffers[base+1], &autoclearBufferSizes[base+1],
                        &autoclearBuffers[base+2], &autoclearBufferSizes[base+2],
                        &autoclearBuffers[base+3], &autoclearBufferSizes[base+3],
                        &autoclearBuffers[base+4], &autoclearBufferSizes[base+4]};
        executeKernel(clearFiveBuffersKernel, args, max(max(max(max(autoclearBufferSizes[base], autoclearBufferSizes[base+1]), autoclearBufferSizes[base+2]), autoclearBufferSizes[base+3]), autoclearBufferSizes[base+4]), 128);
    }
    else if (total-base == 4) {
        void* args[] = {&autoclearBuffers[base], &autoclearBufferSizes[base],
                        &autoclearBuffers[base+1], &autoclearBufferSizes[base+1],
                        &autoclearBuffers[base+2], &autoclearBufferSizes[base+2],
                        &autoclearBuffers[base+3], &autoclearBufferSizes[base+3]};
        executeKernel(clearFourBuffersKernel, args, max(max(max(autoclearBufferSizes[base], autoclearBufferSizes[base+1]), autoclearBufferSizes[base+2]), autoclearBufferSizes[base+3]), 128);
    }
    else if (total-base == 3) {
        void* args[] = {&autoclearBuffers[base], &autoclearBufferSizes[base],
                        &autoclearBuffers[base+1], &autoclearBufferSizes[base+1],
                        &autoclearBuffers[base+2], &autoclearBufferSizes[base+2]};
        executeKernel(clearThreeBuffersKernel, args, max(max(autoclearBufferSizes[base], autoclearBufferSizes[base+1]), autoclearBufferSizes[base+2]), 128);
    }
    else if (total-base == 2) {
        void* args[] = {&autoclearBuffers[base], &autoclearBufferSizes[base],
                        &autoclearBuffers[base+1], &autoclearBufferSizes[base+1]};
        executeKernel(clearTwoBuffersKernel, args, max(autoclearBufferSizes[base], autoclearBufferSizes[base+1]), 128);
    }
    else if (total-base == 1) {
        clearBuffer(autoclearBuffers[base], autoclearBufferSizes[base]*4);
    }
}
788

Peter Eastman's avatar
Peter Eastman committed
789
double CudaContext::reduceEnergy() {
Peter Eastman's avatar
Peter Eastman committed
790
    int bufferSize = energyBuffer.getSize();
Peter Eastman's avatar
Peter Eastman committed
791
    int workGroupSize  = 512;
Peter Eastman's avatar
Peter Eastman committed
792
    void* args[] = {&energyBuffer.getDevicePointer(), &energySum.getDevicePointer(), &bufferSize, &workGroupSize};
793
    executeKernel(reduceEnergyKernel, args, workGroupSize*energySum.getSize(), workGroupSize, workGroupSize*energyBuffer.getElementSize());
Peter Eastman's avatar
Peter Eastman committed
794
    energySum.download(pinnedBuffer);
795
796
797
798
799
800
801
802
803
804
    double result = 0;
    if (getUseDoublePrecision() || getUseMixedPrecision()) {
        for (int i = 0; i < energySum.getSize(); i++)
            result += ((double*) pinnedBuffer)[i];
    }
    else {
        for (int i = 0; i < energySum.getSize(); i++)
            result += ((float*) pinnedBuffer)[i];
    }
    return result;
Peter Eastman's avatar
Peter Eastman committed
805
806
}

807
void CudaContext::setCharges(const vector<double>& charges) {
Peter Eastman's avatar
Peter Eastman committed
808
809
    if (!chargeBuffer.isInitialized())
        chargeBuffer.initialize(*this, numAtoms, useDoublePrecision ? sizeof(double) : sizeof(float), "chargeBuffer");
Peter Eastman's avatar
Peter Eastman committed
810
811
812
813
    vector<double> c(numAtoms);
    for (int i = 0; i < numAtoms; i++)
        c[i] = charges[i];
    chargeBuffer.upload(c, true);
Peter Eastman's avatar
Peter Eastman committed
814
    void* args[] = {&chargeBuffer.getDevicePointer(), &posq.getDevicePointer(), &atomIndexDevice.getDevicePointer(), &numAtoms};
815
816
817
    executeKernel(setChargesKernel, args, numAtoms);
}

818
819
820
821
822
823
bool CudaContext::requestPosqCharges() {
    bool allow = !hasAssignedPosqCharges;
    hasAssignedPosqCharges = true;
    return allow;
}

824
825
826
827
828
829
830
831
832
void CudaContext::addEnergyParameterDerivative(const string& param) {
    // See if this parameter has already been registered.
    
    for (int i = 0; i < energyParamDerivNames.size(); i++)
        if (param == energyParamDerivNames[i])
            return;
    energyParamDerivNames.push_back(param);
}

833
834
void CudaContext::flushQueue() {
    cuStreamSynchronize(getCurrentStream());
835
836
}

837
838
839
840
841
842
843
844
845
846
vector<int> CudaContext::getDevicePrecedence() {
    int numDevices;
    CUdevice thisDevice;
    string errorMessage = "Error initializing Context";
    vector<pair<pair<int, int>, int> > devices;

    CHECK_RESULT(cuDeviceGetCount(&numDevices));
    for (int i = 0; i < numDevices; i++) {
        CHECK_RESULT(cuDeviceGet(&thisDevice, i));
        int major, minor, clock, multiprocessors, speed;
847
848
        CHECK_RESULT(cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, thisDevice));
        CHECK_RESULT(cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, thisDevice));
849
850
851
        if (major == 1 && minor < 2)
            continue;

Robert T. McGibbon's avatar
Robert T. McGibbon committed
852
        if ((useDoublePrecision || useMixedPrecision) && (major+0.1*minor < 1.3))
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
            continue;

        CHECK_RESULT(cuDeviceGetAttribute(&clock, CU_DEVICE_ATTRIBUTE_CLOCK_RATE, thisDevice));
        CHECK_RESULT(cuDeviceGetAttribute(&multiprocessors, CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, thisDevice));
        speed = clock*multiprocessors;
        pair<int, int> deviceProperties = std::make_pair(major, speed);
        devices.push_back(std::make_pair(deviceProperties, -i));
    }

    // sort first by compute capability (higher is better), then speed
    // (higher is better), and finally device index (lower is better)
    std::sort(devices.begin(), devices.end());
    std::reverse(devices.begin(), devices.end());

    vector<int> precedence;
    for (int i = 0; i < static_cast<int>(devices.size()); i++) {
        precedence.push_back(-devices[i].second);
    }

    return precedence;
}
874
875
876
877
878
879
880

unsigned int CudaContext::getEventFlags() {
    unsigned int flags = CU_EVENT_DISABLE_TIMING;
    if (useBlockingSync)
        flags += CU_EVENT_BLOCKING_SYNC;
    return flags;
}