/*
 * 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 MD runner routine calling all integrators.
 *
 * \author David van der Spoel <david.vanderspoel@icm.uu.se>
 * \ingroup module_mdrun
 */
#include "gmxpre.h"

#include "runner.h"

#include "config.h"

#include <cassert>
#include <cinttypes>
#include <cmath>
#include <csignal>
#include <cstdlib>
#include <cstring>

#include <algorithm>
#include <bitset>
#include <filesystem>
#include <memory>
#include <optional>
#include <utility>
#include <variant>
#include <vector>

#include "gromacs/commandline/filenm.h"
#include "gromacs/domdec/builder.h"
#include "gromacs/domdec/domdec.h"
#include "gromacs/domdec/domdec_struct.h"
#include "gromacs/domdec/gpuhaloexchange.h"
#include "gromacs/domdec/localatomsetmanager.h"
#include "gromacs/domdec/makebondedlinks.h"
#include "gromacs/domdec/partition.h"
#include "gromacs/domdec/reversetopology.h"
#include "gromacs/ewald/ewald_utils.h"
#include "gromacs/ewald/pme.h"
#include "gromacs/ewald/pme_gpu_program.h"
#include "gromacs/ewald/pme_only.h"
#include "gromacs/ewald/pme_pp_comm_gpu.h"
#include "gromacs/fileio/checkpoint.h"
#include "gromacs/fileio/filetypes.h"
#include "gromacs/fileio/gmxfio.h"
#include "gromacs/fileio/oenv.h"
#include "gromacs/fileio/tpxio.h"
#include "gromacs/fileio/trrio.h"
#include "gromacs/gmxlib/network.h"
#include "gromacs/gmxlib/nrnb.h"
#include "gromacs/gpu_utils/capabilities.h"
#include "gromacs/gpu_utils/device_stream_manager.h"
#include "gromacs/gpu_utils/gpu_utils.h"
#include "gromacs/gpu_utils/gpueventsynchronizer_helpers.h"
#include "gromacs/gpu_utils/hostallocator.h"
#include "gromacs/gpu_utils/nvshmem_manager.h"
#include "gromacs/hardware/detecthardware.h"
#include "gromacs/hardware/device_management.h"
#include "gromacs/hardware/hardwaretopology.h"
#include "gromacs/hardware/printhardware.h"
#include "gromacs/imd/imd.h"
#include "gromacs/listed_forces/disre.h"
#include "gromacs/listed_forces/listed_forces.h"
#include "gromacs/listed_forces/listed_forces_gpu.h"
#include "gromacs/listed_forces/orires.h"
#include "gromacs/math/functions.h"
#include "gromacs/math/matrix.h"
#include "gromacs/math/utilities.h"
#include "gromacs/mdlib/boxdeformation.h"
#include "gromacs/mdlib/broadcaststructs.h"
#include "gromacs/mdlib/calc_verletbuf.h"
#include "gromacs/mdlib/constr.h"
#include "gromacs/mdlib/dispersioncorrection.h"
#include "gromacs/mdlib/enerdata_utils.h"
#include "gromacs/mdlib/force.h"
#include "gromacs/mdlib/forcerec.h"
#include "gromacs/mdlib/gmx_omp_nthreads.h"
#include "gromacs/mdlib/gpuforcereduction.h"
#include "gromacs/mdlib/makeconstraints.h"
#include "gromacs/mdlib/md_support.h"
#include "gromacs/mdlib/mdatoms.h"
#include "gromacs/mdlib/mdgraph_gpu.h"
#include "gromacs/mdlib/sighandler.h"
#include "gromacs/mdlib/stophandler.h"
#include "gromacs/mdlib/tgroup.h"
#include "gromacs/mdlib/updategroups.h"
#include "gromacs/mdlib/vsite.h"
#include "gromacs/mdrun/mdmodules.h"
#include "gromacs/mdrun/simulationcontext.h"
#include "gromacs/mdrun/simulationinput.h"
#include "gromacs/mdrun/simulationinputhandle.h"
#include "gromacs/mdrunutility/handlerestart.h"
#include "gromacs/mdrunutility/logging.h"
#include "gromacs/mdrunutility/mdmodulesnotifiers.h"
#include "gromacs/mdrunutility/multisim.h"
#include "gromacs/mdrunutility/plainpairlistranges.h"
#include "gromacs/mdrunutility/print_validation.h"
#include "gromacs/mdrunutility/printtime.h"
#include "gromacs/mdrunutility/threadaffinity.h"
#include "gromacs/mdtypes/atominfo.h"
#include "gromacs/mdtypes/checkpointdata.h"
#include "gromacs/mdtypes/commrec.h"
#include "gromacs/mdtypes/enerdata.h"
#include "gromacs/mdtypes/fcdata.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/locality.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/pull_params.h"
#include "gromacs/mdtypes/simulation_workload.h"
#include "gromacs/mdtypes/state.h"
#include "gromacs/mdtypes/state_propagator_data_gpu.h"
#include "gromacs/modularsimulator/modularsimulator.h"
#include "gromacs/nbnxm/gpu_data_mgmt.h"
#include "gromacs/nbnxm/nbnxm.h"
#include "gromacs/nbnxm/nbnxm_enums.h"
#include "gromacs/nbnxm/pairlist_tuning.h"
#include "gromacs/pbcutil/pbc.h"
#include "gromacs/pulling/output.h"
#include "gromacs/pulling/pull.h"
#include "gromacs/pulling/pull_rotation.h"
#include "gromacs/restraint/manager.h"
#include "gromacs/restraint/restraintmdmodule.h"
#include "gromacs/restraint/restraintpotential.h"
#include "gromacs/swap/swapcoords.h"
#include "gromacs/taskassignment/decidegpuusage.h"
#include "gromacs/taskassignment/decidesimulationworkload.h"
#include "gromacs/taskassignment/resourcedivision.h"
#include "gromacs/taskassignment/taskassignment.h"
#include "gromacs/taskassignment/usergpuids.h"
#include "gromacs/timing/gpu_timing.h"
#include "gromacs/timing/wallcycle.h"
#include "gromacs/timing/wallcyclereporting.h"
#include "gromacs/timing/walltime_accounting.h"
#include "gromacs/topology/block.h"
#include "gromacs/topology/ifunc.h"
#include "gromacs/topology/mtop_util.h"
#include "gromacs/topology/topology.h"
#include "gromacs/topology/topology_enums.h"
#include "gromacs/trajectory/trajectoryframe.h"
#include "gromacs/utility/basenetwork.h"
#include "gromacs/utility/cstringutil.h"
#include "gromacs/utility/enumerationhelpers.h"
#include "gromacs/utility/exceptions.h"
#include "gromacs/utility/fatalerror.h"
#include "gromacs/utility/filestream.h"
#include "gromacs/utility/gmxassert.h"
#include "gromacs/utility/gmxmpi.h"
#include "gromacs/utility/keyvaluetree.h"
#include "gromacs/utility/logger.h"
#include "gromacs/utility/loggerbuilder.h"
#include "gromacs/utility/mpiinfo.h"
#include "gromacs/utility/physicalnodecommunicator.h"
#include "gromacs/utility/pleasecite.h"
#include "gromacs/utility/programcontext.h"
#include "gromacs/utility/smalloc.h"
#include "gromacs/utility/stringutil.h"
#include "gromacs/utility/vec.h"
#include "gromacs/utility/vectypes.h"

#include "isimulator.h"
#include "membedholder.h"
#include "replicaexchange.h"
#include "simulatorbuilder.h"

class DeviceContext;
class DeviceStream;
struct DeviceInformation;
struct gmx_pme_t;
struct pull_t;

namespace gmx
{

/*! \brief Manage any development feature flag variables encountered
 *
 * The use of dev features indicated by environment variables is
 * logged in order to ensure that runs with such features enabled can
 * be identified from their log and standard output. Any cross
 * dependencies are also checked, and if unsatisfied, a fatal error
 * issued.
 *
 * Note that some development features overrides are applied already here:
 * the GPU communication flags are set to false in non-tMPI and non-CUDA builds.
 *
 * \param[in]  mdlog                Logger object.
 * \param[in]  pmeRunMode   Run mode indicating what resource is PME executed on.
 * \param[in]  numRanksPerSimulation   The number of ranks in each simulation.
 * \param[in]  numPmeRanksPerSimulation   The number of PME ranks in each simulation, can be -1
 * \param[in]  canUseGpuAwareMpi   Whether GPU-aware MPI can be used.
 * \returns                         The object populated with development feature flags.
 */
static DevelopmentFeatureFlags manageDevelopmentFeatures(const gmx::MDLogger& mdlog,
                                                         const PmeRunMode     pmeRunMode,
                                                         const int            numRanksPerSimulation,
                                                         const int  numPmeRanksPerSimulation,
                                                         const bool canUseGpuAwareMpi)
{
    DevelopmentFeatureFlags devFlags;

    if (std::getenv("GMX_CUDA_GRAPH") != nullptr)
    {
        if (GpuConfigurationCapabilities::GpuGraph)
        {
            devFlags.enableCudaGraphs = true;
            GMX_LOG(mdlog.warning)
                    .asParagraph()
                    .appendText(
                            "GMX_CUDA_GRAPH environment variable is detected. "
                            "The experimental CUDA Graphs feature will be used if run conditions "
                            "allow.");
        }
        else
        {
            devFlags.enableCudaGraphs = false;
            std::string errorReason;
            if (GMX_GPU_CUDA)
            {
                errorReason = "the CUDA version in use is below the minimum requirement (11.1)";
            }
            else if (GMX_GPU_SYCL)
            {
                if (GMX_SYCL_DPCPP)
                {
                    errorReason = "GROMACS is built without GMX_SYCL_ENABLE_GRAPHS";
                }
                else
                {
                    errorReason = "SYCL Graph extension is only supported in oneAPI DPC++";
                }
            }
            else
            {
                errorReason = "CUDA or SYCL build is required";
            }
            GMX_LOG(mdlog.warning)
                    .asParagraph()
                    .appendTextFormatted(
                            "GMX_CUDA_GRAPH environment variable is detected, but %s. GPU Graphs "
                            "will be disabled.",
                            errorReason.c_str());
        }
    }

    if (std::getenv("GMX_ENABLE_NVSHMEM") != nullptr)
    {
        if (GMX_LIB_MPI && GMX_NVSHMEM)
        {
            devFlags.enableNvshmem = true;
            GMX_LOG(mdlog.warning)
                    .asParagraph()
                    .appendText(
                            "GMX_ENABLE_NVSHMEM environment variable is detected. "
                            "The experimental NVSHMEM feature will be used if run conditions "
                            "allow.");
        }
        else
        {
            devFlags.enableNvshmem = false;
            GMX_LOG(mdlog.warning)
                    .asParagraph()
                    .appendText(
                            "GMX_ENABLE_NVSHMEM environment variable is detected, "
                            "but GROMACS was built without NVSHMEM support. "
                            "Direct use of NVSHMEM will be disabled. "
                            "NVSHMEM may still be used indirectly if cuFFTMp is enabled. ");
        }
    }

    // PME decomposition is supported only with CUDA or SYCL and also
    // needs GPU-aware MPI support for it to work.
    const bool pmeGpuDecompositionRequested =
            (pmeRunMode == PmeRunMode::GPU || pmeRunMode == PmeRunMode::Mixed)
            && ((numRanksPerSimulation > 1 && numPmeRanksPerSimulation == 0)
                || numPmeRanksPerSimulation > 1);
    const bool pmeGpuDecompositionSupported =
            canUseGpuAwareMpi && GpuConfigurationCapabilities::PpPmeDirectComm
            && ((GpuConfigurationCapabilities::PmeDecomposition && pmeRunMode == PmeRunMode::GPU)
                || (GpuConfigurationCapabilities::Pme && pmeRunMode == PmeRunMode::Mixed));

    const bool forcePmeGpuDecomposition = std::getenv("GMX_GPU_PME_DECOMPOSITION") != nullptr;

    if (pmeGpuDecompositionSupported && pmeGpuDecompositionRequested)
    {
        // PME decomposition is supported only when it is forced using GMX_GPU_PME_DECOMPOSITION
        if (forcePmeGpuDecomposition)
        {
            GMX_LOG(mdlog.warning)
                    .asParagraph()
                    .appendTextFormatted(
                            "This run has requested the 'GPU PME decomposition' feature, enabled "
                            "by the GMX_GPU_PME_DECOMPOSITION environment variable. "
                            "PME decomposition lacks substantial testing "
                            "and should be used with caution.");
        }
        else
        {
            gmx_fatal(FARGS,
                      "Multiple PME tasks were required to run on GPUs, "
                      "but that is not supported. "
                      "Use GMX_GPU_PME_DECOMPOSITION environment variable to enable it.");
        }
    }

    if (!pmeGpuDecompositionSupported && pmeGpuDecompositionRequested)
    {
        if (GMX_GPU_CUDA)
        {
            gmx_fatal(FARGS,
                      "PME tasks were required to run on more than one CUDA-devices. To enable "
                      "this feature, "
                      "use MPI with CUDA-aware support and build GROMACS with cuFFTMp support.");
        }
        else
        {
            gmx_fatal(
                    FARGS,
                    "PME tasks were required to run on GPUs, but that is not implemented with "
                    "more than one PME rank. Use a single rank simulation, or a separate PME rank, "
                    "or permit PME tasks to be assigned to the CPU.");
        }
    }

    devFlags.enableGpuPmeDecomposition =
            forcePmeGpuDecomposition && pmeGpuDecompositionRequested && pmeGpuDecompositionSupported;

    return devFlags;
}

/*! \brief Barrier for safe simultaneous thread access to mdrunner data
 *
 * Used to ensure that the main thread does not modify mdrunner during copy
 * on the spawned threads. */
static void threadMpiMdrunnerAccessBarrier()
{
#if GMX_THREAD_MPI
    MPI_Barrier(MPI_COMM_WORLD);
#endif
}

Mdrunner Mdrunner::cloneOnSpawnedThread() const
{
    auto newRunner = Mdrunner(std::make_unique<MDModules>());

    // All runners in the same process share a restraint manager resource because it is
    // part of the interface to the client code, which is associated only with the
    // original thread. Handles to the same resources can be obtained by copy.
    {
        newRunner.restraintManager_ = std::make_unique<RestraintManager>(*restraintManager_);
    }

    // Copy members of main runner.
    // \todo Replace with builder when Simulation context and/or runner phases are better defined.
    // Ref https://gitlab.com/gromacs/gromacs/-/issues/2587 and https://gitlab.com/gromacs/gromacs/-/issues/2375
    newRunner.hw_opt    = hw_opt;
    newRunner.filenames = filenames;

    newRunner.hwinfo_         = hwinfo_;
    newRunner.oenv            = oenv;
    newRunner.mdrunOptions    = mdrunOptions;
    newRunner.domdecOptions   = domdecOptions;
    newRunner.nbpu_opt        = nbpu_opt;
    newRunner.nbfe_opt        = nbfe_opt;
    newRunner.pme_opt         = pme_opt;
    newRunner.pme_fft_opt     = pme_fft_opt;
    newRunner.bonded_opt      = bonded_opt;
    newRunner.update_opt      = update_opt;
    newRunner.nstlist_cmdline = nstlist_cmdline;
    newRunner.replExParams    = replExParams;
    newRunner.pforce          = pforce;
    // Give the spawned thread the newly created valid communicator
    // for the simulation.
    newRunner.libraryWorldCommunicator = MPI_COMM_WORLD;
    newRunner.simulationCommunicator   = MPI_COMM_WORLD;
    newRunner.ms                       = ms;
    newRunner.startingBehavior         = startingBehavior;
    newRunner.stopHandlerBuilder_      = std::make_unique<StopHandlerBuilder>(*stopHandlerBuilder_);
    newRunner.inputHolder_             = inputHolder_;

    threadMpiMdrunnerAccessBarrier();

    return newRunner;
}

/*! \brief The callback used for running on spawned threads.
 *
 * Obtains the pointer to the main mdrunner object from the one
 * argument permitted to the thread-launch API call, copies it to make
 * a new runner for this thread, reinitializes necessary data, and
 * proceeds to the simulation. */
static void mdrunner_start_fn(const void* arg)
{
    try
    {
        const auto* mainMdrunner = reinterpret_cast<const gmx::Mdrunner*>(arg);
        /* copy the arg list to make sure that it's thread-local. This
           doesn't copy pointed-to items, of course; fnm, cr and fplog
           are reset in the call below, all others should be const. */
        gmx::Mdrunner mdrunner = mainMdrunner->cloneOnSpawnedThread();
        mdrunner.mdrunner();
    }
    GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR
}


void Mdrunner::spawnThreads(int numThreadsToLaunch)
{
#if GMX_THREAD_MPI
    /* now spawn new threads that start mdrunner_start_fn(), while
       the main thread returns. Thread affinity is handled later. */
    if (tMPI_Init_fn(TRUE, numThreadsToLaunch, TMPI_AFFINITY_NONE, mdrunner_start_fn, static_cast<const void*>(this))
        != TMPI_SUCCESS)
    {
        GMX_THROW(gmx::InternalError("Failed to spawn thread-MPI threads"));
    }

    // Give the main thread the newly created valid communicator for
    // the simulation.
    libraryWorldCommunicator = MPI_COMM_WORLD;
    simulationCommunicator   = MPI_COMM_WORLD;
    threadMpiMdrunnerAccessBarrier();
#else
    GMX_UNUSED_VALUE(numThreadsToLaunch);
    GMX_UNUSED_VALUE(mdrunner_start_fn);
#endif
}

} // namespace gmx

