Began implementing multi-GPU support

platforms/opencl/include/OpenCLPlatform.h

@@ -28,6 +28,7 @@
  * -------------------------------------------------------------------------- */
 
 #include "openmm/Platform.h"
+#include "openmm/System.h"
 
 namespace OpenMM {
     
@@ -66,7 +67,8 @@ class OpenCLPlatform::PlatformData {
 public:
     PlatformData(int numParticles, int deviceIndex);
     ~PlatformData();
-    OpenCLContext* context;
+    void initializeContexts(const System& system);
+    std::vector<OpenCLContext*> contexts;
     bool removeCM;
     int cmMotionFrequency;
     int stepCount, computeForceCount;

platforms/opencl/src/OpenCLArray.h

@@ -127,8 +127,14 @@ public:
      * Copy the values in a vector to the Buffer.
      */
     void upload(std::vector<T>& data) {
+        upload(&data[0]);
+    }
+    /**
+     * Copy the values in an array to the Buffer.
+     */
+    void upload(T* data) {
         try {
-            context.getQueue().enqueueWriteBuffer(*buffer, CL_TRUE, 0, size*sizeof(T), &data[0]);
+            context.getQueue().enqueueWriteBuffer(*buffer, CL_TRUE, 0, size*sizeof(T), data);
         }
         catch (cl::Error err) {
             std::stringstream str;

platforms/opencl/src/OpenCLContext.cpp

@@ -52,9 +52,11 @@ using namespace std;
 const int OpenCLContext::ThreadBlockSize = 64;
 const int OpenCLContext::TileSize = 32;
 
-OpenCLContext::OpenCLContext(int numParticles, int deviceIndex) : time(0.0), stepCount(0), computeForceCount(0), posq(NULL), velm(NULL),
+OpenCLContext::OpenCLContext(int numParticles, int deviceIndex, OpenCLPlatform::PlatformData& platformData) :
+        time(0.0), platformData(platformData), stepCount(0), computeForceCount(0), posq(NULL), velm(NULL),
         forceBuffers(NULL), energyBuffer(NULL), atomIndex(NULL), integration(NULL), nonbonded(NULL) {
     try {
+        contextIndex = platformData.contexts.size();
         std::vector<cl::Platform> platforms;
         cl::Platform::get(&platforms);
         cl_context_properties cprops[] = {CL_CONTEXT_PLATFORM, (cl_context_properties) platforms[0](), 0};

platforms/opencl/src/OpenCLContext.h

@@ -36,6 +36,7 @@
 #endif
 #include <cl.hpp>
 #include "openmm/internal/windowsExport.h"
+#include "OpenCLPlatform.h"
 
 namespace OpenMM {
 
@@ -125,14 +126,17 @@ struct mm_int16 {
 };
 
 /**
- * This class contains the information associated with a Context by the OpenCL Platform.
+ * This class contains the information associated with a Context by the OpenCL Platform.  Each OpenCLContext is
+ * specific to a particular device, and manages data structures and kernels for that device.  When running a simulation
+ * in parallel on multiple devices, there is a separate OpenCLContext for each one.  The list of all contexts is
+ * stored in the OpenCLPlatform::PlatformData.
  */
 
 class OPENMM_EXPORT OpenCLContext {
 public:
     static const int ThreadBlockSize;
     static const int TileSize;
-    OpenCLContext(int numParticles, int deviceIndex);
+    OpenCLContext(int numParticles, int deviceIndex, OpenCLPlatform::PlatformData& platformData);
     ~OpenCLContext();
     /**
      * This is called to initialize internal data structures after all Forces in the system
@@ -161,6 +165,18 @@ public:
     int getDeviceIndex() {
         return deviceIndex;
     }
+    /**
+     * Get the PlatformData object this context is part of.
+     */
+    OpenCLPlatform::PlatformData& getPlatformData() {
+        return platformData;
+    }
+    /**
+     * Get the index of this context in the list stored in the PlatformData.
+     */
+    int getContextIndex() const {
+        return contextIndex;
+    }
     /**
      * Get the cl::CommandQueue associated with this object.
      */
@@ -402,7 +418,9 @@ private:
     void findMoleculeGroups(const System& system);
     static void tagAtomsInMolecule(int atom, int molecule, std::vector<int>& atomMolecule, std::vector<std::vector<int> >& atomBonds);
     double time;
+    OpenCLPlatform::PlatformData& platformData;
     int deviceIndex;
+    int contextIndex;
     int stepCount;
     int computeForceCount;
     int numAtoms;

platforms/opencl/src/OpenCLKernelFactory.cpp

@@ -25,7 +25,7 @@
  * -------------------------------------------------------------------------- */
 
 #include "OpenCLKernelFactory.h"
-#include "OpenCLKernels.h"
+#include "OpenCLParallelKernels.h"
 #include "openmm/internal/ContextImpl.h"
 #include "openmm/OpenMMException.h"
 
@@ -33,7 +33,15 @@ using namespace OpenMM;
 
 KernelImpl* OpenCLKernelFactory::createKernelImpl(std::string name, const Platform& platform, ContextImpl& context) const {
     OpenCLPlatform::PlatformData& data = *static_cast<OpenCLPlatform::PlatformData*>(context.getPlatformData());
-    OpenCLContext& cl = *data.context;
+    if (data.contexts.size() > 0) {
+        // We are running in parallel on multiple devices, so we may want to create a parallel kernel.
+        
+        if (name == CalcForcesAndEnergyKernel::Name())
+            return new OpenCLParallelCalcForcesAndEnergyKernel(name, platform, data);
+        if (name == CalcHarmonicBondForceKernel::Name())
+            return new OpenCLParallelCalcHarmonicBondForceKernel(name, platform, data, context.getSystem());
+    }
+    OpenCLContext& cl = *data.contexts[0];
     if (name == CalcForcesAndEnergyKernel::Name())
         return new OpenCLCalcForcesAndEnergyKernel(name, platform, cl);
     if (name == UpdateStateDataKernel::Name())

platforms/opencl/src/OpenCLKernels.cpp

@@ -103,7 +103,9 @@ double OpenCLUpdateStateDataKernel::getTime(const ContextImpl& context) const {
 }
 
 void OpenCLUpdateStateDataKernel::setTime(ContextImpl& context, double time) {
-    cl.setTime(time);
+    vector<OpenCLContext*>& contexts = cl.getPlatformData().contexts;
+    for (int i = 0; i < (int) contexts.size(); i++)
+        contexts[i]->setTime(time);
 }
 
 void OpenCLUpdateStateDataKernel::getPositions(ContextImpl& context, std::vector<Vec3>& positions) {
@@ -186,7 +188,9 @@ void OpenCLUpdateStateDataKernel::getPeriodicBoxVectors(ContextImpl& context, Ve
 }
 
 void OpenCLUpdateStateDataKernel::setPeriodicBoxVectors(ContextImpl& context, const Vec3& a, const Vec3& b, const Vec3& c) const {
-    cl.setPeriodicBoxSize(a[0], b[1], c[2]);
+    vector<OpenCLContext*>& contexts = cl.getPlatformData().contexts;
+    for (int i = 0; i < (int) contexts.size(); i++)
+        contexts[i]->setPeriodicBoxSize(a[0], b[1], c[2]);
 }
 
 void OpenCLApplyConstraintsKernel::initialize(const System& system) {
@@ -230,7 +234,10 @@ OpenCLCalcHarmonicBondForceKernel::~OpenCLCalcHarmonicBondForceKernel() {
 }
 
 void OpenCLCalcHarmonicBondForceKernel::initialize(const System& system, const HarmonicBondForce& force) {
-    numBonds = force.getNumBonds();
+    int numContexts = cl.getPlatformData().contexts.size();
+    int startIndex = cl.getContextIndex()*force.getNumBonds()/numContexts;
+    int endIndex = (cl.getContextIndex()+1)*force.getNumBonds()/numContexts;
+    numBonds = endIndex-startIndex;
     if (numBonds == 0)
         return;
     params = new OpenCLArray<mm_float2>(cl, numBonds, "bondParams");
@@ -241,7 +248,7 @@ void OpenCLCalcHarmonicBondForceKernel::initialize(const System& system, const H
     for (int i = 0; i < numBonds; i++) {
         int particle1, particle2;
         double length, k;
-        force.getBondParameters(i, particle1, particle2, length, k);
+        force.getBondParameters(startIndex+i, particle1, particle2, length, k);
         paramVector[i] = mm_float2((cl_float) length, (cl_float) k);
         indicesVector[i] = mm_int4(particle1, particle2, forceBufferCounter[particle1]++, forceBufferCounter[particle2]++);
     }
@@ -3235,7 +3242,7 @@ OpenCLIntegrateVerletStepKernel::~OpenCLIntegrateVerletStepKernel() {
 }
 
 void OpenCLIntegrateVerletStepKernel::initialize(const System& system, const VerletIntegrator& integrator) {
-    cl.initialize(system);
+    cl.getPlatformData().initializeContexts(system);
     cl::Program program = cl.createProgram(OpenCLKernelSources::verlet, "");
     kernel1 = cl::Kernel(program, "integrateVerletPart1");
     kernel2 = cl::Kernel(program, "integrateVerletPart2");
@@ -3291,7 +3298,7 @@ OpenCLIntegrateLangevinStepKernel::~OpenCLIntegrateLangevinStepKernel() {
 }
 
 void OpenCLIntegrateLangevinStepKernel::initialize(const System& system, const LangevinIntegrator& integrator) {
-    cl.initialize(system);
+    cl.getPlatformData().initializeContexts(system);
     cl.getIntegrationUtilities().initRandomNumberGenerator(integrator.getRandomNumberSeed());
     map<string, string> defines;
     defines["NUM_ATOMS"] = intToString(cl.getNumAtoms());
@@ -3365,7 +3372,7 @@ OpenCLIntegrateBrownianStepKernel::~OpenCLIntegrateBrownianStepKernel() {
 }
 
 void OpenCLIntegrateBrownianStepKernel::initialize(const System& system, const BrownianIntegrator& integrator) {
-    cl.initialize(system);
+    cl.getPlatformData().initializeContexts(system);
     cl.getIntegrationUtilities().initRandomNumberGenerator(integrator.getRandomNumberSeed());
     map<string, string> defines;
     defines["NUM_ATOMS"] = intToString(cl.getNumAtoms());
@@ -3424,7 +3431,7 @@ OpenCLIntegrateVariableVerletStepKernel::~OpenCLIntegrateVariableVerletStepKerne
 }
 
 void OpenCLIntegrateVariableVerletStepKernel::initialize(const System& system, const VariableVerletIntegrator& integrator) {
-    cl.initialize(system);
+    cl.getPlatformData().initializeContexts(system);
     cl::Program program = cl.createProgram(OpenCLKernelSources::verlet, "");
     kernel1 = cl::Kernel(program, "integrateVerletPart1");
     kernel2 = cl::Kernel(program, "integrateVerletPart2");
@@ -3491,7 +3498,7 @@ OpenCLIntegrateVariableLangevinStepKernel::~OpenCLIntegrateVariableLangevinStepK
 }
 
 void OpenCLIntegrateVariableLangevinStepKernel::initialize(const System& system, const VariableLangevinIntegrator& integrator) {
-    cl.initialize(system);
+    cl.getPlatformData().initializeContexts(system);
     cl.getIntegrationUtilities().initRandomNumberGenerator(integrator.getRandomNumberSeed());
     map<string, string> defines;
     defines["NUM_ATOMS"] = intToString(cl.getNumAtoms());

platforms/opencl/src/OpenCLParallelKernels.cpp

+/* -------------------------------------------------------------------------- *
+ *                                   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.               *
+ *                                                                            *
+ * Portions copyright (c) 2011 Stanford University and the Authors.           *
+ * 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/>.      *
+ * -------------------------------------------------------------------------- */
+
+#include "OpenCLParallelKernels.h"
+
+using namespace OpenMM;
+using namespace std;
+
+OpenCLParallelCalcForcesAndEnergyKernel::OpenCLParallelCalcForcesAndEnergyKernel(string name, const Platform& platform, OpenCLPlatform::PlatformData& data) :
+        CalcForcesAndEnergyKernel(name, platform), data(data) {
+    for (int i = 0; i < (int) data.contexts.size(); i++)
+        kernels.push_back(Kernel(new OpenCLCalcForcesAndEnergyKernel(name, platform, *data.contexts[i])));
+}
+
+void OpenCLParallelCalcForcesAndEnergyKernel::initialize(const System& system) {
+    for (int i = 0; i < (int) kernels.size(); i++)
+        getKernel(i).initialize(system);
+}
+
+void OpenCLParallelCalcForcesAndEnergyKernel::beginComputation(ContextImpl& context, bool includeForce, bool includeEnergy) {
+    // Copy coordinates over to each device.
+    
+    OpenCLContext& mainContext = *data.contexts[0];
+    mainContext.getPosq().download();
+    for (int i = 1; i < (int) data.contexts.size(); i++)
+        data.contexts[i]->getPosq().upload(mainContext.getPosq().getHostBuffer());
+    
+    // Execute the kernel on each device.
+    
+    for (int i = 0; i < (int) kernels.size(); i++)
+        getKernel(i).beginComputation(context, includeForce, includeEnergy);
+}
+
+double OpenCLParallelCalcForcesAndEnergyKernel::finishComputation(ContextImpl& context, bool includeForce, bool includeEnergy) {
+    double energy = 0.0;
+    for (int i = 0; i < (int) kernels.size(); i++)
+        energy += getKernel(i).finishComputation(context, includeForce, includeEnergy);
+    if (includeForce) {
+        // Sum the forces from all devices.
+        
+        for (int i = 0; i < (int) data.contexts.size(); i++)
+            data.contexts[i]->getForce().download();
+        OpenCLArray<mm_float4>& forces = data.contexts[0]->getForce();
+        for (int i = 1; i < (int) data.contexts.size(); i++) {
+            OpenCLArray<mm_float4>& contextForces = data.contexts[i]->getForce();
+            for (int j = 0; j < forces.getSize(); j++) {
+                mm_float4& f1 = forces[j];
+                const mm_float4& f2 = contextForces[j];
+                f1.x += f2.x;
+                f1.y += f2.y;
+                f1.z += f2.z;
+            }
+        }
+        forces.upload();
+    }
+    return energy;
+}
+
+OpenCLParallelCalcHarmonicBondForceKernel::OpenCLParallelCalcHarmonicBondForceKernel(std::string name, const Platform& platform, OpenCLPlatform::PlatformData& data, System& system) :
+        CalcHarmonicBondForceKernel(name, platform), data(data) {
+    for (int i = 0; i < (int) data.contexts.size(); i++)
+        kernels.push_back(Kernel(new OpenCLCalcHarmonicBondForceKernel(name, platform, *data.contexts[i], system)));
+}
+
+void OpenCLParallelCalcHarmonicBondForceKernel::initialize(const System& system, const HarmonicBondForce& force) {
+    for (int i = 0; i < (int) kernels.size(); i++)
+        getKernel(i).initialize(system, force);
+}
+
+double OpenCLParallelCalcHarmonicBondForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
+    for (int i = 0; i < (int) kernels.size(); i++)
+        getKernel(i).execute(context, includeForces, includeEnergy);
+}

platforms/opencl/src/OpenCLPlatform.cpp

@@ -106,12 +106,18 @@ void OpenCLPlatform::contextDestroyed(ContextImpl& context) const {
 }
 
 OpenCLPlatform::PlatformData::PlatformData(int numParticles, int deviceIndex) : removeCM(false), stepCount(0), computeForceCount(0), time(0.0)  {
-    context = new OpenCLContext(numParticles, deviceIndex);
+    contexts.push_back(new OpenCLContext(numParticles, deviceIndex, *this));
     stringstream device;
-    device << context->getDeviceIndex();
+    device << contexts[0]->getDeviceIndex();
     propertyValues[OpenCLPlatform::OpenCLDeviceIndex()] = device.str();
 }
 
 OpenCLPlatform::PlatformData::~PlatformData() {
-    delete context;
+    for (int i = 0; i < (int) contexts.size(); i++)
+        delete contexts[i];
+}
+
+void OpenCLPlatform::PlatformData::initializeContexts(const System& system) {
+    for (int i = 0; i < (int) contexts.size(); i++)
+        contexts[i]->initialize(system);
 }

platforms/opencl/tests/TestOpenCLNonbondedForce.cpp

@@ -468,7 +468,7 @@ void testBlockInteractions(bool periodic) {
     context.setPositions(positions);
     ContextImpl* contextImpl = *reinterpret_cast<ContextImpl**>(&context);
     OpenCLPlatform::PlatformData& data = *static_cast<OpenCLPlatform::PlatformData*>(contextImpl->getPlatformData());
-    OpenCLContext& clcontext = *data.context;
+    OpenCLContext& clcontext = *data.contexts[0];
     OpenCLNonbondedUtilities& nb = clcontext.getNonbondedUtilities();
     State state = context.getState(State::Positions | State::Velocities | State::Forces);
     nb.updateNeighborListSize();

platforms/opencl/tests/TestOpenCLRandom.cpp

@@ -48,7 +48,8 @@ void testGaussian() {
     System system;
     for (int i = 0; i < numAtoms; i++)
         system.addParticle(1.0);
-    OpenCLContext context(numAtoms, -1);
+    OpenCLPlatform::PlatformData platformData(numAtoms, -1);
+    OpenCLContext& context = *platformData.contexts[0];
     context.initialize(system);
     context.getIntegrationUtilities().initRandomNumberGenerator(0);
     OpenCLArray<mm_float4>& random = context.getIntegrationUtilities().getRandom();

platforms/opencl/tests/TestOpenCLSort.cpp

@@ -51,7 +51,8 @@ void verifySorting(vector<float> array) {
 
     System system;
     system.addParticle(0.0);
-    OpenCLContext context(1, -1);
+    OpenCLPlatform::PlatformData platformData(1, -1);
+    OpenCLContext& context = *platformData.contexts[0];
     context.initialize(system);
     OpenCLArray<float> data(context, array.size(), "sortData");
     data.upload(array);