From 905880aec24da8cdc62d0d8978b91d78307b570b Mon Sep 17 00:00:00 2001 From: Artem Zhmurov Date: Thu, 13 Feb 2020 16:05:51 +0100 Subject: [PATCH] Combine initialization routines for GPUs assigned to PME and Nonbonded The subroutines that do the initialization share a lot of logic. Also, the current setup does not allow to have more than one GPU assigned to a single rank. This make it natural to keep one handle to the GPU device, if there is a GPU assigned to the rank. Change-Id: I9177e1a6c00cf6efa7b58370b8fa6220c8fd4496 --- src/gromacs/mdrun/runner.cpp | 29 +++++++-------- src/gromacs/taskassignment/taskassignment.cpp | 53 +++++++++++++-------------- src/gromacs/taskassignment/taskassignment.h | 25 +++---------- 3 files changed, 44 insertions(+), 63 deletions(-) diff --git a/src/gromacs/mdrun/runner.cpp b/src/gromacs/mdrun/runner.cpp index eddbe5efa7..65312657d8 100644 --- a/src/gromacs/mdrun/runner.cpp +++ b/src/gromacs/mdrun/runner.cpp @@ -1141,18 +1141,16 @@ int Mdrunner::mdrunner() EEL_PME(inputrec->coulombtype) && thisRankHasDuty(cr, DUTY_PME)); // Get the device handles for the modules, nullptr when no task is assigned. - // TODO: There should be only one DeviceInformation per rank. - DeviceInformation* nonbondedDeviceInfo = gpuTaskAssignments.initNonbondedDevice(cr); - DeviceInformation* pmeDeviceInfo = gpuTaskAssignments.initPmeDevice(); - + int deviceId = -1; + DeviceInformation* deviceInfo = gpuTaskAssignments.initDevice(&deviceId); std::unique_ptr deviceContext = nullptr; - if (pmeDeviceInfo) - { - deviceContext = std::make_unique(*pmeDeviceInfo); - } - else if (nonbondedDeviceInfo) + if (deviceInfo != nullptr) { - deviceContext = std::make_unique(*nonbondedDeviceInfo); + if (DOMAINDECOMP(cr) && thisRankHasDuty(cr, DUTY_PP)) + { + dd_setup_dlb_resource_sharing(cr, deviceId); + } + deviceContext = std::make_unique(*deviceInfo); } // TODO Initialize GPU streams here. @@ -1361,7 +1359,7 @@ int Mdrunner::mdrunner() cr->mpi_comm_mysim, cr->dd->pme_nodeid, *deviceContext); } - fr->nbv = Nbnxm::init_nb_verlet(mdlog, inputrec, fr, cr, *hwinfo, nonbondedDeviceInfo, + fr->nbv = Nbnxm::init_nb_verlet(mdlog, inputrec, fr, cr, *hwinfo, deviceInfo, fr->deviceContext, &mtop, box, wcycle); if (useGpuForBonded) { @@ -1452,12 +1450,12 @@ int Mdrunner::mdrunner() if (thisRankHasPmeGpuTask) { GMX_RELEASE_ASSERT( - pmeDeviceInfo != nullptr, + deviceInfo != nullptr, "Device information can not be nullptr when building PME GPU program object."); GMX_RELEASE_ASSERT( deviceContext != nullptr, "Device context can not be nullptr when building PME GPU program object."); - pmeGpuProgram = buildPmeGpuProgram(*pmeDeviceInfo, *deviceContext); + pmeGpuProgram = buildPmeGpuProgram(*deviceInfo, *deviceContext); } /* Initiate PME if necessary, @@ -1486,7 +1484,7 @@ int Mdrunner::mdrunner() pmedata = gmx_pme_init(cr, getNumPmeDomains(cr->dd), inputrec, nChargePerturbed != 0, nTypePerturbed != 0, mdrunOptions.reproducible, ewaldcoeff_q, ewaldcoeff_lj, gmx_omp_nthreads_get(emntPME), pmeRunMode, - nullptr, pmeDeviceInfo, pmeGpuProgram.get(), mdlog); + nullptr, deviceInfo, pmeGpuProgram.get(), mdlog); } GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR } @@ -1699,8 +1697,7 @@ int Mdrunner::mdrunner() physicalNodeComm.barrier(); } - free_gpu(nonbondedDeviceInfo); - free_gpu(pmeDeviceInfo); + free_gpu(deviceInfo); deviceContext.reset(nullptr); sfree(fcd); diff --git a/src/gromacs/taskassignment/taskassignment.cpp b/src/gromacs/taskassignment/taskassignment.cpp index 5a62972d61..8e0ace91cf 100644 --- a/src/gromacs/taskassignment/taskassignment.cpp +++ b/src/gromacs/taskassignment/taskassignment.cpp @@ -400,43 +400,40 @@ void GpuTaskAssignments::reportGpuUsage(const MDLogger& mdlog, numRanksOnThisNode_, printHostName, useGpuForBonded, pmeRunMode, useGpuForUpdate); } -DeviceInformation* GpuTaskAssignments::initNonbondedDevice(const t_commrec* cr) const +/*! \brief Function for whether the task of \c mapping has value \c TaskType. + * + * \param[in] mapping Current GPU task mapping. + * \returns If \c TaskType task was assigned to the \c mapping. + */ +template +static bool hasTaskType(const GpuTaskMapping& mapping) { - DeviceInformation* deviceInfo = nullptr; - const GpuTaskAssignment& gpuTaskAssignment = assignmentForAllRanksOnThisNode_[indexOfThisRank_]; - - // This works because only one task of each type per rank is currently permitted. - auto nbGpuTaskMapping = std::find_if(gpuTaskAssignment.begin(), gpuTaskAssignment.end(), - hasTaskType); - if (nbGpuTaskMapping != gpuTaskAssignment.end()) - { - int deviceId = nbGpuTaskMapping->deviceId_; - deviceInfo = getDeviceInfo(hardwareInfo_.gpu_info, deviceId); - init_gpu(deviceInfo); + return mapping.task_ == TaskType; +} - // TODO Setting up this sharing should probably part of - // init_domain_decomposition after further refactoring. - if (DOMAINDECOMP(cr)) - { - /* When we share GPUs over ranks, we need to know this for the DLB */ - dd_setup_dlb_resource_sharing(cr, deviceId); - } - } - return deviceInfo; +/*! \brief Function for whether the \c mapping has the GPU PME or Nonbonded task. + * + * \param[in] mapping Current GPU task mapping. + * \returns If PME on Nonbonded GPU task was assigned to this mapping. + */ +static bool hasPmeOrNonbondedTask(const GpuTaskMapping& mapping) +{ + return hasTaskType(mapping) || hasTaskType(mapping); } -DeviceInformation* GpuTaskAssignments::initPmeDevice() const +DeviceInformation* GpuTaskAssignments::initDevice(int* deviceId) const { DeviceInformation* deviceInfo = nullptr; const GpuTaskAssignment& gpuTaskAssignment = assignmentForAllRanksOnThisNode_[indexOfThisRank_]; - // This works because only one task of each type is currently permitted. - auto pmeGpuTaskMapping = std::find_if(gpuTaskAssignment.begin(), gpuTaskAssignment.end(), - hasTaskType); - const bool thisRankHasPmeGpuTask = (pmeGpuTaskMapping != gpuTaskAssignment.end()); - if (thisRankHasPmeGpuTask) + // This works because only one task of each type per rank is currently permitted. + auto gpuTaskMapping = + std::find_if(gpuTaskAssignment.begin(), gpuTaskAssignment.end(), hasPmeOrNonbondedTask); + + if (gpuTaskMapping != gpuTaskAssignment.end()) { - deviceInfo = getDeviceInfo(hardwareInfo_.gpu_info, pmeGpuTaskMapping->deviceId_); + *deviceId = gpuTaskMapping->deviceId_; + deviceInfo = getDeviceInfo(hardwareInfo_.gpu_info, *deviceId); init_gpu(deviceInfo); } return deviceInfo; diff --git a/src/gromacs/taskassignment/taskassignment.h b/src/gromacs/taskassignment/taskassignment.h index c6ac9ac9de..7e661e2166 100644 --- a/src/gromacs/taskassignment/taskassignment.h +++ b/src/gromacs/taskassignment/taskassignment.h @@ -239,33 +239,20 @@ public: * \param[in] numCompatibleGpusOnThisNode The number of compatible GPUs on this node. * */ void logPerformanceHints(const MDLogger& mdlog, size_t numCompatibleGpusOnThisNode); - /*! \brief Return handle to the initialized GPU to use for the - * nonbonded task on this rank, if any. + /*! \brief Return handle to the initialized GPU to use in the this rank. * - * Returns nullptr if no such task is assigned to this rank. + * \param[out] deviceId Index of the assigned device. * - * \todo This also sets up DLB for device sharing, where - * appropriate, but that responsbility should move - * elsewhere. */ - DeviceInformation* initNonbondedDevice(const t_commrec* cr) const; - /*! \brief Return handle to the initialized GPU to use for the - * PME task on this rank, if any. - * - * Returns nullptr if no such task is assigned to this rank. */ - DeviceInformation* initPmeDevice() const; + * \returns Device information on the selected devicce. Returns nullptr if no GPU task + * is assigned to this rank. + */ + DeviceInformation* initDevice(int* deviceId) const; //! Return whether this rank has a PME task running on a GPU bool thisRankHasPmeGpuTask() const; //! Return whether this rank has any task running on a GPU bool thisRankHasAnyGpuTask() const; }; -//! Function for whether the task of \c mapping has value \c TaskType. -template -bool hasTaskType(const GpuTaskMapping& mapping) -{ - return mapping.task_ == TaskType; -} - } // namespace gmx #endif -- 2.11.4.GIT