/*! \brief Initialize variables for Verlet scheme simulation */
static void prepare_verlet_scheme(FILE*                          fplog,
                                  const gmx::MpiComm&            mpiCommSimulation,
                                  t_inputrec*                    ir,
                                  int                            nstlist_cmdline,
                                  const gmx_mtop_t&              mtop,
                                  gmx::ArrayRef<const gmx::RVec> coordinates,
                                  const matrix                   box,
                                  bool                           makeGpuPairList)
{
    // We checked the cut-offs in grompp, but double-check here.
    // We have PME+LJcutoff kernels for rcoulomb>rvdw.
    if (usingPmeOrEwald(ir->coulombtype) && ir->vdwtype == VanDerWaalsType::Cut)
    {
        GMX_RELEASE_ASSERT(ir->rcoulomb >= ir->rvdw,
                           "With Verlet lists and PME we should have rcoulomb>=rvdw");
    }
    else
    {
        GMX_RELEASE_ASSERT(ir->rcoulomb == ir->rvdw,
                           "With Verlet lists and no PME rcoulomb and rvdw should be identical");
    }

    std::optional<real> effectiveAtomDensity;
    if (EI_DYNAMICS(ir->eI))
    {
        effectiveAtomDensity = computeEffectiveAtomDensity(
                coordinates, box, std::max(ir->rcoulomb, ir->rvdw), mpiCommSimulation.comm());
    }

    /* For NVE simulations, we will retain the initial list buffer */
    if (EI_DYNAMICS(ir->eI) && ir->verletbuf_tol > 0
        && !(EI_MD(ir->eI) && ir->etc == TemperatureCoupling::No))
    {
        /* Update the Verlet buffer size for the current run setup */

        /* Here we assume SIMD-enabled kernels are being used. But as currently
         * calc_verlet_buffer_size gives the same results for 4x8 and 4x4
         * and 4x2 gives a larger buffer than 4x4, this is ok.
         */
        ListSetupType listType =
                (makeGpuPairList ? ListSetupType::Gpu : ListSetupType::CpuSimdWhenSupported);
        VerletbufListSetup listSetup = verletbufGetSafeListSetup(listType);

        const real rlist_new = calcVerletBufferSize(mtop,
                                                    effectiveAtomDensity.value(),
                                                    *ir,
                                                    ir->verletBufferPressureTolerance,
                                                    ir->nstlist,
                                                    ir->nstlist - 1,
                                                    -1,
                                                    listSetup);

        if (rlist_new != ir->rlist)
        {
            if (fplog != nullptr)
            {
                fprintf(fplog,
                        "\nChanging rlist from %g to %g for non-bonded %dx%d atom kernels\n\n",
                        ir->rlist,
                        rlist_new,
                        listSetup.cluster_size_i,
                        listSetup.cluster_size_j);
            }
            ir->rlist = rlist_new;
        }
    }

    if (nstlist_cmdline > 0 && (!EI_DYNAMICS(ir->eI) || ir->verletbuf_tol <= 0))
    {
        gmx_fatal(FARGS,
                  "Can not set nstlist without %s",
                  !EI_DYNAMICS(ir->eI) ? "dynamics" : "verlet-buffer-tolerance");
    }

    if (EI_DYNAMICS(ir->eI))
    {
        /* Set or try nstlist values */
        gmx::increaseNstlist(
                fplog, mpiCommSimulation, ir, nstlist_cmdline, &mtop, box, effectiveAtomDensity.value(), makeGpuPairList);
    }
}

/*! \brief Override the nslist value in inputrec
 *
 * with value passed on the command line (if any)
 */
static void override_nsteps_cmdline(const gmx::MDLogger& mdlog, int64_t nsteps_cmdline, t_inputrec* ir)
{
    assert(ir);

    /* override with anything else than the default -2 */
    if (nsteps_cmdline > -2)
    {
        char sbuf_steps[STEPSTRSIZE];
        char sbuf_msg[STRLEN];

        ir->nsteps = nsteps_cmdline;
        if (EI_DYNAMICS(ir->eI) && nsteps_cmdline != -1)
        {
            sprintf(sbuf_msg,
                    "Overriding nsteps with value passed on the command line: %s steps, %.3g ps",
                    gmx_step_str(nsteps_cmdline, sbuf_steps),
                    std::fabs(nsteps_cmdline * ir->delta_t));
        }
        else
        {
            sprintf(sbuf_msg,
                    "Overriding nsteps with value passed on the command line: %s steps",
                    gmx_step_str(nsteps_cmdline, sbuf_steps));
        }

        GMX_LOG(mdlog.warning).asParagraph().appendText(sbuf_msg);
    }
    else if (nsteps_cmdline < -2)
    {
        gmx_fatal(FARGS, "Invalid nsteps value passed on the command line: %" PRId64, nsteps_cmdline);
    }
    /* Do nothing if nsteps_cmdline == -2 */
}

