CustomNonbondedForce with interaction groups uses neighbor list in OpenCL

platforms/opencl/include/OpenCLKernels.h

@@ -742,15 +742,15 @@ private:
     ForceInfo* info;
     OpenCLParameterSet* params;
     OpenCLArray globals;
-    OpenCLArray interactionGroupData;
-    cl::Kernel interactionGroupKernel;
+    OpenCLArray interactionGroupData, filteredGroupData, numGroupTiles;
+    cl::Kernel interactionGroupKernel, prepareNeighborListKernel, buildNeighborListKernel;
     std::vector<void*> interactionGroupArgs;
     std::vector<std::string> globalParamNames;
     std::vector<cl_float> globalParamValues;
     std::vector<OpenCLArray> tabulatedFunctions;
     double longRangeCoefficient;
     std::vector<double> longRangeCoefficientDerivs;
-    bool hasInitializedLongRangeCorrection, hasInitializedKernel, hasParamDerivs;
+    bool hasInitializedLongRangeCorrection, hasInitializedKernel, hasParamDerivs, useNeighborList;
     int numGroupThreadBlocks;
     CustomNonbondedForce* forceCopy;
     const System& system;

platforms/opencl/include/OpenCLNonbondedUtilities.h

@@ -153,6 +153,11 @@ public:
     bool getHasInteractions() {
         return (groupCutoff.size() > 0);
     }
+    /**
+     * Given a nonbonded cutoff, get the padded cutoff distance used in computing
+     * the neighbor list.
+     */
+    double padCutoff(double cutoff);
     /**
      * Prepare to compute interactions.  This updates the neighbor list.
      */
@@ -225,6 +230,13 @@ public:
     OpenCLArray& getExclusionRowIndices() {
         return exclusionRowIndices;
     }
+    /**
+     * Get the array containing a flag for whether the neighbor list was rebuilt
+     * on the most recent call to prepareInteractions().
+     */
+    OpenCLArray& getRebuildNeighborList() {
+        return rebuildNeighborList;
+    }
     /**
      * Get the index of the first tile this context is responsible for processing.
      */

platforms/opencl/src/OpenCLKernels.cpp

@@ -2713,6 +2713,16 @@ void OpenCLCalcCustomNonbondedForceKernel::initInteractionGroups(const CustomNon
     }
     interactionGroupData.initialize<mm_int4>(cl, groupData.size(), "interactionGroupData");
     interactionGroupData.upload(groupData);
+    numGroupTiles.initialize<cl_int>(cl, 1, "numGroupTiles");
+
+    // Allocate space for a neighbor list, if necessary.
+
+    if (force.getNonbondedMethod() != CustomNonbondedForce::NoCutoff && groupData.size() > cl.getNumThreadBlocks()) {
+        filteredGroupData.initialize<mm_int4>(cl, groupData.size(), "filteredGroupData");
+        interactionGroupData.copyTo(filteredGroupData);
+        int numTiles = groupData.size()/32;
+        numGroupTiles.upload(&numTiles);
+    }
     
     // Create the kernel.
     
