Add D2H copies for the checkpointing

If the update is offloaded, the current coordinates and velocities
are stored on the device and have to be copied for the checkpointing.

Change-Id: I391bb83e1c70a17c1679f479234a157ce606a908

diff --git a/src/gromacs/mdrun/md.cpp b/src/gromacs/mdrun/md.cpp
index b0b006b8bc785224c54251d7605f2184d1cb0e61..e054d2c9d44c8d0d2bd4adbb911a1a3c21c13f82 100644 (file)
--- a/src/gromacs/mdrun/md.cpp
+++ b/src/gromacs/mdrun/md.cpp
@@ -1107,17 +1107,19 @@ void gmx::LegacySimulator::do_md()
             }
         }
 
-        // Copy coordinate from the GPU for the output if the update is offloaded and
+        // Copy coordinate from the GPU for the output/checkpointing if the update is offloaded and
         // coordinates have not already been copied for i) search or ii) CPU force tasks.
         if (useGpuForUpdate && !bNS && !runScheduleWork->domainWork.haveCpuLocalForceWork
-            && (do_per_step(step, ir->nstxout) || do_per_step(step, ir->nstxout_compressed)))
+            && (do_per_step(step, ir->nstxout) || do_per_step(step, ir->nstxout_compressed)
+                || checkpointHandler->isCheckpointingStep()))
         {
             stateGpu->copyCoordinatesFromGpu(state->x, AtomLocality::Local);
             stateGpu->waitCoordinatesReadyOnHost(AtomLocality::Local);
         }
-        // Copy velocities if needed for the output.
+        // Copy velocities if needed for the output/checkpointing.
         // NOTE: Copy on the search steps is done at the beginning of the step.
-        if (useGpuForUpdate && !bNS && do_per_step(step, ir->nstvout))
+        if (useGpuForUpdate && !bNS
+            && (do_per_step(step, ir->nstvout) || checkpointHandler->isCheckpointingStep()))
         {
             stateGpu->copyVelocitiesFromGpu(state->v, AtomLocality::Local);
             stateGpu->waitVelocitiesReadyOnHost(AtomLocality::Local);