Move free-energy kernel dispatch into nbnxm module

Change-Id: I57a76dedff567ad946e970a06d779f9cf8ce475b

diff --git a/src/gromacs/mdlib/sim_util.cpp b/src/gromacs/mdlib/sim_util.cpp
index b89f57fd5ba1228926147b65673d158e5d9f385d..6f3f2e3150de955cc8cad4a1eff14544cfbf06bb 100644 (file)
--- a/src/gromacs/mdlib/sim_util.cpp
+++ b/src/gromacs/mdlib/sim_util.cpp
@@ -438,120 +438,6 @@ static void do_nb_verlet(t_forcerec                       *fr,
     }
 }
 
-static void do_nb_verlet_fep(const nonbonded_verlet_t         &nbv,
-                             const Nbnxm::InteractionLocality  iLocality,
-                             t_forcerec                       *fr,
-                             rvec                              x[],
-                             rvec                              f[],
-                             const t_mdatoms                  *mdatoms,
-                             t_lambda                         *fepvals,
-                             real                             *lambda,
-                             gmx_enerdata_t                   *enerd,
-                             int                               flags,
-                             t_nrnb                           *nrnb,
-                             gmx_wallcycle_t                   wcycle)
-{
-    int              donb_flags;
-    nb_kernel_data_t kernel_data;
-    real             lam_i[efptNR];
-    real             dvdl_nb[efptNR];
-    int              i, j;
-
-    donb_flags = 0;
-    /* Add short-range interactions */
-    donb_flags |= GMX_NONBONDED_DO_SR;
-
-    /* Currently all group scheme kernels always calculate (shift-)forces */
-    if (flags & GMX_FORCE_FORCES)
-    {
-        donb_flags |= GMX_NONBONDED_DO_FORCE;
-    }
-    if (flags & GMX_FORCE_VIRIAL)
-    {
-        donb_flags |= GMX_NONBONDED_DO_SHIFTFORCE;
-    }
-    if (flags & GMX_FORCE_ENERGY)
-    {
-        donb_flags |= GMX_NONBONDED_DO_POTENTIAL;
-    }
-
-    kernel_data.flags  = donb_flags;
-    kernel_data.lambda = lambda;
-    kernel_data.dvdl   = dvdl_nb;
-
-    kernel_data.energygrp_elec = enerd->grpp.ener[egCOULSR];
-    kernel_data.energygrp_vdw  = enerd->grpp.ener[egLJSR];
-
-    /* reset free energy components */
-    for (i = 0; i < efptNR; i++)
-    {
-        dvdl_nb[i]  = 0;
-    }
-
-    const gmx::ArrayRef<t_nblist const * const > nbl_fep = nbv.freeEnergyPairlistSet(iLocality);
-
-    GMX_ASSERT(gmx_omp_nthreads_get(emntNonbonded) == nbl_fep.ssize(), "Number of lists should be same as number of NB threads");
-
-    wallcycle_sub_start(wcycle, ewcsNONBONDED);
-#pragma omp parallel for schedule(static) num_threads(nbl_fep.ssize())
-    for (int th = 0; th < nbl_fep.ssize(); th++)
-    {
-        try
-        {
-            gmx_nb_free_energy_kernel(nbl_fep[th],
-                                      x, f, fr, mdatoms, &kernel_data, nrnb);
-        }
-        GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR;
-    }
-
-    if (fepvals->sc_alpha != 0)
-    {
-        enerd->dvdl_nonlin[efptVDW]  += dvdl_nb[efptVDW];
-        enerd->dvdl_nonlin[efptCOUL] += dvdl_nb[efptCOUL];
-    }
-    else
-    {
-        enerd->dvdl_lin[efptVDW]  += dvdl_nb[efptVDW];
-        enerd->dvdl_lin[efptCOUL] += dvdl_nb[efptCOUL];
-    }
-
-    /* If we do foreign lambda and we have soft-core interactions
-     * we have to recalculate the (non-linear) energies contributions.
-     */
-    if (fepvals->n_lambda > 0 && (flags & GMX_FORCE_DHDL) && fepvals->sc_alpha != 0)
-    {
-        kernel_data.flags          = (donb_flags & ~(GMX_NONBONDED_DO_FORCE | GMX_NONBONDED_DO_SHIFTFORCE)) | GMX_NONBONDED_DO_FOREIGNLAMBDA;
-        kernel_data.lambda         = lam_i;
-        kernel_data.energygrp_elec = enerd->foreign_grpp.ener[egCOULSR];
-        kernel_data.energygrp_vdw  = enerd->foreign_grpp.ener[egLJSR];
-        /* Note that we add to kernel_data.dvdl, but ignore the result */
-
-        for (i = 0; i < enerd->n_lambda; i++)
-        {
-            for (j = 0; j < efptNR; j++)
-            {
-                lam_i[j] = (i == 0 ? lambda[j] : fepvals->all_lambda[j][i-1]);
-            }
-            reset_foreign_enerdata(enerd);
-#pragma omp parallel for schedule(static) num_threads(nbl_fep.ssize())
-            for (int th = 0; th < nbl_fep.ssize(); th++)
-            {
-                try
-                {
-                    gmx_nb_free_energy_kernel(nbl_fep[th],
-                                              x, f, fr, mdatoms, &kernel_data, nrnb);
-                }
-                GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR;
-            }
-
-            sum_epot(&(enerd->foreign_grpp), enerd->foreign_term);
-            enerd->enerpart_lambda[i] += enerd->foreign_term[F_EPOT];
-        }
-    }
-
-    wallcycle_sub_stop(wcycle, ewcsNONBONDED);
-}
-
 gmx_bool use_GPU(const nonbonded_verlet_t *nbv)
 {
     return nbv != nullptr && nbv->useGpu();
@@ -1364,22 +1250,20 @@ static void do_force_cutsVERLET(FILE *fplog,
         /* Calculate the local and non-local free energy interactions here.
          * Happens here on the CPU both with and without GPU.
          */
-        if (fr->nbv->freeEnergyPairlistSet(Nbnxm::InteractionLocality::Local)[0]->nrj > 0)
-        {
-            do_nb_verlet_fep(*nbv, Nbnxm::InteractionLocality::Local,
-                             fr, as_rvec_array(x.unpaddedArrayRef().data()), f, mdatoms,
-                             inputrec->fepvals, lambda,
-                             enerd, flags, nrnb, wcycle);
-        }
+        wallcycle_sub_start(wcycle, ewcsNONBONDED);
+        nbv->dispatchFreeEnergyKernel(Nbnxm::InteractionLocality::Local,
+                                      fr, as_rvec_array(x.unpaddedArrayRef().data()), f, *mdatoms,
+                                      inputrec->fepvals, lambda,
+                                      enerd, flags, nrnb);
 
-        if (havePPDomainDecomposition(cr) &&
-            fr->nbv->freeEnergyPairlistSet(Nbnxm::InteractionLocality::NonLocal)[0]->nrj > 0)
+        if (havePPDomainDecomposition(cr))
         {
-            do_nb_verlet_fep(*nbv, Nbnxm::InteractionLocality::NonLocal,
-                             fr, as_rvec_array(x.unpaddedArrayRef().data()), f, mdatoms,
-                             inputrec->fepvals, lambda,
-                             enerd, flags, nrnb, wcycle);
+            nbv->dispatchFreeEnergyKernel(Nbnxm::InteractionLocality::NonLocal,
+                                          fr, as_rvec_array(x.unpaddedArrayRef().data()), f, *mdatoms,
+                                          inputrec->fepvals, lambda,
+                                          enerd, flags, nrnb);
         }
+        wallcycle_sub_stop(wcycle, ewcsNONBONDED);
     }
 
     if (!bUseOrEmulGPU)

diff --git a/src/gromacs/nbnxm/kerneldispatch.cpp b/src/gromacs/nbnxm/kerneldispatch.cpp
index 5c302d7c8849e18f8ceeec4aa718addb97d2f748..fcac007c027352930d409970db694e842855c081 100644 (file)
--- a/src/gromacs/nbnxm/kerneldispatch.cpp
+++ b/src/gromacs/nbnxm/kerneldispatch.cpp
@@ -36,12 +36,18 @@
 #include "gmxpre.h"
 
 #include "gromacs/gmxlib/nrnb.h"
+#include "gromacs/gmxlib/nonbonded/nb_free_energy.h"
+#include "gromacs/gmxlib/nonbonded/nb_kernel.h"
+#include "gromacs/gmxlib/nonbonded/nonbonded.h"
 #include "gromacs/math/vectypes.h"
+#include "gromacs/mdlib/force.h"
 #include "gromacs/mdlib/force_flags.h"
 #include "gromacs/mdlib/gmx_omp_nthreads.h"
 #include "gromacs/mdtypes/enerdata.h"
+#include "gromacs/mdtypes/inputrec.h"
 #include "gromacs/mdtypes/interaction_const.h"
 #include "gromacs/mdtypes/md_enums.h"
+#include "gromacs/mdtypes/mdatom.h"
 #include "gromacs/nbnxm/gpu_data_mgmt.h"
 #include "gromacs/nbnxm/nbnxm.h"
 #include "gromacs/nbnxm/nbnxm_simd.h"
@@ -530,3 +536,110 @@ void NbnxnDispatchKernel(nonbonded_verlet_t        *nbv,
 
     accountFlops(nrnb, *nbv, iLocality, ic, forceFlags);
 }
+
+void
+nonbonded_verlet_t::dispatchFreeEnergyKernel(Nbnxm::InteractionLocality  iLocality,
+                                             t_forcerec                 *fr,
+                                             rvec                        x[],
+                                             rvec                        f[],
+                                             const t_mdatoms            &mdatoms,
+                                             t_lambda                   *fepvals,
+                                             real                       *lambda,
+                                             gmx_enerdata_t             *enerd,
+                                             const int                   forceFlags,
+                                             t_nrnb                     *nrnb)
+{
+    const gmx::ArrayRef<t_nblist const * const > nbl_fep = pairlistSet(iLocality).nbl_fep;
+
+    /* When the first list is empty, all are empty and there is nothing to do */
+    if (nbl_fep[0]->nrj == 0)
+    {
+        return;
+    }
+
+    int donb_flags = 0;
+    /* Add short-range interactions */
+    donb_flags |= GMX_NONBONDED_DO_SR;
+
+    /* Currently all group scheme kernels always calculate (shift-)forces */
+    if (forceFlags & GMX_FORCE_FORCES)
+    {
+        donb_flags |= GMX_NONBONDED_DO_FORCE;
+    }
+    if (forceFlags & GMX_FORCE_VIRIAL)
+    {
+        donb_flags |= GMX_NONBONDED_DO_SHIFTFORCE;
+    }
+    if (forceFlags & GMX_FORCE_ENERGY)
+    {
+        donb_flags |= GMX_NONBONDED_DO_POTENTIAL;
+    }
+
+    nb_kernel_data_t kernel_data;
+    real             dvdl_nb[efptNR] = { 0 };
+    kernel_data.flags  = donb_flags;
+    kernel_data.lambda = lambda;
+    kernel_data.dvdl   = dvdl_nb;
+
+    kernel_data.energygrp_elec = enerd->grpp.ener[egCOULSR];
+    kernel_data.energygrp_vdw  = enerd->grpp.ener[egLJSR];
+
+    GMX_ASSERT(gmx_omp_nthreads_get(emntNonbonded) == nbl_fep.ssize(), "Number of lists should be same as number of NB threads");
+
+#pragma omp parallel for schedule(static) num_threads(nbl_fep.ssize())
+    for (int th = 0; th < nbl_fep.ssize(); th++)
+    {
+        try
+        {
+            gmx_nb_free_energy_kernel(nbl_fep[th],
+                                      x, f, fr, &mdatoms, &kernel_data, nrnb);
+        }
+        GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR;
+    }
+
+    if (fepvals->sc_alpha != 0)
+    {
+        enerd->dvdl_nonlin[efptVDW]  += dvdl_nb[efptVDW];
+        enerd->dvdl_nonlin[efptCOUL] += dvdl_nb[efptCOUL];
+    }
+    else
+    {
+        enerd->dvdl_lin[efptVDW]  += dvdl_nb[efptVDW];
+        enerd->dvdl_lin[efptCOUL] += dvdl_nb[efptCOUL];
+    }
+
+    /* If we do foreign lambda and we have soft-core interactions
+     * we have to recalculate the (non-linear) energies contributions.
+     */
+    if (fepvals->n_lambda > 0 && (forceFlags & GMX_FORCE_DHDL) && fepvals->sc_alpha != 0)
+    {
+        real lam_i[efptNR];
+        kernel_data.flags          = (donb_flags & ~(GMX_NONBONDED_DO_FORCE | GMX_NONBONDED_DO_SHIFTFORCE)) | GMX_NONBONDED_DO_FOREIGNLAMBDA;
+        kernel_data.lambda         = lam_i;
+        kernel_data.energygrp_elec = enerd->foreign_grpp.ener[egCOULSR];
+        kernel_data.energygrp_vdw  = enerd->foreign_grpp.ener[egLJSR];
+        /* Note that we add to kernel_data.dvdl, but ignore the result */
+
+        for (int i = 0; i < enerd->n_lambda; i++)
+        {
+            for (int j = 0; j < efptNR; j++)
+            {
+                lam_i[j] = (i == 0 ? lambda[j] : fepvals->all_lambda[j][i-1]);
+            }
+            reset_foreign_enerdata(enerd);
+#pragma omp parallel for schedule(static) num_threads(nbl_fep.ssize())
+            for (int th = 0; th < nbl_fep.ssize(); th++)
+            {
+                try
+                {
+                    gmx_nb_free_energy_kernel(nbl_fep[th],
+                                              x, f, fr, &mdatoms, &kernel_data, nrnb);
+                }
+                GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR;
+            }
+
+            sum_epot(&(enerd->foreign_grpp), enerd->foreign_term);
+            enerd->enerpart_lambda[i] += enerd->foreign_term[F_EPOT];
+        }
+    }
+}

diff --git a/src/gromacs/nbnxm/nbnxm.h b/src/gromacs/nbnxm/nbnxm.h
index e0103c849d0b0ea7d386ab9a298f3114fa7619a0..1b020d5916d256fc35034849fb51f69cb6a1940b 100644 (file)
--- a/src/gromacs/nbnxm/nbnxm.h
+++ b/src/gromacs/nbnxm/nbnxm.h
@@ -122,6 +122,8 @@ struct nbnxn_search;
 struct nonbonded_verlet_t;
 struct t_blocka;
 struct t_commrec;
+struct t_lambda;
+struct t_mdatoms;
 struct t_nrnb;
 struct t_forcerec;
 struct t_inputrec;
@@ -262,6 +264,18 @@ struct nonbonded_verlet_t
                                      kernelSetup_.kernelType, nbat, shift_vec);
         }
 
+        //! Dispatches the non-bonded free-energy kernel, always runs on the CPU
+        void dispatchFreeEnergyKernel(Nbnxm::InteractionLocality  iLocality,
+                                      t_forcerec                 *fr,
+                                      rvec                        x[],
+                                      rvec                        f[],
+                                      const t_mdatoms            &mdatoms,
+                                      t_lambda                   *fepvals,
+                                      real                       *lambda,
+                                      gmx_enerdata_t             *enerd,
+                                      int                         forceFlags,
+                                      t_nrnb                     *nrnb);
+
         //! Return the kernel setup
         const Nbnxm::KernelSetup &kernelSetup() const
         {
@@ -274,13 +288,6 @@ struct nonbonded_verlet_t
             kernelSetup_ = kernelSetup;
         }
 
-        //! Returns the a list of free-energy pairlists for the given locality
-        const gmx::ArrayRef<t_nblist const * const>
-        freeEnergyPairlistSet(Nbnxm::InteractionLocality iLocality) const
-        {
-            return pairlistSet(iLocality).nbl_fep;
-        }
-
         //! Parameters for the search and list pruning setup
         std::unique_ptr<NbnxnListParameters>  listParams;
         //! Working data for constructing the pairlists