namespace gmx
{

//! Initializes the logger for mdrun.
static gmx::LoggerOwner buildLogger(FILE* fplog, const bool isSimulationMainRank)
{
    gmx::LoggerBuilder builder;
    if (fplog != nullptr)
    {
        builder.addTargetFile(gmx::MDLogger::LogLevel::Info, fplog);
    }
    if (isSimulationMainRank)
    {
        builder.addTargetStream(gmx::MDLogger::LogLevel::Warning, &gmx::TextOutputFile::standardError());
    }
    return builder.build();
}

//! Make a TaskTarget from an mdrun argument string.
static TaskTarget findTaskTarget(const char* optionString)
{
    TaskTarget returnValue = TaskTarget::Auto;

    if (std::strncmp(optionString, "auto", 3) == 0)
    {
        returnValue = TaskTarget::Auto;
    }
    else if (std::strncmp(optionString, "cpu", 3) == 0)
    {
        returnValue = TaskTarget::Cpu;
    }
    else if (std::strncmp(optionString, "gpu", 3) == 0)
    {
        returnValue = TaskTarget::Gpu;
    }
    else
    {
        GMX_ASSERT(false, "Option string should have been checked for sanity already");
    }

    return returnValue;
}

//! Finish run, aggregate data to print performance info.
static void finish_run(FILE*                     fplog,
                       const gmx::MDLogger&      mdlog,
                       const t_commrec*          cr,
                       const t_inputrec&         inputrec,
                       t_nrnb                    nrnb[],
                       gmx_wallcycle*            wcycle,
                       gmx_walltime_accounting_t walltime_accounting,
                       nonbonded_verlet_t*       nbv,
                       const gmx_pme_t*          pme,
                       const int                 nratoms,
                       gmx_bool                  bWriteStat)
{
    double delta_t = 0;
    double nbfs = 0, mflop = 0;
    double elapsed_time, elapsed_time_over_all_ranks, elapsed_time_over_all_threads,
            elapsed_time_over_all_threads_over_all_ranks;
    /* Control whether it is valid to print a report. Only the
       simulation main may print, but it should not do so if the run
       terminated e.g. before a scheduled reset step. This is
       complicated by the fact that PME ranks are unaware of the
       reason why they were sent a pmerecvqxFINISH. To avoid
       communication deadlocks, we always do the communication for the
       report, even if we've decided not to write the report, because
       how long it takes to finish the run is not important when we've
       decided not to report on the simulation performance.

       Further, we only report performance for dynamical integrators,
       because those are the only ones for which we plan to
       consider doing any optimizations. */
    bool printReport = EI_DYNAMICS(inputrec.eI) && cr->isSimulationMainRank();

    if (printReport && !walltime_accounting_get_valid_finish(walltime_accounting))
    {
        GMX_LOG(mdlog.warning)
                .asParagraph()
                .appendText("Simulation ended prematurely, no performance report will be written.");
        printReport = false;
    }

    t_nrnb*                 nrnb_tot;
    std::unique_ptr<t_nrnb> nrnbTotalStorage;
    if (cr->commMySim.isParallel())
    {
        nrnbTotalStorage = std::make_unique<t_nrnb>();
        nrnb_tot         = nrnbTotalStorage.get();
#if GMX_MPI
        MPI_Allreduce(nrnb->n.data(), nrnb_tot->n.data(), eNRNB, MPI_DOUBLE, MPI_SUM, cr->commMySim.comm());
#endif
    }
    else
    {
        nrnb_tot = nrnb;
    }

    elapsed_time = walltime_accounting_get_time_since_reset(walltime_accounting);
    elapsed_time_over_all_threads =
            walltime_accounting_get_time_since_reset_over_all_threads(walltime_accounting);
    if (GMX_MPI && cr->commMySim.isParallel())
    {
#if GMX_MPI
        /* reduce elapsed_time over all MPI ranks in the current simulation */
        MPI_Allreduce(
                &elapsed_time, &elapsed_time_over_all_ranks, 1, MPI_DOUBLE, MPI_SUM, cr->commMySim.comm());
        elapsed_time_over_all_ranks /= cr->commMySim.size();
        /* Reduce elapsed_time_over_all_threads over all MPI ranks in the
         * current simulation. */
        MPI_Allreduce(&elapsed_time_over_all_threads,
                      &elapsed_time_over_all_threads_over_all_ranks,
                      1,
                      MPI_DOUBLE,
                      MPI_SUM,
                      cr->commMySim.comm());
#endif
    }
    else
    {
        elapsed_time_over_all_ranks                  = elapsed_time;
        elapsed_time_over_all_threads_over_all_ranks = elapsed_time_over_all_threads;
    }

    if (printReport)
    {
        print_flop(fplog, nrnb_tot, &nbfs, &mflop);
    }

    if (thisRankHasPPDuty(cr->dd) && haveDDAtomOrdering(*cr))
    {
        print_dd_statistics(cr->dd, inputrec, fplog);
    }

    /* TODO Move the responsibility for any scaling by thread counts
     * to the code that handled the thread region, so that there's a
     * mechanism to keep cycle counting working during the transition
     * to task parallelism. */
    int nthreads_pp  = gmx_omp_nthreads_get(ModuleMultiThread::Nonbonded);
    int nthreads_pme = gmx_omp_nthreads_get(ModuleMultiThread::Pme);
    wallcycle_scale_by_num_threads(
            wcycle, thisRankHasPmeDuty(cr->dd) && !thisRankHasPPDuty(cr->dd), nthreads_pp, nthreads_pme);
    auto cycle_sum(wallcycle_sum(cr, wcycle));

    if (printReport)
    {
        auto* nbnxn_gpu_timings =
                (nbv != nullptr && nbv->useGpu()) ? gpu_get_timings(nbv->gpuNbv()) : nullptr;
        gmx_wallclock_gpu_pme_t pme_gpu_timings = {};

        if (pme_gpu_task_enabled(pme))
        {
            pme_gpu_get_timings(pme, &pme_gpu_timings);
        }
        wallcycle_print(fplog,
                        mdlog,
                        cr->commMySim.size(),
                        cr->dd ? cr->dd->numPmeOnlyRanks : 0,
                        nthreads_pp,
                        nthreads_pme,
                        elapsed_time_over_all_ranks,
                        wcycle,
                        cycle_sum,
                        nbnxn_gpu_timings,
                        &pme_gpu_timings);

        if (EI_DYNAMICS(inputrec.eI))
        {
            delta_t = inputrec.delta_t;
        }

        if (fplog)
        {
            print_perf(fplog,
                       elapsed_time_over_all_threads_over_all_ranks,
                       elapsed_time_over_all_ranks,
                       walltime_accounting_get_nsteps_done_since_reset(walltime_accounting),
                       delta_t,
                       nbfs,
                       mflop,
                       nratoms);
        }
        if (bWriteStat)
        {
            print_perf(stderr,
                       elapsed_time_over_all_threads_over_all_ranks,
                       elapsed_time_over_all_ranks,
                       walltime_accounting_get_nsteps_done_since_reset(walltime_accounting),
                       delta_t,
                       nbfs,
                       mflop,
                       nratoms);
        }
    }
}

//! Returns whether the run conditions permit the local state to have filler particles
static bool localStateHasFillerParticles(const gmx_mtop_t& mtop,
                                         const t_inputrec& inputrec,
                                         const bool        useDomainDecomposition,
                                         const bool        haveSinglePPRank,
                                         const bool        useGpuDirectHalo)
{
    // Having filler particles in the local states is supported with DD atom ordering.
    // WholeMoleculeTransform and GPU-direct does not support filler particles.
    const bool useEwaldSurfaceCorrection =
            (usingPmeOrEwald(inputrec.coulombtype) && inputrec.epsilon_surface != 0);
    const bool haveOrientationRestraints =
            (gmx_mtop_ftype_count(mtop, InteractionFunction::OrientationRestraints) > 0);
    const bool needWholeMolecules = useEwaldSurfaceCorrection || haveOrientationRestraints;
    const bool canHaveFillerParticlesInLocalState =
            useDomainDecomposition
            && ((haveSinglePPRank && !needWholeMolecules) || (!haveSinglePPRank && !useGpuDirectHalo));
    bool haveFillerParticlesInLocalState = false;
    if (const char* env = getenv("GMX_FILLERS_IN_LOCAL_STATE"))
    {
        int value;
        if (sscanf(env, "%d", &value) == 0 || (value < 0 || value > 2))
        {
            GMX_THROW(gmx::InvalidInputError(
                    "Env.var. GMX_FILLERS_IN_LOCAL_STATE should have value 0, 1 or 2"));
        }
        if (value == 2 && !canHaveFillerParticlesInLocalState)
        {
            GMX_THROW(
                    gmx::InvalidInputError("Fillers in local state requested, but not supported"
                                           " because DD is not used or because GPU-direct comm."
                                           " does not support it or because an algorithm requires"
                                           " whole molecules"));
        }
        haveFillerParticlesInLocalState = canHaveFillerParticlesInLocalState && (value != 0);
    }

    return haveFillerParticlesInLocalState;
}

int Mdrunner::mdrunner()
{
    std::unique_ptr<t_forcerec> fr;
    real                        ewaldcoeff_q     = 0;
    real                        ewaldcoeff_lj    = 0;
    int                         nChargePerturbed = -1, nTypePerturbed = 0;
    gmx_walltime_accounting_t   walltime_accounting = nullptr;
    MembedHolder                membedHolder(filenames.size(), filenames.data());

    /* CAUTION: threads may be started later on in this function, so
       cr doesn't reflect the final parallel state right now */
    gmx_mtop_t mtop;

    /* TODO: inputrec should tell us whether we use an algorithm, not a file option */
    const bool doEssentialDynamics = opt2bSet("-ei", filenames.size(), filenames.data());
    const bool doRerun             = mdrunOptions.rerun;

    // Handle task-assignment related user options.
    EmulateGpuNonbonded emulateGpuNonbonded =
            (std::getenv("GMX_EMULATE_GPU") != nullptr ? EmulateGpuNonbonded::Yes
                                                       : EmulateGpuNonbonded::No);

    std::vector<int> userGpuTaskAssignment;
    try
    {
        userGpuTaskAssignment = parseUserTaskAssignmentString(hw_opt.userGpuTaskAssignment);
    }
    GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR
    auto nonbondedTarget   = findTaskTarget(nbpu_opt);
    auto nonbondedFeTarget = findTaskTarget(nbfe_opt);
    auto pmeTarget         = findTaskTarget(pme_opt);
    auto pmeFftTarget      = findTaskTarget(pme_fft_opt);
    auto bondedTarget      = findTaskTarget(bonded_opt);
    auto updateTarget      = findTaskTarget(update_opt);

    FILE* fplog = nullptr;
    // If we are appending, we don't write log output because we need
    // to check that the old log file matches what the checkpoint file
    // expects. Otherwise, we should start to write log output now if
    // there is a file ready for it.
    if (logFileHandle != nullptr && startingBehavior != StartingBehavior::RestartWithAppending)
    {
        fplog = gmx_fio_getfp(logFileHandle);
    }
    const bool       isSimulationMainRank = findIsSimulationMainRank(ms, simulationCommunicator);
    gmx::LoggerOwner logOwner(buildLogger(fplog, isSimulationMainRank));
    gmx::MDLogger    mdlog(logOwner.logger());

    gmx_print_detected_hardware(fplog, isSimulationMainRank && isMainSim(ms), mdlog, hwinfo_);

    std::vector<int> availableDevices =
            makeListOfAvailableDevices(hwinfo_->deviceInfoList, hw_opt.devicesSelectedByUser);
    const int numAvailableDevices = gmx::ssize(availableDevices);

    // Print citation requests after all software/hardware printing
    pleaseCiteGromacs(fplog);

    // Note: legacy program logic relies on checking whether these pointers are assigned.
    // Objects may or may not be allocated later.
    std::unique_ptr<t_inputrec> inputrec;
    std::unique_ptr<t_state>    globalState;

    auto partialDeserializedTpr = std::make_unique<PartialDeserializedTprFile>();

    if (isSimulationMainRank)
    {
        // Allocate objects to be initialized by later function calls.
        /* Only the main rank has the global state */
        globalState = std::make_unique<t_state>();
        inputrec    = std::make_unique<t_inputrec>();

        /* Read (nearly) all data required for the simulation
         * and keep the partly serialized tpr contents to send to other ranks later
         */
        applyGlobalSimulationState(
                *inputHolder_.get(), partialDeserializedTpr.get(), globalState.get(), inputrec.get(), &mtop);

        static_assert(sc_trrMaxAtomCount == sc_checkpointMaxAtomCount);
        if (mtop.natoms > sc_checkpointMaxAtomCount)
        {
            gmx_fatal(FARGS,
                      "System has %d atoms, which is more than can be stored in checkpoint and trr "
                      "files (max %" PRId64 ")",
                      mtop.natoms,
                      sc_checkpointMaxAtomCount);
        }

        // The XTC format has been updated to support up to 2^31-1 atoms, which is anyway the
        // largest supported by GROMACS, so no need for any particular check here.
    }

    /* Check and update the hardware options for internal consistency */
    checkAndUpdateHardwareOptions(
            mdlog, &hw_opt, isSimulationMainRank, domdecOptions.numPmeRanks, inputrec.get());

    if (GMX_THREAD_MPI && isSimulationMainRank)
    {
        bool useGpuForNonbonded = false;
        bool useGpuForPme       = false;
        try
        {
            GMX_RELEASE_ASSERT(inputrec != nullptr, "Keep the compiler happy");

            // If the user specified the number of ranks, then we must
            // respect that, but in default mode, we need to allow for
            // the number of GPUs to choose the number of ranks.
            std::string gpuNonbondedNotSupportedReason;
            bool        canUseGpuForNonbonded =
                    canUseGpusForNonbonded(*inputrec, doRerun, &gpuNonbondedNotSupportedReason);
            if (GMX_GPU && !canUseGpuForNonbonded)
            {
                GMX_RELEASE_ASSERT(!gpuNonbondedNotSupportedReason.empty(),
                                   "A reason should have been given");
                GMX_LOG(mdlog.warning).asParagraph().appendText(gpuNonbondedNotSupportedReason);
            }
            useGpuForNonbonded =
                    decideWhetherToUseGpusForNonbondedWithThreadMpi(nonbondedTarget,
                                                                    numAvailableDevices > 0,
                                                                    userGpuTaskAssignment,
                                                                    emulateGpuNonbonded,
                                                                    canUseGpuForNonbonded,
                                                                    mdrunOptions.reproducible,
                                                                    hw_opt.nthreads_tmpi);
            useGpuForPme = decideWhetherToUseGpusForPmeWithThreadMpi(useGpuForNonbonded,
                                                                     pmeTarget,
                                                                     pmeFftTarget,
                                                                     numAvailableDevices,
                                                                     userGpuTaskAssignment,
                                                                     *inputrec,
                                                                     hw_opt.nthreads_tmpi,
                                                                     domdecOptions.numPmeRanks);
        }
        GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR

        /* Determine how many thread-MPI ranks to start.
         *
         * TODO Over-writing the user-supplied value here does
         * prevent any possible subsequent checks from working
         * correctly. */
        hw_opt.nthreads_tmpi = get_nthreads_mpi(hwinfo_,
                                                &hw_opt,
                                                numAvailableDevices,
                                                useGpuForNonbonded,
                                                useGpuForPme,
                                                inputrec.get(),
                                                mtop,
                                                mdlog,
                                                membedHolder.doMembed());

        // Now start the threads for thread MPI.
        spawnThreads(hw_opt.nthreads_tmpi);
        // The spawned threads enter mdrunner() and execution of
        // main and spawned threads joins at the end of this block.
    }

    GMX_RELEASE_ASSERT(!GMX_MPI || ms || simulationCommunicator != MPI_COMM_NULL,
                       "Must have valid communicator unless running a multi-simulation");
    MpiComm mpiCommSimulation(simulationCommunicator);

    const bool simulationIsParallel = mpiCommSimulation.isParallel();
    const bool isMainSimulationRank = mpiCommSimulation.isMainRank();

    PhysicalNodeCommunicator physicalNodeComm(libraryWorldCommunicator, gmx_physicalnode_id_hash());

    if (simulationIsParallel)
    {
        /* now broadcast everything to the non-main nodes/threads: */
        if (!isSimulationMainRank)
        {
            // Until now, only the main rank has a non-null pointer.
            // On non-main ranks, allocate the object that will receive data in the following call.
            inputrec = std::make_unique<t_inputrec>();
        }
        init_parallel(mpiCommSimulation.comm(),
                      isMainSimulationRank,
                      inputrec.get(),
                      &mtop,
                      partialDeserializedTpr.get());
    }
    GMX_RELEASE_ASSERT(inputrec != nullptr, "All ranks should have a valid inputrec now");
    partialDeserializedTpr.reset(nullptr);

    // Note that these variables describe only their own node.
    //
    // Note that when bonded interactions run on a GPU they always run
    // alongside a nonbonded task, so do not influence task assignment
    // even though they affect the force calculation workload.
    bool useGpuForNonbonded   = false;
    bool useGpuForNonbondedFE = false;
    bool useGpuForPme         = false;
    bool useGpuForBonded      = false;
    bool useGpuForUpdate      = false;
    bool gpusWereDetected     = hwinfo_->ngpu_compatible_tot > 0;
    try
    {
        // It's possible that there are different numbers of GPUs on
        // different nodes, which is the user's responsibility to
        // handle. If unsuitable, we will notice that during task
        // assignment.
        std::string gpuNonbondedNotSupportedReason;
        bool        canUseGpuForNonbonded =
                canUseGpusForNonbonded(*inputrec, doRerun, &gpuNonbondedNotSupportedReason);
        if (GMX_GPU && !GMX_THREAD_MPI && !canUseGpuForNonbonded) // We already warned for threadMPI earlier
        {
            GMX_RELEASE_ASSERT(!gpuNonbondedNotSupportedReason.empty(),
                               "A reason should have been given");
            GMX_LOG(mdlog.warning).asParagraph().appendText(gpuNonbondedNotSupportedReason);
        }
        useGpuForNonbonded = decideWhetherToUseGpusForNonbonded(nonbondedTarget,
                                                                userGpuTaskAssignment,
                                                                emulateGpuNonbonded,
                                                                canUseGpuForNonbonded,
                                                                mdrunOptions.reproducible,
                                                                gpusWereDetected);
        useGpuForNonbondedFE =
                decideWhetherToUseGpusForNonbondedFE(useGpuForNonbonded, nonbondedFeTarget, *inputrec);
        useGpuForPme    = decideWhetherToUseGpusForPme(useGpuForNonbonded,
                                                    pmeTarget,
                                                    pmeFftTarget,
                                                    userGpuTaskAssignment,
                                                    *inputrec,
                                                    mpiCommSimulation.size(),
                                                    domdecOptions.numPmeRanks,
                                                    gpusWereDetected);
        useGpuForBonded = decideWhetherToUseGpusForBonded(
                useGpuForNonbonded, useGpuForPme, bondedTarget, *inputrec, mtop, domdecOptions.numPmeRanks, gpusWereDetected);
    }
    GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR

    const PmeRunMode pmeRunMode = determinePmeRunMode(useGpuForPme, pmeFftTarget, *inputrec);

    const bool useModularSimulator = ModularSimulator::isInputCompatible(mdlog,
                                                                         inputrec.get(),
                                                                         doRerun,
                                                                         mtop,
                                                                         ms,
                                                                         replExParams,
                                                                         nullptr,
                                                                         doEssentialDynamics,
                                                                         membedHolder.doMembed(),
                                                                         updateTarget == TaskTarget::Gpu);

    // Now the number of ranks is known to all ranks, and each knows
    // the inputrec read by the main rank. The ranks can now all run
    // the task-deciding functions and will agree on the result
    // without needing to communicate.
    // The LBFGS minimizer, test-particle insertion, normal modes and shell dynamics don't support DD
    const bool hasCustomParallelization =
            (EI_TPI(inputrec->eI) || inputrec->eI == IntegrationAlgorithm::NM);
    const bool canUseDomainDecomposition =
            (inputrec->eI != IntegrationAlgorithm::LBFGS && !hasCustomParallelization
             && gmx_mtop_particletype_count(mtop)[ParticleType::Shell] == 0);
    GMX_RELEASE_ASSERT(!simulationIsParallel || hasCustomParallelization || canUseDomainDecomposition,
                       "A parallel run should not arrive here without DD support");

    int useDDWithSingleRank = -1;
    if (const char* ddSingleRankEnv = std::getenv("GMX_DD_SINGLE_RANK"))
    {
        useDDWithSingleRank = std::strtol(ddSingleRankEnv, nullptr, 10);
    }

    // The overhead of DD partitioning is only compensated when we have both non-bondeds and PME on the CPU
    const bool useDomainDecomposition =
            canUseDomainDecomposition
            && (simulationIsParallel
                || (!useGpuForNonbonded && usingFullElectrostatics(inputrec->coulombtype)
                    && useDDWithSingleRank != 0)
                || useDDWithSingleRank == 1);

    ObservablesReducerBuilder observablesReducerBuilder;

    // Build restraints.
    // TODO: hide restraint implementation details from Mdrunner.
    // There is nothing unique about restraints at this point as far as the
    // Mdrunner is concerned. The Mdrunner should just be getting a sequence of
    // factory functions from the SimulationContext on which to call mdModules_->add().
    // TODO: capture all restraints into a single RestraintModule, passed to the runner builder.
    for (auto&& restraint : restraintManager_->getRestraints())
    {
        mdModules_->add(RestraintMDModule::sc_name,
                        RestraintMDModule::create(restraint, restraint->sites()));
    }

    // TODO: Error handling
    mdModules_->assignOptionsToModules(*inputrec->params, nullptr);
    // now that the MDModules know their options, they know which callbacks to sign up to
    mdModules_->subscribeToSimulationSetupNotifications();
    const auto& setupNotifier = mdModules_->notifiers().simulationSetupNotifier_;

    // Notify MdModules of existing logger
    setupNotifier.notify(mdlog);

    // Notify MdModules of internal parameters, saved into KVT
    if (inputrec->internalParameters != nullptr)
    {
        setupNotifier.notify(*inputrec->internalParameters);
    }

    // Let MdModules know the .tpr input and .edr output filenames
    {
        gmx::MdRunInputFilename mdRunInputFilename = { ftp2fn(efTPR, filenames.size(), filenames.data()) };
        setupNotifier.notify(mdRunInputFilename);
        gmx::EdrOutputFilename edrOutputFilename = { ftp2fn(efEDR, filenames.size(), filenames.data()) };
        setupNotifier.notify(edrOutputFilename);

        constexpr const char*    plumedOptionName = "-plumed";
        gmx::PlumedInputFilename plumedFilename;
        if (opt2bSet(plumedOptionName, filenames.size(), filenames.data()))
        {
            plumedFilename.plumedFilename_ =
                    std::string(opt2fn(plumedOptionName, filenames.size(), filenames.data()));
        }

        setupNotifier.notify(plumedFilename);
    }

    if (fplog != nullptr)
    {
        pr_inputrec(fplog, 0, "Input Parameters", inputrec.get(), FALSE);
        fprintf(fplog, "\n");
    }

    if (isSimulationMainRank)
    {
        /* In rerun, set velocities to zero if present */
        if (doRerun && globalState->hasEntry(StateEntry::V))
        {
            // rerun does not use velocities
            GMX_LOG(mdlog.info)
                    .asParagraph()
                    .appendText(
                            "Rerun trajectory contains velocities. Rerun does only evaluate "
                            "potential energy and forces. The velocities will be ignored.");
            for (int i = 0; i < globalState->numAtoms(); i++)
            {
                clear_rvec(globalState->v[i]);
            }
            globalState->setFlags(globalState->flags() & ~enumValueToBitMask(StateEntry::V));
        }

        /* now make sure the state is initialized and propagated */
        set_state_entries(globalState.get(), inputrec.get(), useModularSimulator);
    }

    /* NM and TPI parallelize over force/energy calculations, not atoms,
     * so we need to initialize and broadcast the global state.
     */
    if (inputrec->eI == IntegrationAlgorithm::NM || inputrec->eI == IntegrationAlgorithm::TPI)
    {
        if (!isSimulationMainRank)
        {
            globalState = std::make_unique<t_state>();
        }
        broadcastStateWithoutDynamics(
                mpiCommSimulation.comm(), useDomainDecomposition, simulationIsParallel, globalState.get());
    }

    /* A parallel command line option consistency check that we can
       only do after any threads have started. */
    if (!simulationIsParallel
        && (domdecOptions.numCells[XX] > 1 || domdecOptions.numCells[YY] > 1
            || domdecOptions.numCells[ZZ] > 1 || domdecOptions.numPmeRanks > 0))
    {
        gmx_fatal(FARGS,
                  "The -dd or -npme option request a parallel simulation, "
#if !GMX_MPI
                  "but %s was compiled without threads or MPI enabled",
                  output_env_get_program_display_name(oenv));
#elif GMX_THREAD_MPI
                  "but the number of MPI-threads (option -ntmpi) is not set or is 1");
#else
                  "but %s was not started through mpirun/mpiexec or only one rank was requested "
                  "through mpirun/mpiexec",
                  output_env_get_program_display_name(oenv));
#endif
    }

    if (doRerun && (EI_ENERGY_MINIMIZATION(inputrec->eI) || IntegrationAlgorithm::NM == inputrec->eI))
    {
        gmx_fatal(FARGS,
                  "The .mdp file specified an energy minimization or normal mode algorithm, and "
                  "these are not compatible with mdrun -rerun");
    }

    /* NMR restraints must be initialized before load_checkpoint,
     * since with time averaging the history is added to t_state.
     * For proper consistency check we therefore need to extend
     * t_state here.
     * So the PME-only nodes (if present) will also initialize
     * the distance restraints.
     */

    /* This needs to be called before read_checkpoint to extend the state */
    t_disresdata* disresdata;
    snew(disresdata, 1);
    init_disres(fplog,
                mtop,
                inputrec.get(),
                DisResRunMode::MDRun,
                isSimulationMainRank ? DDRole::Main : DDRole::Agent,
                simulationIsParallel ? NumRanks::Multiple : NumRanks::Single,
                mpiCommSimulation.comm(),
                ms,
                disresdata,
                globalState.get(),
                replExParams.exchangeInterval > 0);

    if (gmx_mtop_ftype_count(mtop, InteractionFunction::OrientationRestraints) > 0 && isSimulationMainRank)
    {
        extendStateWithOriresHistory(mtop, *inputrec, globalState.get());
    }

