Optimize updateParametersInContext() (#4610)

* Optimize CustomNonbondedForce.updateParametersInContext()

* Optimized uploading changed values to GPU

* Optimized updateParametersInContext() for lots of bonded forces

* Optimized updateParametersInContext() for CustomExternalForce

* Optimized updateParametersInContext() for NonbondedForce

* Code changes for HIP platform

olla/include/openmm/kernels.h

@@ -323,8 +323,10 @@ public:
      *
      * @param context    the context to copy parameters to
      * @param force      the HarmonicBondForce to copy the parameters from
+     * @param firstBond  the index of the first bond whose parameters might have changed
+     * @param lastBond   the index of the last bond whose parameters might have changed
      */
-    virtual void copyParametersToContext(ContextImpl& context, const HarmonicBondForce& force) = 0;
+    virtual void copyParametersToContext(ContextImpl& context, const HarmonicBondForce& force, int firstBond, int lastBond) = 0;
 };
 
 /**
@@ -358,8 +360,10 @@ public:
      *
      * @param context    the context to copy parameters to
      * @param force      the CustomBondForce to copy the parameters from
+     * @param firstBond  the index of the first bond whose parameters might have changed
+     * @param lastBond   the index of the last bond whose parameters might have changed
      */
-    virtual void copyParametersToContext(ContextImpl& context, const CustomBondForce& force) = 0;
+    virtual void copyParametersToContext(ContextImpl& context, const CustomBondForce& force, int firstBond, int lastBond) = 0;
 };
 
 /**
@@ -393,8 +397,10 @@ public:
      *
      * @param context    the context to copy parameters to
      * @param force      the HarmonicAngleForce to copy the parameters from
+     * @param firstAngle the index of the first bond whose parameters might have changed
+     * @param lastAngle  the index of the last bond whose parameters might have changed
      */
-    virtual void copyParametersToContext(ContextImpl& context, const HarmonicAngleForce& force) = 0;
+    virtual void copyParametersToContext(ContextImpl& context, const HarmonicAngleForce& force, int firstAngle, int lastAngle) = 0;
 };
 
 /**
@@ -428,8 +434,10 @@ public:
      *
      * @param context    the context to copy parameters to
      * @param force      the CustomAngleForce to copy the parameters from
+     * @param firstAngle the index of the first bond whose parameters might have changed
+     * @param lastAngle  the index of the last bond whose parameters might have changed
      */
-    virtual void copyParametersToContext(ContextImpl& context, const CustomAngleForce& force) = 0;
+    virtual void copyParametersToContext(ContextImpl& context, const CustomAngleForce& force, int firstAngle, int lastAngle) = 0;
 };
 
 /**
@@ -463,8 +471,10 @@ public:
      *
      * @param context      the context to copy parameters to
      * @param force        the PeriodicTorsionForce to copy the parameters from
+     * @param firstTorsion the index of the first torsion whose parameters might have changed
+     * @param lastTorsion  the index of the last torsion whose parameters might have changed
      */
-    virtual void copyParametersToContext(ContextImpl& context, const PeriodicTorsionForce& force) = 0;
+    virtual void copyParametersToContext(ContextImpl& context, const PeriodicTorsionForce& force, int firstTorsion, int lastTorsion) = 0;
 };
 
 /**
@@ -568,8 +578,10 @@ public:
      *
      * @param context      the context to copy parameters to
      * @param force        the CustomTorsionForce to copy the parameters from
+     * @param firstTorsion the index of the first torsion whose parameters might have changed
+     * @param lastTorsion  the index of the last torsion whose parameters might have changed
      */
-    virtual void copyParametersToContext(ContextImpl& context, const CustomTorsionForce& force) = 0;
+    virtual void copyParametersToContext(ContextImpl& context, const CustomTorsionForce& force, int firstTorsion, int lastTorsion) = 0;
 };
 
 /**
@@ -613,8 +625,12 @@ public:
      *
      * @param context        the context to copy parameters to
      * @param force          the NonbondedForce to copy the parameters from
+     * @param firstParticle  the index of the first particle whose parameters might have changed
+     * @param lastParticle   the index of the last particle whose parameters might have changed
+     * @param firstException the index of the first exception whose parameters might have changed
+     * @param lastException  the index of the last exception whose parameters might have changed
      */
