Merge pull request #1150 from peastman/cpuopt

Optimizations to CPU platform

platforms/cpu/include/CpuKernels.h

@@ -92,6 +92,46 @@ private:
     Kernel referenceKernel;
 };
 
+/**
+ * This kernel is invoked by HarmonicAngleForce to calculate the forces acting on the system and the energy of the system.
+ */
+class CpuCalcHarmonicAngleForceKernel : public CalcHarmonicAngleForceKernel {
+public:
+    CpuCalcHarmonicAngleForceKernel(std::string name, const Platform& platform, CpuPlatform::PlatformData& data) :
+            CalcHarmonicAngleForceKernel(name, platform), data(data), angleIndexArray(NULL), angleParamArray(NULL) {
+    }
+    ~CpuCalcHarmonicAngleForceKernel();
+    /**
+     * Initialize the kernel.
+     * 
+     * @param system     the System this kernel will be applied to
+     * @param force      the HarmonicAngleForce this kernel will be used for
+     */
+    void initialize(const System& system, const HarmonicAngleForce& force);
+    /**
+     * Execute the kernel to calculate the forces and/or energy.
+     *
+     * @param context        the context in which to execute this kernel
+     * @param includeForces  true if forces should be calculated
+     * @param includeEnergy  true if the energy should be calculated
+     * @return the potential energy due to the force
+     */
+    double execute(ContextImpl& context, bool includeForces, bool includeEnergy);
+    /**
+     * Copy changed parameters over to a context.
+     *
+     * @param context    the context to copy parameters to
+     * @param force      the HarmonicAngleForce to copy the parameters from
+     */
+    void copyParametersToContext(ContextImpl& context, const HarmonicAngleForce& force);
+private:
+    CpuPlatform::PlatformData& data;
+    int numAngles;
+    int **angleIndexArray;
+    RealOpenMM **angleParamArray;
+    CpuBondForce bondForce;
+};
+
 /**
  * This kernel is invoked by PeriodicTorsionForce to calculate the forces acting on the system and the energy of the system.
  */
@@ -229,6 +269,7 @@ private:
     CpuNeighborList* neighborList;
     CpuNonbondedForce* nonbonded;
     Kernel optimizedPme;
+    CpuBondForce bondForce;
 };
 
 /**

platforms/cpu/include/CpuLangevinDynamics.h

 
-/* Portions copyright (c) 2013 Stanford University and Simbios.
+/* Portions copyright (c) 2013-2015 Stanford University and Simbios.
  * Authors: Peter Eastman
  * Contributors: 
  *
@@ -37,6 +37,7 @@ class CpuLangevinDynamics : public ReferenceStochasticDynamics {
 public:
     class Update1Task;
     class Update2Task;
+    class Update3Task;
     /**
      * Constructor.
      *
@@ -80,9 +81,21 @@ public:
     void updatePart2(int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates, std::vector<OpenMM::RealVec>& velocities,
                      std::vector<OpenMM::RealVec>& forces, std::vector<RealOpenMM>& inverseMasses, std::vector<OpenMM::RealVec>& xPrime);
 
+    /**  
+     * Third update
+     * 
+     * @param numberOfAtoms       number of atoms
+     * @param atomCoordinates     atom coordinates
+     * @param velocities          velocities
+     * @param inverseMasses       inverse atom masses
+     */
+    void updatePart3(int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates, std::vector<OpenMM::RealVec>& velocities,
+                     std::vector<RealOpenMM>& inverseMasses, std::vector<OpenMM::RealVec>& xPrime);
+
 private:
     void threadUpdate1(int threadIndex);
     void threadUpdate2(int threadIndex);
+    void threadUpdate3(int threadIndex);
     OpenMM::ThreadPool& threads;
     OpenMM::CpuRandom& random;
     std::vector<OpenMM_SFMT::SFMT> threadRandom;

platforms/cpu/src/CpuKernelFactory.cpp