#if GMX_FAHCORE
    /* We have to remember the generation's first step before reading checkpoint.
       This way, we can report to the F@H core both the generation's first step
       and the restored first step, thus making it able to distinguish between
       an interruption/resume and start of the n-th generation simulation.
       Having this information, the F@H core can correctly calculate and report
       the progress.
     */
    int gen_first_step = 0;
    if (cr->commMySim.isMainRank())
    {
        gen_first_step = inputrec->init_step;
    }
#endif

    ObservablesHistory observablesHistory = {};

    auto modularSimulatorCheckpointData = std::make_unique<ReadCheckpointDataHolder>();
    if (startingBehavior != StartingBehavior::NewSimulation)
    {
        /* Check if checkpoint file exists before doing continuation.
         * This way we can use identical input options for the first and subsequent runs...
         */
        if (mdrunOptions.numStepsCommandline > -2)
        {
            /* Temporarily set the number of steps to unlimited to avoid
             * triggering the nsteps check in load_checkpoint().
             * This hack will go away soon when the -nsteps option is removed.
             */
            inputrec->nsteps = -1;
        }

        // Finish applying initial simulation state information from external sources on all ranks.
        // Reconcile checkpoint file data with Mdrunner state established up to this point.
        applyLocalState(*inputHolder_.get(),
                        logFileHandle,
                        mpiCommSimulation,
                        inputrec.get(),
                        globalState.get(),
                        &observablesHistory,
                        mdrunOptions.reproducible,
                        mdModules_->notifiers(),
                        modularSimulatorCheckpointData.get(),
                        useModularSimulator);
        // TODO: (#3652) Synchronize filesystem state, SimulationInput contents, and program
        //  invariants
        //  on all code paths.
        // Write checkpoint or provide hook to update SimulationInput.
        // If there was a checkpoint file, SimulationInput contains more information
        // than if there wasn't. At this point, we have synchronized the in-memory
        // state with the filesystem state only for restarted simulations. We should
        // be calling applyLocalState unconditionally and expect that the completeness
        // of SimulationInput is not dependent on its creation method.

        if (startingBehavior == StartingBehavior::RestartWithAppending && logFileHandle)
        {
            // Now we can start normal logging to the truncated log file.
            fplog = gmx_fio_getfp(logFileHandle);
            prepareLogAppending(fplog);
            logOwner = buildLogger(fplog, isSimulationMainRank);
            mdlog    = logOwner.logger();
            // Provide the log file handle to the fatal error handler
            gmx_fatal_set_log_file(fplog);
            // A scope guard in mdrun handles clearing this when the log file is closed.
        }
    }

#if GMX_FAHCORE
    if (cr->commMySim.isMainRank())
    {
        fcRegisterSteps(inputrec->nsteps + inputrec->init_step, gen_first_step);
    }
