Restructured how force and energy computations are done to make it easier to...

Restructured how force and energy computations are done to make it easier to write plugins for the CUDA platform

olla/include/openmm/kernels.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008 Stanford University and the Authors.           *
+ * Portions copyright (c) 2008-2009 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -57,14 +57,16 @@
 namespace OpenMM {
 
 /**
- * This kernel is invoked at the start of each force evaluation to clear the forces.
+ * This kernel is invoked at the beginning and end of force and energy computations.  It gives the
+ * Platform a chance to clear buffers and do other initialization at the beginning, and to do any
+ * necessary work at the end to determine the final results.
  */
-class InitializeForcesKernel : public KernelImpl {
+class CalcForcesAndEnergyKernel : public KernelImpl {
 public:
     static std::string Name() {
-        return "InitializeForces";
+        return "CalcForcesAndEnergyKernel";
     }
-    InitializeForcesKernel(std::string name, const Platform& platform) : KernelImpl(name, platform) {
+    CalcForcesAndEnergyKernel(std::string name, const Platform& platform) : KernelImpl(name, platform) {
     }
     /**
      * Initialize the kernel.
@@ -73,11 +75,36 @@ public:
      */
     virtual void initialize(const System& system) = 0;
     /**
-     * Execute the kernel.
+     * This is called at the beginning of each force computation, before calcForces() has been called on
+     * any ForceImpl.
      * 
      * @param context    the context in which to execute this kernel
      */
-    virtual void execute(ContextImpl& context) = 0;
+    virtual void beginForceComputation(ContextImpl& context) = 0;
+    /**
+     * This is called at the end of each force computation, after calcForces() has been called on
+     * every ForceImpl.
+     *
+     * @param context    the context in which to execute this kernel
+     */
+    virtual void finishForceComputation(ContextImpl& context) = 0;
+    /**
+     * This is called at the beginning of each energy computation, before calcEnergy() has been called on
+     * any ForceImpl.
+     *
+     * @param context    the context in which to execute this kernel
+     */
+    virtual void beginEnergyComputation(ContextImpl& context) = 0;
+    /**
+     * This is called at the end of each energy computation, after calcEnergy() has been called on
+     * every ForceImpl.
+     *
+     * @param context    the context in which to execute this kernel
+     * @return the potential energy of the system.  This value is added to all values returned by ForceImpls'
+     * calcEnergy() methods.  That is, each force kernel may <i>either</i> return its contribution to the
+     * energy directly, <i>or</i> add it to an internal buffer so that it will be included here.
+     */
+    virtual double finishEnergyComputation(ContextImpl& context) = 0;
 };
 
 /**

openmmapi/include/openmm/State.h

@@ -88,15 +88,6 @@ public:
 private:
     friend class Context;
     State(double time, int numParticles, DataType types);
-
-// Retarded visual studio compiler complains about being unable to 
-// export private stl class members.
-// This stanza explains that it should temporarily shut up.
-#if defined(_MSC_VER)
-#pragma warning(push)
-#pragma warning(disable:4251)
-#endif
-
     std::vector<Vec3>& updPositions();
     std::vector<Vec3>& updVelocities();
     std::vector<Vec3>& updForces();
@@ -108,11 +99,6 @@ private:
     std::vector<Vec3> velocities;
     std::vector<Vec3> forces;
     std::map<std::string, double> parameters;
-
-#if defined(_MSC_VER)
-#pragma warning(pop)
-#endif
-    
 };
 
 } // namespace OpenMM

openmmapi/src/ContextImpl.cpp

@@ -51,7 +51,7 @@ ContextImpl::ContextImpl(Context& owner, System& system, Integrator& integrator,
 {
     vector<string> kernelNames;
     kernelNames.push_back(CalcKineticEnergyKernel::Name());
-    kernelNames.push_back(InitializeForcesKernel::Name());
+    kernelNames.push_back(CalcForcesAndEnergyKernel::Name());
     kernelNames.push_back(UpdateTimeKernel::Name());
     for (int i = 0; i < system.getNumForces(); ++i) {
         forceImpls.push_back(system.getForce(i).createImpl());
@@ -67,8 +67,8 @@ ContextImpl::ContextImpl(Context& owner, System& system, Integrator& integrator,
     else if (!platform->supportsKernels(kernelNames))
         throw OpenMMException("Specified a Platform for a Context which does not support all required kernels");
     platform->contextCreated(*this);
-    initializeForcesKernel = platform->createKernel(InitializeForcesKernel::Name(), *this);
-    dynamic_cast<InitializeForcesKernel&>(initializeForcesKernel.getImpl()).initialize(system);
+    initializeForcesKernel = platform->createKernel(CalcForcesAndEnergyKernel::Name(), *this);
+    dynamic_cast<CalcForcesAndEnergyKernel&>(initializeForcesKernel.getImpl()).initialize(system);
     kineticEnergyKernel = platform->createKernel(CalcKineticEnergyKernel::Name(), *this);
     dynamic_cast<CalcKineticEnergyKernel&>(kineticEnergyKernel.getImpl()).initialize(system);
     updateTimeKernel = platform->createKernel(UpdateTimeKernel::Name(), *this);
@@ -110,9 +110,11 @@ void ContextImpl::setParameter(std::string name, double value) {
 }
 
 void ContextImpl::calcForces() {
-    dynamic_cast<InitializeForcesKernel&>(initializeForcesKernel.getImpl()).execute(*this);
+    CalcForcesAndEnergyKernel& kernel = dynamic_cast<CalcForcesAndEnergyKernel&>(initializeForcesKernel.getImpl());
+    kernel.beginForceComputation(*this);
     for (int i = 0; i < (int) forceImpls.size(); ++i)
         forceImpls[i]->calcForces(*this, forces);
+    kernel.finishForceComputation(*this);
 }
 
 double ContextImpl::calcKineticEnergy() {
@@ -120,9 +122,12 @@ double ContextImpl::calcKineticEnergy() {
 }
 
 double ContextImpl::calcPotentialEnergy() {
+    CalcForcesAndEnergyKernel& kernel = dynamic_cast<CalcForcesAndEnergyKernel&>(initializeForcesKernel.getImpl());
+    kernel.beginEnergyComputation(*this);
     double energy = 0.0;
     for (int i = 0; i < (int) forceImpls.size(); ++i)
         energy += forceImpls[i]->calcEnergy(*this);
+    energy += kernel.finishEnergyComputation(*this);
     return energy;
 }
 

platforms/brook/src/BrookInitializeForcesKernel.cpp

@@ -47,7 +47,7 @@ using namespace std;
 
 BrookInitializeForcesKernel::BrookInitializeForcesKernel( std::string name, const Platform& platform,
                                                           OpenMMBrookInterface& openMMBrookInterface, System& system ) :
-                     InitializeForcesKernel( name, platform ), _openMMBrookInterface( openMMBrookInterface ), _system( system ){
+                     CalcForcesAndEnergyKernel( name, platform ), _openMMBrookInterface( openMMBrookInterface ), _system( system ){
 
 // ---------------------------------------------------------------------------------------
 

platforms/brook/src/BrookInitializeForcesKernel.h

@@ -35,7 +35,7 @@ namespace OpenMM {
 /**
  * This kernel initializes the forces
  */
-class BrookInitializeForcesKernel : public InitializeForcesKernel {
+class BrookInitializeForcesKernel : public CalcForcesAndEnergyKernel {
 
    public:
   

platforms/brook/src/BrookKernelFactory.cpp

@@ -54,7 +54,7 @@ KernelImpl* BrookKernelFactory::createKernelImpl( std::string name, const Platfo
 
    // initialize forces
 
-	if( name == InitializeForcesKernel::Name() ){
+	if( name == CalcForcesAndEnergyKernel::Name() ){
 
       return new BrookInitializeForcesKernel( name, platform, openMMBrookInterface, context.getSystem() );
 

platforms/brook/src/BrookPlatform.cpp

@@ -279,7 +279,7 @@ void BrookPlatform::_initializeKernelFactory( void ){
 
    BrookKernelFactory* factory = new BrookKernelFactory();
 
-   registerKernelFactory( InitializeForcesKernel::Name(),         factory );
+   registerKernelFactory( CalcForcesAndEnergyKernel::Name(),         factory );
    registerKernelFactory( CalcHarmonicBondForceKernel::Name(),    factory );
    registerKernelFactory( CalcHarmonicAngleForceKernel::Name(),   factory );
    registerKernelFactory( CalcPeriodicTorsionForceKernel::Name(), factory );

platforms/cuda/include/CudaPlatform.h

@@ -33,8 +33,6 @@
 struct _gpuContext;
 
 namespace OpenMM {
-    
-class KernelImpl;
 
 /**
  * This Platform subclass uses CUDA implementations of the OpenMM kernels to run on NVidia GPUs.
@@ -79,7 +77,6 @@ class CudaPlatform::PlatformData {
 public:
     PlatformData(_gpuContext* gpu);
     _gpuContext* gpu;
-    KernelImpl* primaryKernel;
     bool removeCM;
     bool hasBonds, hasAngles, hasPeriodicTorsions, hasRB, hasNonbonded, hasCustomNonbonded;
     int nonbondedMethod, customNonbondedMethod;

platforms/cuda/src/CudaKernelFactory.cpp

@@ -33,8 +33,8 @@ using namespace OpenMM;
 
 KernelImpl* CudaKernelFactory::createKernelImpl(std::string name, const Platform& platform, ContextImpl& context) const {
     CudaPlatform::PlatformData& data = *static_cast<CudaPlatform::PlatformData*>(context.getPlatformData());
-    if (name == InitializeForcesKernel::Name())
-        return new CudaInitializeForcesKernel(name, platform);
+    if (name == CalcForcesAndEnergyKernel::Name())
+        return new CudaCalcForcesAndEnergyKernel(name, platform, data);
     if (name == UpdateTimeKernel::Name())
         return new CudaUpdateTimeKernel(name, platform, data);
     if (name == CalcHarmonicBondForceKernel::Name())

platforms/cuda/src/CudaKernels.cpp

@@ -28,7 +28,6 @@
 #include "CudaStreamImpl.h"
 #include "openmm/LangevinIntegrator.h"
 #include "openmm/Context.h"
-#include "ReferencePlatform.h"
 #include "openmm/internal/ContextImpl.h"
 #include "kernels/gputypes.h"
 #include "kernels/cudaKernels.h"
@@ -39,12 +38,19 @@ extern "C" int gpuSetConstants( gpuContext gpu );
 using namespace OpenMM;
 using namespace std;
 
-static void calcForces(ContextImpl& context, CudaPlatform::PlatformData& data) {
+void CudaCalcForcesAndEnergyKernel::initialize(const System& system) {
+}
+
+void CudaCalcForcesAndEnergyKernel::beginForceComputation(ContextImpl& context) {
     _gpuContext* gpu = data.gpu;
     if (data.nonbondedMethod != NO_CUTOFF && data.computeForceCount%100 == 0)
         gpuReorderAtoms(gpu);
     data.computeForceCount++;
     kClearForces(gpu);
+}
+
+void CudaCalcForcesAndEnergyKernel::finishForceComputation(ContextImpl& context) {
+    _gpuContext* gpu = data.gpu;
     if (gpu->bIncludeGBSA) {
         gpu->bRecalculateBornRadii = true;
         kCalculateCDLJObcGbsaForces1(gpu);
@@ -59,12 +65,16 @@ static void calcForces(ContextImpl& context, CudaPlatform::PlatformData& data) {
     kReduceForces(gpu);
 }
 
-static double calcEnergy(ContextImpl& context, CudaPlatform::PlatformData& data, System& system) {
+void CudaCalcForcesAndEnergyKernel::beginEnergyComputation(ContextImpl& context) {
     _gpuContext* gpu = data.gpu;
     if (data.nonbondedMethod != NO_CUTOFF && data.stepCount%100 == 0)
         gpuReorderAtoms(gpu);
     data.stepCount++;
     kClearEnergy(gpu);
+}
+
+double CudaCalcForcesAndEnergyKernel::finishEnergyComputation(ContextImpl& context) {
+    _gpuContext* gpu = data.gpu;
     if (gpu->bIncludeGBSA) {
         gpu->bRecalculateBornRadii = true;
         kCalculateCDLJObcGbsaForces1(gpu);
@@ -81,12 +91,6 @@ static double calcEnergy(ContextImpl& context, CudaPlatform::PlatformData& data,
     return kReduceEnergy(gpu)+data.ewaldSelfEnergy;
 }
 
-void CudaInitializeForcesKernel::initialize(const System& system) {
-}
-
-void CudaInitializeForcesKernel::execute(ContextImpl& context) {
-}
-
 void CudaUpdateTimeKernel::initialize(const System& system) {
 }
 
@@ -102,8 +106,6 @@ CudaCalcHarmonicBondForceKernel::~CudaCalcHarmonicBondForceKernel() {
 }
 
 void CudaCalcHarmonicBondForceKernel::initialize(const System& system, const HarmonicBondForce& force) {
-    if (data.primaryKernel == NULL)
-        data.primaryKernel = this;
     data.hasBonds = true;
     numBonds = force.getNumBonds();
     vector<int> particle1(numBonds);
@@ -120,13 +122,9 @@ void CudaCalcHarmonicBondForceKernel::initialize(const System& system, const Har
 }
 
 void CudaCalcHarmonicBondForceKernel::executeForces(ContextImpl& context) {
-    if (data.primaryKernel == this)
-        calcForces(context, data);
 }
 
 double CudaCalcHarmonicBondForceKernel::executeEnergy(ContextImpl& context) {
-    if (data.primaryKernel == this)
-        return calcEnergy(context, data, system);
     return 0.0;
 }
 
@@ -134,8 +132,6 @@ CudaCalcHarmonicAngleForceKernel::~CudaCalcHarmonicAngleForceKernel() {
 }
 
 void CudaCalcHarmonicAngleForceKernel::initialize(const System& system, const HarmonicAngleForce& force) {
-    if (data.primaryKernel == NULL)
-        data.primaryKernel = this;
     data.hasAngles = true;
     numAngles = force.getNumAngles();
     const float RadiansToDegrees = (float) (180.0/3.14159265);
@@ -154,13 +150,9 @@ void CudaCalcHarmonicAngleForceKernel::initialize(const System& system, const Ha
 }
 
 void CudaCalcHarmonicAngleForceKernel::executeForces(ContextImpl& context) {
-    if (data.primaryKernel == this)
-        calcForces(context, data);
 }
 
 double CudaCalcHarmonicAngleForceKernel::executeEnergy(ContextImpl& context) {
-    if (data.primaryKernel == this)
-        return calcEnergy(context, data, system);
     return 0.0;
 }
 
@@ -168,8 +160,6 @@ CudaCalcPeriodicTorsionForceKernel::~CudaCalcPeriodicTorsionForceKernel() {
 }
 
 void CudaCalcPeriodicTorsionForceKernel::initialize(const System& system, const PeriodicTorsionForce& force) {
-    if (data.primaryKernel == NULL)
-        data.primaryKernel = this;
     data.hasPeriodicTorsions = true;
     numTorsions = force.getNumTorsions();
     const float RadiansToDegrees = (float)(180.0/3.14159265);
@@ -190,13 +180,9 @@ void CudaCalcPeriodicTorsionForceKernel::initialize(const System& system, const
 }
 
 void CudaCalcPeriodicTorsionForceKernel::executeForces(ContextImpl& context) {
-    if (data.primaryKernel == this)
-        calcForces(context, data);
 }
 
 double CudaCalcPeriodicTorsionForceKernel::executeEnergy(ContextImpl& context) {
-    if (data.primaryKernel == this)
-        return calcEnergy(context, data, system);
     return 0.0;
 }
 
@@ -204,8 +190,6 @@ CudaCalcRBTorsionForceKernel::~CudaCalcRBTorsionForceKernel() {
 }
 
 void CudaCalcRBTorsionForceKernel::initialize(const System& system, const RBTorsionForce& force) {
-    if (data.primaryKernel == NULL)
-        data.primaryKernel = this;
     data.hasRB = true;
     numTorsions = force.getNumTorsions();
     vector<int> particle1(numTorsions);
@@ -232,13 +216,9 @@ void CudaCalcRBTorsionForceKernel::initialize(const System& system, const RBTors
 }
 
 void CudaCalcRBTorsionForceKernel::executeForces(ContextImpl& context) {
-    if (data.primaryKernel == this)
-        calcForces(context, data);
 }
 
 double CudaCalcRBTorsionForceKernel::executeEnergy(ContextImpl& context) {
-    if (data.primaryKernel == this)
-        return calcEnergy(context, data, system);
     return 0.0;
 }
 
@@ -246,8 +226,6 @@ CudaCalcNonbondedForceKernel::~CudaCalcNonbondedForceKernel() {
 }
 
 void CudaCalcNonbondedForceKernel::initialize(const System& system, const NonbondedForce& force) {
-    if (data.primaryKernel == NULL)
-        data.primaryKernel = this;
     data.hasNonbonded = true;
     numParticles = force.getNumParticles();
     _gpuContext* gpu = data.gpu;
@@ -362,13 +340,9 @@ void CudaCalcNonbondedForceKernel::initialize(const System& system, const Nonbon
 }
 
 void CudaCalcNonbondedForceKernel::executeForces(ContextImpl& context) {
-    if (data.primaryKernel == this)
-        calcForces(context, data);
 }
 
 double CudaCalcNonbondedForceKernel::executeEnergy(ContextImpl& context) {
-    if (data.primaryKernel == this)
-        return calcEnergy(context, data, system);
     return 0.0;
 }
 
@@ -376,7 +350,6 @@ CudaCalcCustomNonbondedForceKernel::~CudaCalcCustomNonbondedForceKernel() {
 }
 
 void CudaCalcCustomNonbondedForceKernel::initialize(const System& system, const CustomNonbondedForce& force) {
-    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;
@@ -462,17 +435,11 @@ void CudaCalcCustomNonbondedForceKernel::initialize(const System& system, const
 }
 
 void CudaCalcCustomNonbondedForceKernel::executeForces(ContextImpl& context) {
-    if (data.primaryKernel == this) {
-        updateGlobalParams(context);
-        calcForces(context, data);
-    }
+    updateGlobalParams(context);
 }
 
 double CudaCalcCustomNonbondedForceKernel::executeEnergy(ContextImpl& context) {
-    if (data.primaryKernel == this) {
-        updateGlobalParams(context);
-        return calcEnergy(context, data, system);
-    }
+    updateGlobalParams(context);
     return 0.0;
 }
 

platforms/cuda/src/CudaKernels.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008 Stanford University and the Authors.           *
+ * Portions copyright (c) 2008-2009 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -41,11 +41,13 @@ class CudaVerletDynamics;
 namespace OpenMM {
 
 /**
- * This kernel is invoked at the start of each force evaluation to clear the forces.
+ * This kernel is invoked at the beginning and end of force and energy computations.  It gives the
+ * Platform a chance to clear buffers and do other initialization at the beginning, and to do any
+ * necessary work at the end to determine the final results.
  */
-class CudaInitializeForcesKernel : public InitializeForcesKernel {
+class CudaCalcForcesAndEnergyKernel : public CalcForcesAndEnergyKernel {
 public:
-    CudaInitializeForcesKernel(std::string name, const Platform& platform) : InitializeForcesKernel(name, platform) {
+    CudaCalcForcesAndEnergyKernel(std::string name, const Platform& platform, CudaPlatform::PlatformData& data) : CalcForcesAndEnergyKernel(name, platform), data(data) {
     }
     /**
      * Initialize the kernel.
@@ -54,11 +56,38 @@ public:
      */
     void initialize(const System& system);
     /**
-     * Execute the kernel.
-     * 
+     * This is called at the beginning of each force computation, before calcForces() has been called on
+     * any ForceImpl.
+     *
      * @param context    the context in which to execute this kernel
      */
-    void execute(ContextImpl& context);
+    void beginForceComputation(ContextImpl& context);
+    /**
+     * This is called at the end of each force computation, after calcForces() has been called on
+     * every ForceImpl.
+     *
+     * @param context    the context in which to execute this kernel
+     */
+    void finishForceComputation(ContextImpl& context);
+    /**
+     * This is called at the beginning of each energy computation, before calcEnergy() has been called on
+     * any ForceImpl.
+     *
+     * @param context    the context in which to execute this kernel
+     */
+    void beginEnergyComputation(ContextImpl& context);
+    /**
+     * This is called at the end of each energy computation, after calcEnergy() has been called on
+     * every ForceImpl.
+     *
+     * @param context    the context in which to execute this kernel
+     * @return the potential energy of the system.  This value is added to all values returned by ForceImpls'
+     * calcEnergy() methods.  That is, each force kernel may <i>either</i> return its contribution to the
+     * energy directly, <i>or</i> add it to an internal buffer so that it will be included here.
+     */
+    double finishEnergyComputation(ContextImpl& context);
+private:
+    CudaPlatform::PlatformData& data;
 };
 
 /**

platforms/cuda/src/CudaPlatform.cpp

@@ -46,7 +46,7 @@ extern "C" void initOpenMMPlugin() {
 
 CudaPlatform::CudaPlatform() {
     CudaKernelFactory* factory = new CudaKernelFactory();
-    registerKernelFactory(InitializeForcesKernel::Name(), factory);
+    registerKernelFactory(CalcForcesAndEnergyKernel::Name(), factory);
     registerKernelFactory(UpdateTimeKernel::Name(), factory);
     registerKernelFactory(CalcHarmonicBondForceKernel::Name(), factory);
     registerKernelFactory(CalcHarmonicAngleForceKernel::Name(), factory);
@@ -106,7 +106,7 @@ void CudaPlatform::contextDestroyed(ContextImpl& context) const {
 }
 
 CudaPlatform::PlatformData::PlatformData(_gpuContext* gpu) : gpu(gpu), removeCM(false), nonbondedMethod(0), customNonbondedMethod(0), hasBonds(false), hasAngles(false),
-        hasPeriodicTorsions(false), hasRB(false), hasNonbonded(false), hasCustomNonbonded(false), primaryKernel(NULL), stepCount(0), computeForceCount(0), time(0.0),
+        hasPeriodicTorsions(false), hasRB(false), hasNonbonded(false), hasCustomNonbonded(false), stepCount(0), computeForceCount(0), time(0.0),
         ewaldSelfEnergy(0.0) {
     stringstream device;
     device << gpu->device;

platforms/opencl/src/OpenCLKernelFactory.cpp

@@ -33,8 +33,8 @@ using namespace OpenMM;
 
 KernelImpl* OpenCLKernelFactory::createKernelImpl(std::string name, const Platform& platform, ContextImpl& context) const {
     OpenCLPlatform::PlatformData& data = *static_cast<OpenCLPlatform::PlatformData*>(context.getPlatformData());
-    if (name == InitializeForcesKernel::Name())
-        return new OpenCLInitializeForcesKernel(name, platform, data);
+    if (name == CalcForcesAndEnergyKernel::Name())
+        return new OpenCLCalcForcesAndEnergyKernel(name, platform, data);
     if (name == UpdateTimeKernel::Name())
         return new OpenCLUpdateTimeKernel(name, platform, data);
 //    if (name == CalcHarmonicBondForceKernel::Name())

platforms/opencl/src/OpenCLKernels.cpp

@@ -34,13 +34,23 @@
 using namespace OpenMM;
 using namespace std;
 
-void OpenCLInitializeForcesKernel::initialize(const System& system) {
+void OpenCLCalcForcesAndEnergyKernel::initialize(const System& system) {
 }
 
-void OpenCLInitializeForcesKernel::execute(ContextImpl& context) {
+void OpenCLCalcForcesAndEnergyKernel::beginForceComputation(ContextImpl& context) {
     data.context->clearBuffer(data.context->getForceBuffers());
 }
 
+void OpenCLCalcForcesAndEnergyKernel::finishForceComputation(ContextImpl& context) {
+}
+
+void OpenCLCalcForcesAndEnergyKernel::beginEnergyComputation(ContextImpl& context) {
+}
+
+double OpenCLCalcForcesAndEnergyKernel::finishEnergyComputation(ContextImpl& context) {
+    return 0.0;
+}
+
 void OpenCLUpdateTimeKernel::initialize(const System& system) {
 }
 

platforms/opencl/src/OpenCLKernels.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008 Stanford University and the Authors.           *
+ * Portions copyright (c) 2008-2009 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -34,11 +34,13 @@
 namespace OpenMM {
 
 /**
- * This kernel is invoked at the start of each force evaluation to clear the forces.
+ * This kernel is invoked at the beginning and end of force and energy computations.  It gives the
+ * Platform a chance to clear buffers and do other initialization at the beginning, and to do any
+ * necessary work at the end to determine the final results.
  */
-class OpenCLInitializeForcesKernel : public InitializeForcesKernel {
+class OpenCLCalcForcesAndEnergyKernel : public CalcForcesAndEnergyKernel {
 public:
-    OpenCLInitializeForcesKernel(std::string name, const Platform& platform, OpenCLPlatform::PlatformData& data) : InitializeForcesKernel(name, platform), data(data) {
+    OpenCLCalcForcesAndEnergyKernel(std::string name, const Platform& platform, OpenCLPlatform::PlatformData& data) : CalcForcesAndEnergyKernel(name, platform), data(data) {
     }
     /**
      * Initialize the kernel.
@@ -47,11 +49,36 @@ public:
      */
     void initialize(const System& system);
     /**
-     * Execute the kernel.
+     * This is called at the beginning of each force computation, before calcForces() has been called on
+     * any ForceImpl.
+     *
+     * @param context    the context in which to execute this kernel
+     */
+    void beginForceComputation(ContextImpl& context);
+    /**
+     * This is called at the end of each force computation, after calcForces() has been called on
+     * every ForceImpl.
+     *
+     * @param context    the context in which to execute this kernel
+     */
+    void finishForceComputation(ContextImpl& context);
+    /**
+     * This is called at the beginning of each energy computation, before calcEnergy() has been called on
+     * any ForceImpl.
+     *
+     * @param context    the context in which to execute this kernel
+     */
+    void beginEnergyComputation(ContextImpl& context);
+    /**
+     * This is called at the end of each energy computation, after calcEnergy() has been called on
+     * every ForceImpl.
      *
      * @param context    the context in which to execute this kernel
+     * @return the potential energy of the system.  This value is added to all values returned by ForceImpls'
+     * calcEnergy() methods.  That is, each force kernel may <i>either</i> return its contribution to the
+     * energy directly, <i>or</i> add it to an internal buffer so that it will be included here.
      */
-    void execute(ContextImpl& context);
+    double finishEnergyComputation(ContextImpl& context);
 private:
    OpenCLPlatform::PlatformData& data;
 };

platforms/opencl/src/OpenCLPlatform.cpp

@@ -45,7 +45,7 @@ extern "C" void initOpenMMPlugin() {
 
 OpenCLPlatform::OpenCLPlatform() {
     OpenCLKernelFactory* factory = new OpenCLKernelFactory();
-    registerKernelFactory(InitializeForcesKernel::Name(), factory);
+    registerKernelFactory(CalcForcesAndEnergyKernel::Name(), factory);
     registerKernelFactory(UpdateTimeKernel::Name(), factory);
 //    registerKernelFactory(CalcHarmonicBondForceKernel::Name(), factory);
 //    registerKernelFactory(CalcHarmonicAngleForceKernel::Name(), factory);

platforms/reference/src/ReferenceKernelFactory.cpp

@@ -38,8 +38,8 @@ using namespace OpenMM;
 
 KernelImpl* ReferenceKernelFactory::createKernelImpl(std::string name, const Platform& platform, ContextImpl& context) const {
     ReferencePlatform::PlatformData& data = *static_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
-    if (name == InitializeForcesKernel::Name())
-        return new ReferenceInitializeForcesKernel(name, platform);
+    if (name == CalcForcesAndEnergyKernel::Name())
+        return new ReferenceCalcForcesAndEnergyKernel(name, platform);
     if (name == UpdateTimeKernel::Name())
         return new ReferenceUpdateTimeKernel(name, platform, data);
     if (name == CalcNonbondedForceKernel::Name())

platforms/reference/src/ReferenceKernels.cpp

@@ -126,14 +126,24 @@ static void findAnglesForCCMA(const System& system, vector<ReferenceCCMAAlgorith
     }
 }
 
-void ReferenceInitializeForcesKernel::initialize(const System& system) {
+void ReferenceCalcForcesAndEnergyKernel::initialize(const System& system) {
 }
 
-void ReferenceInitializeForcesKernel::execute(ContextImpl& context) {
+void ReferenceCalcForcesAndEnergyKernel::beginForceComputation(ContextImpl& context) {
     double zero[] = {0.0, 0.0, 0.0};
     context.getForces().fillWithValue(zero);
 }
 
+void ReferenceCalcForcesAndEnergyKernel::finishForceComputation(ContextImpl& context) {
+}
+
+void ReferenceCalcForcesAndEnergyKernel::beginEnergyComputation(ContextImpl& context) {
+}
+
+double ReferenceCalcForcesAndEnergyKernel::finishEnergyComputation(ContextImpl& context) {
+    return 0.0;
+}
+
 void ReferenceUpdateTimeKernel::initialize(const System& system) {
 }
 

platforms/reference/src/ReferenceKernels.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008 Stanford University and the Authors.           *
+ * Portions copyright (c) 2008-2009 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -51,11 +51,13 @@ class ReferenceVerletDynamics;
 namespace OpenMM {
 
 /**
- * This kernel is invoked at the start of each force evaluation to clear the forces.
+ * This kernel is invoked at the beginning and end of force and energy computations.  It gives the
+ * Platform a chance to clear buffers and do other initialization at the beginning, and to do any
+ * necessary work at the end to determine the final results.
  */
-class ReferenceInitializeForcesKernel : public InitializeForcesKernel {
+class ReferenceCalcForcesAndEnergyKernel : public CalcForcesAndEnergyKernel {
 public:
-    ReferenceInitializeForcesKernel(std::string name, const Platform& platform) : InitializeForcesKernel(name, platform) {
+    ReferenceCalcForcesAndEnergyKernel(std::string name, const Platform& platform) : CalcForcesAndEnergyKernel(name, platform) {
     }
     /**
      * Initialize the kernel.
@@ -64,11 +66,36 @@ public:
      */
     void initialize(const System& system);
     /**
-     * Execute the kernel.
-     * 
+     * This is called at the beginning of each force computation, before calcForces() has been called on
+     * any ForceImpl.
+     *
      * @param context    the context in which to execute this kernel
      */
-    void execute(ContextImpl& context);
+    void beginForceComputation(ContextImpl& context);
+    /**
+     * This is called at the end of each force computation, after calcForces() has been called on
+     * every ForceImpl.
+     *
+     * @param context    the context in which to execute this kernel
+     */
+    void finishForceComputation(ContextImpl& context);
+    /**
+     * This is called at the beginning of each energy computation, before calcEnergy() has been called on
+     * any ForceImpl.
+     *
+     * @param context    the context in which to execute this kernel
+     */
+    void beginEnergyComputation(ContextImpl& context);
+    /**
+     * This is called at the end of each energy computation, after calcEnergy() has been called on
+     * every ForceImpl.
+     *
+     * @param context    the context in which to execute this kernel
+     * @return the potential energy of the system.  This value is added to all values returned by ForceImpls'
+     * calcEnergy() methods.  That is, each force kernel may <i>either</i> return its contribution to the
+     * energy directly, <i>or</i> add it to an internal buffer so that it will be included here.
+     */
+    double finishEnergyComputation(ContextImpl& context);
 };
 
 /**

platforms/reference/src/ReferencePlatform.cpp

@@ -39,7 +39,7 @@ using namespace OpenMM;
 
 ReferencePlatform::ReferencePlatform() {
     ReferenceKernelFactory* factory = new ReferenceKernelFactory();
-    registerKernelFactory(InitializeForcesKernel::Name(), factory);
+    registerKernelFactory(CalcForcesAndEnergyKernel::Name(), factory);
     registerKernelFactory(UpdateTimeKernel::Name(), factory);
     registerKernelFactory(CalcHarmonicBondForceKernel::Name(), factory);
     registerKernelFactory(CalcHarmonicAngleForceKernel::Name(), factory);