Optimization to constraint algorithm: skip the last iteration as soon as we...

Optimization to constraint algorithm: skip the last iteration as soon as we see that all constraints have converged.

platforms/cuda/src/kernels/kCShake.cu

@@ -101,7 +101,7 @@ __global__ void kApplyCShake_kernel(float4* atomPositions, bool addOldPosition)
     unsigned int maxIterations = 150;
     float lowerTol = 1.0f-2.0f*cSim.shakeTolerance+cSim.shakeTolerance*cSim.shakeTolerance;
     float upperTol = 1.0f+2.0f*cSim.shakeTolerance+cSim.shakeTolerance*cSim.shakeTolerance;
-    for (unsigned int iteration = 0; iteration < maxIterations && iteration == requiredIterations; iteration++)
+    for (unsigned int iteration = 0; iteration < maxIterations; iteration++)
     {
         // Calculate the constraint force for each constraint.
 
@@ -130,9 +130,12 @@ __global__ void kApplyCShake_kernel(float4* atomPositions, bool addOldPosition)
                 requiredIterations = iteration+1;
             pos += blockDim.x * gridDim.x;
         }
-        kSyncAllThreads_kernel(cSim.pSyncCounter, iteration);
+        __syncthreads();
         if (threadIdx.x == 0 && requiredIterations > iteration)
             cSim.pRequiredIterations[0] = requiredIterations;
+        kSyncAllThreads_kernel(cSim.pSyncCounter, iteration);
+        if (iteration == cSim.pRequiredIterations[0])
+            break; // All constraints have converged.
 
         // Multiply by the inverse constraint matrix.
 
@@ -178,8 +181,9 @@ __global__ void kApplyCShake_kernel(float4* atomPositions, bool addOldPosition)
             atomPositions[pos] = atomPos;
             pos += blockDim.x*gridDim.x;
         }
+        if (threadIdx.x == 0)
+            requiredIterations = iteration+1;
         kSyncAllThreads_kernel(&cSim.pSyncCounter[2*gridDim.x], iteration);
-        requiredIterations = cSim.pRequiredIterations[0];
     }
 
     // Reset the initial sync counter to be ready for the next call.

platforms/reference/src/SimTKReference/ReferenceRigidShakeAlgorithm.cpp

@@ -568,12 +568,11 @@ int ReferenceRigidShakeAlgorithm::apply( int numberOfAtoms, RealOpenMM** atomCoo
 
    // main loop
 
-   int done                 = 0;
    int iterations           = 0;
    int numberConverged      = 0;
    vector<RealOpenMM> constraintForce(_numberOfConstraints);
    vector<RealOpenMM> tempForce(_numberOfConstraints);
-   while( !done && iterations++ < getMaximumNumberOfIterations() ){
+   while (iterations < getMaximumNumberOfIterations()) {
       numberConverged  = 0;
       for( int ii = 0; ii < _numberOfConstraints; ii++ ){
 
@@ -600,9 +599,9 @@ int ReferenceRigidShakeAlgorithm::apply( int numberOfAtoms, RealOpenMM** atomCoo
             numberConverged++;
          }
       }
-     if( numberConverged == _numberOfConstraints ){
-         done = true;
-      }
+      if( numberConverged == _numberOfConstraints )
+         break;
+      iterations++;
 
       if (_matrix.size() > 0) {
           for (int i = 0; i < _numberOfConstraints; i++) {
@@ -646,11 +645,11 @@ RealOpenMM damping = one;//(RealOpenMM) (iterations < 2 ? 0.5 : 1.0);
 
    // diagnostics
 
-   if( debug || !done ){
+   if( debug || numberConverged < _numberOfConstraints ){
       std::stringstream message;
       message << methodName;
       message << " iterations=" << iterations << " no. converged=" << numberConverged << " out of " << _numberOfConstraints;
-      if( done ){
+      if( numberConverged == _numberOfConstraints ){
          message << " SUCCESS";
       } else {
          message << " FAILED";
@@ -664,7 +663,7 @@ RealOpenMM damping = one;//(RealOpenMM) (iterations < 2 ? 0.5 : 1.0);
       SimTKOpenMMLog::printMessage( message );
    }
 
-   return (done ? ReferenceDynamics::DefaultReturn : ReferenceDynamics::ErrorReturn);
+   return (numberConverged == _numberOfConstraints ? ReferenceDynamics::DefaultReturn : ReferenceDynamics::ErrorReturn);
 
 }
 

platforms/reference/src/SimTKReference/ReferenceShakeAlgorithm.cpp

@@ -65,7 +65,7 @@ ReferenceShakeAlgorithm::ReferenceShakeAlgorithm( int numberOfConstraints,
    _atomIndices                = atomIndices;
    _distance                   = distance;
 
-   _maximumNumberOfIterations  = 15;
+   _maximumNumberOfIterations  = 150;
    _tolerance                  = tolerance;
    _hasInitializedMasses       = false;