#endif

    if (mdrunOptions.numStepsCommandline > -2)
    {
        GMX_LOG(mdlog.info)
                .asParagraph()
                .appendText(
                        "The -nsteps functionality is deprecated, and may be removed in a future "
                        "version. "
                        "Consider using gmx convert-tpr -nsteps or changing the appropriate .mdp "
                        "file field.");
    }
    /* override nsteps with value set on the commandline */
    override_nsteps_cmdline(mdlog, mdrunOptions.numStepsCommandline, inputrec.get());

    matrix box;
    if (isSimulationMainRank)
    {
        copy_mat(globalState->box, box);
    }

    if (simulationIsParallel)
    {
        gmx_bcast(sizeof(box), box, mpiCommSimulation.comm());
    }

    if (inputrec->cutoff_scheme != CutoffScheme::Verlet)
    {
        gmx_fatal(FARGS,
                  "This group-scheme .tpr file can no longer be run by mdrun. Please update to the "
                  "Verlet scheme, or use an earlier version of GROMACS if necessary.");
    }
    /* Update rlist and nstlist. */
    /* Note: prepare_verlet_scheme is calling increaseNstlist(...), which (while attempting to
     * increase rlist) tries to check if the newly chosen value fits with the DD scheme. As this is
     * run before any DD scheme is set up, this check is never executed. See #3334 for more details.
     */
    prepare_verlet_scheme(
            fplog,
            mpiCommSimulation,
            inputrec.get(),
            nstlist_cmdline,
            mtop,
            mpiCommSimulation.isMainRank() ? globalState->x : gmx::ArrayRef<const gmx::RVec>{},
            box,
            useGpuForNonbonded || (emulateGpuNonbonded == EmulateGpuNonbonded::Yes));

    // We need to decide on update groups early, as this affects
    // inter-domain communication distances.
    auto         updateGroupingsPerMoleculeTypeResult = makeUpdateGroupingsPerMoleculeType(mtop);
    UpdateGroups updateGroups;
    if (std::holds_alternative<std::string>(updateGroupingsPerMoleculeTypeResult))
    {
        GMX_LOG(mdlog.warning)
                .asParagraph()
                .appendTextFormatted("Update groups can not be used for this system because %s",
                                     std::get<std::string>(updateGroupingsPerMoleculeTypeResult).c_str());
    }
    else
    {
        auto updateGroupingsPerMoleculeType =
                std::get<std::vector<RangePartitioning>>(updateGroupingsPerMoleculeTypeResult);
        const real maxUpdateGroupRadius = computeMaxUpdateGroupRadius(
                mtop, updateGroupingsPerMoleculeType, maxReferenceTemperature(*inputrec));
        const real cutoffMargin = std::sqrt(max_cutoff2(inputrec->pbcType, box)) - inputrec->rlist;
        updateGroups            = makeUpdateGroups(mdlog,
                                        std::move(updateGroupingsPerMoleculeType),
                                        maxUpdateGroupRadius,
                                        doRerun,
                                        useDomainDecomposition,
                                        systemHasConstraintsOrVsites(mtop),
                                        cutoffMargin);
    }

    try
    {
        const bool haveFrozenAtoms = inputrecFrozenAtoms(inputrec.get());

        useGpuForUpdate = decideWhetherToUseGpuForUpdate(
                useDomainDecomposition,
                updateGroups.useUpdateGroups(),
                pmeRunMode,
                domdecOptions.numPmeRanks > 0,
                useGpuForNonbonded,
                updateTarget,
                gpusWereDetected,
                *inputrec,
                mtop,
                doEssentialDynamics,
                gmx_mtop_ftype_count(mtop, InteractionFunction::OrientationRestraints) > 0,
                haveFrozenAtoms,
                useModularSimulator,
                doRerun,
                mdlog);
    }
    GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR

    // We are using the minimal supported level of GPU-aware MPI
    // support as an approximation of whether such communication
    // will work. It likely would not work in cases where ranks
    // have heterogeneous device types or vendors unless the MPI
    // library supported that.
    const bool canUseDirectGpuComm =
            decideWhetherDirectGpuCommunicationCanBeUsed(hwinfo_->minGpuAwareMpiStatus,
                                                         inputrec->useMts,
                                                         replExParams.exchangeInterval > 0,
                                                         (inputrec->eSwapCoords != SwapType::No),
                                                         gpusWereDetected,
                                                         mdlog);

    // Initialize development feature flags that enabled by environment variable
    // and report those features that are enabled.
    const DevelopmentFeatureFlags devFlags = manageDevelopmentFeatures(
            mdlog, pmeRunMode, mpiCommSimulation.size(), domdecOptions.numPmeRanks, canUseDirectGpuComm);

    bool useGpuDirectHalo = false;

    if (useGpuForNonbonded)
    {
        // domdecOptions.numPmeRanks == -1 results in 0 separate PME ranks when useGpuForNonbonded is true.
        // Todo: remove this assumption later once auto mode has support for separate PME rank
        const int numPmeRanks = domdecOptions.numPmeRanks > 0 ? domdecOptions.numPmeRanks : 0;
        bool      havePPDomainDecomposition = (mpiCommSimulation.size() - numPmeRanks) > 1;
        useGpuDirectHalo = decideWhetherToUseGpuForHalo(havePPDomainDecomposition,
                                                        useGpuForNonbonded,
                                                        canUseDirectGpuComm,
                                                        useModularSimulator,
                                                        doRerun,
                                                        EI_ENERGY_MINIMIZATION(inputrec->eI),
                                                        mdlog);
    }

    const bool haveFillerParticlesInLocalState = localStateHasFillerParticles(
            mtop,
            *inputrec,
            useDomainDecomposition,
            !simulationIsParallel || mpiCommSimulation.size() - domdecOptions.numPmeRanks == 1,
            useGpuDirectHalo);

    if (haveFillerParticlesInLocalState && useDomainDecomposition
        && mpiCommSimulation.size() - domdecOptions.numPmeRanks > 1
        && domdecOptions.dlbOption != gmx::DlbOption::no)
    {
        GMX_LOG(mdlog.info)
                .asParagraph()
                .appendText(
                        "Turning off DLB, as the is not (yet) supported with direct halo exchange");
        domdecOptions.dlbOption = gmx::DlbOption::no;
    }

    // This builder is necessary while we have multi-part construction
    // of DD. Before DD is constructed, we use the existence of
    // the builder object to indicate that further construction of DD
    // is needed.
    std::unique_ptr<DomainDecompositionBuilder> ddBuilder;
    if (useDomainDecomposition)
    {
        // The DD builder will disable useGpuDirectHalo if the Y or Z component of any domain is
        // smaller than twice the communication distance, since GPU-direct communication presently
        // only works with a single pulse in these dimensions, and we want to avoid box scaling
        // resulting in fatal errors far into the simulation. Such small systems will not
        // perform well on multiple GPUs in any case, but it is important that our core functionality
        // (in particular for testing) does not break depending on GPU direct communication being enabled.
        ddBuilder = std::make_unique<DomainDecompositionBuilder>(
                mdlog,
                mpiCommSimulation,
                domdecOptions,
                mdrunOptions,
                mtop,
                *inputrec,
                mdModules_->notifiers(),
                box,
                updateGroups.updateGroupingPerMoleculeType(),
                updateGroups.useUpdateGroups(),
                updateGroups.maxUpdateGroupRadius(),
                positionsFromStatePointer(globalState.get()),
                useGpuForNonbonded,
                useGpuForPme,
                useGpuForUpdate,
                useGpuDirectHalo,
                devFlags.enableGpuPmeDecomposition);
    }
    else
    {
        if (inputrec->pbcType == PbcType::Screw)
        {
            gmx_fatal(FARGS, "pbc=screw is only implemented with domain decomposition");
        }
    }

    // Produce the task assignment for all ranks on this node - done after DD is constructed
    GpuTaskAssignments gpuTaskAssignments = GpuTaskAssignmentsBuilder::build(
            availableDevices,
            userGpuTaskAssignment,
            *hwinfo_,
            simulationCommunicator,
            physicalNodeComm,
            nonbondedTarget,
            nonbondedFeTarget,
            pmeTarget,
            bondedTarget,
            updateTarget,
            useGpuForNonbonded,
            useGpuForPme,
            ddBuilder ? ddBuilder->thisRankHasPPDuty() : true,
            // TODO ddBuilder->thisRankHasPmeDuty should imply that a PME
            // algorithm is active, but currently does not.
            usingPme(inputrec->coulombtype) && (ddBuilder ? ddBuilder->thisRankHasPmeDuty() : true));

    // Get the device handle for the modules on this rank, nullptr
    // when no task is assigned.
    DeviceInformation* deviceInfo = gpuTaskAssignments.initDevice();

    // We will later get the pairlist type from the device information here, for now we hard code it.
    const auto pairlistType = PairlistType::Hierarchical8x8x8;
    GMX_RELEASE_ASSERT(sc_gpuClusterSize(pairlistType) >= 4,
                       "The verlet scheme setup relies on the GPU cluster size to be at least 4");

    // TODO Currently this is always built, yet DD partition code
    // checks if it is built before using it. Probably it should
    // become an MDModule that is made only when another module
    // requires it (e.g. pull, CompEl, density fitting), so that we
    // don't update the local atom sets unilaterally every step.
    LocalAtomSetManager atomSets;

    // Local state and topology are declared (and perhaps constructed)
    // now, because DD needs them for the LocalTopologyChecker, but
    // they do not contain valid data until after the first DD
    // partition.
    std::unique_ptr<t_state> localStateInstance;
    t_state*                 localState;
    gmx_localtop_t           localTopology(mtop.ffparams);

    std::unique_ptr<gmx_domdec_t> domdec;
    std::unique_ptr<t_commrec>    commRec;

    if (ddBuilder)
    {
        localStateInstance = std::make_unique<t_state>();
        localState         = localStateInstance.get();
        // TODO Pass the GPU streams to ddBuilder to use in buffer
        // transfers (e.g. halo exchange)
        domdec = ddBuilder->build(&atomSets, haveFillerParticlesInLocalState, &observablesReducerBuilder);
        // The builder's job is done, so destruct it
        ddBuilder.reset(nullptr);

        // Report on hierarchical reductions, when active
        if (const auto mesg = domdec->mpiComm().getHierarchicalReductionsReport(); mesg)
        {
            GMX_LOG(mdlog.info).asParagraph().appendText(*mesg);
        }

        commRec = std::make_unique<t_commrec>(domdec->mpiCommMySim(), domdec->mpiComm_, domdec.get());

        // Note that local state still does not exist yet.
    }
    else
    {
        // Without DD, the local state is merely an alias to the global state,
        // so we don't need to allocate anything.
        localState = globalState.get();

        // DD is inactive, so all ranks have all PP and PME duties
        commRec = std::make_unique<t_commrec>(mpiCommSimulation, mpiCommSimulation, nullptr);
    }

    t_commrec* cr = commRec.get();

    // Ensure that all atoms within the same update group are in the
    // same periodic image. Otherwise, a simulation that did not use
    // update groups (e.g. a single-rank simulation) cannot always be
    // correctly restarted in a way that does use update groups
    // (e.g. a multi-rank simulation).
    if (isSimulationMainRank)
    {
        const bool useUpdateGroups = cr->dd ? ddUsesUpdateGroups(*cr->dd) : false;
        if (useUpdateGroups)
        {
            putUpdateGroupAtomsInSamePeriodicImage(*cr->dd, mtop, globalState->box, globalState->x);
        }
    }

    const bool disableNonbondedCalculation = (std::getenv("GMX_NO_NONBONDED") != nullptr);
    if (disableNonbondedCalculation)
    {
        /* turn off non-bonded calculations */
        GMX_LOG(mdlog.warning)
                .asParagraph()
                .appendText(
                        "Found environment variable GMX_NO_NONBONDED.\n"
                        "Disabling nonbonded calculations.");
    }

    const NumPmeDomains numPmeDomains = getNumPmeDomains(cr->dd);
    const bool useGpuPmeDecomposition = numPmeDomains.x * numPmeDomains.y > 1 && useGpuForPme;
    GMX_RELEASE_ASSERT(!useGpuPmeDecomposition || devFlags.enableGpuPmeDecomposition,
                       "GPU PME decomposition works only in the cases where it is supported");

    MdrunScheduleWorkload runScheduleWork;

    // Also populates the simulation constant workload description.
    // Note: currently the default duty is DUTY_PP | DUTY_PME for all simulations, including those without PME,
    // so this boolean is sufficient on all ranks to determine whether separate PME ranks are used,
    // but this will no longer be the case if cr->duty is changed for !usingPme(fr->ic->eeltype).
    const bool haveSeparatePmeRank = (thisRankHasPPDuty(cr->dd) != thisRankHasPmeDuty(cr->dd));
    runScheduleWork.simulationWork = createSimulationWorkload(
            mdlog,
            *inputrec,
            inputrecDynamicBox(inputrec.get()) || doRerun || inputrec->eI == IntegrationAlgorithm::Mimic,
            replExParams.exchangeInterval > 0,
            disableNonbondedCalculation,
            devFlags,
            haveFillerParticlesInLocalState,
            havePPDomainDecomposition(cr->dd),
            haveSeparatePmeRank,
            useGpuForNonbonded,
            useGpuForNonbondedFE,
            pmeRunMode,
            useGpuForBonded,
            useGpuForUpdate,
            useGpuDirectHalo,
            canUseDirectGpuComm,
            useGpuPmeDecomposition);

    if (GMX_LIB_MPI && deviceInfo
        && (runScheduleWork.simulationWork.useGpuDirectCommunication
            || runScheduleWork.simulationWork.useGpuPmeDecomposition
            || runScheduleWork.simulationWork.useGpuPmePpCommunication
            || runScheduleWork.simulationWork.useGpuHaloExchange))
    {
        doubleCheckGpuAwareMpiWillWork(*deviceInfo);
    }

    GMX_LOG(mdlog.info)
            .asParagraph()
            .appendTextFormatted("Local state %s filler particles",
                                 runScheduleWork.simulationWork.haveFillerParticlesInLocalState
                                         ? "uses"
                                         : "does not use");

    if (runScheduleWork.simulationWork.useGpuDirectCommunication
        && gmx::GpuConfigurationCapabilities::DisableEventCounting)
    {
        // Don't enable event counting with GPU Direct comm, see #3988.
        gmx::internal::disableGpuEventConsumptionCounting();
    }

    const bool printHostName = (cr->commMySim.isParallel());
    gpuTaskAssignments.reportGpuUsage(mdlog, printHostName, pmeRunMode, runScheduleWork.simulationWork);

    std::unique_ptr<DeviceStreamManager> deviceStreamManager = nullptr;

    if (deviceInfo != nullptr)
    {
        if (runScheduleWork.simulationWork.havePpDomainDecomposition && thisRankHasPPDuty(cr->dd))
        {
            dd_setup_dlb_resource_sharing(cr, uniqueDeviceId(*deviceInfo));
        }
        const bool useGpuTiming = decideGpuTimingsUsage();
        deviceStreamManager     = std::make_unique<DeviceStreamManager>(
                *deviceInfo, runScheduleWork.simulationWork, useGpuTiming);
    }

    std::unique_ptr<NvshmemManager> nvshmemManager;
    if (runScheduleWork.simulationWork.useNvshmem)
    {
        nvshmemManager = std::make_unique<NvshmemManager>(mdlog, cr->commMySim.comm());
    }

    // If the user chose a task assignment, give them some hints
    // where appropriate.
    if (!userGpuTaskAssignment.empty())
    {
        gpuTaskAssignments.logPerformanceHints(mdlog, numAvailableDevices);
    }

#if GMX_MPI
    if (isMultiSim(ms))
    {
        GMX_LOG(mdlog.warning)
                .asParagraph()
                .appendTextFormatted(
                        "This is simulation %d out of %d running as a composite GROMACS\n"
                        "multi-simulation job. Setup for this simulation:\n",
                        ms->simulationIndex_,
                        ms->numSimulations_);
    }
    GMX_LOG(mdlog.warning)
            .appendTextFormatted("Using %d MPI %s\n",
                                 cr->commMySim.size(),
#    if GMX_THREAD_MPI
                                 cr->commMySim.isSerial() ? "thread" : "threads"
#    else
                                 cr->commMySim.isSerial() ? "process" : "processes"
#    endif
            );
    std::fflush(stderr);
