hillclimber 0.1.5a8__cp314-cp314-macosx_15_0_x86_64.whl

plumed/_lib/lib/plumed/patches/gromacs-2023.5.diff/src/gromacs/mdrun/md.cpp

@@ -0,0 +1,2419 @@
+/*
+ * This file is part of the GROMACS molecular simulation package.
+ *
+ * Copyright 1991- The GROMACS Authors
+ * and the project initiators Erik Lindahl, Berk Hess and David van der Spoel.
+ * Consult the AUTHORS/COPYING files and https://www.gromacs.org for details.
+ *
+ * GROMACS is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public License
+ * as published by the Free Software Foundation; either version 2.1
+ * of the License, or (at your option) any later version.
+ *
+ * GROMACS is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with GROMACS; if not, see
+ * https://www.gnu.org/licenses, or write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
+ *
+ * If you want to redistribute modifications to GROMACS, please
+ * consider that scientific software is very special. Version
+ * control is crucial - bugs must be traceable. We will be happy to
+ * consider code for inclusion in the official distribution, but
+ * derived work must not be called official GROMACS. Details are found
+ * in the README & COPYING files - if they are missing, get the
+ * official version at https://www.gromacs.org.
+ *
+ * To help us fund GROMACS development, we humbly ask that you cite
+ * the research papers on the package. Check out https://www.gromacs.org.
+ */
+/*! \internal \file
+ *
+ * \brief Implements the integrator for normal molecular dynamics simulations
+ *
+ * \author David van der Spoel <david.vanderspoel@icm.uu.se>
+ * \ingroup module_mdrun
+ */
+#include "gmxpre.h"
+
+#include <cinttypes>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+
+#include <algorithm>
+#include <memory>
+#include <numeric>
+
+#include "gromacs/applied_forces/awh/awh.h"
+#include "gromacs/applied_forces/awh/read_params.h"
+#include "gromacs/commandline/filenm.h"
+#include "gromacs/domdec/collect.h"
+#include "gromacs/domdec/dlbtiming.h"
+#include "gromacs/domdec/domdec.h"
+#include "gromacs/domdec/domdec_network.h"
+#include "gromacs/domdec/domdec_struct.h"
+#include "gromacs/domdec/gpuhaloexchange.h"
+#include "gromacs/domdec/localtopologychecker.h"
+#include "gromacs/domdec/mdsetup.h"
+#include "gromacs/domdec/partition.h"
+#include "gromacs/essentialdynamics/edsam.h"
+#include "gromacs/ewald/pme_load_balancing.h"
+#include "gromacs/ewald/pme_pp.h"
+#include "gromacs/fileio/trxio.h"
+#include "gromacs/gmxlib/network.h"
+#include "gromacs/gmxlib/nrnb.h"
+#include "gromacs/gpu_utils/device_stream_manager.h"
+#include "gromacs/gpu_utils/gpu_utils.h"
+#include "gromacs/math/units.h"
+#include "gromacs/imd/imd.h"
+#include "gromacs/listed_forces/listed_forces.h"
+#include "gromacs/math/boxmatrix.h"
+#include "gromacs/math/functions.h"
+#include "gromacs/math/vec.h"
+#include "gromacs/math/vectypes.h"
+#include "gromacs/mdlib/checkpointhandler.h"
+#include "gromacs/mdlib/compute_io.h"
+#include "gromacs/mdlib/constr.h"
+#include "gromacs/mdlib/coupling.h"
+#include "gromacs/mdlib/ebin.h"
+#include "gromacs/mdlib/enerdata_utils.h"
+#include "gromacs/mdlib/energyoutput.h"
+#include "gromacs/mdlib/expanded.h"
+#include "gromacs/mdlib/force.h"
+#include "gromacs/mdlib/force_flags.h"
+#include "gromacs/mdlib/forcerec.h"
+#include "gromacs/mdlib/freeenergyparameters.h"
+#include "gromacs/mdlib/md_support.h"
+#include "gromacs/mdlib/mdatoms.h"
+#include "gromacs/mdlib/mdgraph_gpu.h"
+#include "gromacs/mdlib/mdoutf.h"
+#include "gromacs/mdlib/membed.h"
+#include "gromacs/mdlib/resethandler.h"
+#include "gromacs/mdlib/sighandler.h"
+#include "gromacs/mdlib/simulationsignal.h"
+#include "gromacs/mdlib/stat.h"
+#include "gromacs/mdlib/stophandler.h"
+#include "gromacs/mdlib/tgroup.h"
+#include "gromacs/mdlib/trajectory_writing.h"
+#include "gromacs/mdlib/update.h"
+#include "gromacs/mdlib/update_constrain_gpu.h"
+#include "gromacs/mdlib/update_vv.h"
+#include "gromacs/mdlib/vcm.h"
+#include "gromacs/mdlib/vsite.h"
+#include "gromacs/mdrunutility/freeenergy.h"
+#include "gromacs/mdrunutility/handlerestart.h"
+#include "gromacs/mdrunutility/multisim.h"
+#include "gromacs/mdrunutility/printtime.h"
+#include "gromacs/mdtypes/awh_history.h"
+#include "gromacs/mdtypes/awh_params.h"
+#include "gromacs/mdtypes/commrec.h"
+#include "gromacs/mdtypes/df_history.h"
+#include "gromacs/mdtypes/energyhistory.h"
+#include "gromacs/mdtypes/fcdata.h"
+#include "gromacs/mdtypes/forcebuffers.h"
+#include "gromacs/mdtypes/forcerec.h"
+#include "gromacs/mdtypes/group.h"
+#include "gromacs/mdtypes/inputrec.h"
+#include "gromacs/mdtypes/interaction_const.h"
+#include "gromacs/mdtypes/md_enums.h"
+#include "gromacs/mdtypes/mdatom.h"
+#include "gromacs/mdtypes/mdrunoptions.h"
+#include "gromacs/mdtypes/multipletimestepping.h"
+#include "gromacs/mdtypes/observableshistory.h"
+#include "gromacs/mdtypes/observablesreducer.h"
+#include "gromacs/mdtypes/pullhistory.h"
+#include "gromacs/mdtypes/simulation_workload.h"
+#include "gromacs/mdtypes/state.h"
+#include "gromacs/mdtypes/state_propagator_data_gpu.h"
+#include "gromacs/modularsimulator/energydata.h"
+#include "gromacs/nbnxm/gpu_data_mgmt.h"
+#include "gromacs/nbnxm/nbnxm.h"
+#include "gromacs/pbcutil/pbc.h"
+#include "gromacs/pulling/output.h"
+#include "gromacs/pulling/pull.h"
+#include "gromacs/swap/swapcoords.h"
+#include "gromacs/timing/wallcycle.h"
+#include "gromacs/timing/walltime_accounting.h"
+#include "gromacs/topology/atoms.h"
+#include "gromacs/topology/idef.h"
+#include "gromacs/topology/mtop_util.h"
+#include "gromacs/topology/topology.h"
+#include "gromacs/trajectory/trajectoryframe.h"
+#include "gromacs/utility/basedefinitions.h"
+#include "gromacs/utility/cstringutil.h"
+#include "gromacs/utility/fatalerror.h"
+#include "gromacs/utility/logger.h"
+#include "gromacs/utility/real.h"
+#include "gromacs/utility/smalloc.h"
+
+#include "legacysimulator.h"
+#include "replicaexchange.h"
+#include "shellfc.h"
+
+/* PLUMED */
+#include "../../../Plumed.h"
+extern int    plumedswitch;
+extern plumed plumedmain;
+/* END PLUMED */
+
+/* PLUMED HREX */
+extern int plumed_hrex;
+/* END PLUMED HREX */
+
+using gmx::SimulationSignaller;
+
+void gmx::LegacySimulator::do_md()
+{
+    // TODO Historically, the EM and MD "integrators" used different
+    // names for the t_inputrec *parameter, but these must have the
+    // same name, now that it's a member of a struct. We use this ir
+    // alias to avoid a large ripple of nearly useless changes.
+    // t_inputrec is being replaced by IMdpOptionsProvider, so this
+    // will go away eventually.
+    const t_inputrec* ir = inputrec;
+
+    double       t, t0 = ir->init_t;
+    gmx_bool     bGStatEveryStep, bGStat, bCalcVir, bCalcEnerStep, bCalcEner;
+    gmx_bool     bNS = FALSE, bNStList, bStopCM, bFirstStep, bInitStep, bLastStep = FALSE;
+    gmx_bool     bDoExpanded = FALSE;
+    gmx_bool     do_ene, do_log, do_verbose;
+    gmx_bool     bMainState;
+    unsigned int force_flags;
+    tensor    force_vir = { { 0 } }, shake_vir = { { 0 } }, total_vir = { { 0 } }, pres = { { 0 } };
+    int       i, m;
+    rvec      mu_tot;
+    Matrix3x3 pressureCouplingMu{ { 0. } }, parrinelloRahmanM{ { 0. } };
+    gmx_repl_ex_t     repl_ex = nullptr;
+    gmx_global_stat_t gstat;
+    gmx_shellfc_t*    shellfc;
+    gmx_bool          bSumEkinhOld, bDoReplEx, bDoReplExPrev, bExchanged, bNeedRepartition;
+    gmx_bool          bTrotter;
+    real              dvdl_constr;
+    std::vector<RVec> cbuf;
+    matrix            lastbox;
+    int               lamnew = 0;
+    /* for FEP */
+    double    cycles;
+    real      saved_conserved_quantity = 0;
+    real      last_ekin                = 0;
+    t_extmass MassQ;
+    char      sbuf[STEPSTRSIZE], sbuf2[STEPSTRSIZE];
+
+    /* PME load balancing data for GPU kernels */
+    gmx_bool bPMETune         = FALSE;
+    gmx_bool bPMETunePrinting = FALSE;
+
+    bool bInteractiveMDstep = false;
+
+    SimulationSignals signals;
+    /* PLUMED */
+    int plumedNeedsEnergy=0;
+    int plumedWantsToStop=0;
+    matrix plumed_vir;
+    real lambdaForce=0;
+    real realFepState=0;
+    /* END PLUMED */
+
+    // Most global communnication stages don't propagate mdrun
+    // signals, and will use this object to achieve that.
+    SimulationSignaller nullSignaller(nullptr, nullptr, nullptr, false, false);
+
+    if (!mdrunOptions.writeConfout)
+    {
+        // This is on by default, and the main known use case for
+        // turning it off is for convenience in benchmarking, which is
+        // something that should not show up in the general user
+        // interface.
+        GMX_LOG(mdlog.info)
+                .asParagraph()
+                .appendText(
+                        "The -noconfout functionality is deprecated, and may be removed in a "
+                        "future version.");
+    }
+
+    /* md-vv uses averaged full step velocities for T-control
+       md-vv-avek uses averaged half step velocities for T-control (but full step ekin for P control)
+       md uses averaged half step kinetic energies to determine temperature unless defined otherwise by GMX_EKIN_AVE_VEL; */
+    bTrotter = (EI_VV(ir->eI)
+                && (inputrecNptTrotter(ir) || inputrecNphTrotter(ir) || inputrecNvtTrotter(ir)));
+
+    const bool bRerunMD = false;
+
+    int nstglobalcomm = computeGlobalCommunicationPeriod(mdlog, ir, cr);
+    bGStatEveryStep   = (nstglobalcomm == 1);
+
+    const SimulationGroups* groups = &top_global.groups;
+
+    std::unique_ptr<EssentialDynamics> ed = nullptr;
+    if (opt2bSet("-ei", nfile, fnm))
+    {
+        /* Initialize essential dynamics sampling */
+        ed = init_edsam(mdlog,
+                        opt2fn_null("-ei", nfile, fnm),
+                        opt2fn("-eo", nfile, fnm),
+                        top_global,
+                        *ir,
+                        cr,
+                        constr,
+                        state_global,
+                        observablesHistory,
+                        oenv,
+                        startingBehavior);
+    }
+    else if (observablesHistory->edsamHistory)
+    {
+        gmx_fatal(FARGS,
+                  "The checkpoint is from a run with essential dynamics sampling, "
+                  "but the current run did not specify the -ei option. "
+                  "Either specify the -ei option to mdrun, or do not use this checkpoint file.");
+    }
+
+    int*                fep_state = MAIN(cr) ? &state_global->fep_state : nullptr;
+    gmx::ArrayRef<real> lambda    = MAIN(cr) ? state_global->lambda : gmx::ArrayRef<real>();
+    initialize_lambdas(
+            fplog, ir->efep, ir->bSimTemp, *ir->fepvals, ir->simtempvals->temperatures, ekind, MAIN(cr), fep_state, lambda);
+    Update upd(*ir, *ekind, deform);
+
+    // Simulated annealing updates the reference temperature.
+    const bool doSimulatedAnnealing = initSimulatedAnnealing(*ir, ekind, &upd);
+
+    const bool useReplicaExchange = (replExParams.exchangeInterval > 0);
+
+    t_fcdata& fcdata = *fr->fcdata;
+
+    bool simulationsShareState       = false;
+    bool simulationsShareHamiltonian = false;
+    int  nstSignalComm               = nstglobalcomm;
+    {
+        // TODO This implementation of ensemble orientation restraints is nasty because
+        // a user can't just do multi-sim with single-sim orientation restraints.
+        bool usingEnsembleRestraints = (fcdata.disres->nsystems > 1) || ((ms != nullptr) && fcdata.orires);
+        bool awhUsesMultiSim = (ir->bDoAwh && ir->awhParams->shareBiasMultisim() && (ms != nullptr));
+
+        // Replica exchange, ensemble restraints and AWH need all
+        // simulations to remain synchronized, so they need
+        // checkpoints and stop conditions to act on the same step, so
+        // the propagation of such signals must take place between
+        // simulations, not just within simulations.
+        // TODO: Make algorithm initializers set these flags.
+        simulationsShareState = useReplicaExchange || usingEnsembleRestraints || awhUsesMultiSim || (plumedswitch && ms); // PLUMED hack, if we have multiple sim and plumed we usually want them to be in sync 
+
+        // With AWH with bias sharing each simulation uses an non-shared, but identical, Hamiltonian
+        simulationsShareHamiltonian = useReplicaExchange || usingEnsembleRestraints;
+
+        if (simulationsShareState)
+        {
+            // Inter-simulation signal communication does not need to happen
+            // often, so we use a minimum of 200 steps to reduce overhead.
+            const int c_minimumInterSimulationSignallingInterval = 200;
+            nstSignalComm = ((c_minimumInterSimulationSignallingInterval + nstglobalcomm - 1) / nstglobalcomm)
+                            * nstglobalcomm;
+        }
+    }
+
+    if (startingBehavior != StartingBehavior::RestartWithAppending)
+    {
+        pleaseCiteCouplingAlgorithms(fplog, *ir);
+    }
+    gmx_mdoutf*       outf = init_mdoutf(fplog,
+                                   nfile,
+                                   fnm,
+                                   mdrunOptions,
+                                   cr,
+                                   outputProvider,
+                                   mdModulesNotifiers,
+                                   ir,
+                                   top_global,
+                                   oenv,
+                                   wcycle,
+                                   startingBehavior,
+                                   simulationsShareState,
+                                   ms);
+    gmx::EnergyOutput energyOutput(mdoutf_get_fp_ene(outf),
+                                   top_global,
+                                   *ir,
+                                   pull_work,
+                                   mdoutf_get_fp_dhdl(outf),
+                                   false,
+                                   startingBehavior,
+                                   simulationsShareHamiltonian,
+                                   mdModulesNotifiers);
+
+    gstat = global_stat_init(ir);
+
+    const auto& simulationWork     = runScheduleWork->simulationWork;
+    const bool  useGpuForPme       = simulationWork.useGpuPme;
+    const bool  useGpuForNonbonded = simulationWork.useGpuNonbonded;
+    const bool  useGpuForUpdate    = simulationWork.useGpuUpdate;
+
+    /* Check for polarizable models and flexible constraints */
+    shellfc = init_shell_flexcon(fplog,
+                                 top_global,
+                                 constr ? constr->numFlexibleConstraints() : 0,
+                                 ir->nstcalcenergy,
+                                 haveDDAtomOrdering(*cr),
+                                 useGpuForPme);
+
+    {
+        double io = compute_io(ir, top_global.natoms, *groups, energyOutput.numEnergyTerms(), 1);
+        if ((io > 2000) && MAIN(cr))
+        {
+            fprintf(stderr, "\nWARNING: This run will generate roughly %.0f Mb of data\n\n", io);
+        }
+    }
+
+    ObservablesReducer observablesReducer = observablesReducerBuilder->build();
+
+    ForceBuffers     f(simulationWork.useMts,
+                   (simulationWork.useGpuFBufferOps || useGpuForUpdate) ? PinningPolicy::PinnedIfSupported
+                                                                            : PinningPolicy::CannotBePinned);
+    const t_mdatoms* md = mdAtoms->mdatoms();
+    if (haveDDAtomOrdering(*cr))
+    {
+        // Local state only becomes valid now.
+        dd_init_local_state(*cr->dd, state_global, state);
+
+        /* Distribute the charge groups over the nodes from the main node */
+        dd_partition_system(fplog,
+                            mdlog,
+                            ir->init_step,
+                            cr,
+                            TRUE,
+                            state_global,
+                            top_global,
+                            *ir,
+                            imdSession,
+                            pull_work,
+                            state,
+                            &f,
+                            mdAtoms,
+                            top,
+                            fr,
+                            vsite,
+                            constr,
+                            nrnb,
+                            nullptr,
+                            FALSE);
+        upd.updateAfterPartition(state->natoms, md->cFREEZE, md->cTC, md->cACC);
+        fr->longRangeNonbondeds->updateAfterPartition(*md);
+    }
+    else
+    {
+        state_change_natoms(state_global, state_global->natoms);
+
+        /* Generate and initialize new topology */
+        mdAlgorithmsSetupAtomData(cr, *ir, top_global, top, fr, &f, mdAtoms, constr, vsite, shellfc);
+
+        upd.updateAfterPartition(state->natoms, md->cFREEZE, md->cTC, md->cACC);
+        fr->longRangeNonbondeds->updateAfterPartition(*md);
+    }
+
+    // Now that the state is valid we can set up Parrinello-Rahman
+    init_parrinellorahman(ir->pressureCouplingOptions,
+                          ir->deform,
+                          ir->delta_t * ir->pressureCouplingOptions.nstpcouple,
+                          state->box,
+                          state->box_rel,
+                          state->boxv,
+                          &parrinelloRahmanM,
+                          &pressureCouplingMu);
+
+    std::unique_ptr<UpdateConstrainGpu> integrator;
+
+    StatePropagatorDataGpu* stateGpu = fr->stateGpu;
+
+    // TODO: the assertions below should be handled by UpdateConstraintsBuilder.
+    if (useGpuForUpdate)
+    {
+        GMX_RELEASE_ASSERT(!haveDDAtomOrdering(*cr) || ddUsesUpdateGroups(*cr->dd)
+                                   || constr == nullptr || constr->numConstraintsTotal() == 0,
+                           "Constraints in domain decomposition are only supported with update "
+                           "groups if using GPU update.\n");
+        GMX_RELEASE_ASSERT(ir->eConstrAlg != ConstraintAlgorithm::Shake || constr == nullptr
+                                   || constr->numConstraintsTotal() == 0,
+                           "SHAKE is not supported with GPU update.");
+        GMX_RELEASE_ASSERT(useGpuForPme || (useGpuForNonbonded && simulationWork.useGpuXBufferOps),
+                           "Either PME or short-ranged non-bonded interaction tasks must run on "
+                           "the GPU to use GPU update.\n");
+        GMX_RELEASE_ASSERT(ir->eI == IntegrationAlgorithm::MD,
+                           "Only the md integrator is supported with the GPU update.\n");
+        GMX_RELEASE_ASSERT(
+                ir->etc != TemperatureCoupling::NoseHoover,
+                "Nose-Hoover temperature coupling is not supported with the GPU update.\n");
+        GMX_RELEASE_ASSERT(
+                ir->pressureCouplingOptions.epc == PressureCoupling::No
+                        || ir->pressureCouplingOptions.epc == PressureCoupling::ParrinelloRahman
+                        || ir->pressureCouplingOptions.epc == PressureCoupling::Berendsen
+                        || ir->pressureCouplingOptions.epc == PressureCoupling::CRescale,
+                "Only Parrinello-Rahman, Berendsen, and C-rescale pressure coupling are supported "
+                "with the GPU update.\n");
+        GMX_RELEASE_ASSERT(!md->haveVsites,
+                           "Virtual sites are not supported with the GPU update.\n");
+        GMX_RELEASE_ASSERT(ed == nullptr,
+                           "Essential dynamics is not supported with the GPU update.\n");
+        GMX_RELEASE_ASSERT(!ir->bPull || !pull_have_constraint(*ir->pull),
+                           "Constraints pulling is not supported with the GPU update.\n");
+        GMX_RELEASE_ASSERT(fcdata.orires == nullptr,
+                           "Orientation restraints are not supported with the GPU update.\n");
+        GMX_RELEASE_ASSERT(
+                ir->efep == FreeEnergyPerturbationType::No
+                        || (!haveFepPerturbedMasses(top_global) && !havePerturbedConstraints(top_global)),
+                "Free energy perturbation of masses and constraints are not supported with the GPU "
+                "update.");
+
+        if (constr != nullptr && constr->numConstraintsTotal() > 0)
+        {
+            GMX_LOG(mdlog.info)
+                    .asParagraph()
+                    .appendText("Updating coordinates and applying constraints on the GPU.");
+        }
+        else
+        {
+            GMX_LOG(mdlog.info).asParagraph().appendText("Updating coordinates on the GPU.");
+        }
+        GMX_RELEASE_ASSERT(fr->deviceStreamManager != nullptr,
+                           "Device stream manager should be initialized in order to use GPU "
+                           "update-constraints.");
+        GMX_RELEASE_ASSERT(
+                fr->deviceStreamManager->streamIsValid(gmx::DeviceStreamType::UpdateAndConstraints),
+                "Update stream should be initialized in order to use GPU "
+                "update-constraints.");
+        integrator = std::make_unique<UpdateConstrainGpu>(
+                *ir,
+                top_global,
+                ekind->numTemperatureCouplingGroups(),
+                fr->deviceStreamManager->context(),
+                fr->deviceStreamManager->stream(gmx::DeviceStreamType::UpdateAndConstraints),
+                wcycle);
+
+        stateGpu->setXUpdatedOnDeviceEvent(integrator->xUpdatedOnDeviceEvent());
+
+        integrator->setPbc(PbcType::Xyz, state->box);
+    }
+
+    if (useGpuForPme || simulationWork.useGpuXBufferOps || useGpuForUpdate)
+    {
+        changePinningPolicy(&state->x, PinningPolicy::PinnedIfSupported);
+    }
+    if (useGpuForUpdate)
+    {
+        changePinningPolicy(&state->v, PinningPolicy::PinnedIfSupported);
+    }
+
+    // NOTE: The global state is no longer used at this point.
+    // But state_global is still used as temporary storage space for writing
+    // the global state to file and potentially for replica exchange.
+    // (Global topology should persist.)
+
+    update_mdatoms(mdAtoms->mdatoms(), state->lambda[FreeEnergyPerturbationCouplingType::Mass]);
+
+    if (ir->bExpanded)
+    {
+        /* Check nstexpanded here, because the grompp check was broken */
+        if (ir->expandedvals->nstexpanded % ir->nstcalcenergy != 0)
+        {
+            gmx_fatal(FARGS,
+                      "With expanded ensemble, nstexpanded should be a multiple of nstcalcenergy");
+        }
+        init_expanded_ensemble(startingBehavior != StartingBehavior::NewSimulation, ir, state->dfhist);
+    }
+
+    if (MAIN(cr))
+    {
+        EnergyData::initializeEnergyHistory(startingBehavior, observablesHistory, &energyOutput);
+    }
+
+    preparePrevStepPullCom(
+            ir, pull_work, md->massT, state, state_global, cr, startingBehavior != StartingBehavior::NewSimulation);
+
+    // TODO: Remove this by converting AWH into a ForceProvider
+    auto awh = prepareAwhModule(fplog,
+                                *ir,
+                                state_global,
+                                cr,
+                                ms,
+                                startingBehavior != StartingBehavior::NewSimulation,
+                                shellfc != nullptr,
+                                opt2fn("-awh", nfile, fnm),
+                                pull_work);
+
+    if (useReplicaExchange && MAIN(cr))
+    {
+        repl_ex = init_replica_exchange(fplog, ms, top_global.natoms, ir, replExParams);
+    }
+    /* PME tuning is only supported in the Verlet scheme, with PME for
+     * Coulomb. It is not supported with only LJ PME.
+     * Disable PME tuning with GPU PME decomposition */
+    bPMETune = (mdrunOptions.tunePme && usingPme(fr->ic->eeltype) && !mdrunOptions.reproducible
+                && ir->cutoff_scheme != CutoffScheme::Group && !simulationWork.useGpuPmeDecomposition);
+
+    pme_load_balancing_t* pme_loadbal = nullptr;
+    if (bPMETune)
+    {
+        pme_loadbal_init(
+                &pme_loadbal, cr, mdlog, *ir, state->box, *fr->ic, *fr->nbv, fr->pmedata, fr->nbv->useGpu());
+    }
+
+    if (!ir->bContinuation)
+    {
+        if (state->flags & enumValueToBitMask(StateEntry::V))
+        {
+            auto v = makeArrayRef(state->v);
+            /* Set the velocities of vsites, shells and frozen atoms to zero */
+            for (i = 0; i < md->homenr; i++)
+            {
+                if (md->ptype[i] == ParticleType::Shell)
+                {
+                    clear_rvec(v[i]);
+                }
+                else if (!md->cFREEZE.empty())
+                {
+                    for (m = 0; m < DIM; m++)
+                    {
+                        if (ir->opts.nFreeze[md->cFREEZE[i]][m])
+                        {
+                            v[i][m] = 0;
+                        }
+                    }
+                }
+            }
+        }
+
+        if (constr)
+        {
+            /* Constrain the initial coordinates and velocities */
+            do_constrain_first(fplog,
+                               constr,
+                               ir,
+                               md->nr,
+                               md->homenr,
+                               state->x.arrayRefWithPadding(),
+                               state->v.arrayRefWithPadding(),
+                               state->box,
+                               state->lambda[FreeEnergyPerturbationCouplingType::Bonded]);
+        }
+    }
+
+    const int nstfep = computeFepPeriod(*ir, replExParams);
+
+    /* Be REALLY careful about what flags you set here. You CANNOT assume
+     * this is the first step, since we might be restarting from a checkpoint,
+     * and in that case we should not do any modifications to the state.
+     */
+    bStopCM = (ir->comm_mode != ComRemovalAlgorithm::No && !ir->bContinuation);
+
+    // When restarting from a checkpoint, it can be appropriate to
+    // initialize ekind from quantities in the checkpoint. Otherwise,
+    // compute_globals must initialize ekind before the simulation
+    // starts/restarts. However, only the main rank knows what was
+    // found in the checkpoint file, so we have to communicate in
+    // order to coordinate the restart.
+    //
+    // TODO Consider removing this communication if/when checkpoint
+    // reading directly follows .tpr reading, because all ranks can
+    // agree on hasReadEkinState at that time.
+    bool hasReadEkinState = MAIN(cr) ? state_global->ekinstate.hasReadEkinState : false;
+    if (PAR(cr))
+    {
+        gmx_bcast(sizeof(hasReadEkinState), &hasReadEkinState, cr->mpi_comm_mygroup);
+    }
+    if (hasReadEkinState)
+    {
+        restore_ekinstate_from_state(cr, ekind, MAIN(cr) ? &state_global->ekinstate : nullptr);
+    }
+
+    unsigned int cglo_flags =
+            (CGLO_TEMPERATURE | CGLO_GSTAT | (EI_VV(ir->eI) ? CGLO_PRESSURE : 0)
+             | (EI_VV(ir->eI) ? CGLO_CONSTRAINT : 0) | (hasReadEkinState ? CGLO_READEKIN : 0));
+
+    bSumEkinhOld = FALSE;
+
+    t_vcm vcm(top_global.groups, *ir);
+    reportComRemovalInfo(fplog, vcm);
+
+    int64_t step     = ir->init_step;
+    int64_t step_rel = 0;
+
+    /* To minimize communication, compute_globals computes the COM velocity
+     * and the kinetic energy for the velocities without COM motion removed.
+     * Thus to get the kinetic energy without the COM contribution, we need
+     * to call compute_globals twice.
+     */
+    for (int cgloIteration = 0; cgloIteration < (bStopCM ? 2 : 1); cgloIteration++)
+    {
+        unsigned int cglo_flags_iteration = cglo_flags;
+        if (bStopCM && cgloIteration == 0)
+        {
+            cglo_flags_iteration |= CGLO_STOPCM;
+            cglo_flags_iteration &= ~CGLO_TEMPERATURE;
+        }
+        compute_globals(gstat,
+                        cr,
+                        ir,
+                        fr,
+                        ekind,
+                        makeConstArrayRef(state->x),
+                        makeConstArrayRef(state->v),
+                        state->box,
+                        md,
+                        nrnb,
+                        &vcm,
+                        nullptr,
+                        enerd,
+                        force_vir,
+                        shake_vir,
+                        total_vir,
+                        pres,
+                        &nullSignaller,
+                        state->box,
+                        &bSumEkinhOld,
+                        cglo_flags_iteration,
+                        step,
+                        &observablesReducer);
+        // Clean up after pre-step use of compute_globals()
+        observablesReducer.markAsReadyToReduce();
+
+        if (cglo_flags_iteration & CGLO_STOPCM)
+        {
+            /* At initialization, do not pass x with acceleration-correction mode
+             * to avoid (incorrect) correction of the initial coordinates.
+             */
+            auto x = (vcm.mode == ComRemovalAlgorithm::LinearAccelerationCorrection)
+                             ? ArrayRef<RVec>()
+                             : makeArrayRef(state->x);
+            process_and_stopcm_grp(fplog, &vcm, *md, x, makeArrayRef(state->v));
+            inc_nrnb(nrnb, eNR_STOPCM, md->homenr);
+        }
+    }
+    if (ir->eI == IntegrationAlgorithm::VVAK)
+    {
+        /* a second call to get the half step temperature initialized as well */
+        /* we do the same call as above, but turn the pressure off -- internally to
+           compute_globals, this is recognized as a velocity verlet half-step
+           kinetic energy calculation.  This minimized excess variables, but
+           perhaps loses some logic?*/
+
+        compute_globals(gstat,
+                        cr,
+                        ir,
+                        fr,
+                        ekind,
+                        makeConstArrayRef(state->x),
+                        makeConstArrayRef(state->v),
+                        state->box,
+                        md,
+                        nrnb,
+                        &vcm,
+                        nullptr,
+                        enerd,
+                        force_vir,
+                        shake_vir,
+                        total_vir,
+                        pres,
+                        &nullSignaller,
+                        state->box,
+                        &bSumEkinhOld,
+                        cglo_flags & ~CGLO_PRESSURE,
+                        step,
+                        &observablesReducer);
+        // Clean up after pre-step use of compute_globals()
+        observablesReducer.markAsReadyToReduce();
+    }
+
+    /* Calculate the initial half step temperature, and save the ekinh_old */
+    if (startingBehavior == StartingBehavior::NewSimulation)
+    {
+        for (i = 0; (i < ir->opts.ngtc); i++)
+        {
+            copy_mat(ekind->tcstat[i].ekinh, ekind->tcstat[i].ekinh_old);
+        }
+    }
+
+    /* need to make an initiation call to get the Trotter variables set, as well as other constants
+       for non-trotter temperature control */
+    auto trotter_seq = init_npt_vars(ir, *ekind, state, &MassQ, bTrotter);
+
+    if (MAIN(cr))
+    {
+        if (!ir->bContinuation)
+        {
+            if (constr && ir->eConstrAlg == ConstraintAlgorithm::Lincs)
+            {
+                fprintf(fplog,
+                        "RMS relative constraint deviation after constraining: %.2e\n",
+                        constr->rmsd());
+            }
+            if (EI_STATE_VELOCITY(ir->eI))
+            {
+                real temp = enerd->term[F_TEMP];
+                if (ir->eI != IntegrationAlgorithm::VV)
+                {
+                    /* Result of Ekin averaged over velocities of -half
+                     * and +half step, while we only have -half step here.
+                     */
+                    temp *= 2;
+                }
+                fprintf(fplog, "Initial temperature: %g K\n", temp);
+            }
+        }
+
+        char tbuf[20];
+        fprintf(stderr, "starting mdrun '%s'\n", *(top_global.name));
+        if (ir->nsteps >= 0)
+        {
+            sprintf(tbuf, "%8.1f", (ir->init_step + ir->nsteps) * ir->delta_t);
+        }
+        else
+        {
+            sprintf(tbuf, "%s", "infinite");
+        }
+        if (ir->init_step > 0)
+        {
+            fprintf(stderr,
+                    "%s steps, %s ps (continuing from step %s, %8.1f ps).\n",
+                    gmx_step_str(ir->init_step + ir->nsteps, sbuf),
+                    tbuf,
+                    gmx_step_str(ir->init_step, sbuf2),
+                    ir->init_step * ir->delta_t);
+        }
+        else
+        {
+            fprintf(stderr, "%s steps, %s ps.\n", gmx_step_str(ir->nsteps, sbuf), tbuf);
+        }
+        fprintf(fplog, "\n");
+    }
+
+    /* PLUMED */
+    if(plumedswitch){
+      /* detect plumed API version */
+      int pversion=0;
+      plumed_cmd(plumedmain,"getApiVersion",&pversion);
+      /* setting kbT is only implemented with api>1) */
+      real kbT=ir->opts.ref_t[0]*gmx::c_boltz;
+      if(pversion>1) plumed_cmd(plumedmain,"setKbT",&kbT);
+      if(pversion>2){
+        int res=1;
+        if( (startingBehavior != StartingBehavior::NewSimulation) ) plumed_cmd(plumedmain,"setRestart",&res);
+      }
+
+      if(isMultiSim(ms)) {
+        if(MAIN(cr)) plumed_cmd(plumedmain,"GREX setMPIIntercomm",&ms->mainRanksComm_);
+        if(PAR(cr)){
+          if(haveDDAtomOrdering(*cr)) {
+            plumed_cmd(plumedmain,"GREX setMPIIntracomm",&cr->dd->mpi_comm_all);
+          }else{
+            plumed_cmd(plumedmain,"GREX setMPIIntracomm",&cr->mpi_comm_mysim);
+          }
+        }
+        plumed_cmd(plumedmain,"GREX init",nullptr);
+      }
+      if(PAR(cr)){
+        if(haveDDAtomOrdering(*cr)) {
+          plumed_cmd(plumedmain,"setMPIComm",&cr->dd->mpi_comm_all);
+        }
+      }
+      plumed_cmd(plumedmain,"setNatoms",top_global.natoms);
+      plumed_cmd(plumedmain,"setMDEngine","gromacs");
+      plumed_cmd(plumedmain,"setLog",fplog);
+      real real_delta_t=ir->delta_t;
+      plumed_cmd(plumedmain,"setTimestep",&real_delta_t);
+      plumed_cmd(plumedmain,"init",nullptr);
+
+      if(haveDDAtomOrdering(*cr)) {
+        int nat_home = dd_numHomeAtoms(*cr->dd);
+        plumed_cmd(plumedmain,"setAtomsNlocal",&nat_home);
+        plumed_cmd(plumedmain,"setAtomsGatindex",cr->dd->globalAtomIndices.data());
+      }
+      realFepState = state->fep_state;
+      plumed_cmd(plumedmain, "setExtraCV lambda", &realFepState);
+      plumed_cmd(plumedmain, "setExtraCVForce lambda", &lambdaForce);
+    }
+    /* END PLUMED */
+
+    walltime_accounting_start_time(walltime_accounting);
+    wallcycle_start(wcycle, WallCycleCounter::Run);
+    print_start(fplog, cr, walltime_accounting, "mdrun");
+
+    /***********************************************************
+     *
+     *             Loop over MD steps
+     *
+     ************************************************************/
+
+    bFirstStep = TRUE;
+    /* Skip the first Nose-Hoover integration when we get the state from tpx */
+    bInitStep        = startingBehavior == StartingBehavior::NewSimulation || EI_VV(ir->eI);
+    bSumEkinhOld     = FALSE;
+    bExchanged       = FALSE;
+    bNeedRepartition = FALSE;
+    bDoReplEx        = FALSE;
+
+    auto stopHandler = stopHandlerBuilder->getStopHandlerMD(
+            compat::not_null<SimulationSignal*>(&signals[eglsSTOPCOND]),
+            simulationsShareState,
+            MAIN(cr),
+            ir->nstlist,
+            mdrunOptions.reproducible,
+            nstSignalComm,
+            mdrunOptions.maximumHoursToRun,
+            ir->nstlist == 0,
+            fplog,
+            step,
+            bNS,
+            walltime_accounting);
+
+    real checkpointPeriod = mdrunOptions.checkpointOptions.period;
+    if (ir->bExpanded)
+    {
+        GMX_LOG(mdlog.info)
+                .asParagraph()
+                .appendText(
+                        "Expanded ensemble with the legacy simulator does not always "
+                        "checkpoint correctly, so checkpointing is disabled. You will "
+                        "not be able to do a checkpoint restart of this simulation. "
+                        "If you use the modular simulator (e.g. by choosing md-vv integrator) "
+                        "then checkpointing is enabled. See "
+                        "https://gitlab.com/gromacs/gromacs/-/issues/4629 for details.");
+        // Use a negative period to disable checkpointing.
+        checkpointPeriod = -1;
+    }
+    auto checkpointHandler = std::make_unique<CheckpointHandler>(
+            compat::make_not_null<SimulationSignal*>(&signals[eglsCHKPT]),
+            simulationsShareState,
+            ir->nstlist == 0,
+            MAIN(cr),
+            mdrunOptions.writeConfout,
+            checkpointPeriod);
+
+    const bool resetCountersIsLocal = true;
+    auto       resetHandler         = std::make_unique<ResetHandler>(
+            compat::make_not_null<SimulationSignal*>(&signals[eglsRESETCOUNTERS]),
+            !resetCountersIsLocal,
+            ir->nsteps,
+            MAIN(cr),
+            mdrunOptions.timingOptions.resetHalfway,
+            mdrunOptions.maximumHoursToRun,
+            mdlog,
+            wcycle,
+            walltime_accounting);
+
+    const DDBalanceRegionHandler ddBalanceRegionHandler(cr);
+
+    if (MAIN(cr) && isMultiSim(ms) && !useReplicaExchange)
+    {
+        logInitialMultisimStatus(ms, cr, mdlog, simulationsShareState, ir->nsteps, ir->init_step);
+    }
+
+    bool usedMdGpuGraphLastStep = false;
+    /* and stop now if we should */
+    bLastStep = (bLastStep || (ir->nsteps >= 0 && step_rel > ir->nsteps));
+    while (!bLastStep)
+    {
+
+        /* Determine if this is a neighbor search step */
+        bNStList = (ir->nstlist > 0 && step % ir->nstlist == 0);
+
+        if (bPMETune && bNStList)
+        {
+            // This has to be here because PME load balancing is called so early.
+            // TODO: Move to after all booleans are defined.
+            if (useGpuForUpdate && !bFirstStep)
+            {
+                stateGpu->copyCoordinatesFromGpu(state->x, AtomLocality::Local);
+                stateGpu->waitCoordinatesReadyOnHost(AtomLocality::Local);
+            }
+            /* PME grid + cut-off optimization with GPUs or PME nodes */
+            pme_loadbal_do(pme_loadbal,
+                           cr,
+                           (mdrunOptions.verbose && MAIN(cr)) ? stderr : nullptr,
+                           fplog,
+                           mdlog,
+                           *ir,
+                           fr,
+                           state->box,
+                           state->x,
+                           wcycle,
+                           step,
+                           step_rel,
+                           &bPMETunePrinting,
+                           simulationWork.useGpuPmePpCommunication);
+        }
+
+        wallcycle_start(wcycle, WallCycleCounter::Step);
+
+        bLastStep = (step_rel == ir->nsteps);
+        t         = t0 + step * ir->delta_t;
+
+        // TODO Refactor this, so that nstfep does not need a default value of zero
+        if (ir->efep != FreeEnergyPerturbationType::No || ir->bSimTemp)
+        {
+            /* find and set the current lambdas */
+            state->lambda = currentLambdas(step, *(ir->fepvals), state->fep_state);
+
+            bDoExpanded = (do_per_step(step, ir->expandedvals->nstexpanded) && (ir->bExpanded)
+                           && (!bFirstStep));
+        }
+
+        bDoReplExPrev = bDoReplEx;
+        bDoReplEx = (useReplicaExchange && (step > 0) && !bLastStep
+                     && do_per_step(step, replExParams.exchangeInterval));
+
+        if (doSimulatedAnnealing)
+        {
+            // Simulated annealing updates the reference temperature.
+            update_annealing_target_temp(*ir, t, ekind, &upd);
+        }
+
+        /* Stop Center of Mass motion */
+        bStopCM = (ir->comm_mode != ComRemovalAlgorithm::No && do_per_step(step, ir->nstcomm));
+
+        /* Determine whether or not to do Neighbour Searching */
+        bNS = (bFirstStep || bNStList || bExchanged || bNeedRepartition);
+
+        /* Note that the stopHandler will cause termination at nstglobalcomm
+         * steps. Since this concides with nstcalcenergy, nsttcouple and/or
+         * nstpcouple steps, we have computed the half-step kinetic energy
+         * of the previous step and can always output energies at the last step.
+         */
+        bLastStep = bLastStep || stopHandler->stoppingAfterCurrentStep(bNS);
+
+        /* do_log triggers energy and virial calculation. Because this leads
+         * to different code paths, forces can be different. Thus for exact
+         * continuation we should avoid extra log output.
+         * Note that the || bLastStep can result in non-exact continuation
+         * beyond the last step. But we don't consider that to be an issue.
+         */
+        do_log     = (do_per_step(step, ir->nstlog)
+                  || (bFirstStep && startingBehavior == StartingBehavior::NewSimulation) || bLastStep);
+        do_verbose = mdrunOptions.verbose
+                     && (step % mdrunOptions.verboseStepPrintInterval == 0 || bFirstStep || bLastStep);
+
+        // On search steps, when doing the update on the GPU, copy
+        // the coordinates and velocities to the host unless they are
+        // already there (ie on the first step and after replica
+        // exchange).
+        if (useGpuForUpdate && bNS && !bFirstStep && !bExchanged)
+        {
+            if (usedMdGpuGraphLastStep)
+            {
+                // Wait on coordinates produced from GPU graph
+                stateGpu->waitCoordinatesUpdatedOnDevice();
+            }
+            stateGpu->copyVelocitiesFromGpu(state->v, AtomLocality::Local);
+            stateGpu->copyCoordinatesFromGpu(state->x, AtomLocality::Local);
+            stateGpu->waitVelocitiesReadyOnHost(AtomLocality::Local);
+            stateGpu->waitCoordinatesReadyOnHost(AtomLocality::Local);
+        }
+
+        // We need to calculate virtual velocities if we are writing them in the current step.
+        // They also need to be periodically updated. Every 1000 steps is arbitrary, but a reasonable number.
+        // The reason why the velocities need to be updated regularly is that the virtual site coordinates
+        // are updated using these velocities during integration. Those coordinates are used for, e.g., domain
+        // decomposition. Before computing any forces the positions of the virtual sites are recalculated.
+        // This fixes a bug, #4879, which was introduced in MR !979.
+        const int  c_virtualSiteVelocityUpdateInterval = 1000;
+        const bool needVirtualVelocitiesThisStep =
+                (vsite != nullptr)
+                && (do_per_step(step, ir->nstvout) || checkpointHandler->isCheckpointingStep()
+                    || do_per_step(step, c_virtualSiteVelocityUpdateInterval));
+
+        if (vsite != nullptr)
+        {
+            // Virtual sites need to be updated before domain decomposition and forces are calculated
+            wallcycle_start(wcycle, WallCycleCounter::VsiteConstr);
+            // md-vv calculates virtual velocities once it has full-step real velocities
+            vsite->construct(state->x,
+                             state->v,
+                             state->box,
+                             (!EI_VV(inputrec->eI) && needVirtualVelocitiesThisStep)
+                                     ? VSiteOperation::PositionsAndVelocities
+                                     : VSiteOperation::Positions);
+            wallcycle_stop(wcycle, WallCycleCounter::VsiteConstr);
+        }
+
+        if (bNS && !(bFirstStep && ir->bContinuation))
+        {
+            bMainState = FALSE;
+            /* Correct the new box if it is too skewed */
+            if (inputrecDynamicBox(ir))
+            {
+                if (correct_box(fplog, step, state->box))
+                {
+                    bMainState = TRUE;
+                }
+            }
+            // If update is offloaded, and the box was changed either
+            // above or in a replica exchange on the previous step,
+            // the GPU Update object should be informed
+            if (useGpuForUpdate && (bMainState || bExchanged))
+            {
+                integrator->setPbc(PbcType::Xyz, state->box);
+            }
+            if (haveDDAtomOrdering(*cr) && bMainState)
+            {
+                dd_collect_state(cr->dd, state, state_global);
+            }
+
+            if (haveDDAtomOrdering(*cr))
+            {
+                /* Repartition the domain decomposition */
+                dd_partition_system(fplog,
+                                    mdlog,
+                                    step,
+                                    cr,
+                                    bMainState,
+                                    state_global,
+                                    top_global,
+                                    *ir,
+                                    imdSession,
+                                    pull_work,
+                                    state,
+                                    &f,
+                                    mdAtoms,
+                                    top,
+                                    fr,
+                                    vsite,
+                                    constr,
+                                    nrnb,
+                                    wcycle,
+                                    do_verbose && !bPMETunePrinting);
+                upd.updateAfterPartition(state->natoms, md->cFREEZE, md->cTC, md->cACC);
+                fr->longRangeNonbondeds->updateAfterPartition(*md);
+                /* PLUMED */
+                if(plumedswitch){
+                  int nat_home = dd_numHomeAtoms(*cr->dd);
+                  plumed_cmd(plumedmain,"setAtomsNlocal",&nat_home);
+                  plumed_cmd(plumedmain,"setAtomsGatindex",cr->dd->globalAtomIndices.data());
+                }
+                /* END PLUMED */
+            }
+        }
+
+        // Allocate or re-size GPU halo exchange object, if necessary
+        if (bNS && simulationWork.havePpDomainDecomposition && simulationWork.useGpuHaloExchange)
+        {
+            GMX_RELEASE_ASSERT(fr->deviceStreamManager != nullptr,
+                               "GPU device manager has to be initialized to use GPU "
+                               "version of halo exchange.");
+            constructGpuHaloExchange(*cr, *fr->deviceStreamManager, wcycle);
+        }
+
+        if (MAIN(cr) && do_log)
+        {
+            gmx::EnergyOutput::printHeader(
+                    fplog, step, t); /* can we improve the information printed here? */
+        }
+
+        if (ir->efep != FreeEnergyPerturbationType::No)
+        {
+            update_mdatoms(mdAtoms->mdatoms(), state->lambda[FreeEnergyPerturbationCouplingType::Mass]);
+        }
+
+        if (bExchanged)
+        {
+            /* We need the kinetic energy at minus the half step for determining
+             * the full step kinetic energy and possibly for T-coupling.*/
+            /* This may not be quite working correctly yet . . . . */
+            int cglo_flags = CGLO_GSTAT | CGLO_TEMPERATURE;
+            compute_globals(gstat,
+                            cr,
+                            ir,
+                            fr,
+                            ekind,
+                            makeConstArrayRef(state->x),
+                            makeConstArrayRef(state->v),
+                            state->box,
+                            md,
+                            nrnb,
+                            &vcm,
+                            wcycle,
+                            enerd,
+                            nullptr,
+                            nullptr,
+                            nullptr,
+                            nullptr,
+                            &nullSignaller,
+                            state->box,
+                            &bSumEkinhOld,
+                            cglo_flags,
+                            step,
+                            &observablesReducer);
+        }
+        clear_mat(force_vir);
+
+        /* PLUMED HREX */
+        gmx_bool bHREX = bDoReplEx && plumed_hrex;
+
+        if (plumedswitch && bHREX) {
+          // gmx_enerdata_t *hrex_enerd;
+          int nlambda = enerd->foreignLambdaTerms.numLambdas();
+          gmx_enerdata_t hrex_enerd(enerd->grpp.nener, nlambda == 0 ? nullptr: &inputrec->fepvals->all_lambda);
+          int repl  = -1;
+          int nrepl = -1;
+          if (MAIN(cr)){
+            repl  = replica_exchange_get_repl(repl_ex);
+            nrepl = replica_exchange_get_nrepl(repl_ex);
+          }
+
+          if (haveDDAtomOrdering(*cr)) {
+            dd_collect_state(cr->dd,state,state_global);
+          } else {
+            copy_state_serial(state, state_global);
+          }
+
+          if(MAIN(cr)){
+            if(repl%2==step/replExParams.exchangeInterval%2){
+              if(repl-1>=0) exchange_state(ms,repl-1,state_global);
+            }else{
+              if(repl+1<nrepl) exchange_state(ms,repl+1,state_global);
+            }
+          }
+          if (!haveDDAtomOrdering(*cr)) {
+            copy_state_serial(state_global, state);
+          }
+          if (haveDDAtomOrdering(*cr)) {
+              dd_partition_system(fplog,
+                                  mdlog,
+                                  step,
+                                  cr,
+                                  TRUE,
+                                  state_global,
+                                  top_global,
+                                  *ir,
+                                  imdSession,
+                                  pull_work,
+                                  state,
+                                  &f,
+                                  mdAtoms,
+                                  top,
+                                  fr,
+                                  vsite,
+                                  constr,
+                                  nrnb,
+                                  wcycle,
+                                  do_verbose && !bPMETunePrinting);
+          }
+          do_force(fplog,
+                   cr,
+                   ms,
+                   *ir,
+                   awh.get(),
+                   enforcedRotation,
+                   imdSession,
+                   pull_work,
+                   step,
+                   nrnb,
+                   wcycle,
+                   top,
+                   state->box,
+                   state->x.arrayRefWithPadding(),
+                   &state->hist,
+                   &f.view(),
+                   force_vir,
+                   md,
+                   &hrex_enerd,
+                   state->lambda,
+                   fr,
+                   runScheduleWork,
+                   vsite,
+                   mu_tot,
+                   t,
+                   ed ? ed->getLegacyED() : nullptr,
+                   fr->longRangeNonbondeds.get(),
+                   GMX_FORCE_STATECHANGED |
+                   GMX_FORCE_DYNAMICBOX |
+                   GMX_FORCE_ALLFORCES |
+                   GMX_FORCE_VIRIAL |
+                   GMX_FORCE_ENERGY |
+                   GMX_FORCE_DHDL |
+                   GMX_FORCE_NS,
+                   ddBalanceRegionHandler);
+
+          plumed_cmd(plumedmain,"GREX cacheLocalUSwap",&(&hrex_enerd)->term[F_EPOT]);
+
+          /* exchange back */
+          if (haveDDAtomOrdering(*cr)) {
+            dd_collect_state(cr->dd,state,state_global);
+          } else {
+            copy_state_serial(state, state_global);
+          }
+
+          if(MAIN(cr)){
+            if(repl%2==step/replExParams.exchangeInterval%2){
+              if(repl-1>=0) exchange_state(ms,repl-1,state_global);
+            }else{
+              if(repl+1<nrepl) exchange_state(ms,repl+1,state_global);
+            }
+          }
+
+          if (!haveDDAtomOrdering(*cr)) {
+            copy_state_serial(state_global, state);
+          }
+          if (haveDDAtomOrdering(*cr)) {
+              dd_partition_system(fplog,
+                                  mdlog,
+                                  step,
+                                  cr,
+                                  TRUE,
+                                  state_global,
+                                  top_global,
+                                  *ir,
+                                  imdSession,
+                                  pull_work,
+                                  state,
+                                  &f,
+                                  mdAtoms,
+                                  top,
+                                  fr,
+                                  vsite,
+                                  constr,
+                                  nrnb,
+                                  wcycle,
+                                  do_verbose && !bPMETunePrinting);
+            int nat_home = dd_numHomeAtoms(*cr->dd);
+            plumed_cmd(plumedmain,"setAtomsNlocal",&nat_home);
+            plumed_cmd(plumedmain,"setAtomsGatindex",cr->dd->globalAtomIndices.data());
+          }
+          bNS=true;
+        }
+        /* END PLUMED HREX */
+
+        checkpointHandler->decideIfCheckpointingThisStep(bNS||bDoReplExPrev, bFirstStep, bLastStep);
+
+        /* Determine the energy and pressure:
+         * at nstcalcenergy steps and at energy output steps (set below).
+         */
+        if (EI_VV(ir->eI) && (!bInitStep))
+        {
+            bCalcEnerStep = do_per_step(step, ir->nstcalcenergy);
+            bCalcVir      = bCalcEnerStep
+                       || (ir->pressureCouplingOptions.epc != PressureCoupling::No
+                           && (do_per_step(step, ir->pressureCouplingOptions.nstpcouple)
+                               || do_per_step(step - 1, ir->pressureCouplingOptions.nstpcouple)));
+        }
+        else
+        {
+            bCalcEnerStep = do_per_step(step, ir->nstcalcenergy);
+            bCalcVir      = bCalcEnerStep
+                       || (ir->pressureCouplingOptions.epc != PressureCoupling::No
+                           && do_per_step(step, ir->pressureCouplingOptions.nstpcouple));
+        }
+        bCalcEner = bCalcEnerStep;
+
+        do_ene = (do_per_step(step, ir->nstenergy) || bLastStep);
+
+        if (do_ene || do_log || bDoReplEx)
+        {
+            bCalcVir  = TRUE;
+            bCalcEner = TRUE;
+        }
+
+        // bCalcEner is only here for when the last step is not a mulitple of nstfep
+        const bool computeDHDL = ((ir->efep != FreeEnergyPerturbationType::No || ir->bSimTemp)
+                                  && (do_per_step(step, nstfep) || bCalcEner));
+
+        /* Do we need global communication ? */
+        bGStat = (bCalcVir || bCalcEner || bStopCM || do_per_step(step, nstglobalcomm)
+                  || (EI_VV(ir->eI) && inputrecNvtTrotter(ir) && do_per_step(step - 1, nstglobalcomm)));
+
+        force_flags = (GMX_FORCE_STATECHANGED | ((inputrecDynamicBox(ir)) ? GMX_FORCE_DYNAMICBOX : 0)
+                       | GMX_FORCE_ALLFORCES | (bCalcVir ? GMX_FORCE_VIRIAL : 0)
+                       | (bCalcEner ? GMX_FORCE_ENERGY : 0) | (computeDHDL ? GMX_FORCE_DHDL : 0));
+        if (simulationWork.useMts && !do_per_step(step, ir->nstfout))
+        {
+            // TODO: merge this with stepWork.useOnlyMtsCombinedForceBuffer
+            force_flags |= GMX_FORCE_DO_NOT_NEED_NORMAL_FORCE;
+        }
+
+        const bool doTemperatureScaling = (ir->etc != TemperatureCoupling::No
+                                           && do_per_step(step + ir->nsttcouple - 1, ir->nsttcouple));
+
+        /* With leap-frog type integrators we compute the kinetic energy
+         * at a whole time step as the average of the half-time step kinetic
+         * energies of two subsequent steps. Therefore we need to compute the
+         * half step kinetic energy also if we need energies at the next step.
+         */
+        const bool needHalfStepKineticEnergy =
+                (!EI_VV(ir->eI) && (do_per_step(step + 1, nstglobalcomm) || step_rel + 1 == ir->nsteps));
+
+        // Parrinello-Rahman requires the pressure to be availible before the update to compute
+        // the velocity scaling matrix. Hence, it runs one step after the nstpcouple step.
+        const bool doParrinelloRahman =
+                (ir->pressureCouplingOptions.epc == PressureCoupling::ParrinelloRahman
+                 && do_per_step(step + ir->pressureCouplingOptions.nstpcouple - 1,
+                                ir->pressureCouplingOptions.nstpcouple));
+
+        MdGpuGraph* mdGraph = simulationWork.useMdGpuGraph ? fr->mdGraph[step % 2].get() : nullptr;
+
+        if (simulationWork.useMdGpuGraph)
+        {
+            // Reset graph on search step (due to changing neighbour list etc)
+            // or virial step (due to changing shifts and box).
+            if (bNS || bCalcVir)
+            {
+                fr->mdGraph[MdGraphEvenOrOddStep::EvenStep]->reset();
+                fr->mdGraph[MdGraphEvenOrOddStep::OddStep]->reset();
+            }
+            else
+            {
+                mdGraph->setUsedGraphLastStep(usedMdGpuGraphLastStep);
+                bool canUseMdGpuGraphThisStep =
+                        !bNS && !bCalcVir && !doTemperatureScaling && !doParrinelloRahman && !bGStat
+                        && !needHalfStepKineticEnergy && !do_per_step(step, ir->nstxout)
+                        && !do_per_step(step, ir->nstxout_compressed)
+                        && !do_per_step(step, ir->nstvout) && !do_per_step(step, ir->nstfout)
+                        && !checkpointHandler->isCheckpointingStep();
+                if (mdGraph->captureThisStep(canUseMdGpuGraphThisStep))
+                {
+                    // getCoordinatesReadyOnDeviceEvent() uses stepWork.doNeighborSearch but this is not set until
+                    // later in do_force(), so preset this part here.
+                    // TODO remove this when stepWork initialization has been refactored to start of step.
+                    runScheduleWork->stepWork.doNeighborSearch = false;
+                    mdGraph->startRecord(stateGpu->getCoordinatesReadyOnDeviceEvent(
+                            AtomLocality::Local, simulationWork, runScheduleWork->stepWork));
+                }
+            }
+        }
+
+        if (!simulationWork.useMdGpuGraph || mdGraph->graphIsCapturingThisStep()
+            || !mdGraph->useGraphThisStep())
+        {
+
+            if (shellfc)
+            {
+                /* Now is the time to relax the shells */
+                relax_shell_flexcon(fplog,
+                                    cr,
+                                    ms,
+                                    mdrunOptions.verbose,
+                                    enforcedRotation,
+                                    step,
+                                    ir,
+                                    imdSession,
+                                    pull_work,
+                                    bNS,
+                                    force_flags,
+                                    top,
+                                    constr,
+                                    enerd,
+                                    state->natoms,
+                                    state->x.arrayRefWithPadding(),
+                                    state->v.arrayRefWithPadding(),
+                                    state->box,
+                                    state->lambda,
+                                    &state->hist,
+                                    &f.view(),
+                                    force_vir,
+                                    *md,
+                                    fr->longRangeNonbondeds.get(),
+                                    nrnb,
+                                    wcycle,
+                                    shellfc,
+                                    fr,
+                                    runScheduleWork,
+                                    t,
+                                    mu_tot,
+                                    vsite,
+                                    ddBalanceRegionHandler);
+            }
+            else
+            {
+                /* The AWH history need to be saved _before_ doing force calculations where the AWH bias
+                   is updated (or the AWH update will be performed twice for one step when continuing).
+                   It would be best to call this update function from do_md_trajectory_writing but that
+                   would occur after do_force. One would have to divide the update_awh function into one
+                   function applying the AWH force and one doing the AWH bias update. The update AWH
+                   bias function could then be called after do_md_trajectory_writing (then containing
+                   update_awh_history). The checkpointing will in the future probably moved to the start
+                   of the md loop which will rid of this issue. */
+                if (awh && checkpointHandler->isCheckpointingStep() && MAIN(cr))
+                {
+                    awh->updateHistory(state_global->awhHistory.get());
+                }
+
+                /* The coordinates (x) are shifted (to get whole molecules)
+                 * in do_force.
+                 * This is parallellized as well, and does communication too.
+                 * Check comments in sim_util.c
+                 */
+
+             /* PLUMED */
+            plumedNeedsEnergy=0;
+            if(plumedswitch){
+              int pversion=0;
+              plumed_cmd(plumedmain,"getApiVersion",&pversion);
+              long int lstep=step; plumed_cmd(plumedmain,"setStepLong",&lstep);
+              plumed_cmd(plumedmain,"setPositions",&state->x[0][0]);
+              plumed_cmd(plumedmain,"setMasses",&md->massT[0]);
+              plumed_cmd(plumedmain,"setCharges",&md->chargeA[0]);
+              plumed_cmd(plumedmain,"setBox",&state->box[0][0]);
+              plumed_cmd(plumedmain,"prepareCalc",nullptr);
+              plumed_cmd(plumedmain,"setStopFlag",&plumedWantsToStop);
+              int checkp=0; if(checkpointHandler->isCheckpointingStep()) checkp=1;
+              if(pversion>3) plumed_cmd(plumedmain,"doCheckPoint",&checkp);
+              plumed_cmd(plumedmain,"setForces",&f.view().force()[0][0]);
+              plumed_cmd(plumedmain,"isEnergyNeeded",&plumedNeedsEnergy);
+              if(plumedNeedsEnergy) force_flags |= GMX_FORCE_ENERGY | GMX_FORCE_VIRIAL;
+              clear_mat(plumed_vir);
+              plumed_cmd(plumedmain,"setVirial",&plumed_vir[0][0]);
+            }
+            /* END PLUMED */
+                do_force(fplog,
+                         cr,
+                         ms,
+                         *ir,
+                         awh.get(),
+                         enforcedRotation,
+                         imdSession,
+                         pull_work,
+                         step,
+                         nrnb,
+                         wcycle,
+                         top,
+                         state->box,
+                         state->x.arrayRefWithPadding(),
+                         &state->hist,
+                         &f.view(),
+                         force_vir,
+                         md,
+                         enerd,
+                         state->lambda,
+                         fr,
+                         runScheduleWork,
+                         vsite,
+                         mu_tot,
+                         t,
+                         ed ? ed->getLegacyED() : nullptr,
+                         fr->longRangeNonbondeds.get(),
+                         (bNS ? GMX_FORCE_NS : 0) | force_flags,
+                         ddBalanceRegionHandler);
+            /* PLUMED */
+            if(plumedswitch){
+              if(plumedNeedsEnergy){
+                msmul(force_vir,2.0,plumed_vir);
+                plumed_cmd(plumedmain,"setEnergy",&enerd->term[F_EPOT]);
+                plumed_cmd(plumedmain,"performCalc",nullptr);
+                msmul(plumed_vir,0.5,force_vir);
+              } else {
+                msmul(plumed_vir,0.5,plumed_vir);
+                m_add(force_vir,plumed_vir,force_vir);
+              }
+              if(bDoReplEx) plumed_cmd(plumedmain,"GREX savePositions",nullptr);
+              if(plumedWantsToStop) bLastStep = true;
+              if(bHREX) plumed_cmd(plumedmain,"GREX cacheLocalUNow",&enerd->term[F_EPOT]);
+            }
+            /* END PLUMED */
+            }
+
+            // VV integrators do not need the following velocity half step
+            // if it is the first step after starting from a checkpoint.
+            // That is, the half step is needed on all other steps, and
+            // also the first step when starting from a .tpr file.
+            if (EI_VV(ir->eI))
+            {
+                integrateVVFirstStep(step,
+                                     bFirstStep,
+                                     bInitStep,
+                                     startingBehavior,
+                                     nstglobalcomm,
+                                     ir,
+                                     fr,
+                                     cr,
+                                     state,
+                                     mdAtoms->mdatoms(),
+                                     &fcdata,
+                                     &MassQ,
+                                     &vcm,
+                                     enerd,
+                                     &observablesReducer,
+                                     ekind,
+                                     gstat,
+                                     &last_ekin,
+                                     bCalcVir,
+                                     total_vir,
+                                     shake_vir,
+                                     force_vir,
+                                     pres,
+                                     do_log,
+                                     do_ene,
+                                     bCalcEner,
+                                     bGStat,
+                                     bStopCM,
+                                     bTrotter,
+                                     bExchanged,
+                                     &bSumEkinhOld,
+                                     &saved_conserved_quantity,
+                                     &f,
+                                     &upd,
+                                     constr,
+                                     &nullSignaller,
+                                     trotter_seq,
+                                     nrnb,
+                                     fplog,
+                                     wcycle);
+                if (vsite != nullptr && needVirtualVelocitiesThisStep)
+                {
+                    // Positions were calculated earlier
+                    wallcycle_start(wcycle, WallCycleCounter::VsiteConstr);
+                    vsite->construct(state->x, state->v, state->box, VSiteOperation::Velocities);
+                    wallcycle_stop(wcycle, WallCycleCounter::VsiteConstr);
+                }
+            }
+
+            /* ########  END FIRST UPDATE STEP  ############## */
+            /* ########  If doing VV, we now have v(dt) ###### */
+            if (bDoExpanded)
+            {
+                /* perform extended ensemble sampling in lambda - we don't
+                   actually move to the new state before outputting
+                   statistics, but if performing simulated tempering, we
+                   do update the velocities and the tau_t. */
+                lamnew = ExpandedEnsembleDynamics(fplog,
+                                                  *inputrec,
+                                                  *enerd,
+                                                  ekind,
+                                                  state,
+                                                  &MassQ,
+                                                  state->fep_state,
+                                                  state->dfhist,
+                                                  step,
+                                                  state->v.rvec_array(),
+                                              &realFepState,
+                                                  md->homenr,
+                                                  md->cTC);
+                /* history is maintained in state->dfhist, but state_global is what is sent to trajectory and log output */
+                if (MAIN(cr))
+                {
+                    copy_df_history(state_global->dfhist, state->dfhist);
+                }
+            }
+
+            // 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)
+                    || checkpointHandler->isCheckpointingStep()))
+            {
+                stateGpu->copyCoordinatesFromGpu(state->x, AtomLocality::Local);
+                stateGpu->waitCoordinatesReadyOnHost(AtomLocality::Local);
+            }
+            // 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) || checkpointHandler->isCheckpointingStep()))
+            {
+                stateGpu->copyVelocitiesFromGpu(state->v, AtomLocality::Local);
+                stateGpu->waitVelocitiesReadyOnHost(AtomLocality::Local);
+            }
+            // Copy forces for the output if the forces were reduced on the GPU (not the case on virial steps)
+            // and update is offloaded hence forces are kept on the GPU for update and have not been
+            // already transferred in do_force().
+            // TODO: There should be an improved, explicit mechanism that ensures this copy is only executed
+            //       when the forces are ready on the GPU -- the same synchronizer should be used as the one
+            //       prior to GPU update.
+            // TODO: When the output flags will be included in step workload, this copy can be combined with the
+            //       copy call in do_force(...).
+            // NOTE: The forces should not be copied here if the vsites are present, since they were modified
+            //       on host after the D2H copy in do_force(...).
+            if (runScheduleWork->stepWork.useGpuFBufferOps
+                && (simulationWork.useGpuUpdate && !vsite) && do_per_step(step, ir->nstfout))
+            {
+                stateGpu->copyForcesFromGpu(f.view().force(), AtomLocality::Local);
+                stateGpu->waitForcesReadyOnHost(AtomLocality::Local);
+            }
+            /* Now we have the energies and forces corresponding to the
+             * coordinates at time t. We must output all of this before
+             * the update.
+             */
+            const EkindataState ekindataState =
+                    bGStat ? (bSumEkinhOld ? EkindataState::UsedNeedToReduce
+                                           : EkindataState::UsedDoNotNeedToReduce)
+                           : EkindataState::NotUsed;
+            do_md_trajectory_writing(fplog,
+                                     cr,
+                                     nfile,
+                                     fnm,
+                                     step,
+                                     step_rel,
+                                     t,
+                                     ir,
+                                     state,
+                                     state_global,
+                                     observablesHistory,
+                                     top_global,
+                                     fr,
+                                     outf,
+                                     energyOutput,
+                                     ekind,
+                                     f.view().force(),
+                                     checkpointHandler->isCheckpointingStep(),
+                                     bRerunMD,
+                                     bLastStep,
+                                     mdrunOptions.writeConfout,
+                                     ekindataState);
+            /* Check if IMD step and do IMD communication, if bIMD is TRUE. */
+            bInteractiveMDstep = imdSession->run(step, bNS, state->box, state->x, t);
+
+            /* kludge -- virial is lost with restart for MTTK NPT control. Must reload (saved earlier). */
+            if (startingBehavior != StartingBehavior::NewSimulation && bFirstStep
+                && (inputrecNptTrotter(ir) || inputrecNphTrotter(ir)))
+            {
+                copy_mat(state->svir_prev, shake_vir);
+                copy_mat(state->fvir_prev, force_vir);
+            }
+
+            stopHandler->setSignal();
+            resetHandler->setSignal(walltime_accounting);
+
+            if (bGStat || !PAR(cr))
+            {
+                /* In parallel we only have to check for checkpointing in steps
+                 * where we do global communication,
+                 *  otherwise the other nodes don't know.
+                 */
+                checkpointHandler->setSignal(walltime_accounting);
+            }
+
+            /* #########   START SECOND UPDATE STEP ################# */
+
+            /* at the start of step, randomize or scale the velocities ((if vv. Restriction of
+               Andersen controlled in preprocessing */
+
+            if (ETC_ANDERSEN(ir->etc)) /* keep this outside of update_tcouple because of the extra info required to pass */
+            {
+                gmx_bool bIfRandomize;
+                bIfRandomize = update_randomize_velocities(
+                        ir, step, cr, md->homenr, md->cTC, md->invmass, state->v, &upd, constr);
+                /* if we have constraints, we have to remove the kinetic energy parallel to the bonds */
+                if (constr && bIfRandomize)
+                {
+                    constrain_velocities(constr, do_log, do_ene, step, state, nullptr, false, nullptr);
+                }
+            }
+            /* Box is changed in update() when we do pressure coupling,
+             * but we should still use the old box for energy corrections and when
+             * writing it to the energy file, so it matches the trajectory files for
+             * the same timestep above. Make a copy in a separate array.
+             */
+            copy_mat(state->box, lastbox);
+
+            dvdl_constr = 0;
+
+            if (!useGpuForUpdate)
+            {
+                wallcycle_start(wcycle, WallCycleCounter::Update);
+            }
+            /* UPDATE PRESSURE VARIABLES IN TROTTER FORMULATION WITH CONSTRAINTS */
+            if (bTrotter)
+            {
+                trotter_update(ir,
+                               step,
+                               ekind,
+                               enerd,
+                               state,
+                               total_vir,
+                               md->homenr,
+                               md->cTC,
+                               md->invmass,
+                               &MassQ,
+                               trotter_seq,
+                               TrotterSequence::Three);
+                /* We can only do Berendsen coupling after we have summed
+                 * the kinetic energy or virial. Since the happens
+                 * in global_state after update, we should only do it at
+                 * step % nstlist = 1 with bGStatEveryStep=FALSE.
+                 */
+            }
+            else
+            {
+                update_tcouple(step, ir, state, ekind, &MassQ, md->homenr, md->cTC);
+                update_pcouple_before_coordinates(mdlog,
+                                                  step,
+                                                  ir->pressureCouplingOptions,
+                                                  ir->deform,
+                                                  ir->delta_t,
+                                                  state,
+                                                  &pressureCouplingMu,
+                                                  &parrinelloRahmanM);
+            }
+
+            if (EI_VV(ir->eI))
+            {
+                GMX_ASSERT(!useGpuForUpdate, "GPU update is not supported with VVAK integrator.");
+
+                integrateVVSecondStep(step,
+                                      ir,
+                                      fr,
+                                      cr,
+                                      state,
+                                      mdAtoms->mdatoms(),
+                                      &fcdata,
+                                      &MassQ,
+                                      &vcm,
+                                      pull_work,
+                                      enerd,
+                                      &observablesReducer,
+                                      ekind,
+                                      gstat,
+                                      &dvdl_constr,
+                                      bCalcVir,
+                                      total_vir,
+                                      shake_vir,
+                                      force_vir,
+                                      pres,
+                                      lastbox,
+                                      do_log,
+                                      do_ene,
+                                      bGStat,
+                                      &bSumEkinhOld,
+                                      &f,
+                                      &cbuf,
+                                      &upd,
+                                      constr,
+                                      &nullSignaller,
+                                      trotter_seq,
+                                      nrnb,
+                                      wcycle);
+            }
+            else
+            {
+                if (useGpuForUpdate)
+                {
+                    // On search steps, update handles to device vectors
+                    // TODO: this condition has redundant / unnecessary clauses
+                    if (bNS && (bFirstStep || haveDDAtomOrdering(*cr) || bExchanged))
+                    {
+                        integrator->set(stateGpu->getCoordinates(),
+                                        stateGpu->getVelocities(),
+                                        stateGpu->getForces(),
+                                        top->idef,
+                                        *md);
+
+                        // Copy data to the GPU after buffers might have been reinitialized
+                        /* The velocity copy is redundant if we had Center-of-Mass motion removed on
+                         * the previous step. We don't check that now. */
+                        stateGpu->copyVelocitiesToGpu(state->v, AtomLocality::Local);
+                    }
+
+                    // Copy x to the GPU unless we have already transferred in do_force().
+                    // We transfer in do_force() if a GPU force task requires x (PME or x buffer ops).
+                    if (!(runScheduleWork->stepWork.haveGpuPmeOnThisRank
+                          || runScheduleWork->stepWork.useGpuXBufferOps))
+                    {
+                        stateGpu->copyCoordinatesToGpu(state->x, AtomLocality::Local);
+                        // Coordinates are later used by the integrator running in the same stream.
+                        stateGpu->consumeCoordinatesCopiedToDeviceEvent(AtomLocality::Local);
+                    }
+
+                    if ((simulationWork.useGpuPme && simulationWork.useCpuPmePpCommunication)
+                        || (!runScheduleWork->stepWork.useGpuFBufferOps))
+                    {
+                        // The PME forces were recieved to the host, and reduced on the CPU with the
+                        // rest of the forces computed on the GPU, so the final forces have to be
+                        // copied back to the GPU. Or the buffer ops were not offloaded this step,
+                        // so the forces are on the host and have to be copied
+                        stateGpu->copyForcesToGpu(f.view().force(), AtomLocality::Local);
+                    }
+                    const bool doTemperatureScaling =
+                            (ir->etc != TemperatureCoupling::No
+                             && do_per_step(step + ir->nsttcouple - 1, ir->nsttcouple));
+
+                    // This applies Leap-Frog, LINCS and SETTLE in succession
+                    integrator->integrate(stateGpu->getLocalForcesReadyOnDeviceEvent(
+                                                  runScheduleWork->stepWork, runScheduleWork->simulationWork),
+                                          ir->delta_t,
+                                          true,
+                                          bCalcVir,
+                                          shake_vir,
+                                          doTemperatureScaling,
+                                          ekind->tcstat,
+                                          doParrinelloRahman,
+                                          ir->pressureCouplingOptions.nstpcouple * ir->delta_t,
+                                          parrinelloRahmanM);
+                }
+                else
+                {
+                    /* With multiple time stepping we need to do an additional normal
+                     * update step to obtain the virial and dH/dl, as the actual MTS integration
+                     * using an acceleration where the slow forces are multiplied by mtsFactor.
+                     * Using that acceleration would result in a virial with the slow
+                     * force contribution would be a factor mtsFactor too large.
+                     */
+                    const bool separateVirialConstraining =
+                            (simulationWork.useMts && (bCalcVir || computeDHDL) && constr != nullptr);
+                    if (separateVirialConstraining)
+                    {
+                        upd.update_for_constraint_virial(*ir,
+                                                         md->homenr,
+                                                         md->havePartiallyFrozenAtoms,
+                                                         md->invmass,
+                                                         md->invMassPerDim,
+                                                         *state,
+                                                         f.view().forceWithPadding(),
+                                                         *ekind);
+
+                        constrain_coordinates(constr,
+                                              do_log,
+                                              do_ene,
+                                              step,
+                                              state,
+                                              upd.xp()->arrayRefWithPadding(),
+                                              &dvdl_constr,
+                                              bCalcVir,
+                                              shake_vir);
+                    }
+
+                    ArrayRefWithPadding<const RVec> forceCombined =
+                            (simulationWork.useMts && step % ir->mtsLevels[1].stepFactor == 0)
+                                    ? f.view().forceMtsCombinedWithPadding()
+                                    : f.view().forceWithPadding();
+                    upd.update_coords(*ir,
+                                      step,
+                                      md->homenr,
+                                      md->havePartiallyFrozenAtoms,
+                                      md->ptype,
+                                      md->invmass,
+                                      md->invMassPerDim,
+                                      state,
+                                      forceCombined,
+                                      &fcdata,
+                                      ekind,
+                                      parrinelloRahmanM,
+                                      etrtPOSITION,
+                                      cr,
+                                      constr != nullptr);
+
+                    wallcycle_stop(wcycle, WallCycleCounter::Update);
+
+                    constrain_coordinates(constr,
+                                          do_log,
+                                          do_ene,
+                                          step,
+                                          state,
+                                          upd.xp()->arrayRefWithPadding(),
+                                          separateVirialConstraining ? nullptr : &dvdl_constr,
+                                          bCalcVir && !separateVirialConstraining,
+                                          shake_vir);
+
+                    upd.update_sd_second_half(*ir,
+                                              step,
+                                              &dvdl_constr,
+                                              md->homenr,
+                                              md->ptype,
+                                              md->invmass,
+                                              state,
+                                              cr,
+                                              nrnb,
+                                              wcycle,
+                                              constr,
+                                              do_log,
+                                              do_ene);
+                    upd.finish_update(
+                            *ir, md->havePartiallyFrozenAtoms, md->homenr, state, wcycle, constr != nullptr);
+                }
+
+                if (ir->bPull && ir->pull->bSetPbcRefToPrevStepCOM)
+                {
+                    updatePrevStepPullCom(pull_work, state->pull_com_prev_step);
+                }
+
+                enerd->term[F_DVDL_CONSTR] += dvdl_constr;
+            }
+        }
+
+        if (simulationWork.useMdGpuGraph)
+        {
+            GMX_ASSERT((mdGraph != nullptr), "MD GPU graph does not exist.");
+            if (mdGraph->graphIsCapturingThisStep())
+            {
+                mdGraph->endRecord();
+                // Force graph reinstantiation (instead of graph exec
+                // update): with PME tuning, since the GPU kernels
+                // chosen by the FFT library can vary with grid size;
+                // or with an odd nstlist, since the odd/even step
+                // pruning pattern will change
+                bool forceGraphReinstantiation =
+                        pme_loadbal_is_active(pme_loadbal) || ((ir->nstlist % 2) == 1);
+                mdGraph->createExecutableGraph(forceGraphReinstantiation);
+            }
+            if (mdGraph->useGraphThisStep())
+            {
+                mdGraph->launchGraphMdStep(integrator->xUpdatedOnDeviceEvent());
+            }
+            if (bNS)
+            {
+                // TODO: merge disableForDomainIfAnyPpRankHasCpuForces() back into reset() when
+                // domainWork initialization is moved out of do_force().
+                fr->mdGraph[MdGraphEvenOrOddStep::EvenStep]->disableForDomainIfAnyPpRankHasCpuForces(
+                        runScheduleWork->domainWork.haveCpuLocalForceWork);
+                fr->mdGraph[MdGraphEvenOrOddStep::OddStep]->disableForDomainIfAnyPpRankHasCpuForces(
+                        runScheduleWork->domainWork.haveCpuLocalForceWork);
+            }
+            usedMdGpuGraphLastStep = mdGraph->useGraphThisStep();
+        }
+
+        /* ############## IF NOT VV, Calculate globals HERE  ############ */
+        /* With Leap-Frog we can skip compute_globals at
+         * non-communication steps, but we need to calculate
+         * the kinetic energy one step before communication.
+         */
+        {
+            // Organize to do inter-simulation signalling on steps if
+            // and when algorithms require it.
+            const bool doInterSimSignal = (simulationsShareState && do_per_step(step, nstSignalComm));
+
+            if (useGpuForUpdate)
+            {
+                const bool coordinatesRequiredForStopCM =
+                        bStopCM && (bGStat || needHalfStepKineticEnergy || doInterSimSignal)
+                        && !EI_VV(ir->eI);
+
+                // Copy coordinates when needed to stop the CM motion or for replica exchange
+                if (coordinatesRequiredForStopCM || bDoReplEx)
+                {
+                    stateGpu->copyCoordinatesFromGpu(state->x, AtomLocality::Local);
+                    stateGpu->waitCoordinatesReadyOnHost(AtomLocality::Local);
+                }
+
+                // Copy velocities back to the host if:
+                // - Globals are computed this step (includes the energy output steps).
+                // - Temperature is needed for the next step.
+                // - This is a replica exchange step (even though we will only need
+                //     the velocities if an exchange succeeds)
+                if (bGStat || needHalfStepKineticEnergy || bDoReplEx)
+                {
+                    stateGpu->copyVelocitiesFromGpu(state->v, AtomLocality::Local);
+                    stateGpu->waitVelocitiesReadyOnHost(AtomLocality::Local);
+                }
+            }
+
+            if (bGStat || needHalfStepKineticEnergy || doInterSimSignal)
+            {
+                // Since we're already communicating at this step, we
+                // can propagate intra-simulation signals. Note that
+                // check_nstglobalcomm has the responsibility for
+                // choosing the value of nstglobalcomm that is one way
+                // bGStat becomes true, so we can't get into a
+                // situation where e.g. checkpointing can't be
+                // signalled.
+                bool                doIntraSimSignal = true;
+                SimulationSignaller signaller(&signals, cr, ms, doInterSimSignal, doIntraSimSignal);
+
+                compute_globals(gstat,
+                                cr,
+                                ir,
+                                fr,
+                                ekind,
+                                makeConstArrayRef(state->x),
+                                makeConstArrayRef(state->v),
+                                state->box,
+                                md,
+                                nrnb,
+                                &vcm,
+                                wcycle,
+                                enerd,
+                                force_vir,
+                                shake_vir,
+                                total_vir,
+                                pres,
+                                &signaller,
+                                lastbox,
+                                &bSumEkinhOld,
+                                (bGStat ? CGLO_GSTAT : 0) | (!EI_VV(ir->eI) && bCalcEner ? CGLO_ENERGY : 0)
+                                        | (!EI_VV(ir->eI) && bStopCM ? CGLO_STOPCM : 0)
+                                        | (!EI_VV(ir->eI) ? CGLO_TEMPERATURE : 0)
+                                        | (!EI_VV(ir->eI) ? CGLO_PRESSURE : 0) | CGLO_CONSTRAINT,
+                                step,
+                                &observablesReducer);
+                if (!EI_VV(ir->eI) && bStopCM)
+                {
+                    process_and_stopcm_grp(
+                            fplog, &vcm, *md, makeArrayRef(state->x), makeArrayRef(state->v));
+                    inc_nrnb(nrnb, eNR_STOPCM, md->homenr);
+
+                    // TODO: The special case of removing CM motion should be dealt more gracefully
+                    if (useGpuForUpdate)
+                    {
+                        // Issue #3988, #4106.
+                        stateGpu->resetCoordinatesCopiedToDeviceEvent(AtomLocality::Local);
+                        stateGpu->copyCoordinatesToGpu(state->x, AtomLocality::Local);
+                        // Here we block until the H2D copy completes because event sync with the
+                        // force kernels that use the coordinates on the next steps is not implemented
+                        // (not because of a race on state->x being modified on the CPU while H2D is in progress).
+                        stateGpu->waitCoordinatesCopiedToDevice(AtomLocality::Local);
+                        // If the COM removal changed the velocities on the CPU, this has to be accounted for.
+                        if (vcm.mode != ComRemovalAlgorithm::No)
+                        {
+                            stateGpu->copyVelocitiesToGpu(state->v, AtomLocality::Local);
+                        }
+                    }
+                }
+            }
+        }
+
+        /* #############  END CALC EKIN AND PRESSURE ################# */
+
+        /* Note: this is OK, but there are some numerical precision issues with using the convergence of
+           the virial that should probably be addressed eventually. state->veta has better properies,
+           but what we actually need entering the new cycle is the new shake_vir value. Ideally, we could
+           generate the new shake_vir, but test the veta value for convergence.  This will take some thought. */
+
+        if (ir->efep != FreeEnergyPerturbationType::No && !EI_VV(ir->eI))
+        {
+            /* Sum up the foreign energy and dK/dl terms for md and sd.
+               Currently done every step so that dH/dl is correct in the .edr */
+            accumulateKineticLambdaComponents(enerd, state->lambda, *ir->fepvals);
+        }
+
+        const real currentSystemRefT =
+                (haveEnsembleTemperature(*ir) ? ekind->currentEnsembleTemperature() : 0.0_real);
+        const bool scaleCoordinates = !useGpuForUpdate || bDoReplEx;
+        update_pcouple_after_coordinates(fplog,
+                                         step,
+                                         ir->pressureCouplingOptions,
+                                         ir->ld_seed,
+                                         currentSystemRefT,
+                                         ir->opts.nFreeze,
+                                         ir->deform,
+                                         ir->delta_t,
+                                         md->homenr,
+                                         md->cFREEZE,
+                                         pres,
+                                         force_vir,
+                                         shake_vir,
+                                         &pressureCouplingMu,
+                                         state,
+                                         nrnb,
+                                         upd.deform(),
+                                         scaleCoordinates);
+
+        const bool doBerendsenPressureCoupling =
+                (inputrec->pressureCouplingOptions.epc == PressureCoupling::Berendsen
+                 && do_per_step(step, inputrec->pressureCouplingOptions.nstpcouple));
+        const bool doCRescalePressureCoupling =
+                (inputrec->pressureCouplingOptions.epc == PressureCoupling::CRescale
+                 && do_per_step(step, inputrec->pressureCouplingOptions.nstpcouple));
+        if (useGpuForUpdate
+            && (doBerendsenPressureCoupling || doCRescalePressureCoupling || doParrinelloRahman))
+        {
+            integrator->scaleCoordinates(pressureCouplingMu);
+            if (doCRescalePressureCoupling)
+            {
+                integrator->scaleVelocities(invertBoxMatrix(pressureCouplingMu));
+            }
+            integrator->setPbc(PbcType::Xyz, state->box);
+        }
+
+        /* ################# END UPDATE STEP 2 ################# */
+        /* #### We now have r(t+dt) and v(t+dt/2)  ############# */
+
+        /* The coordinates (x) were unshifted in update */
+        if (!bGStat)
+        {
+            /* We will not sum ekinh_old,
+             * so signal that we still have to do it.
+             */
+            bSumEkinhOld = TRUE;
+        }
+
+        if (bCalcEner)
+        {
+            /* #########  BEGIN PREPARING EDR OUTPUT  ###########  */
+
+            /* use the directly determined last velocity, not actually the averaged half steps */
+            if (bTrotter && ir->eI == IntegrationAlgorithm::VV)
+            {
+                enerd->term[F_EKIN] = last_ekin;
+            }
+            enerd->term[F_ETOT] = enerd->term[F_EPOT] + enerd->term[F_EKIN];
+
+            if (integratorHasConservedEnergyQuantity(ir))
+            {
+                if (EI_VV(ir->eI))
+                {
+                    enerd->term[F_ECONSERVED] = enerd->term[F_ETOT] + saved_conserved_quantity;
+                }
+                else
+                {
+                    enerd->term[F_ECONSERVED] =
+                            enerd->term[F_ETOT]
+                            + NPT_energy(ir->pressureCouplingOptions,
+                                         ir->etc,
+                                         gmx::constArrayRefFromArray(ir->opts.nrdf, ir->opts.ngtc),
+                                         *ekind,
+                                         inputrecNvtTrotter(ir) || inputrecNptTrotter(ir),
+                                         state,
+                                         &MassQ);
+                }
+            }
+            /* #########  END PREPARING EDR OUTPUT  ###########  */
+        }
+
+        /* Output stuff */
+        if (MAIN(cr))
+        {
+            if (fplog && do_log && bDoExpanded)
+            {
+                /* only needed if doing expanded ensemble */
+                PrintFreeEnergyInfoToFile(fplog,
+                                          ir->fepvals.get(),
+                                          ir->expandedvals.get(),
+                                          ir->bSimTemp ? ir->simtempvals.get() : nullptr,
+                                          state_global->dfhist,
+                                          state->fep_state,
+                                          ir->nstlog,
+                                          step);
+            }
+            if (bCalcEner)
+            {
+                const bool outputDHDL = (computeDHDL && do_per_step(step, ir->fepvals->nstdhdl));
+
+                energyOutput.addDataAtEnergyStep(outputDHDL,
+                                                 bCalcEnerStep,
+                                                 t,
+                                                 md->tmass,
+                                                 enerd,
+                                                 ir->fepvals.get(),
+                                                 lastbox,
+                                                 PTCouplingArrays{ state->boxv,
+                                                                   state->nosehoover_xi,
+                                                                   state->nosehoover_vxi,
+                                                                   state->nhpres_xi,
+                                                                   state->nhpres_vxi },
+                                                 state->fep_state,
+                                                 total_vir,
+                                                 pres,
+                                                 ekind,
+                                                 mu_tot,
+                                                 constr);
+            }
+            else
+            {
+                energyOutput.recordNonEnergyStep();
+            }
+
+            gmx_bool do_dr = do_per_step(step, ir->nstdisreout);
+            gmx_bool do_or = do_per_step(step, ir->nstorireout);
+
+            if (doSimulatedAnnealing)
+            {
+                gmx::EnergyOutput::printAnnealingTemperatures(
+                        do_log ? fplog : nullptr, *groups, ir->opts, *ekind);
+            }
+            if (do_log || do_ene || do_dr || do_or)
+            {
+                energyOutput.printStepToEnergyFile(mdoutf_get_fp_ene(outf),
+                                                   do_ene,
+                                                   do_dr,
+                                                   do_or,
+                                                   do_log ? fplog : nullptr,
+                                                   step,
+                                                   t,
+                                                   fr->fcdata.get(),
+                                                   awh.get());
+            }
+            if (do_log && ((ir->bDoAwh && awh->hasFepLambdaDimension()) || ir->fepvals->delta_lambda != 0))
+            {
+                const bool isInitialOutput = false;
+                printLambdaStateToLog(fplog, state->lambda, isInitialOutput);
+            }
+
+            if (ir->bPull)
+            {
+                pull_print_output(pull_work, step, t);
+            }
+
+            if (do_per_step(step, ir->nstlog))
+            {
+                if (fflush(fplog) != 0)
+                {
+                    gmx_fatal(FARGS, "Cannot flush logfile - maybe you are out of disk space?");
+                }
+            }
+        }
+        if (bDoExpanded)
+        {
+            /* Have to do this part _after_ outputting the logfile and the edr file */
+            /* Gets written into the state at the beginning of next loop*/
+            state->fep_state = lamnew;
+            if(plumedswitch)
+            {
+                realFepState = state->fep_state;
+            }
+        }
+        else if (ir->bDoAwh && awh->needForeignEnergyDifferences(step))
+        {
+            state->fep_state = awh->fepLambdaState();
+        }
+        /* Print the remaining wall clock time for the run */
+        if (isMainSimMainRank(ms, MAIN(cr)) && (do_verbose || gmx_got_usr_signal()) && !bPMETunePrinting)
+        {
+            if (shellfc)
+            {
+                fprintf(stderr, "\n");
+            }
+            print_time(stderr, walltime_accounting, step, ir, cr);
+        }
+
+        /* Ion/water position swapping.
+         * Not done in last step since trajectory writing happens before this call
+         * in the MD loop and exchanges would be lost anyway. */
+        bNeedRepartition = FALSE;
+        if ((ir->eSwapCoords != SwapType::No) && (step > 0) && !bLastStep
+            && do_per_step(step, ir->swap->nstswap))
+        {
+            bNeedRepartition = do_swapcoords(cr,
+                                             step,
+                                             t,
+                                             ir,
+                                             swap,
+                                             wcycle,
+                                             as_rvec_array(state->x.data()),
+                                             state->box,
+                                             MAIN(cr) && mdrunOptions.verbose,
+                                             bRerunMD);
+
+            if (bNeedRepartition && haveDDAtomOrdering(*cr))
+            {
+                dd_collect_state(cr->dd, state, state_global);
+            }
+        }
+
+        /* Replica exchange */
+        bExchanged = FALSE;
+        if (bDoReplEx)
+        {
+            bExchanged = replica_exchange(fplog, cr, ms, repl_ex, state_global, enerd, state, step, t);
+        }
+
+        if ((bExchanged || bNeedRepartition) && haveDDAtomOrdering(*cr))
+        {
+            dd_partition_system(fplog,
+                                mdlog,
+                                step,
+                                cr,
+                                TRUE,
+                                state_global,
+                                top_global,
+                                *ir,
+                                imdSession,
+                                pull_work,
+                                state,
+                                &f,
+                                mdAtoms,
+                                top,
+                                fr,
+                                vsite,
+                                constr,
+                                nrnb,
+                                wcycle,
+                                FALSE);
+            upd.updateAfterPartition(state->natoms, md->cFREEZE, md->cTC, md->cACC);
+            fr->longRangeNonbondeds->updateAfterPartition(*md);
+        }
+
+        bFirstStep = FALSE;
+        bInitStep  = FALSE;
+
+        /* #######  SET VARIABLES FOR NEXT ITERATION IF THEY STILL NEED IT ###### */
+        /* With all integrators, except VV, we need to retain the pressure
+         * at the current step for coupling at the next step.
+         */
+        if ((state->flags & enumValueToBitMask(StateEntry::PressurePrevious))
+            && (bGStatEveryStep
+                || (ir->pressureCouplingOptions.nstpcouple > 0
+                    && step % ir->pressureCouplingOptions.nstpcouple == 0)))
+        {
+            /* Store the pressure in t_state for pressure coupling
+             * at the next MD step.
+             */
+            copy_mat(pres, state->pres_prev);
+        }
+
+        /* #######  END SET VARIABLES FOR NEXT ITERATION ###### */
+
+        if ((membed != nullptr) && (!bLastStep))
+        {
+            rescale_membed(step_rel, membed, as_rvec_array(state_global->x.data()));
+        }
+
+        cycles = wallcycle_stop(wcycle, WallCycleCounter::Step);
+        if (haveDDAtomOrdering(*cr) && wcycle)
+        {
+            dd_cycles_add(cr->dd, cycles, ddCyclStep);
+        }
+
+        /* increase the MD step number */
+        step++;
+        step_rel++;
+        observablesReducer.markAsReadyToReduce();
+
+#if GMX_FAHCORE
+        if (MAIN(cr))
+        {
+            fcReportProgress(ir->nsteps + ir->init_step, step);
+        }
+#endif
+
+        resetHandler->resetCounters(
+                step, step_rel, mdlog, fplog, cr, fr->nbv.get(), nrnb, fr->pmedata, pme_loadbal, wcycle, walltime_accounting);
+
+        /* If bIMD is TRUE, the main updates the IMD energy record and sends positions to VMD client */
+        imdSession->updateEnergyRecordAndSendPositionsAndEnergies(bInteractiveMDstep, step, bCalcEner);
+
+        // any run that uses GPUs must be at least offloading nonbondeds
+        const bool usingGpu = simulationWork.useGpuNonbonded;
+        if (usingGpu)
+        {
+            // ensure that GPU errors do not propagate between MD steps
+            checkPendingDeviceErrorBetweenSteps();
+        }
+    }
+    /* End of main MD loop */
+
+    /* Closing TNG files can include compressing data. Therefore it is good to do that
+     * before stopping the time measurements. */
+    mdoutf_tng_close(outf);
+
+    /* Stop measuring walltime */
+    walltime_accounting_end_time(walltime_accounting);
+
+    if (simulationWork.haveSeparatePmeRank)
+    {
+        /* Tell the PME only node to finish */
+        gmx_pme_send_finish(cr);
+    }
+
+    if (MAIN(cr))
+    {
+        if (ir->nstcalcenergy > 0)
+        {
+            energyOutput.printEnergyConservation(fplog, ir->simulation_part, EI_MD(ir->eI));
+
+            gmx::EnergyOutput::printAnnealingTemperatures(fplog, *groups, ir->opts, *ekind);
+            energyOutput.printAverages(fplog, groups);
+        }
+    }
+    done_mdoutf(outf);
+
+    if (bPMETune)
+    {
+        pme_loadbal_done(pme_loadbal, fplog, mdlog, fr->nbv->useGpu());
+    }
+
+    done_shellfc(fplog, shellfc, step_rel);
+
+    if (useReplicaExchange && MAIN(cr))
+    {
+        print_replica_exchange_statistics(fplog, repl_ex);
+    }
+
+    walltime_accounting_set_nsteps_done(walltime_accounting, step_rel);
+
+    global_stat_destroy(gstat);
+}