Continuing CUDA implementation of CustomNonbondedForce

platforms/cuda/src/CudaKernels.cpp

@@ -82,7 +82,6 @@ static double calcEnergy(ContextImpl& context, CudaPlatform::PlatformData& data,
             pos[i] = Vec3(posData[3*i], posData[3*i+1], posData[3*i+2]);
         delete[] posData;
         refContext.setPositions(pos);
- //       printf("Total CPU energies: %f\n", refContext.getState(State::Energy).getPotentialEnergy());
         return refContext.getState(State::Energy).getPotentialEnergy();
     }
     else
@@ -97,9 +96,10 @@ static double calcEnergy(ContextImpl& context, CudaPlatform::PlatformData& data,
             kReduceObcGbsaBornForces(gpu);
             kCalculateObcGbsaForces2(gpu);
         }
-        else if (data.hasNonbonded) {
+        else if (data.hasNonbonded)
             kCalculateCDLJForces(gpu);
-        }
+        if (data.hasCustomNonbonded)
+            kCalculateCustomNonbondedForces(gpu, data.hasNonbonded);
         kCalculateLocalForces(gpu);
         if (gpu->bIncludeGBSA)
             kReduceBornSumAndForces(gpu);
@@ -389,8 +389,7 @@ CudaCalcCustomNonbondedForceKernel::~CudaCalcCustomNonbondedForceKernel() {
 }
 
 void CudaCalcCustomNonbondedForceKernel::initialize(const System& system, const CustomNonbondedForce& force) {
-    if (data.primaryKernel == NULL)
-        data.primaryKernel = this;
+    data.primaryKernel = this; // This must always be the primary kernel so it can update the global parameters
     data.hasCustomNonbonded = true;
     numParticles = force.getNumParticles();
     _gpuContext* gpu = data.gpu;
@@ -452,21 +451,45 @@ void CudaCalcCustomNonbondedForceKernel::initialize(const System& system, const
         paramNames.push_back(force.getParameterName(i));
         combiningRules.push_back(force.getParameterCombiningRule(i));
     }
+    globalParamNames.resize(force.getNumGlobalParameters());
+    globalParamValues.resize(force.getNumGlobalParameters());
+    for (int i = 0; i < force.getNumGlobalParameters(); i++) {
+        globalParamNames[i] = force.getGlobalParameterName(i);
+        globalParamValues[i] = force.getGlobalParameterDefaultValue(i);
+    }
     gpuSetCustomNonbondedParameters(gpu, parameters, exclusionList, exceptionParticle1, exceptionParticle2, exceptionParams, method,
-            (float)force.getCutoffDistance(), force.getEnergyFunction(), combiningRules, paramNames);
+            (float)force.getCutoffDistance(), force.getEnergyFunction(), combiningRules, paramNames, globalParamNames);
+    if (globalParamValues.size() > 0)
+        SetCustomNonbondedGlobalParams(globalParamValues);
 }
 
 void CudaCalcCustomNonbondedForceKernel::executeForces(ContextImpl& context) {
-    if (data.primaryKernel == this)
+    if (data.primaryKernel == this) {
+        updateGlobalParams(context);
         calcForces(context, data);
+    }
 }
 
 double CudaCalcCustomNonbondedForceKernel::executeEnergy(ContextImpl& context) {
-    if (data.primaryKernel == this)
+    if (data.primaryKernel == this) {
+        updateGlobalParams(context);
         return calcEnergy(context, data, system);
+    }
     return 0.0;
 }
 
+void CudaCalcCustomNonbondedForceKernel::updateGlobalParams(ContextImpl& context) {
+    bool changed = false;
+    for (int i = 0; i < globalParamNames.size(); i++) {
+        float value = (float) context.getParameter(globalParamNames[i]);
+        if (value != globalParamValues[i])
+            changed = true;
+        globalParamValues[i] = value;
+    }
+    if (changed)
+        SetCustomNonbondedGlobalParams(globalParamValues);
+}
+
 CudaCalcGBSAOBCForceKernel::~CudaCalcGBSAOBCForceKernel() {
 }
 

platforms/cuda/src/CudaKernels.h

@@ -289,8 +289,11 @@ public:
      */
     double executeEnergy(ContextImpl& context);
 private:
+    void updateGlobalParams(ContextImpl& context);
     CudaPlatform::PlatformData& data;
     int numParticles;
+    std::vector<std::string> globalParamNames;
+    std::vector<float> globalParamValues;
     System& system;
 };
 

platforms/cuda/src/kernels/cudaKernels.h