#endif

    // If mdrun -pin auto honors any affinity setting that already
    // exists. If so, it is nice to provide feedback about whether
    // that existing affinity setting was from OpenMP or something
    // else, so we run this code both before and after we initialize
    // the OpenMP support.
    gmx_check_thread_affinity_set(
            mdlog, &hw_opt, hwinfo_->hardwareTopology->maxThreads(), FALSE, libraryWorldCommunicator);
    /* Check and update the number of OpenMP threads requested */
    checkAndUpdateRequestedNumOpenmpThreads(
            &hw_opt, *hwinfo_, cr->commMySim, ms, physicalNodeComm.size_, pmeRunMode, mtop, *inputrec);

    gmx_omp_nthreads_init(mdlog,
                          cr->commMySim,
                          haveSeparatePmeRank,
                          hwinfo_->hardwareTopology->maxThreads(),
                          physicalNodeComm.size_,
                          hw_opt.nthreads_omp,
                          hw_opt.nthreads_omp_pme,
                          !thisRankHasPPDuty(cr->dd));

    const bool bEnableFPE = gmxShouldEnableFPExceptions();
    // FIXME - reconcile with gmx_feenableexcept() call from CommandLineModuleManager::run()
    if (bEnableFPE)
    {
        gmx_feenableexcept();
    }

    /* Now that we know the setup is consistent, check for efficiency */
    check_resource_division_efficiency(
            hwinfo_, gpuTaskAssignments.thisRankHasAnyGpuTask(), &cr->commMySim, mdlog);

    /* getting number of PP/PME threads on this MPI / tMPI rank.
       PME: env variable should be read only on one node to make sure it is
       identical everywhere;
     */
    const int numThreadsOnThisRank = thisRankHasPPDuty(cr->dd)
                                             ? gmx_omp_nthreads_get(ModuleMultiThread::Nonbonded)
                                             : gmx_omp_nthreads_get(ModuleMultiThread::Pme);
    checkHardwareOversubscription(
            numThreadsOnThisRank, cr->commMyGroup.rank(), *hwinfo_->hardwareTopology, physicalNodeComm, mdlog);

    // Enable Peer access between GPUs where available
    // Only for DD, only main PP rank needs to perform setup, and only if thread MPI plus
    // any of the GPU communication features are active.
    if (haveDDAtomOrdering(*cr) && cr->commMySim.isMainRank() && thisRankHasPPDuty(cr->dd) && GMX_THREAD_MPI
        && (runScheduleWork.simulationWork.useGpuHaloExchange
            || runScheduleWork.simulationWork.useGpuPmePpCommunication))
    {
        setupGpuDevicePeerAccess(gpuTaskAssignments.deviceIdsAssigned(), mdlog);
    }

    if (mdrunOptions.timingOptions.resetStep > -1)
    {
        GMX_LOG(mdlog.info)
                .asParagraph()
                .appendText(
                        "The -resetstep functionality is deprecated, and may be removed in a "
                        "future version.");
    }
    std::unique_ptr<gmx_wallcycle> wcycle =
            wallcycle_init(fplog, mdrunOptions.timingOptions.resetStep, cr);

    if (cr->commMySim.isParallel())
    {
        /* Main synchronizes its value of reset_counters with all nodes
         * including PME only nodes */
        int64_t reset_counters = wcycle_get_reset_counters(wcycle.get());
        gmx_bcast(sizeof(reset_counters), &reset_counters, cr->commMySim.comm());
        wcycle_set_reset_counters(wcycle.get(), reset_counters);
    }

    // Membrane embedding must be initialized before we call init_forcerec()
    membedHolder.initializeMembed(fplog,
                                  filenames.size(),
                                  filenames.data(),
                                  &mtop,
                                  inputrec.get(),
                                  globalState.get(),
                                  cr,
                                  &mdrunOptions.checkpointOptions.period);

    const bool thisRankHasPmeGpuTask = gpuTaskAssignments.thisRankHasPmeGpuTask();
    std::unique_ptr<BoxDeformation>      deform;
    std::unique_ptr<MDAtoms>             mdAtoms;
    std::unique_ptr<VirtualSitesHandler> vsite;

    t_nrnb nrnb;
    if (thisRankHasPPDuty(cr->dd))
    {
        setupNotifier.notify(cr->commMyGroup);
        setupNotifier.notify(ms);
        setupNotifier.notify(&atomSets);
        setupNotifier.notify(mtop);
        setupNotifier.notify(inputrec->pbcType);
        setupNotifier.notify(SimulationTimeStep{ inputrec->delta_t });
        setupNotifier.notify(startingBehavior);
        setupNotifier.notify(EnsembleTemperature{ *inputrec });
        PlainPairlistRanges plainPairlistRanges(mtop, *inputrec);
        setupNotifier.notify(&plainPairlistRanges);
        MDModulesDirectProvider mdModuleCoulombDirectProvider;
        setupNotifier.notify(&mdModuleCoulombDirectProvider);

        /* Initiate forcerecord */
        fr                 = std::make_unique<t_forcerec>();
        fr->forceProviders = mdModules_->initForceProviders(wcycle.get());

        std::optional<bool> anMDModuleProvidesDirectCoulomb = mdModuleCoulombDirectProvider.isDirectProvider;
        if (setupNotifier.haveSubscribers<MDModulesDirectProvider*>()
            && !anMDModuleProvidesDirectCoulomb.has_value())
        {
            GMX_THROW(gmx::APIError(
                    "At least one MD module subscribed to the direct provider notification, "
                    "but none reported whether it provides direct interactions."));
        }
        if (!setupNotifier.haveSubscribers<MDModulesDirectProvider*>()
            && anMDModuleProvidesDirectCoulomb.has_value())
        {
            GMX_THROW(
                    gmx::APIError("No MD module subscribed to the direct provider notification, "
                                  "but direct interactions provider is reported."));
        }

        logValidationMessages(mdlog,
                              runScheduleWork.simulationWork,
                              haveFillerParticlesInLocalState,
                              fn2ftp(ftp2fn(efTRN, filenames.size(), filenames.data())) == efH5MD,
                              useModularSimulator,
                              inputrec->eI,
                              deviceInfo);

        init_forcerec(fplog,
                      mdlog,
                      runScheduleWork.simulationWork,
                      fr.get(),
                      *inputrec,
                      mtop,
                      cr,
                      ms,
                      box,
                      opt2fn("-table", filenames.size(), filenames.data()),
                      opt2fn("-tablep", filenames.size(), filenames.data()),
                      opt2fns("-tableb", filenames.size(), filenames.data()),
                      pforce,
                      anMDModuleProvidesDirectCoulomb);

        // Dirty hack, for fixing disres and orires should be made mdmodules
        fr->fcdata->disres = disresdata;
        if (gmx_mtop_ftype_count(mtop, InteractionFunction::OrientationRestraints) > 0)
        {
            fr->fcdata->orires = std::make_unique<t_oriresdata>(
                    fplog, mtop, *inputrec, ms, globalState.get(), &atomSets);
        }

        // Now that all simulation setup notifications have been emitted,
        // set up the simulation run notification subscriptions
        mdModules_->subscribeToSimulationRunNotifications();

        if (mdModules_->notifiers().simulationRunNotifier_.haveSubscribers<const MDModulesPairlistConstructedSignal&>())
        {
            if (plainPairlistRanges.ranges().empty())
            {
                GMX_THROW(gmx::APIError(
                        "At least one MDModule has subscribed to the pairlist construction signal, "
                        "but no MDModules requested a range for the plain pairlist"));
            }

            fr->plainPairlistRange = *std::max_element(plainPairlistRanges.ranges().begin(),
                                                       plainPairlistRanges.ranges().end());
            if (fr->plainPairlistRange.value() > inputrec->rlist)
            {
                // This is not nice. We could consider increasing rlist instead.
                const std::string mesg = gmx::formatString(
                        "MDModules request a plain pairlist with a range of %f nm, which is larger "
                        "than the normal pairlist range of %f nm",
                        fr->plainPairlistRange.value(),
                        inputrec->rlist);
                GMX_THROW(gmx::APIError(mesg));
            }
        }

        deform = buildBoxDeformation(
                globalState != nullptr ? createMatrix3x3FromLegacyMatrix(globalState->box)
                                       : diagonalMatrix<real>(0.0),
                cr->commMySim.isMainRank() ? DDRole::Main : DDRole::Agent,
                cr->commMySim.isParallel() ? NumRanks::Multiple : NumRanks::Single,
                cr->commMyGroup.comm(),
                *inputrec);

        // Save a handle to device stream manager to use elsewhere in the code
        // TODO: Forcerec is not a correct place to store it.
        fr->deviceStreamManager = deviceStreamManager.get();

        if (runScheduleWork.simulationWork.useGpuPmePpCommunication && !thisRankHasPmeDuty(cr->dd))
        {
            GMX_RELEASE_ASSERT(
                    deviceStreamManager != nullptr,
                    "GPU device stream manager should be valid in order to use PME-PP direct "
                    "communications.");
            GMX_RELEASE_ASSERT(
                    deviceStreamManager->streamIsValid(DeviceStreamType::PmePpTransfer),
                    "GPU PP-PME stream should be valid in order to use GPU PME-PP direct "
                    "communications.");
            fr->pmePpCommGpu = std::make_unique<gmx::PmePpCommGpu>(
                    cr->commMySim.comm(),
                    cr->dd->pme_nodeid,
                    &cr->dd->pmeForceReceiveBuffer,
                    deviceStreamManager->context(),
                    deviceStreamManager->stream(DeviceStreamType::PmePpTransfer),
                    runScheduleWork.simulationWork.useNvshmem);
        }

        fr->nbv = init_nb_verlet(
                mdlog,
                *inputrec,
                *fr,
                cr->commMyGroup,
                cr->dd,
                *hwinfo_,
                runScheduleWork.simulationWork.useGpuNonbonded,
                runScheduleWork.simulationWork.useGpuNonbondedFE,
                deviceStreamManager.get(),
                mtop,
                runScheduleWork.simulationWork.haveFillerParticlesInLocalState,
                (cr->dd && cr->dd->mpiComm().isParallel()) ? &observablesReducerBuilder : nullptr,
                isSimulationMainRank ? globalState->x : gmx::ArrayRef<const gmx::RVec>{},
                box,
                wcycle.get());
        // TODO: Move the logic below to a GPU bonded builder
        if (runScheduleWork.simulationWork.useGpuBonded)
        {
            GMX_RELEASE_ASSERT(deviceStreamManager != nullptr,
                               "GPU device stream manager should be valid in order to use GPU "
                               "version of bonded forces.");
            fr->listedForcesGpu =
                    std::make_unique<ListedForcesGpu>(mtop.ffparams,
                                                      fr->ic->coulomb.epsfac * fr->fudgeQQ,
                                                      inputrec->opts.ngener - inputrec->nwall,
                                                      deviceStreamManager->context(),
                                                      deviceStreamManager->bondedStream(),
                                                      wcycle.get());
        }
        fr->longRangeNonbondeds = std::make_unique<CpuPpLongRangeNonbondeds>(fr->n_tpi,
                                                                             fr->ic->coulomb.ewaldCoeff,
                                                                             fr->ic->coulomb.epsilon_r,
                                                                             fr->qsum,
                                                                             fr->ic->coulomb.type,
                                                                             fr->ic->vdw.type,
                                                                             *inputrec,
                                                                             &nrnb,
                                                                             wcycle.get(),
                                                                             fplog);

        /* Initialize the mdAtoms structure.
         * mdAtoms is not filled with atom data,
         * as this can not be done now with domain decomposition.
         */
        mdAtoms = makeMDAtoms(fplog, mtop, *inputrec, thisRankHasPmeGpuTask);
        if (globalState && thisRankHasPmeGpuTask)
        {
            // The pinning of coordinates in the global state object works, because we only use
            // PME on GPU without DD or on a separate PME rank, and because the local state pointer
            // points to the global state object without DD.
            // FIXME: MD and EM separately set up the local state - this should happen in the same
            // function, which should also perform the pinning.
            changePinningPolicy(&globalState->x, pme_get_pinning_policy());
        }

        /* Initialize the virtual site communication */
        vsite = makeVirtualSitesHandler(
                mtop, cr->dd, fr->pbcType, updateGroups.updateGroupingPerMoleculeType());

        calc_shifts(box, fr->shift_vec);

        /* With periodic molecules the charge groups should be whole at start up
         * and the virtual sites should not be far from their proper positions.
         */
        if (!inputrec->bContinuation && cr->commMySim.isMainRank()
            && !(inputrec->pbcType != PbcType::No && inputrec->bPeriodicMols))
        {
            /* Make molecules whole at start of run */
            if (fr->pbcType != PbcType::No)
            {
                do_pbc_first_mtop(fplog,
                                  inputrec->pbcType,
                                  ir_haveBoxDeformation(*inputrec),
                                  inputrec->deform,
                                  box,
                                  &mtop,
                                  globalState->x,
                                  globalState->v);
            }
            if (vsite)
            {
                /* Correct initial vsite positions are required
                 * for the initial distribution in the domain decomposition
                 * and for the initial shell prediction.
                 */
                constructVirtualSitesGlobal(mtop, globalState->x);
            }
        }
        // Make the DD reverse topology, now that any vsites that are present are available
        if (haveDDAtomOrdering(*cr))
        {
            dd_make_reverse_top(fplog, cr->dd, mtop, vsite.get(), *inputrec, domdecOptions.ddBondedChecking);
        }

        if (usingPme(fr->ic->coulomb.type) || usingLJPme(fr->ic->vdw.type))
        {
            ewaldcoeff_q  = fr->ic->coulomb.ewaldCoeff;
            ewaldcoeff_lj = fr->ic->vdw.ewaldCoeff;
        }
    }
    else
    {
        /* This is a PME only node */

        GMX_ASSERT(globalState == nullptr,
                   "We don't need the state on a PME only rank and expect it to be uninitialized");

        ewaldcoeff_q  = calc_ewaldcoeff_q(inputrec->rcoulomb, inputrec->ewald_rtol);
        ewaldcoeff_lj = calc_ewaldcoeff_lj(inputrec->rvdw, inputrec->ewald_rtol_lj);
    }

    gmx_pme_t* sepPmeData = nullptr;
    // This reference hides the fact that PME data is owned by runner on PME-only ranks and by forcerec on other ranks
    GMX_ASSERT(thisRankHasPPDuty(cr->dd) == (fr != nullptr),
               "Double-checking that only PME-only ranks have no forcerec");
    gmx_pme_t*& pmedata = fr ? fr->pmedata : sepPmeData;

    // TODO should live in ewald module once its testing is improved
    //
    // Later, this program could contain kernels that might be later
    // re-used as auto-tuning progresses, or subsequent simulations
    // are invoked.
    PmeGpuProgramStorage pmeGpuProgram;
    if (thisRankHasPmeGpuTask)
    {
        GMX_RELEASE_ASSERT(
                (deviceStreamManager != nullptr),
                "GPU device stream manager should be initialized in order to use GPU for PME.");
        GMX_RELEASE_ASSERT((deviceInfo != nullptr),
                           "GPU device should be initialized in order to use GPU for PME.");
        pmeGpuProgram = buildPmeGpuProgram(deviceStreamManager->context());
    }

    /* Initiate PME if necessary,
     * either on all nodes or on dedicated PME nodes only. */
    if (usingPme(inputrec->coulombtype) || usingLJPme(inputrec->vdwtype))
    {
        if (mdAtoms && mdAtoms->mdatoms())
        {
            nChargePerturbed = mdAtoms->mdatoms()->nChargePerturbed;
            if (usingLJPme(inputrec->vdwtype))
            {
                nTypePerturbed = mdAtoms->mdatoms()->nTypePerturbed;
            }
        }
        if (cr->dd && cr->dd->numPmeOnlyRanks > 0)
        {
            /* The PME only nodes need to know nChargePerturbed(FEP on Q) and nTypePerturbed(FEP on LJ)*/
            gmx_bcast(sizeof(nChargePerturbed), &nChargePerturbed, cr->commMySim.comm());
            gmx_bcast(sizeof(nTypePerturbed), &nTypePerturbed, cr->commMySim.comm());
        }

        if (thisRankHasPmeDuty(cr->dd))
        {
            try
            {
                // TODO: This should be in the builder.
                GMX_RELEASE_ASSERT(!runScheduleWork.simulationWork.useGpuPme
                                           || (deviceStreamManager != nullptr),
                                   "Device stream manager should be valid in order to use GPU "
                                   "version of PME.");
                GMX_RELEASE_ASSERT(
                        !runScheduleWork.simulationWork.useGpuPme
                                || deviceStreamManager->streamIsValid(DeviceStreamType::Pme),
                        "GPU PME stream should be valid in order to use GPU version of PME.");

                const DeviceContext* deviceContext = runScheduleWork.simulationWork.useGpuPme
                                                             ? &deviceStreamManager->context()
                                                             : nullptr;
                const DeviceStream*  pmeStream =
                        runScheduleWork.simulationWork.useGpuPme
                                 ? &deviceStreamManager->stream(DeviceStreamType::Pme)
                                 : nullptr;

                const t_inputrec* ir = inputrec.get();
                /* For each atom we allow a relative (cut-off) error of up to ewald_rtol.
                 * Thus we can also tolerate an error of an order of magnitude less due to
                 * atoms being slightly outside the halo extent. This will only cause a fraction
                 * of the charge to be missing on the grid. So we pass ewald_rtol as the allowed
                 * chance per atom (ChanceTarget::Atom) to be outside the halo extent.
                 */
                const real haloExtentForAtomDisplacement =
                        updateGroups.maxUpdateGroupRadius()
                        + minCellSizeForAtomDisplacement(mtop,
                                                         *ir,
                                                         updateGroups.updateGroupingPerMoleculeType(),
                                                         ir->ewald_rtol,
                                                         ChanceTarget::Atom);
                pmedata = gmx_pme_init(cr->dd,
                                       getNumPmeDomains(cr->dd),
                                       ir,
                                       box,
                                       haloExtentForAtomDisplacement,
                                       nChargePerturbed != 0,
                                       nTypePerturbed != 0,
                                       mdrunOptions.reproducible,
                                       ewaldcoeff_q,
                                       ewaldcoeff_lj,
                                       gmx_omp_nthreads_get(ModuleMultiThread::Pme),
                                       pmeRunMode,
                                       nullptr,
                                       deviceContext,
                                       pmeStream,
                                       pmeGpuProgram.get(),
                                       mdlog,
                                       nullptr);
            }
            GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR
        }
    }

    /* Set thread affinity after gmx_pme_init(), otherwise with cuFFTMp the NVSHMEM helper thread
     * can be pinned to the same core as the PME thread, causing performance degradation.
     */
    if (hw_opt.threadAffinity != ThreadAffinity::Off)
    {
        /* Before setting affinity, check whether the affinity has changed
         * - which indicates that probably the OpenMP library has changed it
         * since we first checked).
         */
        gmx_check_thread_affinity_set(
                mdlog, &hw_opt, hwinfo_->hardwareTopology->maxThreads(), TRUE, libraryWorldCommunicator);
    }
    {
        /* We always call gmx_set_thread_affinity since it might want to write things to log
         * even with ThreadAffinity::Off */
        gmx_set_thread_affinity(
                mdlog, cr->commMySim, physicalNodeComm, &hw_opt, *hwinfo_->hardwareTopology, numThreadsOnThisRank, nullptr);
    }

    if (EI_DYNAMICS(inputrec->eI))
    {
        /* Turn on signal handling on all nodes */
        /*
         * (A user signal from the PME nodes (if any)
         * is communicated to the PP nodes.
         */
        signal_handler_install();
    }

    try
    {
        pull_t* pull_work = nullptr;
        if (thisRankHasPPDuty(cr->dd))
        {
            /* Assumes uniform use of the number of OpenMP threads */
            walltime_accounting =
                    walltime_accounting_init(gmx_omp_nthreads_get(ModuleMultiThread::Default));

            if (inputrec->bPull)
            {
                /* Initialize pull code */
                pull_work = init_pull(fplog,
                                      inputrec->pull.get(),
                                      inputrec.get(),
                                      mtop,
                                      cr->commMyGroup,
                                      cr->dd,
                                      &atomSets,
                                      inputrec->fepvals->init_lambda_without_states);
                if (inputrec->pull->bXOutAverage || inputrec->pull->bFOutAverage)
                {
                    initPullHistory(pull_work, &observablesHistory);
                }
                if (EI_DYNAMICS(inputrec->eI) && cr->commMySim.isMainRank())
                {
                    init_pull_output_files(
                            pull_work, filenames.size(), filenames.data(), oenv, startingBehavior);
                }
            }

            std::unique_ptr<EnforcedRotation> enforcedRotation;
            if (inputrec->bRot)
            {
                /* Initialize enforced rotation code */
                enforcedRotation = init_rot(fplog,
                                            inputrec.get(),
                                            filenames.size(),
                                            filenames.data(),
                                            cr->commMyGroup,
                                            &atomSets,
                                            globalState.get(),
                                            mtop,
                                            oenv,
                                            mdrunOptions,
                                            startingBehavior);
            }

            std::unique_ptr<SwapCoords> swap;
            if (inputrec->eSwapCoords != SwapType::No)
            {
                /* Initialize ion swapping code */
                swap = init_swapcoords(fplog,
                                       inputrec.get(),
                                       opt2fn_main("-swap", filenames.size(), filenames.data(), cr),
                                       mtop,
                                       globalState.get(),
                                       &observablesHistory,
                                       cr->commMyGroup,
                                       cr->dd,
                                       &atomSets,
                                       oenv,
                                       mdrunOptions,
                                       startingBehavior);
            }

            /* Let makeConstraints know whether we have essential dynamics constraints. */
            auto constr = makeConstraints(
                    mtop,
                    *inputrec,
                    pull_work,
                    pull_work != nullptr ? pull_have_constraint(*pull_work) : false,
                    doEssentialDynamics,
                    fplog,
                    cr->commMyGroup,
                    cr->dd,
                    updateGroups.useUpdateGroups(),
                    ms,
                    &nrnb,
                    wcycle.get(),
                    fr->bMolPBC,
                    (cr->dd && cr->dd->mpiComm().isParallel()) ? &observablesReducerBuilder : nullptr);

            /* Energy terms and groups */
            gmx_enerdata_t enerd(mtop.groups.groups[SimulationAtomGroupType::EnergyOutput].size(),
                                 &inputrec->fepvals->all_lambda);

            // cos acceleration is only supported by md, but older tpr
            // files might still combine it with other integrators
            GMX_RELEASE_ASSERT(inputrec->cos_accel == 0.0 || inputrec->eI == IntegrationAlgorithm::MD,
                               "cos_acceleration is only supported by integrator=md");

            /* Kinetic energy data */
            gmx_ekindata_t ekind(gmx::constArrayRefFromArray(inputrec->opts.ref_t, inputrec->opts.ngtc),
                                 inputrec->ensembleTemperatureSetting,
                                 inputrec->ensembleTemperature,
                                 fr->haveBoxDeformation,
                                 inputrec->cos_accel,
                                 gmx_omp_nthreads_get(ModuleMultiThread::Update));

            /* Set up interactive MD (IMD) */
            auto imdSession = makeImdSession(
                    inputrec.get(),
                    cr->commMyGroup,
                    cr->dd,
                    wcycle.get(),
                    &enerd,
                    ms,
                    mtop,
                    mdlog,
                    cr->commMySim.isMainRank() ? globalState->x : gmx::ArrayRef<gmx::RVec>{},
                    filenames.size(),
                    filenames.data(),
                    oenv,
                    mdrunOptions.imdOptions,
                    startingBehavior);

            if (haveDDAtomOrdering(*cr))
            {
                GMX_RELEASE_ASSERT(fr, "fr was NULL while cr->duty was DUTY_PP");
                /* This call is not included in init_domain_decomposition
                 * because fr->atomInfoForEachMoleculeBlock is set later.
                 */
                makeBondedLinks(cr->dd, mtop, fr->atomInfoForEachMoleculeBlock);
            }

            if (runScheduleWork.simulationWork.useGpuFBufferOpsWhenAllowed)
            {
                fr->gpuForceReduction[gmx::AtomLocality::Local] = std::make_unique<gmx::GpuForceReduction>(
                        deviceStreamManager->context(),
                        deviceStreamManager->stream(gmx::DeviceStreamType::NonBondedLocal),
                        wcycle.get());

                if (runScheduleWork.simulationWork.havePpDomainDecomposition)
                {
                    fr->gpuForceReduction[gmx::AtomLocality::NonLocal] =
                            std::make_unique<gmx::GpuForceReduction>(
                                    deviceStreamManager->context(),
                                    deviceStreamManager->stream(gmx::DeviceStreamType::NonBondedNonLocal),
                                    wcycle.get());
                }

                if (runScheduleWork.simulationWork.useMdGpuGraph)
                {
                    fr->mdGraph[MdGraphEvenOrOddStep::EvenStep] =
                            std::make_unique<gmx::MdGpuGraph>(*fr->deviceStreamManager,
                                                              runScheduleWork.simulationWork,
                                                              cr->commMyGroup.comm(),
                                                              MdGraphEvenOrOddStep::EvenStep,
                                                              wcycle.get());

                    fr->mdGraph[MdGraphEvenOrOddStep::OddStep] =
                            std::make_unique<gmx::MdGpuGraph>(*fr->deviceStreamManager,
                                                              runScheduleWork.simulationWork,
                                                              cr->commMyGroup.comm(),
                                                              MdGraphEvenOrOddStep::OddStep,
                                                              wcycle.get());

                    fr->mdGraph[MdGraphEvenOrOddStep::EvenStep]->setAlternateStepPpTaskCompletionEvent(
                            fr->mdGraph[MdGraphEvenOrOddStep::OddStep]->getPpTaskCompletionEvent());
                    fr->mdGraph[MdGraphEvenOrOddStep::OddStep]->setAlternateStepPpTaskCompletionEvent(
                            fr->mdGraph[MdGraphEvenOrOddStep::EvenStep]->getPpTaskCompletionEvent());
                }
            }

            std::unique_ptr<gmx::StatePropagatorDataGpu> stateGpu;
            if (gpusWereDetected && gmx::needStateGpu(runScheduleWork.simulationWork))
            {
                GpuApiCallBehavior transferKind =
                        (inputrec->eI == IntegrationAlgorithm::MD && !doRerun && !useModularSimulator)
                                ? GpuApiCallBehavior::Async
                                : GpuApiCallBehavior::Sync;
                GMX_RELEASE_ASSERT(deviceStreamManager != nullptr,
                                   "GPU device stream manager should be initialized to use GPU.");
                stateGpu = std::make_unique<gmx::StatePropagatorDataGpu>(
                        *deviceStreamManager,
                        transferKind,
                        pme_gpu_get_block_size(fr->pmedata),
                        runScheduleWork.simulationWork.useNvshmem,
                        runScheduleWork.simulationWork.useGpuFBufferOpsWhenAllowed,
                        wcycle.get());
                fr->stateGpu = stateGpu.get();
            }

            GMX_ASSERT(stopHandlerBuilder_, "Runner must provide StopHandlerBuilder to simulator.");
            SimulatorBuilder simulatorBuilder;

            simulatorBuilder.add(SimulatorStateData(
                    globalState.get(), localState, &observablesHistory, &enerd, &ekind));
            simulatorBuilder.add(std::move(membedHolder));
            simulatorBuilder.add(std::move(stopHandlerBuilder_));
            simulatorBuilder.add(SimulatorConfig(mdrunOptions, startingBehavior, &runScheduleWork));


            simulatorBuilder.add(SimulatorEnv(fplog, cr, ms, mdlog, oenv, &observablesReducerBuilder));
            simulatorBuilder.add(Profiling(&nrnb, walltime_accounting, wcycle.get()));
            simulatorBuilder.add(ConstraintsParam(
                    constr.get(),
                    enforcedRotation ? enforcedRotation->getLegacyEnfrot() : nullptr,
                    vsite.get()));
            // TODO: Separate `fr` to a separate add, and make the `build` handle the coupling sensibly.
            simulatorBuilder.add(LegacyInput(
                    static_cast<int>(filenames.size()), filenames.data(), inputrec.get(), fr.get()));
            simulatorBuilder.add(ReplicaExchangeParameters(replExParams));
            simulatorBuilder.add(InteractiveMD(imdSession.get()));
            simulatorBuilder.add(SimulatorModules(mdModules_->outputProvider(), mdModules_->notifiers()));
            simulatorBuilder.add(CenterOfMassPulling(pull_work));
            // Todo move to an MDModule
            simulatorBuilder.add(IonSwapping(swap.get()));
            simulatorBuilder.add(TopologyData(mtop, &localTopology, mdAtoms.get()));
            simulatorBuilder.add(BoxDeformationHandle(deform.get()));
            simulatorBuilder.add(std::move(modularSimulatorCheckpointData));

            // build and run simulator object based on user-input
            auto simulator = simulatorBuilder.build(useModularSimulator);
            simulator->run();

            if (fr->pmePpCommGpu)
            {
                // destroy object since it is no longer required. (This needs to be done while the GPU context still exists.)
                fr->pmePpCommGpu.reset();
            }

            if (inputrec->bPull)
            {
                finish_pull(pull_work);
            }
        }
        else
        {
            GMX_RELEASE_ASSERT(pmedata, "pmedata was NULL while cr->duty was not DUTY_PP");
            /* do PME only */
            walltime_accounting = walltime_accounting_init(gmx_omp_nthreads_get(ModuleMultiThread::Pme));
            gmx_pmeonly(&pmedata,
                        *cr->dd,
                        &nrnb,
                        wcycle.get(),
                        walltime_accounting,
                        inputrec.get(),
                        pmeRunMode,
                        runScheduleWork.simulationWork.useGpuPmePpCommunication,
                        runScheduleWork.simulationWork.useNvshmem,
                        runScheduleWork.simulationWork.useGpuHaloExchange,
                        deviceStreamManager.get());
        }

        if (!hwinfo_->deviceInfoList.empty())
        {
            /* stop the GPU profiler (only CUDA);
             * Doing so here avoids including a lot of cleanup/freeing API calls in the trace. */
            stopGpuProfiler();
        }

        wallcycle_stop(wcycle.get(), WallCycleCounter::Run);

        /* Finish up, write some stuff
         * if rerunMD, don't write last frame again
         */
        finish_run(fplog,
                   mdlog,
                   cr,
                   *inputrec,
                   &nrnb,
                   wcycle.get(),
                   walltime_accounting,
                   fr ? fr->nbv.get() : nullptr,
                   pmedata,
                   mtop.natoms,
                   EI_DYNAMICS(inputrec->eI) && !isMultiSim(ms));
    }
    GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR

    /* Does what it says */
    print_date_and_time(fplog, cr->commMyGroup.rank(), "Finished mdrun", gmx_gettime());

    try
    {
        // Free PME data
        if (pmedata)
        {
            gmx_pme_destroy(pmedata);
            pmedata = nullptr;
        }

        // FIXME: this is only here to manually unpin mdAtoms->chargeA_ and state->x,
        // before we destroy the GPU context(s)
        // Pinned buffers are associated with contexts in CUDA.
        // As soon as we destroy GPU contexts after mdrunner() exits, these lines should go.
        domdec.reset(nullptr);
        // Destroy commRec, as it has references to domdec
        commRec.reset(nullptr);
        cr = nullptr;
        mdAtoms.reset(nullptr);
        globalState.reset(nullptr);
        localStateInstance.reset(nullptr);
        mdModules_.reset(nullptr); // destruct force providers here as they might also use the GPU
        fr.reset(nullptr);         // destruct forcerec before gpu
        // TODO convert to C++ so we can get rid of these frees
        sfree(disresdata);

        nvshmemManager.reset(nullptr);
        // Destroy streams after all the structures using them
        deviceStreamManager.reset(nullptr);

        /* With tMPI we need to wait for all ranks to finish deallocation before
         * destroying the CUDA context as some tMPI ranks may be sharing
         * GPU and context.
         *
         * This is not a concern in OpenCL where we use one context per rank.
         *
         * Note: it is safe to not call the barrier on the ranks which do not use GPU,
         * but it is easier and more futureproof to call it on the whole node.
         *
         * Note that this function needs to be called even if GPUs are not used
         * in this run because the PME ranks have no knowledge of whether GPUs
         * are used or not, but all ranks need to enter the barrier below.
         * \todo Remove this physical node barrier after making sure
         * that it's not needed anymore (with a shared GPU run).
         */
        if (GMX_THREAD_MPI)
        {
            physicalNodeComm.barrier();
        }

        // Only release GPU when not using torch, as torch keeps some device handles that
        // only get cleaned up on process exit
        if (GMX_GPU && !GMX_TORCH)
        {
            const bool haveDetectedOrForcedCudaAwareMpi =
                    (gmx::checkMpiCudaAwareSupport() == gmx::GpuAwareMpiStatus::Supported
                     || gmx::checkMpiCudaAwareSupport() == gmx::GpuAwareMpiStatus::Forced);
            const bool haveDetectedOrForcedHipAwareMpi =
                    (gmx::checkMpiHipAwareSupport() == gmx::GpuAwareMpiStatus::Supported
                     || gmx::checkMpiHipAwareSupport() == gmx::GpuAwareMpiStatus::Forced);

            if (!haveDetectedOrForcedCudaAwareMpi && !haveDetectedOrForcedHipAwareMpi)
            {
                // Don't reset GPU in case of GPU-AWARE MPI
                // UCX creates GPU buffers which are cleaned-up as part of MPI_Finalize()
                // resetting the device before MPI_Finalize() results in crashes inside UCX
                // This can also cause issues in tests that invoke mdrunner() multiple
                // times in the same process; ref #3952.
                releaseDevice();
            }
        }
    }
    GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR

    walltime_accounting_destroy(walltime_accounting);

    // Ensure log file content is written
    if (logFileHandle)
    {
        gmx_fio_flush(logFileHandle);
    }

    /* Reset FPEs (important for unit tests) by disabling them. Assumes no
     * exceptions were enabled before function was called. */
    if (bEnableFPE)
    {
        gmx_fedisableexcept();
    }

    auto rc = static_cast<int>(gmx_get_stop_condition());