-    virtual void copyParametersToContext(ContextImpl& context, const NonbondedForce& force) = 0;
+    virtual void copyParametersToContext(ContextImpl& context, const NonbondedForce& force, int firstParticle, int lastParticle, int firstException, int lastException) = 0;
     /**
      * Get the parameters being used for PME.
      *
@@ -671,8 +687,10 @@ public:
      *
      * @param context    the context to copy parameters to
      * @param force      the CustomNonbondedForce to copy the parameters from
+     * @param firstParticle  the index of the first particle whose parameters might have changed
+     * @param lastParticle   the index of the last particle whose parameters might have changed
      */
-    virtual void copyParametersToContext(ContextImpl& context, const CustomNonbondedForce& force) = 0;
+    virtual void copyParametersToContext(ContextImpl& context, const CustomNonbondedForce& force, int firstParticle, int lastParticle) = 0;
 };
 
 /**
@@ -781,8 +799,10 @@ public:
      *
      * @param context        the context to copy parameters to
      * @param force          the CustomExternalForce to copy the parameters from
+     * @param firstParticle  the index of the first particle whose parameters might have changed
+     * @param lastParticle   the index of the last particle whose parameters might have changed
      */
-    virtual void copyParametersToContext(ContextImpl& context, const CustomExternalForce& force) = 0;
+    virtual void copyParametersToContext(ContextImpl& context, const CustomExternalForce& force, int firstParticle, int lastParticle) = 0;
 };
 
 /**

openmmapi/include/openmm/CustomAngleForce.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) 2010-2016 Stanford University and the Authors.      *
+ * Portions copyright (c) 2010-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -263,6 +263,7 @@ private:
     std::vector<AngleInfo> angles;
     std::vector<int> energyParameterDerivatives;
     bool usePeriodic;
+    mutable int numContexts, firstChangedAngle, lastChangedAngle;
 };
 
 /**

openmmapi/include/openmm/CustomBondForce.h

@@ -260,6 +260,7 @@ private:
     std::vector<BondInfo> bonds;
     std::vector<int> energyParameterDerivatives;
     bool usePeriodic;
+    mutable int numContexts, firstChangedBond, lastChangedBond;
 };
 
 /**

openmmapi/include/openmm/CustomExternalForce.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-2012 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -240,6 +240,7 @@ private:
     std::vector<ParticleParameterInfo> parameters;
     std::vector<GlobalParameterInfo> globalParameters;
     std::vector<ParticleInfo> particles;
+    mutable int numContexts, firstChangedParticle, lastChangedParticle;
 };
 
 /**

openmmapi/include/openmm/CustomNonbondedForce.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-2022 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -590,6 +590,7 @@ private:
     std::vector<ComputedValueInfo> computedValues;
     std::vector<InteractionGroupInfo> interactionGroups;
     std::vector<int> energyParameterDerivatives;
+    mutable int numContexts, firstChangedParticle, lastChangedParticle;
 };
 
 /**

openmmapi/include/openmm/CustomTorsionForce.h

@@ -276,6 +276,7 @@ private:
     std::vector<TorsionInfo> torsions;
     std::vector<int> energyParameterDerivatives;
     bool usePeriodic;
+    mutable int numContexts, firstChangedTorsion, lastChangedTorsion;
 };
 
 /**

openmmapi/include/openmm/HarmonicAngleForce.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-2016 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -119,6 +119,7 @@ private:
     class AngleInfo;
     std::vector<AngleInfo> angles;
     bool usePeriodic;
+    mutable int numContexts, firstChangedAngle, lastChangedAngle;
 };
 
 /**

openmmapi/include/openmm/HarmonicBondForce.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-2016 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -117,6 +117,7 @@ private:
     class BondInfo;
     std::vector<BondInfo> bonds;
     bool usePeriodic;
+    mutable int numContexts, firstChangedBond, lastChangedBond;
 };
 
 /**

openmmapi/include/openmm/NonbondedForce.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-2022 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -637,6 +637,7 @@ private:
     std::vector<ParticleOffsetInfo> particleOffsets;
     std::vector<ExceptionOffsetInfo> exceptionOffsets;
     std::map<std::pair<int, int>, int> exceptionMap;
+    mutable int numContexts, firstChangedParticle, lastChangedParticle, firstChangedException, lastChangedException;
 };
 
 /**

openmmapi/include/openmm/PeriodicTorsionForce.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-2016 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -126,6 +126,7 @@ private:
     class PeriodicTorsionInfo;
     std::vector<PeriodicTorsionInfo> periodicTorsions;
     bool usePeriodic;
+    mutable int numContexts, firstChangedTorsion, lastChangedTorsion;
 };
 
 /**

openmmapi/include/openmm/internal/CustomAngleForceImpl.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) 2010 Stanford University and the Authors.           *
+ * Portions copyright (c) 2010-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -59,7 +59,7 @@ public:
     double calcForcesAndEnergy(ContextImpl& context, bool includeForces, bool includeEnergy, int groups);
     std::map<std::string, double> getDefaultParameters();
     std::vector<std::string> getKernelNames();
-    void updateParametersInContext(ContextImpl& context);
+    void updateParametersInContext(ContextImpl& context, int firstAngle, int lastAngle);
 private:
     const CustomAngleForce& owner;
     Kernel kernel;

openmmapi/include/openmm/internal/CustomBondForceImpl.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-2009 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -60,7 +60,7 @@ public:
     std::map<std::string, double> getDefaultParameters();
     std::vector<std::string> getKernelNames();
     std::vector<std::pair<int, int> > getBondedParticles() const;
-    void updateParametersInContext(ContextImpl& context);
+    void updateParametersInContext(ContextImpl& context, int firstBond, int lastBond);
 private:
     const CustomBondForce& owner;
     Kernel kernel;

openmmapi/include/openmm/internal/CustomExternalForceImpl.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-2012 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -59,7 +59,7 @@ public:
     double calcForcesAndEnergy(ContextImpl& context, bool includeForces, bool includeEnergy, int groups);
     std::map<std::string, double> getDefaultParameters();
     std::vector<std::string> getKernelNames();
-    void updateParametersInContext(ContextImpl& context);
+    void updateParametersInContext(ContextImpl& context, int firstParticle, int lastParticle);
 private:
     const CustomExternalForce& owner;
     Kernel kernel;

openmmapi/include/openmm/internal/CustomNonbondedForceImpl.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-2022 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -63,7 +63,7 @@ public:
     double calcForcesAndEnergy(ContextImpl& context, bool includeForces, bool includeEnergy, int groups);
     std::map<std::string, double> getDefaultParameters();
     std::vector<std::string> getKernelNames();
-    void updateParametersInContext(ContextImpl& context);
+    void updateParametersInContext(ContextImpl& context, int firstParticle, int lastParticle);
     /**
      * Prepare for computing the long range correction.  This function pre-computes anything
      * that depends only on the Force (such as particle parameters) but not on information in

openmmapi/include/openmm/internal/CustomTorsionForceImpl.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) 2010 Stanford University and the Authors.           *
+ * Portions copyright (c) 2010-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -59,7 +59,7 @@ public:
     double calcForcesAndEnergy(ContextImpl& context, bool includeForces, bool includeEnergy, int groups);
     std::map<std::string, double> getDefaultParameters();
     std::vector<std::string> getKernelNames();
-    void updateParametersInContext(ContextImpl& context);
+    void updateParametersInContext(ContextImpl& context, int firstTorsion, int lastTorsion);
 private:
     const CustomTorsionForce& owner;
     Kernel kernel;

openmmapi/include/openmm/internal/HarmonicAngleForceImpl.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-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -61,7 +61,7 @@ public:
         return std::map<std::string, double>(); // This force field doesn't define any parameters.
     }
     std::vector<std::string> getKernelNames();
-    void updateParametersInContext(ContextImpl& context);
+    void updateParametersInContext(ContextImpl& context, int firstAngle, int lastAngle);
 private:
     const HarmonicAngleForce& owner;
     Kernel kernel;

openmmapi/include/openmm/internal/HarmonicBondForceImpl.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-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -62,7 +62,7 @@ public:
     }
     std::vector<std::string> getKernelNames();
     std::vector<std::pair<int, int> > getBondedParticles() const;
-    void updateParametersInContext(ContextImpl& context);
+    void updateParametersInContext(ContextImpl& context, int firstBond, int lastBond);
 private:
     const HarmonicBondForce& owner;
     Kernel kernel;

openmmapi/include/openmm/internal/NonbondedForceImpl.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-2018 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -61,7 +61,7 @@ public:
     double calcForcesAndEnergy(ContextImpl& context, bool includeForces, bool includeEnergy, int groups);
     std::map<std::string, double> getDefaultParameters();
     std::vector<std::string> getKernelNames();
-    void updateParametersInContext(ContextImpl& context);
+    void updateParametersInContext(ContextImpl& context, int firstParticle, int lastParticle, int firstException, int lastException);
     void getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const;
     void getLJPMEParameters(double& alpha, int& nx, int& ny, int& nz) const;
     /**

openmmapi/include/openmm/internal/PeriodicTorsionForceImpl.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-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -61,7 +61,7 @@ public:
         return std::map<std::string, double>(); // This force field doesn't define any parameters.
     }
     std::vector<std::string> getKernelNames();
-    void updateParametersInContext(ContextImpl& context);
+    void updateParametersInContext(ContextImpl& context, int firstTorsion, int lastTorsion);
 private:
     const PeriodicTorsionForce& owner;
     Kernel kernel;

openmmapi/src/CustomAngleForce.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) 2010-2016 Stanford University and the Authors.      *
+ * Portions copyright (c) 2010-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -40,11 +40,9 @@
 #include <utility>
 
 using namespace OpenMM;
-using std::string;
-using std::stringstream;
-using std::vector;
+using namespace std;
 
-CustomAngleForce::CustomAngleForce(const string& energy) : energyExpression(energy), usePeriodic(false) {
+CustomAngleForce::CustomAngleForce(const string& energy) : energyExpression(energy), usePeriodic(false), numContexts(0) {
 }
 
 const string& CustomAngleForce::getEnergyFunction() const {
@@ -128,14 +126,30 @@ void CustomAngleForce::setAngleParameters(int index, int particle1, int particle
     angles[index].particle1 = particle1;
     angles[index].particle2 = particle2;
     angles[index].particle3 = particle3;
+    if (numContexts > 0) {
+        firstChangedAngle = min(index, firstChangedAngle);
+        lastChangedAngle = max(index, lastChangedAngle);
+    }
 }
 
 ForceImpl* CustomAngleForce::createImpl() const {
+    if (numContexts == 0) {
+        // Begin tracking changes to angles.
+        firstChangedAngle = angles.size();
+        lastChangedAngle = -1;
+    }
+    numContexts++;
     return new CustomAngleForceImpl(*this);
 }
 
 void CustomAngleForce::updateParametersInContext(Context& context) {
-    dynamic_cast<CustomAngleForceImpl&>(getImplInContext(context)).updateParametersInContext(getContextImpl(context));
+    dynamic_cast<CustomAngleForceImpl&>(getImplInContext(context)).updateParametersInContext(getContextImpl(context), firstChangedAngle, lastChangedAngle);
+    if (numContexts == 1) {
+        // We just updated the only existing context for this force, so we can reset
+        // the tracking of changed angles.
+        firstChangedAngle = angles.size();
+        lastChangedAngle = -1;
+    }
 }
 
 void CustomAngleForce::setUsesPeriodicBoundaryConditions(bool periodic) {