@@ -96,3 +96,4 @@ extern void GetRandomSim(gpuContext gpu);
 extern void SetCustomNonbondedForceExpression(const Expression<128>& expression);
 extern void SetCustomNonbondedEnergyExpression(const Expression<128>& expression);
 extern void SetCustomNonbondedCombiningRules(const Expression<64>* expressions);
+extern void SetCustomNonbondedGlobalParams(const std::vector<float>& paramValues);

platforms/cuda/src/kernels/cudatypes.h

@@ -244,7 +244,7 @@ enum CudaNonbondedMethod
 };
 
 enum ExpressionOp {
-    CONSTANT = 0, VARIABLE0, VARIABLE1, VARIABLE2, VARIABLE3, VARIABLE4, VARIABLE5, VARIABLE6, VARIABLE7, VARIABLE8, CUSTOM, ADD, SUBTRACT, MULTIPLY, DIVIDE,
+    CONSTANT = 0, VARIABLE0, VARIABLE1, VARIABLE2, VARIABLE3, VARIABLE4, VARIABLE5, VARIABLE6, VARIABLE7, VARIABLE8, GLOBAL, CUSTOM, ADD, SUBTRACT, MULTIPLY, DIVIDE,
         POWER, NEGATE, SQRT, EXP, LOG, SIN, COS, SEC, CSC, TAN, COT, ASIN, ACOS, ATAN, SQUARE, CUBE, RECIPROCAL, INCREMENT, DECREMENT
 };
 

platforms/cuda/src/kernels/gpu.cpp

@@ -132,7 +132,7 @@ static const float BOLTZ                    =    (RGAS / KILO);            // (k
 #define DUMP_PARAMETERS 0
 
 template <int SIZE>