#if GMX_THREAD_MPI
    /* we need to join all threads. The sub-threads join when they
       exit this function, but the main thread needs to be told to
       wait for that. */
    if (mpiCommSimulation.isMainRank())
    {
        tMPI_Finalize();
    }
#endif

    return rc;
} // Mdrunner::mdrunner

Mdrunner::~Mdrunner()
{
    // Clean up of the Manager.
    // This will end up getting called on every thread-MPI rank, which is unnecessary,
    // but okay as long as threads synchronize some time before adding or accessing
    // a new set of restraints.
    if (restraintManager_)
    {
        restraintManager_->clear();
        GMX_ASSERT(restraintManager_->countRestraints() == 0,
                   "restraints added during runner life time should be cleared at runner "
                   "destruction.");
    }
}

void Mdrunner::addPotential(std::shared_ptr<gmx::IRestraintPotential> puller, const std::string& name)
{
    GMX_ASSERT(restraintManager_, "Mdrunner must have a restraint manager.");
    // Not sure if this should be logged through the md logger or something else,
    // but it is helpful to have some sort of INFO level message sent somewhere.
    //    std::cout << "Registering restraint named " << name << std::endl;

    // When multiple restraints are used, it may be wasteful to register them separately.
    // Maybe instead register an entire Restraint Manager as a force provider.
    restraintManager_->addToSpec(std::move(puller), name);
}

Mdrunner::Mdrunner(std::unique_ptr<MDModules> mdModules) : mdModules_(std::move(mdModules)) {}

Mdrunner::Mdrunner(Mdrunner&&) noexcept = default;

Mdrunner& Mdrunner::operator=(Mdrunner&& /*handle*/) noexcept = default;