@@ -2791,11 +2801,16 @@ void OpenCLCalcCustomNonbondedForceKernel::initInteractionGroups(const CustomNon
         defines["USE_CUTOFF"] = "1";
     if (force.getNonbondedMethod() == CustomNonbondedForce::CutoffPeriodic)
         defines["USE_PERIODIC"] = "1";
-    defines["LOCAL_MEMORY_SIZE"] = cl.intToString(max(32, cl.getNonbondedUtilities().getForceThreadBlockSize()));
+    int localMemorySize = max(32, cl.getNonbondedUtilities().getForceThreadBlockSize());
+    defines["LOCAL_MEMORY_SIZE"] = cl.intToString(localMemorySize);
+    defines["WARPS_IN_BLOCK"] = cl.intToString(localMemorySize/32);
     double cutoff = force.getCutoffDistance();
     defines["CUTOFF_SQUARED"] = cl.doubleToString(cutoff*cutoff);
+    double paddedCutoff = cl.getNonbondedUtilities().padCutoff(cutoff);
+    defines["PADDED_CUTOFF_SQUARED"] = cl.doubleToString(paddedCutoff*paddedCutoff);
     defines["PADDED_NUM_ATOMS"] = cl.intToString(cl.getPaddedNumAtoms());
     defines["TILE_SIZE"] = "32";
+    defines["NUM_TILES"] = cl.intToString(numTileSets);
     int numContexts = cl.getPlatformData().contexts.size();
     int startIndex = cl.getContextIndex()*numTileSets/numContexts;
     int endIndex = (cl.getContextIndex()+1)*numTileSets/numContexts;
@@ -2805,10 +2820,17 @@ void OpenCLCalcCustomNonbondedForceKernel::initInteractionGroups(const CustomNon
         defines["PARAMETER_SIZE_IS_EVEN"] = "1";
     cl::Program program = cl.createProgram(cl.replaceStrings(OpenCLKernelSources::customNonbondedGroups, replacements), defines);
     interactionGroupKernel = cl::Kernel(program, "computeInteractionGroups");
+    prepareNeighborListKernel = cl::Kernel(program, "prepareToBuildNeighborList");
+    buildNeighborListKernel = cl::Kernel(program, "buildNeighborList");
     numGroupThreadBlocks = cl.getNonbondedUtilities().getNumForceThreadBlocks();
 }
 
 double OpenCLCalcCustomNonbondedForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
+    useNeighborList = (filteredGroupData.isInitialized() && cl.getNonbondedUtilities().getUseCutoff());
+    if (useNeighborList && cl.getContextIndex() > 0) {
+        // When using a neighbor list, run the whole calculation on a single device.
+        return 0.0;
+    }
     if (globals.isInitialized()) {
         bool changed = false;
         for (int i = 0; i < (int) globalParamNames.size(); i++) {
@@ -2837,7 +2859,9 @@ double OpenCLCalcCustomNonbondedForceKernel::execute(ContextImpl& context, bool
             interactionGroupKernel.setArg<cl::Buffer>(index++, (useLong ? cl.getLongForceBuffer() : cl.getForceBuffers()).getDeviceBuffer());
             interactionGroupKernel.setArg<cl::Buffer>(index++, cl.getEnergyBuffer().getDeviceBuffer());
             interactionGroupKernel.setArg<cl::Buffer>(index++, cl.getPosq().getDeviceBuffer());
-            interactionGroupKernel.setArg<cl::Buffer>(index++, interactionGroupData.getDeviceBuffer());
+            interactionGroupKernel.setArg<cl::Buffer>(index++, (useNeighborList ? filteredGroupData : interactionGroupData).getDeviceBuffer());
+            interactionGroupKernel.setArg<cl::Buffer>(index++, numGroupTiles.getDeviceBuffer());
+            interactionGroupKernel.setArg<cl_bool>(index++, useNeighborList);
             index += 5;
             for (auto& buffer : params->getBuffers())
                 interactionGroupKernel.setArg<cl::Memory>(index++, buffer.getMemory());
@@ -2847,9 +2871,27 @@ double OpenCLCalcCustomNonbondedForceKernel::execute(ContextImpl& context, bool
                 interactionGroupKernel.setArg<cl::Buffer>(index++, globals.getDeviceBuffer());
             if (hasParamDerivs)
                 interactionGroupKernel.setArg<cl::Memory>(index++, cl.getEnergyParamDerivBuffer().getDeviceBuffer());
+            if (useNeighborList) {
+                // Initialize kernels for building the interaction group neighbor list.
+                
+                prepareNeighborListKernel.setArg<cl::Buffer>(0, cl.getNonbondedUtilities().getRebuildNeighborList().getDeviceBuffer());
+                prepareNeighborListKernel.setArg<cl::Buffer>(1, numGroupTiles.getDeviceBuffer());
+                buildNeighborListKernel.setArg<cl::Buffer>(0, cl.getNonbondedUtilities().getRebuildNeighborList().getDeviceBuffer());
+                buildNeighborListKernel.setArg<cl::Buffer>(1, numGroupTiles.getDeviceBuffer());
+                buildNeighborListKernel.setArg<cl::Buffer>(2, cl.getPosq().getDeviceBuffer());
+                buildNeighborListKernel.setArg<cl::Buffer>(3, interactionGroupData.getDeviceBuffer());
+                buildNeighborListKernel.setArg<cl::Buffer>(4, filteredGroupData.getDeviceBuffer());
+            }
         }
-        setPeriodicBoxArgs(cl, interactionGroupKernel, 4);
         int forceThreadBlockSize = max(32, cl.getNonbondedUtilities().getForceThreadBlockSize());
+        if (useNeighborList) {
+            // Rebuild the neighbor list, if necessary.
+
+            setPeriodicBoxArgs(cl, buildNeighborListKernel, 5);
+            cl.executeKernel(prepareNeighborListKernel, 1, 1);
+            cl.executeKernel(buildNeighborListKernel, numGroupThreadBlocks*forceThreadBlockSize, forceThreadBlockSize);
+        }
+        setPeriodicBoxArgs(cl, interactionGroupKernel, 6);
         cl.executeKernel(interactionGroupKernel, numGroupThreadBlocks*forceThreadBlockSize, forceThreadBlockSize);
     }
     mm_double4 boxSize = cl.getPeriodicBoxSizeDouble();

platforms/opencl/src/OpenCLNonbondedUtilities.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2009-2016 Stanford University and the Authors.      *
+ * Portions copyright (c) 2009-2018 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -323,6 +323,11 @@ double OpenCLNonbondedUtilities::getMaxCutoffDistance() {
     return cutoff;
 }
 
+double OpenCLNonbondedUtilities::padCutoff(double cutoff) {
+    double padding = (usePadding ? 0.1*cutoff : 0.0);
+    return cutoff+padding;
+}
+
 void OpenCLNonbondedUtilities::prepareInteractions(int forceGroups) {
     if ((forceGroups&groupFlags) == 0)
         return;
@@ -464,12 +469,11 @@ void OpenCLNonbondedUtilities::createKernelsForGroups(int groups) {
     kernels.cutoffDistance = cutoff;
     kernels.source = source;
     if (useCutoff) {
-        double padding = (usePadding ? 0.1*cutoff : 0.0);
-        double paddedCutoff = cutoff+padding;
+        double paddedCutoff = padCutoff(cutoff);
         map<string, string> defines;
         defines["TILE_SIZE"] = context.intToString(OpenCLContext::TileSize);
         defines["NUM_ATOMS"] = context.intToString(context.getNumAtoms());
-        defines["PADDING"] = context.doubleToString(padding);
+        defines["PADDING"] = context.doubleToString(paddedCutoff-cutoff);
         defines["PADDED_CUTOFF"] = context.doubleToString(paddedCutoff);
         defines["PADDED_CUTOFF_SQUARED"] = context.doubleToString(paddedCutoff*paddedCutoff);
         defines["NUM_TILES_WITH_EXCLUSIONS"] = context.intToString(exclusionTiles.getSize());

platforms/opencl/src/kernels/customNonbondedGroups.cl

@@ -43,6 +43,7 @@ __kernel void computeInteractionGroups(
         __global real4* restrict forceBuffers,
 #endif
         __global mixed* restrict energyBuffer, __global const real4* restrict posq, __global const int4* restrict groupData,
+        __global int* restrict numGroupTiles, bool useNeighborList,
         real4 periodicBoxSize, real4 invPeriodicBoxSize, real4 periodicBoxVecX, real4 periodicBoxVecY, real4 periodicBoxVecZ
         PARAMETER_ARGUMENTS) {
     const unsigned int totalWarps = get_global_size(0)/TILE_SIZE;
@@ -53,8 +54,8 @@ __kernel void computeInteractionGroups(
     INIT_DERIVATIVES
     __local AtomData localData[LOCAL_MEMORY_SIZE];
 
-    const unsigned int startTile = FIRST_TILE+warp*(LAST_TILE-FIRST_TILE)/totalWarps;
-    const unsigned int endTile = FIRST_TILE+(warp+1)*(LAST_TILE-FIRST_TILE)/totalWarps;
+    const unsigned int startTile = (useNeighborList ? warp*numGroupTiles[0]/totalWarps : FIRST_TILE+warp*(LAST_TILE-FIRST_TILE)/totalWarps);
+    const unsigned int endTile = (useNeighborList ? (warp+1)*numGroupTiles[0]/totalWarps : FIRST_TILE+(warp+1)*(LAST_TILE-FIRST_TILE)/totalWarps);
     for (int tile = startTile; tile < endTile; tile++) {
         const int4 atomData = groupData[TILE_SIZE*tile+tgx];
         const int atom1 = atomData.x;
@@ -129,3 +130,73 @@ __kernel void computeInteractionGroups(
     energyBuffer[get_global_id(0)] += energy;
     SAVE_DERIVATIVES
 }
+
+/**
+ * If the neighbor list needs to be rebuilt, reset the number of tiles to 0.  This is
+ * executed by a single thread.
+ */
+__kernel void prepareToBuildNeighborList(__global int* restrict rebuildNeighborList, __global int* restrict numGroupTiles) {
+    if (rebuildNeighborList[0] == 1)
+        numGroupTiles[0] = 0;
+}
+
+/**
+ * Filter the list of tiles to include only ones that have interactions within the
+ * padded cutoff.
+ */
+__kernel void buildNeighborList(__global int* restrict rebuildNeighborList, __global int* restrict numGroupTiles,
+        __global const real4* restrict posq, __global const int4* restrict groupData, __global int4* restrict filteredGroupData,
+        real4 periodicBoxSize, real4 invPeriodicBoxSize, real4 periodicBoxVecX, real4 periodicBoxVecY, real4 periodicBoxVecZ) {
+    
+    // If the neighbor list doesn't need to be rebuilt on this step, return immediately.
+    
+    if (rebuildNeighborList[0] == 0)
+        return;
+
+    const unsigned int totalWarps = get_global_size(0)/TILE_SIZE;
+    const unsigned int warp = get_global_id(0)/TILE_SIZE; // global warpIndex
+    const unsigned int local_warp = get_local_id(0)/TILE_SIZE; // local warpIndex
+    const unsigned int tgx = get_local_id(0) & (TILE_SIZE-1); // index within the warp
+    const unsigned int tbx = get_local_id(0) - tgx;           // block warpIndex
+    __local real4 localPos[LOCAL_MEMORY_SIZE];
+    __local bool anyInteraction[WARPS_IN_BLOCK];
+    __local int tileIndex[WARPS_IN_BLOCK];
+
+    const unsigned int startTile = warp*NUM_TILES/totalWarps;
+    const unsigned int endTile = (warp+1)*NUM_TILES/totalWarps;
+    for (int tile = startTile; tile < endTile; tile++) {
+        const int4 atomData = groupData[TILE_SIZE*tile+tgx];
+        const int atom1 = atomData.x;
+        const int atom2 = atomData.y;
+        const int rangeStart = atomData.z&0xFFFF;
+        const int rangeEnd = (atomData.z>>16)&0xFFFF;
+        const int exclusions = atomData.w;
+        real4 posq1 = posq[atom1];
+        localPos[get_local_id(0)] = posq[atom2];
+        if (tgx == 0)
+            anyInteraction[local_warp] = false;
+        int tj = tgx;
+        SYNC_WARPS;
+        for (int j = rangeStart; j < rangeEnd && !anyInteraction[local_warp]; j++) {
+            if (tj < rangeEnd) {
+                bool isExcluded = (((exclusions>>tj)&1) == 0);
+                int localIndex = tbx+tj;
+                real4 delta = (real4) (localPos[localIndex].xyz - posq1.xyz, 0);
+#ifdef USE_PERIODIC
+                APPLY_PERIODIC_TO_DELTA(delta)
+#endif
+                real r2 = delta.x*delta.x + delta.y*delta.y + delta.z*delta.z;
+                if (!isExcluded && r2 < PADDED_CUTOFF_SQUARED)
+                    anyInteraction[local_warp] = true;
+            }
+            tj = (tj == rangeEnd-1 ? rangeStart : tj+1);
+            SYNC_WARPS;
+        }
+        if (anyInteraction[local_warp]) {
+            if (tgx == 0)
+                tileIndex[local_warp] = atom_add(numGroupTiles, 1);
+            SYNC_WARPS;
+            filteredGroupData[TILE_SIZE*tileIndex[local_warp]+tgx] = atomData;
+        }
+    }
+}

tests/TestCustomNonbondedForce.h

@@ -7,7 +7,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2016 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2018 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -978,6 +978,68 @@ void testInteractionGroupTabulatedFunction() {
     }
 }
 
+void testInteractionGroupWithCutoff() {
+    const int numParticles = 1000;
+    const double boxSize = 10.0;
+    System system;
+    system.setDefaultPeriodicBoxVectors(Vec3(boxSize, 0, 0), Vec3(0, boxSize, 0), Vec3(0, 0, boxSize));
+    NonbondedForce* standard = new NonbondedForce();
+    CustomNonbondedForce* custom = new CustomNonbondedForce("100/(r+0.1)");
+    system.addForce(standard);
+    system.addForce(custom);
+    standard->setNonbondedMethod(NonbondedForce::CutoffPeriodic);
+    custom->setNonbondedMethod(CustomNonbondedForce::CutoffPeriodic);
+    standard->setCutoffDistance(1.0);
+    custom->setCutoffDistance(1.0);
+    standard->setUseSwitchingFunction(true);
+    custom->setUseSwitchingFunction(true);
+    standard->setSwitchingDistance(0.9);
+    custom->setSwitchingDistance(0.8);
+    vector<Vec3> positions(numParticles);
+    OpenMM_SFMT::SFMT sfmt;
+    init_gen_rand(0, sfmt);
+    for (int i = 0; i < numParticles; i++) {
+        system.addParticle(10.0);
+        standard->addParticle(0.0, 0.2, 0.1);
+        custom->addParticle();
+        while (true) {
+            positions[i] = Vec3(genrand_real2(sfmt), genrand_real2(sfmt), genrand_real2(sfmt))*boxSize;
+            bool tooClose = false;
+            for (int j = 0; j < i; j++) {
+                Vec3 delta = positions[i]-positions[j];
+                if (delta.dot(delta) < 0.5*0.5)
+                    tooClose = true;
+            }
+            if (!tooClose)
+                break;
+        }
+    }
+    set<int> set1, set2;
+    for (int i = 0; i < 10; i++)
+        set1.insert(2*i);
+    for (int i = 0; i < numParticles; i++)
+        set2.insert(i);
+    custom->addInteractionGroup(set1, set2);
+    custom->setForceGroup(1);
+    
+    // Try simulating it and see if energy is conserved (indicating that any optimizations
+    // for combining the cutoff with the interaction group are behaving consistently).
+
+    VerletIntegrator integrator(0.001);
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    context.setVelocitiesToTemperature(100);
+    ASSERT(context.getState(State::Energy, false, 1<<1).getPotentialEnergy() != 0.0);
+    State initialState = context.getState(State::Energy);
+    double initialEnergy = initialState.getPotentialEnergy()+initialState.getKineticEnergy();
+    for (int i = 0; i < 100; i++) {
+        integrator.step(10);
+        State state = context.getState(State::Energy);
+        double energy = state.getPotentialEnergy()+state.getKineticEnergy();
+        ASSERT_EQUAL_TOL(initialEnergy, energy, 0.001);
+    }
+}
+
 void testMultipleCutoffs() {
     System system;
     system.addParticle(1.0);
@@ -1253,6 +1315,7 @@ int main(int argc, char* argv[]) {
         testLargeInteractionGroup();
         testInteractionGroupLongRangeCorrection();
         testInteractionGroupTabulatedFunction();
+        testInteractionGroupWithCutoff();
         testMultipleCutoffs();
         testMultipleSwitches();
         testIllegalVariable();