OpenCL implementation of parameter derivatives for CustomGBForce

platforms/opencl/include/OpenCLKernels.h

@@ -808,13 +808,15 @@ public:
     void copyParametersToContext(ContextImpl& context, const CustomGBForce& force);
 private:
     double cutoff;
-    bool hasInitializedKernels, needParameterGradient;
+    bool hasInitializedKernels, needParameterGradient, needEnergyParamDerivs;
     int maxTiles, numComputedValues;
     OpenCLContext& cl;
     OpenCLParameterSet* params;
     OpenCLParameterSet* computedValues;
     OpenCLParameterSet* energyDerivs;
     OpenCLParameterSet* energyDerivChain;
+    std::vector<OpenCLParameterSet*> dValuedParam;
+    std::vector<OpenCLArray*> dValue0dParam;
     OpenCLArray* longEnergyDerivs;
     OpenCLArray* globals;
     OpenCLArray* valueBuffers;

platforms/opencl/src/OpenCLBondedUtilities.cpp

@@ -204,7 +204,7 @@ void OpenCLBondedUtilities::initialize(const System& system) {
         for (int i = 0; i < (int) arguments.size(); i++)
             s<<", __global "<<argTypes[i]<<"* customArg"<<(i+1);
         if (energyParameterDerivatives.size() > 0)
-            s<<", __global mixed* energyParamDerivs";
+            s<<", __global mixed* restrict energyParamDerivs";
         s<<") {\n";
         s<<"mixed energy = 0;\n";
         for (int i = 0; i < energyParameterDerivatives.size(); i++)
@@ -219,7 +219,7 @@ void OpenCLBondedUtilities::initialize(const System& system) {
         for (int i = 0; i < energyParameterDerivatives.size(); i++)
             for (int index = 0; index < numDerivs; index++)
                 if (allParamDerivNames[index] == energyParameterDerivatives[i])
-                    s<<"energyParamDerivs[get_global_id(0)*"<<numDerivs<<"+"<<i<<"] += energyParamDeriv"<<i<<";\n";
+                    s<<"energyParamDerivs[get_global_id(0)*"<<numDerivs<<"+"<<index<<"] += energyParamDeriv"<<i<<";\n";
         s<<"}\n";
         map<string, string> defines;
         defines["PADDED_NUM_ATOMS"] = context.intToString(context.getPaddedNumAtoms());

platforms/opencl/src/OpenCLNonbondedUtilities.cpp

@@ -605,7 +605,7 @@ cl::Kernel OpenCLNonbondedUtilities::createInteractionKernel(const string& sourc
         }
     }
     if (energyParameterDerivatives.size() > 0)
-        args << ", __global mixed* energyParamDerivs";
+        args << ", __global mixed* restrict energyParamDerivs";
     replacements["PARAMETER_ARGUMENTS"] = args.str();
     stringstream loadLocal1;
     for (int i = 0; i < (int) params.size(); i++) {
@@ -666,7 +666,7 @@ cl::Kernel OpenCLNonbondedUtilities::createInteractionKernel(const string& sourc
     for (int i = 0; i < energyParameterDerivatives.size(); i++)
         for (int index = 0; index < numDerivs; index++)
             if (allParamDerivNames[index] == energyParameterDerivatives[i])
-                saveDerivs<<"energyParamDerivs[get_global_id(0)*"<<numDerivs<<"+"<<i<<"] += energyParamDeriv"<<i<<";\n";
+                saveDerivs<<"energyParamDerivs[get_global_id(0)*"<<numDerivs<<"+"<<index<<"] += energyParamDeriv"<<i<<";\n";
     replacements["SAVE_DERIVATIVES"] = saveDerivs.str();
     map<string, string> defines;
     if (useCutoff)

platforms/opencl/src/kernels/customGBEnergyN2.cl

@@ -32,6 +32,7 @@ __kernel void computeN2Energy(
     const unsigned int tgx = get_local_id(0) & (TILE_SIZE-1);
     const unsigned int tbx = get_local_id(0) - tgx;
     mixed energy = 0;
+    INIT_PARAM_DERIVS
 
     // First loop: process tiles that contain exclusions.
     
@@ -73,6 +74,7 @@ __kernel void computeN2Energy(
                     atom2 = y*TILE_SIZE+j;
                     real dEdR = 0;
                     real tempEnergy = 0;
+                    const real interactionScale = 0.5f;
 #ifdef USE_EXCLUSIONS
                     bool isExcluded = !(excl & 0x1);
 #endif
@@ -123,6 +125,7 @@ __kernel void computeN2Energy(
                     atom2 = y*TILE_SIZE+tj;
                     real dEdR = 0;
                     real tempEnergy = 0;
+                    const real interactionScale = 1.0f;
 #ifdef USE_EXCLUSIONS
                     bool isExcluded = !(excl & 0x1);
 #endif
@@ -281,6 +284,7 @@ __kernel void computeN2Energy(
                         atom2 = atomIndices[tbx+tj];
                         real dEdR = 0;
                         real tempEnergy = 0;
+                        const real interactionScale = 1.0f;
                         if (atom1 < NUM_ATOMS && atom2 < NUM_ATOMS) {
                             COMPUTE_INTERACTION
                             dEdR /= -r;
@@ -319,6 +323,7 @@ __kernel void computeN2Energy(
                         atom2 = atomIndices[tbx+tj];
                         real dEdR = 0;
                         real tempEnergy = 0;
+                        const real interactionScale = 1.0f;
                         if (atom1 < NUM_ATOMS && atom2 < NUM_ATOMS) {
                             COMPUTE_INTERACTION
                             dEdR /= -r;
@@ -373,4 +378,5 @@ __kernel void computeN2Energy(
         pos++;
     }
     energyBuffer[get_global_id(0)] += energy;
+    SAVE_PARAM_DERIVS
 }

platforms/opencl/src/kernels/customGBEnergyPerParticle.cl

@@ -12,6 +12,7 @@
 __kernel void computePerParticleEnergy(int bufferSize, int numBuffers, __global real4* restrict forceBuffers, __global mixed* restrict energyBuffer, __global const real4* restrict posq
         PARAMETER_ARGUMENTS) {
     mixed energy = 0;
+    INIT_PARAM_DERIVS
     unsigned int index = get_global_id(0);
     while (index < NUM_ATOMS) {
         // Reduce the derivatives
@@ -27,4 +28,5 @@ __kernel void computePerParticleEnergy(int bufferSize, int numBuffers, __global
         index += get_global_size(0);
     }
     energyBuffer[get_global_id(0)] += energy;
+    SAVE_PARAM_DERIVS
 }

platforms/opencl/src/kernels/customGBGradientChainRule.cl

@@ -4,6 +4,7 @@
 
 __kernel void computeGradientChainRuleTerms(__global real4* restrict forceBuffers, __global const real4* restrict posq
         PARAMETER_ARGUMENTS) {
+    INIT_PARAM_DERIVS
     unsigned int index = get_global_id(0);
     while (index < NUM_ATOMS) {
         real4 pos = posq[index];
@@ -12,4 +13,5 @@ __kernel void computeGradientChainRuleTerms(__global real4* restrict forceBuffer
         forceBuffers[index] = force;
         index += get_global_size(0);
     }
+    SAVE_PARAM_DERIVS
 }

platforms/opencl/src/kernels/customGBValueN2.cl

@@ -74,6 +74,7 @@ __kernel void computeN2Value(__global const real4* restrict posq, __local real4*
                         COMPUTE_VALUE
                     }
                     value += tempValue1;
+                    ADD_TEMP_DERIVS1
 #ifdef USE_CUTOFF
                 }
 #endif
@@ -123,6 +124,8 @@ __kernel void computeN2Value(__global const real4* restrict posq, __local real4*
                     }
                     value += tempValue1;
                     local_value[tbx+tj] += tempValue2;
+                    ADD_TEMP_DERIVS1
+                    ADD_TEMP_DERIVS2
 #ifdef USE_CUTOFF
                 }
 #endif
@@ -137,18 +140,23 @@ __kernel void computeN2Value(__global const real4* restrict posq, __local real4*
         // Write results.
 
 #ifdef SUPPORTS_64_BIT_ATOMICS
-        unsigned int offset = x*TILE_SIZE + tgx;
-        atom_add(&global_value[offset], (long) (value*0x100000000));
+        unsigned int offset1 = x*TILE_SIZE + tgx;
+        atom_add(&global_value[offset1], (long) (value*0x100000000));
+        STORE_PARAM_DERIVS1
         if (x != y) {
-            offset = y*TILE_SIZE + tgx;
-            atom_add(&global_value[offset], (long) (local_value[get_local_id(0)]*0x100000000));
+            unsigned int offset2 = y*TILE_SIZE + tgx;
+            atom_add(&global_value[offset2], (long) (local_value[get_local_id(0)]*0x100000000));
+            STORE_PARAM_DERIVS2
         }
 #else
         unsigned int offset1 = x*TILE_SIZE + tgx + warp*PADDED_NUM_ATOMS;
         unsigned int offset2 = y*TILE_SIZE + tgx + warp*PADDED_NUM_ATOMS;
         global_value[offset1] += value;
-        if (x != y)
+        STORE_PARAM_DERIVS1
+        if (x != y) {
             global_value[offset2] += local_value[get_local_id(0)];
+            STORE_PARAM_DERIVS2
+        }
 #endif
     }
 
@@ -292,6 +300,8 @@ __kernel void computeN2Value(__global const real4* restrict posq, __local real4*
                         }
                         value += tempValue1;
                         local_value[tbx+tj] += tempValue2;
+                        ADD_TEMP_DERIVS1
+                        ADD_TEMP_DERIVS2
 #ifdef USE_CUTOFF
                     }
 #endif
@@ -308,15 +318,23 @@ __kernel void computeN2Value(__global const real4* restrict posq, __local real4*
             unsigned int atom2 = y*TILE_SIZE + tgx;
 #endif
 #ifdef SUPPORTS_64_BIT_ATOMICS
-            atom_add(&global_value[atom1], (long) (value*0x100000000));
-            if (atom2 < PADDED_NUM_ATOMS)
-                atom_add(&global_value[atom2], (long) (local_value[get_local_id(0)]*0x100000000));
+            unsigned in offset1 = atom1;
+            atom_add(&global_value[offset1], (long) (value*0x100000000));
+            STORE_PARAM_DERIVS1
+            if (atom2 < PADDED_NUM_ATOMS) {
+                unsigned int offset2 = atom2;
+                atom_add(&global_value[offset2], (long) (local_value[get_local_id(0)]*0x100000000));
+                STORE_PARAM_DERIVS2
+            }
 #else
             unsigned int offset1 = atom1 + warp*PADDED_NUM_ATOMS;
-            unsigned int offset2 = atom2 + warp*PADDED_NUM_ATOMS;
             global_value[offset1] += value;
-            if (atom2 < PADDED_NUM_ATOMS)
+            STORE_PARAM_DERIVS1
+            if (atom2 < PADDED_NUM_ATOMS) {
+                unsigned int offset2 = atom2 + warp*PADDED_NUM_ATOMS;
                 global_value[offset2] += local_value[get_local_id(0)];
+                STORE_PARAM_DERIVS2
+            }
 #endif
         }
         pos++;

platforms/opencl/src/kernels/customGBValuePerParticle.cl

@@ -21,6 +21,7 @@ __kernel void computePerParticleValues(int bufferSize, int numBuffers, __global
         for (int i = index+bufferSize; i < totalSize; i += bufferSize)
             sum += valueBuffers[i];
 #endif
+        REDUCE_PARAM0_DERIV
         
         // Now calculate other values
 

tests/TestCustomBondForce.h

@@ -187,8 +187,8 @@ void testEnergyParameterDerivatives() {
     CustomBondForce* bonds = new CustomBondForce("k*(r-r0)^2");
     bonds->addGlobalParameter("r0", 0.0);
     bonds->addGlobalParameter("k", 0.0);
-    bonds->addEnergyParameterDerivative("r0");
     bonds->addEnergyParameterDerivative("k");
+    bonds->addEnergyParameterDerivative("r0");
     vector<double> parameters;
     bonds->addBond(0, 1, parameters);
     bonds->addBond(1, 2, parameters);

tests/TestCustomGBForce.h

@@ -505,10 +505,10 @@ void testEnergyParameterDerivatives() {
     force->addComputedValue("b", "a+B", CustomGBForce::SingleParticle);
     force->addEnergyTerm("C*(a1+b1+a2+b2+r)^0.8", CustomGBForce::ParticlePair);
     force->addEnergyTerm("(D-B)*b", CustomGBForce::SingleParticle);
-    for (int i = 0; i < numParameters; i++) {
+    for (int i = 0; i < numParameters; i++)
         force->addGlobalParameter(paramNames[i], paramValues[i]);
+    for (int i = numParameters-1; i >= 0; i--)
         force->addEnergyParameterDerivative(paramNames[i]);
-    }
     force->setNonbondedMethod(CustomGBForce::CutoffPeriodic);
     force->setCutoffDistance(1.0);
     vector<Vec3> positions;

tests/TestCustomNonbondedForce.h

@@ -1050,8 +1050,8 @@ void testEnergyParameterDerivatives() {
     CustomNonbondedForce* nonbonded = new CustomNonbondedForce("k*(r-r0)^2");
     nonbonded->addGlobalParameter("r0", 0.0);
     nonbonded->addGlobalParameter("k", 0.0);
-    nonbonded->addEnergyParameterDerivative("r0");
     nonbonded->addEnergyParameterDerivative("k");
+    nonbonded->addEnergyParameterDerivative("r0");
     vector<double> parameters;
     nonbonded->addParticle(parameters);
     nonbonded->addParticle(parameters);