@@ -41,6 +41,8 @@ KernelImpl* CpuKernelFactory::createKernelImpl(std::string name, const Platform&
     CpuPlatform::PlatformData& data = CpuPlatform::getPlatformData(context);
     if (name == CalcForcesAndEnergyKernel::Name())
         return new CpuCalcForcesAndEnergyKernel(name, platform, data, context);
+    if (name == CalcHarmonicAngleForceKernel::Name())
+        return new CpuCalcHarmonicAngleForceKernel(name, platform, data);
     if (name == CalcPeriodicTorsionForceKernel::Name())
         return new CpuCalcPeriodicTorsionForceKernel(name, platform, data);
     if (name == CalcRBTorsionForceKernel::Name())

platforms/cpu/src/CpuKernels.cpp

@@ -30,6 +30,7 @@
  * -------------------------------------------------------------------------- */
 
 #include "CpuKernels.h"
+#include "ReferenceAngleBondIxn.h"
 #include "ReferenceBondForce.h"
 #include "ReferenceConstraints.h"
 #include "ReferenceKernelFactory.h"
@@ -252,6 +253,64 @@ double CpuCalcForcesAndEnergyKernel::finishComputation(ContextImpl& context, boo
     return referenceKernel.getAs<ReferenceCalcForcesAndEnergyKernel>().finishComputation(context, includeForce, includeEnergy, groups, valid);
 }
 
+CpuCalcHarmonicAngleForceKernel::~CpuCalcHarmonicAngleForceKernel() {
+    if (angleIndexArray != NULL) {
+        for (int i = 0; i < numAngles; i++) {
+            delete[] angleIndexArray[i];
+            delete[] angleParamArray[i];
+        }
+        delete[] angleIndexArray;
+        delete[] angleParamArray;
+    }
+}
+
+void CpuCalcHarmonicAngleForceKernel::initialize(const System& system, const HarmonicAngleForce& force) {
+    numAngles = force.getNumAngles();
+    angleIndexArray = new int*[numAngles];
+    for (int i = 0; i < numAngles; i++)
+        angleIndexArray[i] = new int[3];
+    angleParamArray = new RealOpenMM*[numAngles];
+    for (int i = 0; i < numAngles; i++)
+        angleParamArray[i] = new RealOpenMM[2];
+    for (int i = 0; i < numAngles; ++i) {
+        int particle1, particle2, particle3;
+        double angle, k;
+        force.getAngleParameters(i, particle1, particle2, particle3, angle, k);
+        angleIndexArray[i][0] = particle1;
+        angleIndexArray[i][1] = particle2;
+        angleIndexArray[i][2] = particle3;
+        angleParamArray[i][0] = (RealOpenMM) angle;
+        angleParamArray[i][1] = (RealOpenMM) k;
+    }
+    bondForce.initialize(system.getNumParticles(), numAngles, 3, angleIndexArray, data.threads);
+}
+
+double CpuCalcHarmonicAngleForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
+    vector<RealVec>& posData = extractPositions(context);
+    vector<RealVec>& forceData = extractForces(context);
+    RealOpenMM energy = 0;
+    ReferenceAngleBondIxn angleBond;
+    bondForce.calculateForce(posData, angleParamArray, forceData, includeEnergy ? &energy : NULL, angleBond);
+    return energy;
+}
+
+void CpuCalcHarmonicAngleForceKernel::copyParametersToContext(ContextImpl& context, const HarmonicAngleForce& force) {
+    if (numAngles != force.getNumAngles())
+        throw OpenMMException("updateParametersInContext: The number of angles has changed");
+
+    // Record the values.
+
+    for (int i = 0; i < numAngles; ++i) {
+        int particle1, particle2, particle3;
+        double angle, k;
+        force.getAngleParameters(i, particle1, particle2, particle3, angle, k);
+        if (particle1 != angleIndexArray[i][0] || particle2 != angleIndexArray[i][1] || particle3 != angleIndexArray[i][2])
+            throw OpenMMException("updateParametersInContext: The set of particles in an angle has changed");
+        angleParamArray[i][0] = (RealOpenMM) angle;
+        angleParamArray[i][1] = (RealOpenMM) k;
+    }
+}
+
 CpuCalcPeriodicTorsionForceKernel::~CpuCalcPeriodicTorsionForceKernel() {
     if (torsionIndexArray != NULL) {
         for (int i = 0; i < numTorsions; i++) {
@@ -479,6 +538,7 @@ void CpuCalcNonbondedForceKernel::initialize(const System& system, const Nonbond
         bonded14ParamArray[i][1] = static_cast<RealOpenMM>(4.0*depth);
         bonded14ParamArray[i][2] = static_cast<RealOpenMM>(charge);
     }
+    bondForce.initialize(system.getNumParticles(), num14, 2, bonded14IndexArray, data.threads);
     
     // Record other parameters.
     
@@ -611,9 +671,8 @@ double CpuCalcNonbondedForceKernel::execute(ContextImpl& context, bool includeFo
     }
     energy += nonbondedEnergy;
     if (includeDirect) {
-        ReferenceBondForce refBondForce;
         ReferenceLJCoulomb14 nonbonded14;
-        refBondForce.calculateForce(num14, bonded14IndexArray, posData, bonded14ParamArray, forceData, includeEnergy ? &energy : NULL, nonbonded14);
+        bondForce.calculateForce(posData, bonded14ParamArray, forceData, includeEnergy ? &energy : NULL, nonbonded14);
         if (data.isPeriodic)
             energy += dispersionCoefficient/(boxVectors[0][0]*boxVectors[1][1]*boxVectors[2][2]);
     }

platforms/cpu/src/CpuLangevinDynamics.cpp

 
-/* Portions copyright (c) 2006-2013 Stanford University and Simbios.
+/* Portions copyright (c) 2006-2015 Stanford University and Simbios.
  * Authors: Peter Eastman
  * Contributors: 
  *
@@ -49,6 +49,16 @@ public:
     CpuLangevinDynamics& owner;
 };
 
+class CpuLangevinDynamics::Update3Task : public ThreadPool::Task {
+public:
+    Update3Task(CpuLangevinDynamics& owner) : owner(owner) {
+    }
+    void execute(ThreadPool& threads, int threadIndex) {
+        owner.threadUpdate3(threadIndex);
+    }
+    CpuLangevinDynamics& owner;
+};
+
 CpuLangevinDynamics::CpuLangevinDynamics(int numberOfAtoms, RealOpenMM deltaT, RealOpenMM tau, RealOpenMM temperature, ThreadPool& threads, CpuRandom& random) : 
            ReferenceStochasticDynamics(numberOfAtoms, deltaT, tau, temperature), threads(threads), random(random) {
 }
@@ -92,6 +102,23 @@ void CpuLangevinDynamics::updatePart2(int numberOfAtoms, vector<RealVec>& atomCo
     threads.waitForThreads();
 }
 
+void CpuLangevinDynamics::updatePart3(int numberOfAtoms, vector<RealVec>& atomCoordinates, vector<RealVec>& velocities,
+                                       vector<RealOpenMM>& inverseMasses, vector<RealVec>& xPrime) {
+    // Record the parameters for the threads.
+    
+    this->numberOfAtoms = numberOfAtoms;
+    this->atomCoordinates = &atomCoordinates[0];
+    this->velocities = &velocities[0];
+    this->inverseMasses = &inverseMasses[0];
+    this->xPrime = &xPrime[0];
+    
+    // Signal the threads to start running and wait for them to finish.
+    
+    Update3Task task(*this);
+    threads.execute(task);
+    threads.waitForThreads();
+}
+
 void CpuLangevinDynamics::threadUpdate1(int threadIndex) {
     const RealOpenMM tau = getTau();
     const RealOpenMM vscale = EXP(-getDeltaT()/tau);
@@ -122,3 +149,16 @@ void CpuLangevinDynamics::threadUpdate2(int threadIndex) {
         }
    }
 }
+
+void CpuLangevinDynamics::threadUpdate3(int threadIndex) {
+   const RealOpenMM invStepSize = 1.0/getDeltaT();
+    int start = threadIndex*numberOfAtoms/threads.getNumThreads();
+    int end = (threadIndex+1)*numberOfAtoms/threads.getNumThreads();
+
+   for (int i = start; i < end; ++i)
+       if (inverseMasses[i] != 0.0) {
+            velocities[i] = (xPrime[i]-atomCoordinates[i])*invStepSize;
+            atomCoordinates[i] = xPrime[i];
+       }
+}
+

platforms/cpu/src/CpuPlatform.cpp

@@ -61,6 +61,7 @@ map<const ContextImpl*, CpuPlatform::PlatformData*> CpuPlatform::contextData;
 CpuPlatform::CpuPlatform() {
     CpuKernelFactory* factory = new CpuKernelFactory();
     registerKernelFactory(CalcForcesAndEnergyKernel::Name(), factory);
+    registerKernelFactory(CalcHarmonicAngleForceKernel::Name(), factory);
     registerKernelFactory(CalcPeriodicTorsionForceKernel::Name(), factory);
     registerKernelFactory(CalcRBTorsionForceKernel::Name(), factory);
     registerKernelFactory(CalcNonbondedForceKernel::Name(), factory);

platforms/cpu/tests/TestCpuHarmonicAngleForce.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) 2008-2015 Stanford University and the Authors.      *
+ * Authors: Peter Eastman                                                     *
+ * Contributors:                                                              *
+ *                                                                            *
+ * Permission is hereby granted, free of charge, to any person obtaining a    *
+ * copy of this software and associated documentation files (the "Software"), *
+ * to deal in the Software without restriction, including without limitation  *
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,   *
+ * and/or sell copies of the Software, and to permit persons to whom the      *
+ * Software is furnished to do so, subject to the following conditions:       *
+ *                                                                            *
+ * The above copyright notice and this permission notice shall be included in *
+ * all copies or substantial portions of the Software.                        *
+ *                                                                            *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR *
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,   *
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL    *
+ * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,    *
+ * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR      *
+ * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE  *
+ * USE OR OTHER DEALINGS IN THE SOFTWARE.                                     *
+ * -------------------------------------------------------------------------- */
+
+#include "CpuTests.h"
+#include "TestHarmonicAngleForce.h"
+
+void testParallelComputation() {
+    System system;
+    const int numParticles = 200;
+    for (int i = 0; i < numParticles; i++)
+        system.addParticle(1.0);
+    HarmonicAngleForce* force = new HarmonicAngleForce();
+    for (int i = 2; i < numParticles; i++)
+        force->addAngle(i-2, i-1, i, 1.1, i);
+    system.addForce(force);
+    vector<Vec3> positions(numParticles);
+    for (int i = 0; i < numParticles; i++)
+        positions[i] = Vec3(i, i%2, 0);
+    VerletIntegrator integrator1(0.01);
+    ReferencePlatform reference;
+    Context context1(system, integrator1, reference);
+    context1.setPositions(positions);
+    State state1 = context1.getState(State::Forces | State::Energy);
+    VerletIntegrator integrator2(0.01);
+    Context context2(system, integrator2, platform);
+    context2.setPositions(positions);
+    State state2 = context2.getState(State::Forces | State::Energy);
+    ASSERT_EQUAL_TOL(state1.getPotentialEnergy(), state2.getPotentialEnergy(), 1e-5);
+    for (int i = 0; i < numParticles; i++)
+        ASSERT_EQUAL_VEC(state1.getForces()[i], state2.getForces()[i], 1e-5);
+}
+
+void runPlatformTests() {
+    testParallelComputation();
+}

platforms/reference/include/ReferenceAngleBondIxn.h

@@ -29,7 +29,7 @@
 
 namespace OpenMM {
 
-class ReferenceAngleBondIxn : public ReferenceBondIxn {
+class OPENMM_EXPORT ReferenceAngleBondIxn : public ReferenceBondIxn {
 
    private:
 

platforms/reference/include/ReferenceStochasticDynamics.h

@@ -111,13 +111,26 @@ class OPENMM_EXPORT ReferenceStochasticDynamics : public ReferenceDynamics {
          @param atomCoordinates     atom coordinates
          @param velocities          velocities
          @param forces              forces
-         @param masses              atom masses
+         @param inverseMasses       inverse atom masses
       
          --------------------------------------------------------------------------------------- */
       
       virtual void updatePart2(int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates, std::vector<OpenMM::RealVec>& velocities,
                        std::vector<OpenMM::RealVec>& forces, std::vector<RealOpenMM>& inverseMasses, std::vector<OpenMM::RealVec>& xPrime);
       
+      /**---------------------------------------------------------------------------------------
+      
+         Third update
+      
+         @param numberOfAtoms       number of atoms
+         @param atomCoordinates     atom coordinates
+         @param velocities          velocities
+         @param inverseMasses       inverse atom masses
+      
+         --------------------------------------------------------------------------------------- */
+      
+      virtual void updatePart3(int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates, std::vector<OpenMM::RealVec>& velocities,
+                       std::vector<RealOpenMM>& inverseMasses, std::vector<OpenMM::RealVec>& xPrime);
 };
 
 } // namespace OpenMM

platforms/reference/src/SimTKReference/ReferenceStochasticDynamics.cpp

@@ -163,11 +163,21 @@ void ReferenceStochasticDynamics::updatePart2(int numberOfAtoms, vector<RealVec>
 
    for (int ii = 0; ii < numberOfAtoms; ii++) {
        if (inverseMasses[ii] != 0.0)
-           for (int jj = 0; jj < 3; jj++)
-               xPrime[ii][jj] = atomCoordinates[ii][jj]+getDeltaT()*velocities[ii][jj];
+            xPrime[ii] = atomCoordinates[ii]+velocities[ii]*getDeltaT();
    }
 }
 
+void ReferenceStochasticDynamics::updatePart3(int numberOfAtoms, vector<RealVec>& atomCoordinates,
+                                              vector<RealVec>& velocities, vector<RealOpenMM>& inverseMasses,
+                                              vector<RealVec>& xPrime) {
+   RealOpenMM invStepSize = 1.0/getDeltaT();
+   for (int i = 0; i < numberOfAtoms; ++i)
+       if (inverseMasses[i] != 0) {
+            velocities[i] = (xPrime[i]-atomCoordinates[i])*invStepSize;
+            atomCoordinates[i] = xPrime[i];
+       }
+}
+
 /**---------------------------------------------------------------------------------------
 
    Update -- driver routine for performing stochastic dynamics update of coordinates
@@ -221,13 +231,7 @@ void ReferenceStochasticDynamics::update(const OpenMM::System& system, vector<Re
 
    // copy xPrime -> atomCoordinates
 
-   RealOpenMM invStepSize = 1.0/getDeltaT();
-   for (int i = 0; i < numberOfAtoms; ++i)
-       if (masses[i] != zero)
-           for (int j = 0; j < 3; ++j) {
-               velocities[i][j] = invStepSize*(xPrime[i][j]-atomCoordinates[i][j]);
-               atomCoordinates[i][j] = xPrime[i][j];
-           }
+   updatePart3(numberOfAtoms, atomCoordinates, velocities, inverseMasses, xPrime);
 
    ReferenceVirtualSites::computePositions(system, atomCoordinates);
    incrementTimeStep();