Refactored nonbonded kernel to allow the same code to be used for other forces

platforms/opencl/src/OpenCLContext.cpp

@@ -121,6 +121,10 @@ void OpenCLContext::addForce(OpenCLForceInfo* force) {
 }
 
 string OpenCLContext::loadSourceFromFile(const string& filename) const {
+    return loadSourceFromFile(filename, map<string, string>());
+}
+
+string OpenCLContext::loadSourceFromFile(const string& filename, const std::map<std::string, std::string>& replacements) const {
     ifstream file((Platform::getDefaultPluginsDirectory()+"/opencl/"+filename).c_str());
     if (!file.is_open())
         throw OpenMMException("Unable to load kernel: "+filename);
@@ -132,6 +136,14 @@ string OpenCLContext::loadSourceFromFile(const string& filename) const {
         kernel += '\n';
     }
     file.close();
+    for (map<string, string>::const_iterator iter = replacements.begin(); iter != replacements.end(); iter++) {
+        int index = -1;
+        do {
+            index = kernel.find(iter->first);
+            if (index != kernel.npos)
+                kernel.replace(index, iter->first.size(), iter->second);
+        } while (index != kernel.npos);
+    }
     return kernel;
 }
 

platforms/opencl/src/OpenCLContext.h

@@ -137,6 +137,13 @@ public:
      * Load OpenCL source code from a file in the kernels directory.
      */
     std::string loadSourceFromFile(const std::string& filename) const;
+    /**
+     * Load OpenCL source code from a file in the kernels directory.
+     *
+     * @param filename     the file to load
+     * @param replacements a set of strings that should be replaced with new strings wherever they appear in the
+     */
+    std::string loadSourceFromFile(const std::string& filename, const std::map<std::string, std::string>& replacements) const;
     /**
      * Create an OpenCL Program from source code.
      */

platforms/opencl/src/OpenCLKernels.cpp