class Mdrunner::BuilderImplementation
{
public:
    BuilderImplementation() = delete;
    BuilderImplementation(std::unique_ptr<MDModules> mdModules, compat::not_null<SimulationContext*> context);
    ~BuilderImplementation();

    BuilderImplementation& setExtraMdrunOptions(const MdrunOptions& options,
                                                real                forceWarningThreshold,
                                                StartingBehavior    startingBehavior);

    void addHardwareDetectionResult(const gmx_hw_info_t* hwinfo);

    void addDomdec(const DomdecOptions& options);

    void addInput(SimulationInputHandle inputHolder);

    void addVerletList(int nstlist);

    void addReplicaExchange(const ReplicaExchangeParameters& params);

    void addNonBonded(const char* nbpu_opt);

    void addNonBondedFETaskAssignment(const char* nbfe_opt);

    void addPME(const char* pme_opt_, const char* pme_fft_opt_);

    void addBondedTaskAssignment(const char* bonded_opt);

    void addUpdateTaskAssignment(const char* update_opt);

    void addHardwareOptions(const gmx_hw_opt_t& hardwareOptions);

    void addFilenames(ArrayRef<const t_filenm> filenames);

    void addOutputEnvironment(gmx_output_env_t* outputEnvironment);

    void addLogFile(t_fileio* logFileHandle);

    void addStopHandlerBuilder(std::unique_ptr<StopHandlerBuilder> builder);

    Mdrunner build();

private:
    // Default parameters copied from runner.h
    // \todo Clarify source(s) of default parameters.

    const char* nbpu_opt_    = nullptr;
    const char* nbfe_opt_    = nullptr;
    const char* pme_opt_     = nullptr;
    const char* pme_fft_opt_ = nullptr;
    const char* bonded_opt_  = nullptr;
    const char* update_opt_  = nullptr;

    MdrunOptions mdrunOptions_;

    DomdecOptions domdecOptions_;

    ReplicaExchangeParameters replicaExchangeParameters_;

    //! Command-line override for the duration of a neighbor list with the Verlet scheme.
    int nstlist_ = 0;

    //! World communicator, used for hardware detection and task assignment
    MPI_Comm libraryWorldCommunicator_ = MPI_COMM_NULL;

    //! Multisim communicator handle.
    gmx_multisim_t* multiSimulation_;

    //! mdrun communicator
    MPI_Comm simulationCommunicator_ = MPI_COMM_NULL;

    //! Print a warning if any force is larger than this (in kJ/mol nm).
    real forceWarningThreshold_ = -1;

    //! Whether the simulation will start afresh, or restart with/without appending.
    StartingBehavior startingBehavior_ = StartingBehavior::NewSimulation;

    //! The modules that comprise the functionality of mdrun.
    std::unique_ptr<MDModules> mdModules_;

    //! Detected hardware.
    const gmx_hw_info_t* hwinfo_ = nullptr;

    //! \brief Parallelism information.
    gmx_hw_opt_t hardwareOptions_;

    //! filename options for simulation.
    ArrayRef<const t_filenm> filenames_;

    /*! \brief Handle to output environment.
     *
     * \todo gmx_output_env_t needs lifetime management.
     */
    gmx_output_env_t* outputEnvironment_ = nullptr;

    /*! \brief Non-owning handle to MD log file.
     *
     * \todo Context should own output facilities for client.
     * \todo Improve log file handle management.
     * \internal
     * Code managing the FILE* relies on the ability to set it to
     * nullptr to check whether the filehandle is valid.
     */
    t_fileio* logFileHandle_ = nullptr;

    /*!
     * \brief Builder for simulation stop signal handler.
     */
    std::unique_ptr<StopHandlerBuilder> stopHandlerBuilder_ = nullptr;

    /*!
     * \brief Sources for initial simulation state.
     *
     * See issue #3652 for near-term refinements to the SimulationInput interface.
     *
     * See issue #3379 for broader discussion on API aspects of simulation inputs and outputs.
     */
    SimulationInputHandle inputHolder_;
};

Mdrunner::BuilderImplementation::BuilderImplementation(std::unique_ptr<MDModules> mdModules,
                                                       compat::not_null<SimulationContext*> context) :
    mdModules_(std::move(mdModules))
{
    libraryWorldCommunicator_ = context->libraryWorldCommunicator_;
    simulationCommunicator_   = context->simulationCommunicator_;
    multiSimulation_          = context->multiSimulation_.get();
}

Mdrunner::BuilderImplementation::~BuilderImplementation() = default;

Mdrunner::BuilderImplementation&
Mdrunner::BuilderImplementation::setExtraMdrunOptions(const MdrunOptions&    options,
                                                      const real             forceWarningThreshold,
                                                      const StartingBehavior startingBehavior)
{
    mdrunOptions_          = options;
    forceWarningThreshold_ = forceWarningThreshold;
    startingBehavior_      = startingBehavior;
    return *this;
}

void Mdrunner::BuilderImplementation::addDomdec(const DomdecOptions& options)
{
    domdecOptions_ = options;
}

void Mdrunner::BuilderImplementation::addVerletList(int nstlist)
{
    nstlist_ = nstlist;
}

void Mdrunner::BuilderImplementation::addReplicaExchange(const ReplicaExchangeParameters& params)
{
    replicaExchangeParameters_ = params;
}

Mdrunner Mdrunner::BuilderImplementation::build()
{
    auto newRunner = Mdrunner(std::move(mdModules_));

    newRunner.mdrunOptions     = mdrunOptions_;
    newRunner.pforce           = forceWarningThreshold_;
    newRunner.startingBehavior = startingBehavior_;
    newRunner.domdecOptions    = domdecOptions_;

    // \todo determine an invariant to check or confirm that all gmx_hw_opt_t objects are valid
    newRunner.hw_opt = hardwareOptions_;

    // No invariant to check. This parameter exists to optionally override other behavior.
    newRunner.nstlist_cmdline = nstlist_;

    newRunner.replExParams = replicaExchangeParameters_;

    newRunner.filenames = filenames_;

    newRunner.libraryWorldCommunicator = libraryWorldCommunicator_;

    newRunner.simulationCommunicator = simulationCommunicator_;

    // nullptr is a valid value for the multisim handle
    newRunner.ms = multiSimulation_;

    if (hwinfo_)
    {
        newRunner.hwinfo_ = hwinfo_;
    }
    else
    {
        GMX_THROW(gmx::APIError(
                "MdrunnerBuilder::addHardwareDetectionResult() is required before build()"));
    }

    if (inputHolder_)
    {
        newRunner.inputHolder_ = std::move(inputHolder_);
    }
    else
    {
        GMX_THROW(gmx::APIError("MdrunnerBuilder::addInput() is required before build()."));
    }

    // \todo Clarify ownership and lifetime management for gmx_output_env_t
    // \todo Update sanity checking when output environment has clearly specified invariants.
    // Initialization and default values for oenv are not well specified in the current version.
    if (outputEnvironment_)
    {
        newRunner.oenv = outputEnvironment_;
    }
    else
    {
        GMX_THROW(gmx::APIError(
                "MdrunnerBuilder::addOutputEnvironment() is required before build()"));
    }

    newRunner.logFileHandle = logFileHandle_;

    if (nbpu_opt_)
    {
        newRunner.nbpu_opt = nbpu_opt_;
    }
    else
    {
        GMX_THROW(gmx::APIError("MdrunnerBuilder::addNonBonded() is required before build()"));
    }

    if (nbfe_opt_)
    {
        newRunner.nbfe_opt = nbfe_opt_;
    }
    else
    {
        GMX_THROW(
                gmx::APIError("MdrunnerBuilder::    void addNonBondedFETaskAssignment(); is "
                              "required before build()"));
    }

    if (pme_opt_ && pme_fft_opt_)
    {
        newRunner.pme_opt     = pme_opt_;
        newRunner.pme_fft_opt = pme_fft_opt_;
    }
    else
    {
        GMX_THROW(gmx::APIError("MdrunnerBuilder::addElectrostatics() is required before build()"));
    }

    if (bonded_opt_)
    {
        newRunner.bonded_opt = bonded_opt_;
    }
    else
    {
        GMX_THROW(gmx::APIError(
                "MdrunnerBuilder::addBondedTaskAssignment() is required before build()"));
    }

    if (update_opt_)
    {
        newRunner.update_opt = update_opt_;
    }
    else
    {
        GMX_THROW(gmx::APIError(
                "MdrunnerBuilder::addUpdateTaskAssignment() is required before build()  "));
    }

    newRunner.restraintManager_ = std::make_unique<gmx::RestraintManager>();

    if (stopHandlerBuilder_)
    {
        newRunner.stopHandlerBuilder_ = std::move(stopHandlerBuilder_);
    }
    else
    {
        newRunner.stopHandlerBuilder_ = std::make_unique<StopHandlerBuilder>();
    }

    return newRunner;
}

void Mdrunner::BuilderImplementation::addHardwareDetectionResult(const gmx_hw_info_t* hwinfo)
{
    hwinfo_ = hwinfo;
}

void Mdrunner::BuilderImplementation::addNonBonded(const char* nbpu_opt)
{
    nbpu_opt_ = nbpu_opt;
}

void Mdrunner::BuilderImplementation::addNonBondedFETaskAssignment(const char* nbfe_opt)
{
    nbfe_opt_ = nbfe_opt;
}

void Mdrunner::BuilderImplementation::addPME(const char* pme_opt, const char* pme_fft_opt)
{
    pme_opt_     = pme_opt;
    pme_fft_opt_ = pme_fft_opt;
}

void Mdrunner::BuilderImplementation::addBondedTaskAssignment(const char* bonded_opt)
{
    bonded_opt_ = bonded_opt;
}

void Mdrunner::BuilderImplementation::addUpdateTaskAssignment(const char* update_opt)
{
    update_opt_ = update_opt;
}

void Mdrunner::BuilderImplementation::addHardwareOptions(const gmx_hw_opt_t& hardwareOptions)
{
    hardwareOptions_ = hardwareOptions;
}

void Mdrunner::BuilderImplementation::addFilenames(ArrayRef<const t_filenm> filenames)
{
    filenames_ = filenames;
}

void Mdrunner::BuilderImplementation::addOutputEnvironment(gmx_output_env_t* outputEnvironment)
{
    outputEnvironment_ = outputEnvironment;
}

void Mdrunner::BuilderImplementation::addLogFile(t_fileio* logFileHandle)
{
    logFileHandle_ = logFileHandle;
}

void Mdrunner::BuilderImplementation::addStopHandlerBuilder(std::unique_ptr<StopHandlerBuilder> builder)
{
    stopHandlerBuilder_ = std::move(builder);
}

void Mdrunner::BuilderImplementation::addInput(SimulationInputHandle inputHolder)
{
    inputHolder_ = std::move(inputHolder);
}

MdrunnerBuilder::MdrunnerBuilder(std::unique_ptr<MDModules>           mdModules,
                                 compat::not_null<SimulationContext*> context) :
    impl_{ std::make_unique<Mdrunner::BuilderImplementation>(std::move(mdModules), context) }
{
}

MdrunnerBuilder::~MdrunnerBuilder() = default;

MdrunnerBuilder& MdrunnerBuilder::addHardwareDetectionResult(const gmx_hw_info_t* hwinfo)
{
    impl_->addHardwareDetectionResult(hwinfo);
    return *this;
}

MdrunnerBuilder& MdrunnerBuilder::addSimulationMethod(const MdrunOptions&    options,
                                                      real                   forceWarningThreshold,
                                                      const StartingBehavior startingBehavior)
{
    impl_->setExtraMdrunOptions(options, forceWarningThreshold, startingBehavior);
    return *this;
}

MdrunnerBuilder& MdrunnerBuilder::addDomainDecomposition(const DomdecOptions& options)
{
    impl_->addDomdec(options);
    return *this;
}

MdrunnerBuilder& MdrunnerBuilder::addNeighborList(int nstlist)
{
    impl_->addVerletList(nstlist);
    return *this;
}

MdrunnerBuilder& MdrunnerBuilder::addReplicaExchange(const ReplicaExchangeParameters& params)
{
    impl_->addReplicaExchange(params);
    return *this;
}

MdrunnerBuilder& MdrunnerBuilder::addNonBonded(const char* nbpu_opt)
{
    impl_->addNonBonded(nbpu_opt);
    return *this;
}

MdrunnerBuilder& MdrunnerBuilder::addNonBondedFETaskAssignment(const char* nbfe_opt)
{
    impl_->addNonBondedFETaskAssignment(nbfe_opt);
    return *this;
}

MdrunnerBuilder& MdrunnerBuilder::addElectrostatics(const char* pme_opt, const char* pme_fft_opt)
{
    // The builder method may become more general in the future, but in this version,
    // parameters for PME electrostatics are both required and the only parameters
    // available.
    if (pme_opt && pme_fft_opt)
    {
        impl_->addPME(pme_opt, pme_fft_opt);
    }
    else
    {
        GMX_THROW(
                gmx::InvalidInputError("addElectrostatics() arguments must be non-null pointers."));
    }
    return *this;
}

MdrunnerBuilder& MdrunnerBuilder::addBondedTaskAssignment(const char* bonded_opt)
{
    impl_->addBondedTaskAssignment(bonded_opt);
    return *this;
}

MdrunnerBuilder& MdrunnerBuilder::addUpdateTaskAssignment(const char* update_opt)
{
    impl_->addUpdateTaskAssignment(update_opt);
    return *this;
}

Mdrunner MdrunnerBuilder::build()
{
    return impl_->build();
}

MdrunnerBuilder& MdrunnerBuilder::addHardwareOptions(const gmx_hw_opt_t& hardwareOptions)
{
    impl_->addHardwareOptions(hardwareOptions);
    return *this;
}

MdrunnerBuilder& MdrunnerBuilder::addFilenames(ArrayRef<const t_filenm> filenames)
{
    impl_->addFilenames(filenames);
    return *this;
}

MdrunnerBuilder& MdrunnerBuilder::addOutputEnvironment(gmx_output_env_t* outputEnvironment)
{
    impl_->addOutputEnvironment(outputEnvironment);
    return *this;
}

MdrunnerBuilder& MdrunnerBuilder::addLogFile(t_fileio* logFileHandle)
{
    impl_->addLogFile(logFileHandle);
    return *this;
}

MdrunnerBuilder& MdrunnerBuilder::addStopHandlerBuilder(std::unique_ptr<StopHandlerBuilder> builder)
{
    impl_->addStopHandlerBuilder(std::move(builder));
    return *this;
}

MdrunnerBuilder& MdrunnerBuilder::addInput(SimulationInputHandle input)
{
    impl_->addInput(std::move(input));
    return *this;
}

MdrunnerBuilder::MdrunnerBuilder(MdrunnerBuilder&&) noexcept = default;

MdrunnerBuilder& MdrunnerBuilder::operator=(MdrunnerBuilder&&) noexcept = default;

} // namespace gmx
