Add Nose-Hoover chain serialization on CUDA

platforms/cuda/include/CudaIntegrationUtilities.h

@@ -120,6 +120,12 @@ public:
      * @param timeShift   the amount by which to shift the velocities in time
      */
     double computeKineticEnergy(double timeShift);
+    /**
+     * Get the data structure that holds the state of all Nose-Hoover chains
+     * 
+     * @return  vector of chain states
+     */
+    std::vector<CudaArray>& getNoseHooverChainState();
 private:
     void applyConstraints(bool constrainVelocities, double tol);
     CudaContext& context;
@@ -168,6 +174,7 @@ private:
     double2 lastStepSize;
     struct ShakeCluster;
     struct ConstraintOrderer;
+    std::vector<CudaArray> noseHooverChainState;
 };
 
 } // namespace OpenMM

platforms/cuda/include/CudaKernels.h

@@ -1762,7 +1762,7 @@ public:
 
 private:
     CudaContext& cu;
-    CUfunction kernel;
+    std::map<int, CUfunction> propagateKernels;
 };
 
 /**

platforms/cuda/src/CudaIntegrationUtilities.cpp

@@ -716,6 +716,15 @@ void CudaIntegrationUtilities::createCheckpoint(ostream& stream) {
     vector<int4> randomSeedVec;
     randomSeed.download(randomSeedVec);
     stream.write((char*) &randomSeedVec[0], sizeof(int4)*randomSeed.getSize());
+    size_t numChains = noseHooverChainState.size();
+    stream.write((char*) &numChains, sizeof(size_t));
+    for (auto &chainState: noseHooverChainState){
+        vector<float2> stateVec;
+        chainState.download(stateVec);
+        size_t vecLength = stateVec.size();
+        stream.write((char*) &vecLength, sizeof(size_t));
+        stream.write((char*) stateVec.data(), sizeof(float2)*stateVec.size());
+    }
 }
 
 void CudaIntegrationUtilities::loadCheckpoint(istream& stream) {
@@ -728,6 +737,17 @@ void CudaIntegrationUtilities::loadCheckpoint(istream& stream) {
     vector<int4> randomSeedVec(randomSeed.getSize());
     stream.read((char*) &randomSeedVec[0], sizeof(int4)*randomSeed.getSize());
     randomSeed.upload(randomSeedVec);
+    size_t numChains, chainLength;
+    stream.read((char*) &numChains, sizeof(size_t));
+    noseHooverChainState.clear();
+    for (size_t i=0; i<numChains; i++){
+        stream.read((char*) &chainLength, sizeof(size_t));
+        std::vector<float2> stateVec(chainLength);
+        stream.read((char*) &stateVec[0], sizeof(float2)*chainLength);
+        CudaArray state;
+        state.upload(stateVec);
+        noseHooverChainState.push_back(state);
+    }
 }
 
 double CudaIntegrationUtilities::computeKineticEnergy(double timeShift) {
@@ -775,3 +795,7 @@ double CudaIntegrationUtilities::computeKineticEnergy(double timeShift) {
         posDelta.copyTo(context.getVelm());
     return 0.5*energy;
 }
+
+std::vector<CudaArray>& CudaIntegrationUtilities::getNoseHooverChainState(){
+        return noseHooverChainState;
+};

platforms/cuda/src/CudaKernels.cpp

@@ -387,7 +387,7 @@ void CudaUpdateStateDataKernel::setPeriodicBoxVectors(ContextImpl& context, cons
 
 void CudaUpdateStateDataKernel::createCheckpoint(ContextImpl& context, ostream& stream) {
     cu.setAsCurrent();
-    int version = 2;
+    int version = 3;
     stream.write((char*) &version, sizeof(int));
     int precision = (cu.getUseDoublePrecision() ? 2 : cu.getUseMixedPrecision() ? 1 : 0);
     stream.write((char*) &precision, sizeof(int));
@@ -419,7 +419,7 @@ void CudaUpdateStateDataKernel::loadCheckpoint(ContextImpl& context, istream& st
     cu.setAsCurrent();
     int version;
     stream.read((char*) &version, sizeof(int));
-    if (version != 2)
+    if (version != 3)
         throw OpenMMException("Checkpoint was created with a different version of OpenMM");
     int precision;
     stream.read((char*) &precision, sizeof(int));
@@ -8350,17 +8350,35 @@ void CudaApplyAndersenThermostatKernel::execute(ContextImpl& context) {
 }
 
 void CudaNoseHooverChainKernel::initialize() {
+    cu.setAsCurrent();
+
+    map<string, string> defines;
+    defines["BEGIN_YS_LOOP"] = "for(const real & ys : {1} {";
+    defines["END_YS_LOOP"] = "}";
+    CUmodule module = cu.createModule(CudaKernelSources::noseHooverChain, defines, "");
+    propagateKernels[1] = cu.getKernel(module, "propagateNoseHooverChain");
+    defines["BEGIN_YS_LOOP"] = "for(const real & ys : {0.828981543588751, -0.657963087177502, 0.828981543588751} {";
+    module = cu.createModule(CudaKernelSources::noseHooverChain, defines, "");
+    propagateKernels[3] = cu.getKernel(module, "propagateNoseHooverChain");
+    defines["BEGIN_YS_LOOP"] = "for(const real & ys : {0.2967324292201065, 0.2967324292201065, -0.186929716880426, 0.2967324292201065, 0.2967324292201065} {";
+    module = cu.createModule(CudaKernelSources::noseHooverChain, defines, "");
+    propagateKernels[5] = cu.getKernel(module, "propagateNoseHooverChain");
+
 }
 
 double CudaNoseHooverChainKernel::propagateChain(ContextImpl& context, const NoseHooverChain &nhc, double kineticEnergy, double timeStep) {
+    cu.setAsCurrent();
+
     return 1;
 }
 
 double CudaNoseHooverChainKernel::computeHeatBathEnergy(ContextImpl& context, const NoseHooverChain &nhc) {
+    cu.setAsCurrent();
     return 1;
 }
 
 double CudaNoseHooverChainKernel::computeMaskedKineticEnergy(ContextImpl& context, const NoseHooverChain &noseHooverChain) {
+    cu.setAsCurrent();
     return 1;
 }