@@ -576,7 +576,8 @@ void OpenCLCalcNonbondedForceKernel::initialize(const System& system, const Nonb
 //    }
 //    data.nonbondedMethod = method;
 //    gpuSetCoulombParameters(gpu, 138.935485f, particle, c6, c12, q, symbol, exclusionList, method);
-    cl.getNonbondedUtilities().addInteraction(useCutoff, usePeriodic, force.getCutoffDistance(), exclusionList, defines);
+    string source = cl.loadSourceFromFile("coulombLennardJones.cl", defines);
+    cl.getNonbondedUtilities().addInteraction(useCutoff, usePeriodic, force.getCutoffDistance(), exclusionList, source);
     cl.getNonbondedUtilities().addParameter("sigmaEpsilon", "float2", 8, sigmaEpsilon->getDeviceBuffer());
     cutoffSquared = force.getCutoffDistance()*force.getCutoffDistance();
 

platforms/opencl/src/OpenCLNonbondedUtilities.cpp

@@ -68,7 +68,7 @@ OpenCLNonbondedUtilities::~OpenCLNonbondedUtilities() {
         delete compact;
 }
 
-void OpenCLNonbondedUtilities::addInteraction(bool usesCutoff, bool usesPeriodic, double cutoffDistance, const vector<vector<int> >& exclusionList, const map<string, string>& defines) {
+void OpenCLNonbondedUtilities::addInteraction(bool usesCutoff, bool usesPeriodic, double cutoffDistance, const vector<vector<int> >& exclusionList, const string& kernel) {
     if (cutoff != -1.0) {
         if (usesCutoff != useCutoff)
             throw OpenMMException("All Forces must agree on whether to use a cutoff");
@@ -86,18 +86,13 @@ void OpenCLNonbondedUtilities::addInteraction(bool usesCutoff, bool usesPeriodic
         }
         if (!sameExclusions)
             throw OpenMMException("All Forces must have identical exceptions");
-        for (map<string, string>::const_iterator iter = defines.begin(); iter != defines.end(); ++iter) {
-            map<string, string>::const_iterator existing = defines.find(iter->first);
-            if (existing != defines.end() && existing->second != iter->second)
-                throw OpenMMException("Two Forces define different values for "+iter->first);
-        }
     }
     else {
         useCutoff = usesCutoff;
         usePeriodic = usesPeriodic;
         cutoff = cutoffDistance;
         atomExclusions = exclusionList;
-        kernelDefines.insert(defines.begin(), defines.end());
+        kernelSource += kernel+"\n";
     }
 }
 
@@ -122,14 +117,58 @@ void OpenCLNonbondedUtilities::initialize(const System& system) {
 
     // Create kernels.
 
-    map<string, string> defines = kernelDefines;
+    map<string, string> replacements;
+    replacements["COMPUTE_INTERACTION"] = kernelSource;
+    stringstream args;
+    for (int i = 0; i < parameters.size(); i++) {
+        args << ", __global ";
+        args << parameters[i].type;
+        args << "* global_";
+        args << parameters[i].name;
+        args << ", __local ";
+        args << parameters[i].type;
+        args << "* local_";
+        args << parameters[i].name;
+    }
+    replacements["PARAMETER_ARGUMENTS"] = args.str();
+    stringstream load1;
+    for (int i = 0; i < parameters.size(); i++) {
+        load1 << parameters[i].type;
+        load1 << " ";
+        load1 << parameters[i].name;
+        load1 << "1 = global_";
+        load1 << parameters[i].name;
+        load1 << "[i];";
+    }
+    replacements["LOAD_ATOM1_PARAMETERS"] = load1.str();
+    stringstream load2j;
+    for (int i = 0; i < parameters.size(); i++) {
+        load2j << parameters[i].type;
+        load2j << " ";
+        load2j << parameters[i].name;
+        load2j << "2 = local_";
+        load2j << parameters[i].name;
+        load2j << "[tbx+j];";
+    }
+    replacements["LOAD_ATOM2_PARAMETERS_J"] = load2j.str();
+    stringstream load2tj;
+    for (int i = 0; i < parameters.size(); i++) {
+        load2tj << parameters[i].type;
+        load2tj << " ";
+        load2tj << parameters[i].name;
+        load2tj << "2 = local_";
+        load2tj << parameters[i].name;
+        load2tj << "[tbx+tj];";
+    }
+    replacements["LOAD_ATOM2_PARAMETERS_TJ"] = load2tj.str();
+    map<string, string> defines;
     if (forceBufferPerAtomBlock)
         defines["USE_OUTPUT_BUFFER_PER_BLOCK"] = "1";
     if (useCutoff)
         defines["USE_CUTOFF"] = "1";
     if (usePeriodic)
         defines["USE_PERIODIC"] = "1";
-    cl::Program forceProgram = context.createProgram(context.loadSourceFromFile("nonbonded.cl"), defines);
+    cl::Program forceProgram = context.createProgram(context.loadSourceFromFile("nonbonded.cl", replacements), defines);
     forceKernel = cl::Kernel(forceProgram, "computeNonbonded");
     cl::Program interactingBlocksProgram = context.createProgram(context.loadSourceFromFile("findInteractingBlocks.cl"), defines);
     findBlockBoundsKernel = cl::Kernel(interactingBlocksProgram, "findBlockBounds");

platforms/opencl/src/OpenCLNonbondedUtilities.h

@@ -52,9 +52,9 @@ public:
      * @param usesPeriodic   specifies whether periodic boundary conditions should be applied to this interaction
      * @param cutoffDistance the cutoff distance for this interaction (ignored if usesCutoff is false)
      * @param exclusionList  for each atom, specifies the list of other atoms whose interactions should be excluded
-     * @param defines        preprocessor macros to define when compiling the kernel
+     * @param kernel         the code to evaluate the interaction
      */
-    void addInteraction(bool usesCutoff, bool usesPeriodic, double cutoffDistance, const std::vector<std::vector<int> >& exclusionList, const std::map<std::string, std::string>& defines);
+    void addInteraction(bool usesCutoff, bool usesPeriodic, double cutoffDistance, const std::vector<std::vector<int> >& exclusionList, const std::string& kernel);
     /**
      * Add a per-atom parameter that interactions may depend on.
      *
@@ -143,6 +143,7 @@ private:
     std::vector<std::vector<int> > atomExclusions;
     std::vector<ParameterInfo> parameters;
     OpenCLCompact* compact;
+    std::string kernelSource;
     std::map<std::string, std::string> kernelDefines;
     double cutoff;
     bool useCutoff, usePeriodic, forceBufferPerAtomBlock, hasComputedInteractions;

platforms/opencl/src/kernels/coulombLennardJones.cl

+{
+    float sig = sigmaEpsilon1.x + sigmaEpsilon2.x;
+    float sig2 = invR*sig;
+    sig2 *= sig2;
+    float sig6 = sig2*sig2*sig2;
+    float eps = sigmaEpsilon1.y*sigmaEpsilon2.y;
+    float tempForce = eps*(12.0f*sig6 - 6.0f)*sig6;
+    tempEnergy += eps*(sig6 - 1.0f)*sig6;
+    const float EpsilonFactor = 138.935485f;
+#ifdef USE_CUTOFF
+    tempForce += EpsilonFactor*posq1.w*posq2.w*(invR - 2.0f*REACTION_FIELD_K*r2);
+    tempEnergy += EpsilonFactor*posq1.w*posq2.w*(invR + REACTION_FIELD_K*r2 - REACTION_FIELD_C);
+#else
+    tempForce += EpsilonFactor*posq1.w*posq2.w*invR;
+    tempEnergy += EpsilonFactor*posq1.w*posq2.w*invR;
+#endif
+    dEdR += tempForce*invR*invR;
+}

platforms/opencl/src/kernels/nonbonded.cl

 const unsigned int TileSize = 32;
-const float EpsilonFactor = 138.935485f;
 
 /**
  * Compute nonbonded interactions.
@@ -7,11 +6,11 @@ const float EpsilonFactor = 138.935485f;
 
 __kernel void computeNonbonded(int numTiles, int paddedNumAtoms,
         __global float4* forceBuffers, __global float* energyBuffer, __global float4* posq, __global unsigned int* tiles,
-        __global unsigned int* exclusions,  __global unsigned int* exclusionIndices, __local float4* local_posq, __local float4* local_force,
+        __global unsigned int* exclusions,  __global unsigned int* exclusionIndices, __local float4* local_posq, __local float4* local_force
 #ifdef USE_CUTOFF
-        float cutoffSquared, float4 periodicBoxSize, __global unsigned int* interactionFlags, __local float4* tempBuffer,
+        , float cutoffSquared, float4 periodicBoxSize, __global unsigned int* interactionFlags, __local float4* tempBuffer
 #endif
-        __global float2* sigmaEpsilon, __local float2* local_sigmaEpsilon) {
+        PARAMETER_ARGUMENTS) {
     unsigned int totalWarps = get_global_size(0)/TileSize;
     unsigned int warp = get_global_id(0)/TileSize;
     unsigned int pos = warp*numTiles/totalWarps;
@@ -20,54 +19,42 @@ __kernel void computeNonbonded(int numTiles, int paddedNumAtoms,
     unsigned int lasty = 0xFFFFFFFF;
 
     while (pos < end) {
-        // Extract tile coordinates from appropriate work unit
+        // Extract the coordinates of this tile
         unsigned int x = tiles[pos];
         unsigned int y = ((x >> 2) & 0x7fff)*TileSize;
         bool hasExclusions = (x & 0x1);
         x = (x>>17)*TileSize;
-        float4 apos;   // Local atom x, y, z, q
-        float4 af = 0.0f;     // Local atom fx, fy, fz
         unsigned int tgx = get_local_id(0) & (TileSize-1);
         unsigned int tbx = get_local_id(0) - tgx;
         unsigned int tj = tgx;
         unsigned int i = x + tgx;
-        apos = posq[i];
-        float2 a = sigmaEpsilon[i];
+        float4 force = 0.0f;
+        float4 posq1 = posq[i];
+        LOAD_ATOM1_PARAMETERS
         if (x == y) {
-            // Handle diagonals uniquely at 50% efficiency
-            // Read fixed atom data into registers and GRF
+            // This tile is on the diagonal.
 
-            local_posq[get_local_id(0)] = apos;
-            local_sigmaEpsilon[get_local_id(0)] = a;
-            apos.w *= EpsilonFactor;
+            local_posq[get_local_id(0)] = posq1;
+            local_sigmaEpsilon[get_local_id(0)] = sigmaEpsilon1;
             unsigned int xi = x/TileSize;
             unsigned int tile = xi+xi*paddedNumAtoms/TileSize-xi*(xi+1)/2;
             unsigned int excl = exclusions[exclusionIndices[tile]+tgx];
             for (unsigned int j = 0; j < TileSize; j++) {
                 bool isExcluded = !(excl & 0x1);
-                float4 delta = (float4) (local_posq[tbx+j].xyz - apos.xyz, 0.0f);
+                float4 delta = (float4) (local_posq[tbx+j].xyz - posq1.xyz, 0.0f);
 #ifdef USE_PERIODIC
                 delta.x -= floor(delta.x/periodicBoxSize.x+0.5f)*periodicBoxSize.x;
                 delta.y -= floor(delta.y/periodicBoxSize.y+0.5f)*periodicBoxSize.y;
                 delta.z -= floor(delta.z/periodicBoxSize.z+0.5f)*periodicBoxSize.z;
 #endif
-                float r2 = delta.x * delta.x + delta.y * delta.y + delta.z * delta.z;
-                float invR = 1.0f / sqrt(r2);
-                float sig = a.x + local_sigmaEpsilon[tbx+j].x;
-                float sig2 = invR * sig;
-                sig2 *= sig2;
-                float sig6 = sig2 * sig2 * sig2;
-                float eps = a.y * local_sigmaEpsilon[tbx+j].y;
-                float dEdR = eps * (12.0f * sig6 - 6.0f) * sig6;
-                float tempEnergy = eps * (sig6 - 1.0f) * sig6;
-#ifdef USE_CUTOFF
-                dEdR += apos.w * local_posq[tbx+j].w * (invR - 2.0f * REACTION_FIELD_K * r2);
-                tempEnergy += apos.w * local_posq[tbx+j].w * (invR + REACTION_FIELD_K * r2 - REACTION_FIELD_C);
-#else
-                dEdR += apos.w * local_posq[tbx+j].w * invR;
-                tempEnergy += apos.w * local_posq[tbx+j].w * invR;
-#endif
-                dEdR *= invR * invR;
+                float r2 = delta.x*delta.x + delta.y*delta.y + delta.z*delta.z;
+                float r = sqrt(r2);
+                float invR = 1.0f / r;
+                float4 posq2 = local_posq[tbx+j];
+                LOAD_ATOM2_PARAMETERS_J
+                float dEdR = 0.0f;
+                float tempEnergy = 0.0f;
+                COMPUTE_INTERACTION
 #ifdef USE_CUTOFF
                 if (isExcluded || r2 > cutoffSquared) {
 #else
@@ -78,35 +65,33 @@ __kernel void computeNonbonded(int numTiles, int paddedNumAtoms,
                 }
                 energy += 0.5f*tempEnergy;
                 delta.xyz *= dEdR;
-                af.xyz -= delta.xyz;
+                force.xyz -= delta.xyz;
                 excl >>= 1;
             }
 
             // Write results
             float4 of;
 #ifdef USE_OUTPUT_BUFFER_PER_BLOCK
-            of.xyz = af.xyz;
+            of.xyz = force.xyz;
             of.w = 0.0f;
-            unsigned int offset = x + tgx + (x/TileSize) * paddedNumAtoms;
+            unsigned int offset = x + tgx + (x/TileSize)*paddedNumAtoms;
             forceBuffers[offset] = of;
 #else
             unsigned int offset = x + tgx + warp*paddedNumAtoms;
             of = forceBuffers[offset];
-            of.xyz += af.xyz;
+            of.xyz += force.xyz;
             forceBuffers[offset] = of;
 #endif
         }
         else {
-            // 100% utilization
-            // Read fixed atom data into registers and GRF
+            // This is an off-diagonal tile.
+
             if (lasty != y) {
                 unsigned int j = y + tgx;
-                float2 temp1 = sigmaEpsilon[j];
                 local_posq[get_local_id(0)] = posq[j];
-                local_sigmaEpsilon[get_local_id(0)] = sigmaEpsilon[j];
+                local_sigmaEpsilon[get_local_id(0)] = global_sigmaEpsilon[j];
             }
             local_force[get_local_id(0)] = 0.0f;
-            apos.w *= EpsilonFactor;
 #ifdef USE_CUTOFF
             unsigned int flags = interactionFlags[pos];
             if (!hasExclusions && flags != 0xFFFFFFFF) {
@@ -119,29 +104,20 @@ __kernel void computeNonbonded(int numTiles, int paddedNumAtoms,
                     for (unsigned int j = 0; j < TileSize; j++) {
                         if ((flags&(1<<j)) != 0) {
                             bool isExcluded = false;
-                            float4 delta = (float4) (local_posq[tbx+j].xyz - apos.xyz, 0.0f);
+                            float4 delta = (float4) (local_posq[tbx+j].xyz - posq1.xyz, 0.0f);
 #ifdef USE_PERIODIC
                             delta.x -= floor(delta.x/periodicBoxSize.x+0.5f)*periodicBoxSize.x;
                             delta.y -= floor(delta.y/periodicBoxSize.y+0.5f)*periodicBoxSize.y;
                             delta.z -= floor(delta.z/periodicBoxSize.z+0.5f)*periodicBoxSize.z;
 #endif
-                            float r2 = delta.x * delta.x + delta.y * delta.y + delta.z * delta.z;
-                            float invR = 1.0f / sqrt(r2);
-                            float sig = a.x + local_sigmaEpsilon[tbx+j].x;
-                            float sig2 = invR * sig;
-                            sig2 *= sig2;
-                            float sig6 = sig2 * sig2 * sig2;
-                            float eps = a.y * local_sigmaEpsilon[tbx+j].y;
-                            float dEdR = eps * (12.0f * sig6 - 6.0f) * sig6;
-			    float tempEnergy = eps * (sig6 - 1.0f) * sig6;
-#ifdef USE_CUTOFF
-                            dEdR += apos.w * local_posq[tbx+j].w * (invR - 2.0f * REACTION_FIELD_K * r2);
-                            tempEnergy += apos.w * local_posq[tbx+j].w * (invR + REACTION_FIELD_K * r2 - REACTION_FIELD_C);
-#else
-                            dEdR += apos.w * local_posq[tbx+j].w * invR;
-                            tempEnergy += apos.w * local_posq[tbx+j].w * invR;
-#endif
-                            dEdR *= invR * invR;
+                            float r2 = delta.x*delta.x + delta.y*delta.y + delta.z*delta.z;
+                            float r = sqrt(r2);
+                            float invR = 1.0f / r;
+                            float4 posq2 = local_posq[tbx+j];
+                            LOAD_ATOM2_PARAMETERS_J
+                            float dEdR = 0.0f;
+                            float tempEnergy = 0.0f;
+                            COMPUTE_INTERACTION
 #ifdef USE_CUTOFF
                             if (r2 > cutoffSquared) {
                                 dEdR = 0.0f;
@@ -150,7 +126,7 @@ __kernel void computeNonbonded(int numTiles, int paddedNumAtoms,
 #endif
 			    energy += tempEnergy;
                             delta.xyz *= dEdR;
-                            af.xyz -= delta.xyz;
+                            force.xyz -= delta.xyz;
                             tempBuffer[get_local_id(0)] = delta;
 
                             // Sum the forces on atom j.
@@ -172,7 +148,8 @@ __kernel void computeNonbonded(int numTiles, int paddedNumAtoms,
             else
 #endif
             {
-                // Read fixed atom data into registers and GRF
+                // Compute the full set of interactions in this tile.
+
                 unsigned int xi = x/TileSize;
                 unsigned int yi = y/TileSize;
                 unsigned int tile = xi+yi*paddedNumAtoms/TileSize-yi*(yi+1)/2;
@@ -180,29 +157,20 @@ __kernel void computeNonbonded(int numTiles, int paddedNumAtoms,
                 excl = (excl >> tgx) | (excl << (TileSize - tgx));
                 for (unsigned int j = 0; j < TileSize; j++) {
                     bool isExcluded = !(excl & 0x1);
-                    float4 delta = (float4) (local_posq[tbx+tj].xyz - apos.xyz, 0.0f);
+                    float4 delta = (float4) (local_posq[tbx+tj].xyz - posq1.xyz, 0.0f);
 #ifdef USE_PERIODIC
                     delta.x -= floor(delta.x/periodicBoxSize.x+0.5f)*periodicBoxSize.x;
                     delta.y -= floor(delta.y/periodicBoxSize.y+0.5f)*periodicBoxSize.y;
                     delta.z -= floor(delta.z/periodicBoxSize.z+0.5f)*periodicBoxSize.z;
 #endif
-                    float r2 = delta.x * delta.x + delta.y * delta.y + delta.z * delta.z;
-                    float invR = 1.0f / sqrt(r2);
-                    float sig = a.x + local_sigmaEpsilon[tbx+tj].x;
-                    float sig2 = invR * sig;
-                    sig2 *= sig2;
-                    float sig6 = sig2 * sig2 * sig2;
-                    float eps = a.y * local_sigmaEpsilon[tbx+tj].y;
-                    float dEdR = eps * (12.0f * sig6 - 6.0f) * sig6;
-		    float tempEnergy = eps * (sig6 - 1.0f) * sig6;
-#ifdef USE_CUTOFF
-                    dEdR += apos.w * local_posq[tbx+tj].w * (invR - 2.0f * REACTION_FIELD_K * r2);
-		    tempEnergy += apos.w * local_posq[tbx+tj].w * (invR + REACTION_FIELD_K * r2 - REACTION_FIELD_C);
-#else
-                    dEdR += apos.w * local_posq[tbx+tj].w * invR;
-                    tempEnergy += apos.w * local_posq[tbx+tj].w * invR;
-#endif
-                    dEdR *= invR * invR;
+                    float r2 = delta.x*delta.x + delta.y*delta.y + delta.z*delta.z;
+                    float r = sqrt(r2);
+                    float invR = 1.0f / r;
+                    float4 posq2 = local_posq[tbx+tj];
+                    LOAD_ATOM2_PARAMETERS_TJ
+                    float dEdR = 0.0f;
+                    float tempEnergy = 0.0f;
+                    COMPUTE_INTERACTION
 #ifdef USE_CUTOFF
                     if (isExcluded || r2 > cutoffSquared) {
 #else
@@ -213,7 +181,7 @@ __kernel void computeNonbonded(int numTiles, int paddedNumAtoms,
                     }
 		    energy += tempEnergy;
                     delta.xyz *= dEdR;
-                    af.xyz -= delta.xyz;
+                    force.xyz -= delta.xyz;
                     local_force[tbx+tj].xyz += delta.xyz;
                     excl >>= 1;
                     tj = (tj + 1) & (TileSize - 1);
@@ -223,17 +191,17 @@ __kernel void computeNonbonded(int numTiles, int paddedNumAtoms,
             // Write results
             float4 of;
 #ifdef USE_OUTPUT_BUFFER_PER_BLOCK
-            of.xyz = af.xyz;
+            of.xyz = force.xyz;
             of.w = 0.0f;
-            unsigned int offset = x + tgx + (y/TileSize) * paddedNumAtoms;
+            unsigned int offset = x + tgx + (y/TileSize)*paddedNumAtoms;
             forceBuffers[offset] = of;
             of = local_force[get_local_id(0)];
-            offset = y + tgx + (x/TileSize) * paddedNumAtoms;
+            offset = y + tgx + (x/TileSize)*paddedNumAtoms;
             forceBuffers[offset] = of;
 #else
             unsigned int offset = x + tgx + warp*paddedNumAtoms;
             of = forceBuffers[offset];
-            of.xyz += af.xyz;
+            of.xyz += force.xyz;
             forceBuffers[offset] = of;
             offset = y + tgx + warp*paddedNumAtoms;
             of = forceBuffers[offset];