-static Expression<SIZE> createExpression(const string& expression, const Lepton::ExpressionProgram& program, const vector<string>& variables) {
+static Expression<SIZE> createExpression(const string& expression, const Lepton::ExpressionProgram& program, const vector<string>& variables, const vector<string>& globalParamNames) {
     Expression<SIZE> exp;
     if (program.getNumOperations() > SIZE)
         throw OpenMMException("Expression contains too many operations: "+expression);
@@ -164,8 +164,14 @@ static Expression<SIZE> createExpression(const string& expression, const Lepton:
                     exp.op[i] = VARIABLE7;
                 else if (variables.size() > 8 && op.getName() == variables[8])
                     exp.op[i] = VARIABLE8;
-                else
+                else {
+                    int j;
+                    for (j = 0; j < globalParamNames.size() && op.getName() != globalParamNames[j]; j++);
+                    if (j == globalParamNames.size())
                         throw OpenMMException("Unknown variable '"+op.getName()+"' in expression: "+expression);
+                    exp.op[i] = GLOBAL;
+                    exp.arg[i] = j;
+                }
                 break;
             case Operation::ADD:
                 exp.op[i] = ADD;
@@ -560,12 +566,15 @@ void gpuSetNonbondedCutoff(gpuContext gpu, float cutoffDistance, float solventDi
 extern "C"
 void gpuSetCustomNonbondedParameters(gpuContext gpu, const vector<vector<double> >& parameters, const vector<vector<int> >& exclusions,
             const vector<int>& exceptionAtom1, const vector<int>& exceptionAtom2, const vector<vector<double> >& exceptionParams,
-            CudaNonbondedMethod method, float cutoffDistance, const string& energyExp, const vector<string>& combiningRules, const vector<string>& paramNames)
+            CudaNonbondedMethod method, float cutoffDistance, const string& energyExp, const vector<string>& combiningRules,
+            const vector<string>& paramNames, const vector<string>& globalParamNames)
 {
     if (gpu->sim.nonbondedCutoff != 0.0f && gpu->sim.nonbondedCutoff != cutoffDistance)
         throw OpenMMException("All nonbonded forces must use the same cutoff");
     if (paramNames.size() > 4)
         throw OpenMMException("CudaPlatform only supports four per-atom parameters for custom nonbonded forces");
+    if (globalParamNames.size() > 8)
+        throw OpenMMException("CudaPlatform only supports eight global parameters for custom nonbonded forces");
     gpu->sim.nonbondedCutoff = cutoffDistance;
     gpu->sim.nonbondedCutoffSqr = cutoffDistance*cutoffDistance;
     gpu->sim.customNonbondedMethod = method;
@@ -612,8 +621,8 @@ void gpuSetCustomNonbondedParameters(gpuContext gpu, const vector<vector<double>
     variables.push_back("r");
     for (int i = 0; i < paramNames.size(); i++)
         variables.push_back(paramNames[i]);
-    SetCustomNonbondedEnergyExpression(createExpression<128>(energyExp, Lepton::Parser::parse(energyExp).optimize().createProgram(), variables));
-    SetCustomNonbondedForceExpression(createExpression<128>(energyExp, Lepton::Parser::parse(energyExp).differentiate("r").optimize().createProgram(), variables));
+    SetCustomNonbondedEnergyExpression(createExpression<128>(energyExp, Lepton::Parser::parse(energyExp).optimize().createProgram(), variables, globalParamNames));
+    SetCustomNonbondedForceExpression(createExpression<128>(energyExp, Lepton::Parser::parse(energyExp).differentiate("r").optimize().createProgram(), variables, globalParamNames));
     Expression<64> paramExpressions[4];
     vector<string> combiningRuleParams;
     combiningRuleParams.push_back("");
@@ -627,7 +636,7 @@ void gpuSetCustomNonbondedParameters(gpuContext gpu, const vector<vector<double>
             combiningRuleParams.push_back("");
     }
     for (int i = 0; i < paramNames.size(); i++)
-        paramExpressions[i] = createExpression<64>(combiningRules[i], Lepton::Parser::parse(combiningRules[i]).optimize().createProgram(), combiningRuleParams);
+        paramExpressions[i] = createExpression<64>(combiningRules[i], Lepton::Parser::parse(combiningRules[i]).optimize().createProgram(), combiningRuleParams, globalParamNames);
     SetCustomNonbondedCombiningRules(paramExpressions);
 }
 

platforms/cuda/src/kernels/gputypes.h

@@ -186,7 +186,7 @@ extern "C"
 void gpuSetCustomNonbondedParameters(gpuContext gpu, const std::vector<std::vector<double> >& parameters, const std::vector<std::vector<int> >& exclusions,
             const std::vector<int>& exceptionAtom1, const std::vector<int>& exceptionAtom2, const std::vector<std::vector<double> >& exceptionParams,
             CudaNonbondedMethod method, float cutoffDistance, const std::string& energyExp, const std::vector<std::string>& combiningRules,
-            const std::vector<std::string>& paramNames);
+            const std::vector<std::string>& paramNames, const std::vector<std::string>& globalParamNames);
 
 extern "C"
 void gpuSetEwaldParameters(gpuContext gpu, float alpha, int kmaxx, int kmaxy, int kmaxz);

platforms/cuda/src/kernels/kCalculateCustomNonbondedForces.cu

@@ -53,6 +53,7 @@ static __constant__ cudaGmxSimulation cSim;
 static __constant__ Expression<128> forceExp;
 static __constant__ Expression<128> energyExp;
 static __constant__ Expression<64> combiningRules[4];
+static __constant__ float globalParams[8];
 
 void SetCalculateCustomNonbondedForcesSim(gpuContext gpu)
 {
@@ -89,6 +90,13 @@ void SetCustomNonbondedCombiningRules(const Expression<64>* expressions)
     RTERROR(status, "SetCustomNonbondedCombiningRules: cudaMemcpyToSymbol failed");
 }
 
+void SetCustomNonbondedGlobalParams(const vector<float>& paramValues)
+{
+    cudaError_t status;
+    status = cudaMemcpyToSymbol(globalParams, &paramValues[0], sizeof(globalParams));
+    RTERROR(status, "SetCustomNonbondedGlobalParams: cudaMemcpyToSymbol failed");
+}
+
 template<int SIZE>
 __device__ float kEvaluateExpression_kernel(Expression<SIZE>* expression, float* stack, float var0, float4 vars1, float4 vars2)
 {
@@ -127,6 +135,9 @@ __device__ float kEvaluateExpression_kernel(Expression<SIZE>* expression, float*
             case VARIABLE8:
                 stack[++stackPointer] = vars2.w;
                 break;
+            case GLOBAL:
+                stack[++stackPointer] = globalParams[(int) expression->arg[i]];
+                break;
             case ADD:
             {
                 float temp = stack[stackPointer];

platforms/cuda/src/kernels/kCalculateCustomNonbondedForces.h

@@ -39,6 +39,7 @@ __global__ void METHOD_NAME(kCalculateCustomNonbonded, Forces_kernel)(unsigned i
     unsigned int numWorkUnits = cSim.pInteractionCount[0];
     unsigned int pos = warp*numWorkUnits/totalWarps;
     unsigned int end = (warp+1)*numWorkUnits/totalWarps;
+    float totalEnergy = 0.0f;
 #ifdef USE_CUTOFF
     float3* tempBuffer = (float3*) &sA[cSim.nonbond_threads_per_block];
 #endif
@@ -112,6 +113,7 @@ __global__ void METHOD_NAME(kCalculateCustomNonbonded, Forces_kernel)(unsigned i
                 float r         = sqrt(dx*dx + dy*dy + dz*dz);
                 float invR      = 1.0f/r;
                 float dEdR      = -kEvaluateExpression_kernel(&forceExp, &stack[MAX_STACK_SIZE*threadIdx.x], r, combinedParams, combinedParams)*invR;
+                float energy    = kEvaluateExpression_kernel(&energyExp, &stack[MAX_STACK_SIZE*threadIdx.x], r, combinedParams, combinedParams);
 #ifdef USE_CUTOFF
                 if (!(excl & 0x1) || r > cSim.nonbondedCutoff)
 #else
@@ -119,7 +121,9 @@ __global__ void METHOD_NAME(kCalculateCustomNonbonded, Forces_kernel)(unsigned i
 #endif
                 {
                     dEdR = 0.0f;
+                    energy = 0.0f;
                 }
+                totalEnergy    += 0.5f*energy;
                 dx             *= dEdR;
                 dy             *= dEdR;
                 dz             *= dEdR;
@@ -213,12 +217,15 @@ __global__ void METHOD_NAME(kCalculateCustomNonbonded, Forces_kernel)(unsigned i
                         float r         = sqrt(dx*dx + dy*dy + dz*dz);
                         float invR      = 1.0f/r;
                         float dEdR      = -kEvaluateExpression_kernel(&forceExp, &stack[MAX_STACK_SIZE*threadIdx.x], r, combinedParams, combinedParams)*invR;
+                        float energy    = kEvaluateExpression_kernel(&energyExp, &stack[MAX_STACK_SIZE*threadIdx.x], r, combinedParams, combinedParams);
 #ifdef USE_CUTOFF
                         if (r > cSim.nonbondedCutoff)
                         {
                             dEdR = 0.0f;
+                            energy = 0.0f;
                         }
 #endif
+                        totalEnergy    += energy;
                         dx             *= dEdR;
                         dy             *= dEdR;
                         dz             *= dEdR;
@@ -276,12 +283,15 @@ __global__ void METHOD_NAME(kCalculateCustomNonbonded, Forces_kernel)(unsigned i
                             float r         = sqrt(dx*dx + dy*dy + dz*dz);
                             float invR      = 1.0f/r;
                             float dEdR      = -kEvaluateExpression_kernel(&forceExp, &stack[MAX_STACK_SIZE*threadIdx.x], r, combinedParams, combinedParams)*invR;
+                            float energy    = kEvaluateExpression_kernel(&energyExp, &stack[MAX_STACK_SIZE*threadIdx.x], r, combinedParams, combinedParams);
 #ifdef USE_CUTOFF
                             if (r > cSim.nonbondedCutoff)
                             {
                                 dEdR = 0.0f;
+                                energy = 0.0f;
                             }
 #endif
+                            totalEnergy    += energy;
                             dx             *= dEdR;
                             dy             *= dEdR;
                             dz             *= dEdR;
@@ -375,6 +385,7 @@ __global__ void METHOD_NAME(kCalculateCustomNonbonded, Forces_kernel)(unsigned i
                     float r         = sqrt(dx*dx + dy*dy + dz*dz);
                     float invR      = 1.0f/r;
                     float dEdR      = -kEvaluateExpression_kernel(&forceExp, &stack[MAX_STACK_SIZE*threadIdx.x], r, combinedParams, combinedParams)*invR;
+                    float energy    = kEvaluateExpression_kernel(&energyExp, &stack[MAX_STACK_SIZE*threadIdx.x], r, combinedParams, combinedParams);
 #ifdef USE_CUTOFF
                     if (!(excl & 0x1) || r > cSim.nonbondedCutoff)
 #else
@@ -382,7 +393,9 @@ __global__ void METHOD_NAME(kCalculateCustomNonbonded, Forces_kernel)(unsigned i
 #endif
                     {
                         dEdR = 0.0f;
+                        energy = 0.0f;
                     }
+                    totalEnergy    += energy;
                     dx             *= dEdR;
                     dy             *= dEdR;
                     dz             *= dEdR;
@@ -430,4 +443,5 @@ __global__ void METHOD_NAME(kCalculateCustomNonbonded, Forces_kernel)(unsigned i
 
         pos++;
     }
+    cSim.pEnergy[blockIdx.x*blockDim.x+threadIdx.x] += totalEnergy;
 }

platforms/cuda/tests/TestCudaCustomNonbondedForce.cpp

@@ -205,7 +205,7 @@ void testPeriodic() {
 int main() {
     try {
         testSimpleExpression();
-//        testParameters();
+        testParameters();
 //        testExceptions();
         testCutoff();
         testPeriodic();