src/gromacs/nbnxm/kerneldispatch.cpp

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2012,2013,2014,2015,2016 by the GROMACS development team.
   5  * Copyright (c) 2017,2018,2019,2020,2021, by the GROMACS development team, led by
   6  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   7  * and including many others, as listed in the AUTHORS file in the
   8  * top-level source directory and at http://www.gromacs.org.
   9  *
  10  * GROMACS is free software; you can redistribute it and/or
  11  * modify it under the terms of the GNU Lesser General Public License
  12  * as published by the Free Software Foundation; either version 2.1
  13  * of the License, or (at your option) any later version.
  14  *
  15  * GROMACS is distributed in the hope that it will be useful,
  16  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  17  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  18  * Lesser General Public License for more details.
  19  *
  20  * You should have received a copy of the GNU Lesser General Public
  21  * License along with GROMACS; if not, see
  22  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
  23  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  24  *
  25  * If you want to redistribute modifications to GROMACS, please
  26  * consider that scientific software is very special. Version
  27  * control is crucial - bugs must be traceable. We will be happy to
  28  * consider code for inclusion in the official distribution, but
  29  * derived work must not be called official GROMACS. Details are found
  30  * in the README & COPYING files - if they are missing, get the
  31  * official version at http://www.gromacs.org.
  32  *
  33  * To help us fund GROMACS development, we humbly ask that you cite
  34  * the research papers on the package. Check out http://www.gromacs.org.
  35  */
  36
  37 #include "gmxpre.h"
  38
  39 #include "gromacs/gmxlib/nrnb.h"
  40 #include "gromacs/gmxlib/nonbonded/nb_free_energy.h"
  41 #include "gromacs/gmxlib/nonbonded/nonbonded.h"
  42 #include "gromacs/math/vectypes.h"
  43 #include "gromacs/mdlib/enerdata_utils.h"
  44 #include "gromacs/mdlib/force.h"
  45 #include "gromacs/mdlib/gmx_omp_nthreads.h"
  46 #include "gromacs/mdtypes/enerdata.h"
  47 #include "gromacs/mdtypes/forceoutput.h"
  48 #include "gromacs/mdtypes/forcerec.h"
  49 #include "gromacs/mdtypes/inputrec.h"
  50 #include "gromacs/mdtypes/interaction_const.h"
  51 #include "gromacs/mdtypes/md_enums.h"
  52 #include "gromacs/mdtypes/mdatom.h"
  53 #include "gromacs/mdtypes/nblist.h"
  54 #include "gromacs/mdtypes/simulation_workload.h"
  55 #include "gromacs/nbnxm/gpu_data_mgmt.h"
  56 #include "gromacs/nbnxm/nbnxm.h"
  57 #include "gromacs/simd/simd.h"
  58 #include "gromacs/timing/wallcycle.h"
  59 #include "gromacs/utility/enumerationhelpers.h"
  60 #include "gromacs/utility/fatalerror.h"
  61 #include "gromacs/utility/gmxassert.h"
  62 #include "gromacs/utility/real.h"
  63
  64 #include "kernel_common.h"
  65 #include "nbnxm_gpu.h"
  66 #include "nbnxm_simd.h"
  67 #include "pairlistset.h"
  68 #include "pairlistsets.h"
  69 #include "kernels_reference/kernel_gpu_ref.h"
  70 #define INCLUDE_KERNELFUNCTION_TABLES
  71 #include "kernels_reference/kernel_ref.h"
  72 #ifdef GMX_NBNXN_SIMD_2XNN
  73 #    include "kernels_simd_2xmm/kernels.h"
  74 #endif
  75 #ifdef GMX_NBNXN_SIMD_4XN
  76 #    include "kernels_simd_4xm/kernels.h"
  77 #endif
  78 #undef INCLUDE_FUNCTION_TABLES
  79
  80 /*! \brief Clears the energy group output buffers
  81  *
  82  * \param[in,out] out  nbnxn kernel output struct
  83  */
  84 static void clearGroupEnergies(nbnxn_atomdata_output_t* out)
  85 {
  86     std::fill(out->Vvdw.begin(), out->Vvdw.end(), 0.0_real);
  87     std::fill(out->Vc.begin(), out->Vc.end(), 0.0_real);
  88     std::fill(out->VSvdw.begin(), out->VSvdw.end(), 0.0_real);
  89     std::fill(out->VSc.begin(), out->VSc.end(), 0.0_real);
  90 }
  91
  92 /*! \brief Reduce the group-pair energy buffers produced by a SIMD kernel
  93  * to single terms in the output buffers.
  94  *
  95  * The SIMD kernels produce a large number of energy buffer in SIMD registers
  96  * to avoid scattered reads and writes.
  97  *
  98  * \tparam        unrollj         The unroll size for j-particles in the SIMD kernel
  99  * \param[in]     numGroups       The number of energy groups
 100  * \param[in]     numGroups_2log  Log2 of numGroups, rounded up
 101  * \param[in,out] out             Struct with energy buffers
 102  */
 103 template<int unrollj>
 104 static void reduceGroupEnergySimdBuffers(int numGroups, int numGroups_2log, nbnxn_atomdata_output_t* out)
 105 {
 106     const int unrollj_half = unrollj / 2;
 107     /* Energies are stored in SIMD registers with size 2^numGroups_2log */
 108     const int numGroupsStorage = (1 << numGroups_2log);
 109
 110     const real* gmx_restrict vVdwSimd     = out->VSvdw.data();
 111     const real* gmx_restrict vCoulombSimd = out->VSc.data();
 112     real* gmx_restrict       vVdw         = out->Vvdw.data();
 113     real* gmx_restrict       vCoulomb     = out->Vc.data();
 114
 115     /* The size of the SIMD energy group buffer array is:
 116      * numGroups*numGroups*numGroupsStorage*unrollj_half*simd_width
 117      */
 118     for (int i = 0; i < numGroups; i++)
 119     {
 120         for (int j1 = 0; j1 < numGroups; j1++)
 121         {
 122             for (int j0 = 0; j0 < numGroups; j0++)
 123             {
 124                 int c = ((i * numGroups + j1) * numGroupsStorage + j0) * unrollj_half * unrollj;
 125                 for (int s = 0; s < unrollj_half; s++)
 126                 {
 127                     vVdw[i * numGroups + j0] += vVdwSimd[c + 0];
 128                     vVdw[i * numGroups + j1] += vVdwSimd[c + 1];
 129                     vCoulomb[i * numGroups + j0] += vCoulombSimd[c + 0];
 130                     vCoulomb[i * numGroups + j1] += vCoulombSimd[c + 1];
 131                     c += unrollj + 2;
 132                 }
 133             }
 134         }
 135     }
 136 }
 137
 138 CoulombKernelType getCoulombKernelType(const Nbnxm::EwaldExclusionType ewaldExclusionType,
 139                                        const CoulombInteractionType    coulombInteractionType,
 140                                        const bool                      haveEqualCoulombVwdRadii)
 141 {
 142
 143     if (EEL_RF(coulombInteractionType) || coulombInteractionType == CoulombInteractionType::Cut)
 144     {
 145         return CoulombKernelType::ReactionField;
 146     }
 147     else
 148     {
 149         if (ewaldExclusionType == Nbnxm::EwaldExclusionType::Table)
 150         {
 151             if (haveEqualCoulombVwdRadii)
 152             {
 153                 return CoulombKernelType::Table;
 154             }
 155             else
 156             {
 157                 return CoulombKernelType::TableTwin;
 158             }
 159         }
 160         else
 161         {
 162             if (haveEqualCoulombVwdRadii)
 163             {
 164                 return CoulombKernelType::Ewald;
 165             }
 166             else
 167             {
 168                 return CoulombKernelType::EwaldTwin;
 169             }
 170         }
 171     }
 172 }
 173
 174 int getVdwKernelType(const Nbnxm::KernelType    kernelType,
 175                      const LJCombinationRule    ljCombinationRule,
 176                      const VanDerWaalsType      vanDerWaalsType,
 177                      const InteractionModifiers interactionModifiers,
 178                      const LongRangeVdW         longRangeVdW)
 179 {
 180     if (vanDerWaalsType == VanDerWaalsType::Cut)
 181     {
 182         switch (interactionModifiers)
 183         {
 184             case InteractionModifiers::None:
 185             case InteractionModifiers::PotShift:
 186                 switch (ljCombinationRule)
 187                 {
 188                     case LJCombinationRule::Geometric: return vdwktLJCUT_COMBGEOM;
 189                     case LJCombinationRule::LorentzBerthelot: return vdwktLJCUT_COMBLB;
 190                     case LJCombinationRule::None: return vdwktLJCUT_COMBNONE;
 191                     default: GMX_THROW(gmx::InvalidInputError("Unknown combination rule"));
 192                 }
 193             case InteractionModifiers::ForceSwitch: return vdwktLJFORCESWITCH;
 194             case InteractionModifiers::PotSwitch: return vdwktLJPOTSWITCH;
 195             default:
 196                 std::string errorMsg =
 197                         gmx::formatString("Unsupported VdW interaction modifier %s (%d)",
 198                                           enumValueToString(interactionModifiers),
 199                                           static_cast<int>(interactionModifiers));
 200                 GMX_THROW(gmx::InvalidInputError(errorMsg));
 201         }
 202     }
 203     else if (vanDerWaalsType == VanDerWaalsType::Pme)
 204     {
 205         if (longRangeVdW == LongRangeVdW::Geom)
 206         {
 207             return vdwktLJEWALDCOMBGEOM;
 208         }
 209         else
 210         {
 211             /* At setup we (should have) selected the C reference kernel */
 212             GMX_RELEASE_ASSERT(kernelType == Nbnxm::KernelType::Cpu4x4_PlainC,
 213                                "Only the C reference nbnxn SIMD kernel supports LJ-PME with LB "
 214                                "combination rules");
 215             return vdwktLJEWALDCOMBLB;
 216         }
 217     }
 218     else
 219     {
 220         std::string errorMsg = gmx::formatString("Unsupported VdW interaction type %s (%d)",
 221                                                  enumValueToString(vanDerWaalsType),
 222                                                  static_cast<int>(vanDerWaalsType));
 223         GMX_THROW(gmx::InvalidInputError(errorMsg));
 224     }
 225 }
 226
 227 /*! \brief Dispatches the non-bonded N versus M atom cluster CPU kernels.
 228  *
 229  * OpenMP parallelization is performed within this function.
 230  * Energy reduction, but not force and shift force reduction, is performed
 231  * within this function.
 232  *
 233  * \param[in]     pairlistSet   Pairlists with local or non-local interactions to compute
 234  * \param[in]     kernelSetup   The non-bonded kernel setup
 235  * \param[in,out] nbat          The atomdata for the interactions
 236  * \param[in]     ic            Non-bonded interaction constants
 237  * \param[in]     shiftVectors  The PBC shift vectors
 238  * \param[in]     stepWork      Flags that tell what to compute
 239  * \param[in]     clearF        Enum that tells if to clear the force output buffer
 240  * \param[out]    vCoulomb      Output buffer for Coulomb energies
 241  * \param[out]    vVdw          Output buffer for Van der Waals energies
 242  * \param[in]     wcycle        Pointer to cycle counting data structure.
 243  */
 244 static void nbnxn_kernel_cpu(const PairlistSet&             pairlistSet,
 245                              const Nbnxm::KernelSetup&      kernelSetup,
 246                              nbnxn_atomdata_t*              nbat,
 247                              const interaction_const_t&     ic,
 248                              gmx::ArrayRef<const gmx::RVec> shiftVectors,
 249                              const gmx::StepWorkload&       stepWork,
 250                              int                            clearF,
 251                              real*                          vCoulomb,
 252                              real*                          vVdw,
 253                              gmx_wallcycle*                 wcycle)
 254 {
 255
 256     const nbnxn_atomdata_t::Params& nbatParams = nbat->params();
 257
 258     const int coulkt = static_cast<int>(getCoulombKernelType(
 259             kernelSetup.ewaldExclusionType, ic.eeltype, (ic.rcoulomb == ic.rvdw)));
 260     const int vdwkt  = getVdwKernelType(
 261             kernelSetup.kernelType, nbatParams.ljCombinationRule, ic.vdwtype, ic.vdw_modifier, ic.ljpme_comb_rule);
 262
 263     gmx::ArrayRef<const NbnxnPairlistCpu> pairlists = pairlistSet.cpuLists();
 264
 265     const auto* shiftVecPointer = as_rvec_array(shiftVectors.data());
 266
 267     int gmx_unused nthreads = gmx_omp_nthreads_get(ModuleMultiThread::Nonbonded);
 268     wallcycle_sub_start(wcycle, WallCycleSubCounter::NonbondedClear);
 269 #pragma omp parallel for schedule(static) num_threads(nthreads)
 270     for (gmx::index nb = 0; nb < pairlists.ssize(); nb++)
 271     {
 272         // Presently, the kernels do not call C++ code that can throw,
 273         // so no need for a try/catch pair in this OpenMP region.
 274         nbnxn_atomdata_output_t* out = &nbat->out[nb];
 275
 276         if (clearF == enbvClearFYes)
 277         {
 278             clearForceBuffer(nbat, nb);
 279
 280             clear_fshift(out->fshift.data());
 281         }
 282
 283         if (nb == 0)
 284         {
 285             wallcycle_sub_stop(wcycle, WallCycleSubCounter::NonbondedClear);
 286             wallcycle_sub_start(wcycle, WallCycleSubCounter::NonbondedKernel);
 287         }
 288
 289         // TODO: Change to reference
 290         const NbnxnPairlistCpu* pairlist = &pairlists[nb];
 291
 292         if (!stepWork.computeEnergy)
 293         {
 294             /* Don't calculate energies */
 295             switch (kernelSetup.kernelType)
 296             {
 297                 case Nbnxm::KernelType::Cpu4x4_PlainC:
 298                     nbnxn_kernel_noener_ref[coulkt][vdwkt](pairlist, nbat, &ic, shiftVecPointer, out);
 299                     break;
 300 #ifdef GMX_NBNXN_SIMD_2XNN
 301                 case Nbnxm::KernelType::Cpu4xN_Simd_2xNN:
 302                     nbnxm_kernel_noener_simd_2xmm[coulkt][vdwkt](pairlist, nbat, &ic, shiftVecPointer, out);
 303                     break;
 304 #endif
 305 #ifdef GMX_NBNXN_SIMD_4XN
 306                 case Nbnxm::KernelType::Cpu4xN_Simd_4xN:
 307                     nbnxm_kernel_noener_simd_4xm[coulkt][vdwkt](pairlist, nbat, &ic, shiftVecPointer, out);
 308                     break;
 309 #endif
 310                 default: GMX_RELEASE_ASSERT(false, "Unsupported kernel architecture");
 311             }
 312         }
 313         else if (out->Vvdw.size() == 1)
 314         {
 315             /* A single energy group (pair) */
 316             out->Vvdw[0] = 0;
 317             out->Vc[0]   = 0;
 318
 319             switch (kernelSetup.kernelType)
 320             {
 321                 case Nbnxm::KernelType::Cpu4x4_PlainC:
 322                     nbnxn_kernel_ener_ref[coulkt][vdwkt](pairlist, nbat, &ic, shiftVecPointer, out);
 323                     break;
 324 #ifdef GMX_NBNXN_SIMD_2XNN
 325                 case Nbnxm::KernelType::Cpu4xN_Simd_2xNN:
 326                     nbnxm_kernel_ener_simd_2xmm[coulkt][vdwkt](pairlist, nbat, &ic, shiftVecPointer, out);
 327                     break;
 328 #endif
 329 #ifdef GMX_NBNXN_SIMD_4XN
 330                 case Nbnxm::KernelType::Cpu4xN_Simd_4xN:
 331                     nbnxm_kernel_ener_simd_4xm[coulkt][vdwkt](pairlist, nbat, &ic, shiftVecPointer, out);
 332                     break;
 333 #endif
 334                 default: GMX_RELEASE_ASSERT(false, "Unsupported kernel architecture");
 335             }
 336         }
 337         else
 338         {
 339             /* Calculate energy group contributions */
 340             clearGroupEnergies(out);
 341
 342             int unrollj = 0;
 343
 344             switch (kernelSetup.kernelType)
 345             {
 346                 case Nbnxm::KernelType::Cpu4x4_PlainC:
 347                     unrollj = c_nbnxnCpuIClusterSize;
 348                     nbnxn_kernel_energrp_ref[coulkt][vdwkt](pairlist, nbat, &ic, shiftVecPointer, out);
 349                     break;
 350 #ifdef GMX_NBNXN_SIMD_2XNN
 351                 case Nbnxm::KernelType::Cpu4xN_Simd_2xNN:
 352                     unrollj = GMX_SIMD_REAL_WIDTH / 2;
 353                     nbnxm_kernel_energrp_simd_2xmm[coulkt][vdwkt](
 354                             pairlist, nbat, &ic, shiftVecPointer, out);
 355                     break;
 356 #endif
 357 #ifdef GMX_NBNXN_SIMD_4XN
 358                 case Nbnxm::KernelType::Cpu4xN_Simd_4xN:
 359                     unrollj = GMX_SIMD_REAL_WIDTH;
 360                     nbnxm_kernel_energrp_simd_4xm[coulkt][vdwkt](pairlist, nbat, &ic, shiftVecPointer, out);
 361                     break;
 362 #endif
 363                 default: GMX_RELEASE_ASSERT(false, "Unsupported kernel architecture");
 364             }
 365
 366             if (kernelSetup.kernelType != Nbnxm::KernelType::Cpu4x4_PlainC)
 367             {
 368                 switch (unrollj)
 369                 {
 370                     case 2:
 371                         reduceGroupEnergySimdBuffers<2>(nbatParams.nenergrp, nbatParams.neg_2log, out);
 372                         break;
 373                     case 4:
 374                         reduceGroupEnergySimdBuffers<4>(nbatParams.nenergrp, nbatParams.neg_2log, out);
 375                         break;
 376                     case 8:
 377                         reduceGroupEnergySimdBuffers<8>(nbatParams.nenergrp, nbatParams.neg_2log, out);
 378                         break;
 379                     default: GMX_RELEASE_ASSERT(false, "Unsupported j-unroll size");
 380                 }
 381             }
 382         }
 383     }
 384     wallcycle_sub_stop(wcycle, WallCycleSubCounter::NonbondedKernel);
 385
 386     if (stepWork.computeEnergy)
 387     {
 388         reduce_energies_over_lists(nbat, pairlists.ssize(), vVdw, vCoulomb);
 389     }
 390 }
 391
 392 static void accountFlops(t_nrnb*                    nrnb,
 393                          const PairlistSet&         pairlistSet,
 394                          const nonbonded_verlet_t&  nbv,
 395                          const interaction_const_t& ic,
 396                          const gmx::StepWorkload&   stepWork)
 397 {
 398     const bool usingGpuKernels = nbv.useGpu();
 399
 400     int enr_nbnxn_kernel_ljc = eNRNB;
 401     if (EEL_RF(ic.eeltype) || ic.eeltype == CoulombInteractionType::Cut)
 402     {
 403         enr_nbnxn_kernel_ljc = eNR_NBNXN_LJ_RF;
 404     }
 405     else if ((!usingGpuKernels && nbv.kernelSetup().ewaldExclusionType == Nbnxm::EwaldExclusionType::Analytical)
 406              || (usingGpuKernels && Nbnxm::gpu_is_kernel_ewald_analytical(nbv.gpu_nbv)))
 407     {
 408         enr_nbnxn_kernel_ljc = eNR_NBNXN_LJ_EWALD;
 409     }
 410     else
 411     {
 412         enr_nbnxn_kernel_ljc = eNR_NBNXN_LJ_TAB;
 413     }
 414     int enr_nbnxn_kernel_lj = eNR_NBNXN_LJ;
 415     if (stepWork.computeEnergy)
 416     {
 417         /* In eNR_??? the nbnxn F+E kernels are always the F kernel + 1 */
 418         enr_nbnxn_kernel_ljc += 1;
 419         enr_nbnxn_kernel_lj += 1;
 420     }
 421
 422     inc_nrnb(nrnb, enr_nbnxn_kernel_ljc, pairlistSet.natpair_ljq_);
 423     inc_nrnb(nrnb, enr_nbnxn_kernel_lj, pairlistSet.natpair_lj_);
 424     /* The Coulomb-only kernels are offset -eNR_NBNXN_LJ_RF+eNR_NBNXN_RF */
 425     inc_nrnb(nrnb, enr_nbnxn_kernel_ljc - eNR_NBNXN_LJ_RF + eNR_NBNXN_RF, pairlistSet.natpair_q_);
 426
 427     if (ic.vdw_modifier == InteractionModifiers::ForceSwitch)
 428     {
 429         /* We add up the switch cost separately */
 430         inc_nrnb(nrnb,
 431                  eNR_NBNXN_ADD_LJ_FSW + (stepWork.computeEnergy ? 1 : 0),
 432                  pairlistSet.natpair_ljq_ + pairlistSet.natpair_lj_);
 433     }
 434     if (ic.vdw_modifier == InteractionModifiers::PotSwitch)
 435     {
 436         /* We add up the switch cost separately */
 437         inc_nrnb(nrnb,
 438                  eNR_NBNXN_ADD_LJ_PSW + (stepWork.computeEnergy ? 1 : 0),
 439                  pairlistSet.natpair_ljq_ + pairlistSet.natpair_lj_);
 440     }
 441     if (ic.vdwtype == VanDerWaalsType::Pme)
 442     {
 443         /* We add up the LJ Ewald cost separately */
 444         inc_nrnb(nrnb,
 445                  eNR_NBNXN_ADD_LJ_EWALD + (stepWork.computeEnergy ? 1 : 0),
 446                  pairlistSet.natpair_ljq_ + pairlistSet.natpair_lj_);
 447     }
 448 }
 449
 450 void nonbonded_verlet_t::dispatchNonbondedKernel(gmx::InteractionLocality       iLocality,
 451                                                  const interaction_const_t&     ic,
 452                                                  const gmx::StepWorkload&       stepWork,
 453                                                  int                            clearF,
 454                                                  gmx::ArrayRef<const gmx::RVec> shiftvec,
 455                                                  gmx::ArrayRef<real> repulsionDispersionSR,
 456                                                  gmx::ArrayRef<real> CoulombSR,
 457                                                  t_nrnb*             nrnb) const
 458 {
 459     const PairlistSet& pairlistSet = pairlistSets().pairlistSet(iLocality);
 460
 461     switch (kernelSetup().kernelType)
 462     {
 463         case Nbnxm::KernelType::Cpu4x4_PlainC:
 464         case Nbnxm::KernelType::Cpu4xN_Simd_4xN:
 465         case Nbnxm::KernelType::Cpu4xN_Simd_2xNN:
 466             nbnxn_kernel_cpu(pairlistSet,
 467                              kernelSetup(),
 468                              nbat.get(),
 469                              ic,
 470                              shiftvec,
 471                              stepWork,
 472                              clearF,
 473                              CoulombSR.data(),
 474                              repulsionDispersionSR.data(),
 475                              wcycle_);
 476             break;
 477
 478         case Nbnxm::KernelType::Gpu8x8x8:
 479             Nbnxm::gpu_launch_kernel(gpu_nbv, stepWork, iLocality);
 480             break;
 481
 482         case Nbnxm::KernelType::Cpu8x8x8_PlainC:
 483             nbnxn_kernel_gpu_ref(pairlistSet.gpuList(),
 484                                  nbat.get(),
 485                                  &ic,
 486                                  shiftvec,
 487                                  stepWork,
 488                                  clearF,
 489                                  nbat->out[0].f,
 490                                  nbat->out[0].fshift.data(),
 491                                  CoulombSR.data(),
 492                                  repulsionDispersionSR.data());
 493             break;
 494
 495         default: GMX_RELEASE_ASSERT(false, "Invalid nonbonded kernel type passed!");
 496     }
 497
 498     accountFlops(nrnb, pairlistSet, *this, ic, stepWork);
 499 }
 500
 501 void nonbonded_verlet_t::dispatchFreeEnergyKernel(gmx::InteractionLocality       iLocality,
 502                                                   gmx::ArrayRef<const gmx::RVec> coords,
 503                                                   gmx::ForceWithShiftForces* forceWithShiftForces,
 504                                                   bool                       useSimd,
 505                                                   int                        ntype,
 506                                                   real                       rlist,
 507                                                   const interaction_const_t& ic,
 508                                                   gmx::ArrayRef<const gmx::RVec> shiftvec,
 509                                                   gmx::ArrayRef<const real>      nbfp,
 510                                                   gmx::ArrayRef<const real>      nbfp_grid,
 511                                                   gmx::ArrayRef<const real>      chargeA,
 512                                                   gmx::ArrayRef<const real>      chargeB,
 513                                                   gmx::ArrayRef<const int>       typeA,
 514                                                   gmx::ArrayRef<const int>       typeB,
 515                                                   t_lambda*                      fepvals,
 516                                                   gmx::ArrayRef<const real>      lambda,
 517                                                   gmx_enerdata_t*                enerd,
 518                                                   const gmx::StepWorkload&       stepWork,
 519                                                   t_nrnb*                        nrnb)
 520 {
 521     const auto nbl_fep = pairlistSets().pairlistSet(iLocality).fepLists();
 522
 523     /* When the first list is empty, all are empty and there is nothing to do */
 524     if (!pairlistSets().params().haveFep || nbl_fep[0]->nrj == 0)
 525     {
 526         return;
 527     }
 528
 529     int donb_flags = 0;
 530     /* Add short-range interactions */
 531     donb_flags |= GMX_NONBONDED_DO_SR;
 532
 533     if (stepWork.computeForces)
 534     {
 535         donb_flags |= GMX_NONBONDED_DO_FORCE;
 536     }
 537     if (stepWork.computeVirial)
 538     {
 539         donb_flags |= GMX_NONBONDED_DO_SHIFTFORCE;
 540     }
 541     if (stepWork.computeEnergy)
 542     {
 543         donb_flags |= GMX_NONBONDED_DO_POTENTIAL;
 544     }
 545
 546     gmx::EnumerationArray<FreeEnergyPerturbationCouplingType, real> dvdl_nb      = { 0 };
 547     int                                                             kernelFlags  = donb_flags;
 548     gmx::ArrayRef<const real>                                       kernelLambda = lambda;
 549     gmx::ArrayRef<real>                                             kernelDvdl   = dvdl_nb;
 550
 551     gmx::ArrayRef<real> energygrp_elec = enerd->grpp.energyGroupPairTerms[NonBondedEnergyTerms::CoulombSR];
 552     gmx::ArrayRef<real> energygrp_vdw = enerd->grpp.energyGroupPairTerms[NonBondedEnergyTerms::LJSR];
 553
 554     GMX_ASSERT(gmx_omp_nthreads_get(ModuleMultiThread::Nonbonded) == nbl_fep.ssize(),
 555                "Number of lists should be same as number of NB threads");
 556
 557     wallcycle_sub_start(wcycle_, WallCycleSubCounter::NonbondedFep);
 558 #pragma omp parallel for schedule(static) num_threads(nbl_fep.ssize())
 559     for (gmx::index th = 0; th < nbl_fep.ssize(); th++)
 560     {
 561         try
 562         {
 563             gmx_nb_free_energy_kernel(*nbl_fep[th],
 564                                       coords,
 565                                       forceWithShiftForces,
 566                                       useSimd,
 567                                       ntype,
 568                                       rlist,
 569                                       ic,
 570                                       shiftvec,
 571                                       nbfp,
 572                                       nbfp_grid,
 573                                       chargeA,
 574                                       chargeB,
 575                                       typeA,
 576                                       typeB,
 577                                       kernelFlags,
 578                                       kernelLambda,
 579                                       kernelDvdl,
 580                                       energygrp_elec,
 581                                       energygrp_vdw,
 582                                       nrnb);
 583         }
 584         GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR
 585     }
 586
 587     if (fepvals->sc_alpha != 0)
 588     {
 589         enerd->dvdl_nonlin[FreeEnergyPerturbationCouplingType::Vdw] +=
 590                 dvdl_nb[FreeEnergyPerturbationCouplingType::Vdw];
 591         enerd->dvdl_nonlin[FreeEnergyPerturbationCouplingType::Coul] +=
 592                 dvdl_nb[FreeEnergyPerturbationCouplingType::Coul];
 593     }
 594     else
 595     {
 596         enerd->dvdl_lin[FreeEnergyPerturbationCouplingType::Vdw] +=
 597                 dvdl_nb[FreeEnergyPerturbationCouplingType::Vdw];
 598         enerd->dvdl_lin[FreeEnergyPerturbationCouplingType::Coul] +=
 599                 dvdl_nb[FreeEnergyPerturbationCouplingType::Coul];
 600     }
 601
 602     /* If we do foreign lambda and we have soft-core interactions
 603      * we have to recalculate the (non-linear) energies contributions.
 604      */
 605     if (fepvals->n_lambda > 0 && stepWork.computeDhdl && fepvals->sc_alpha != 0)
 606     {
 607         gmx::EnumerationArray<FreeEnergyPerturbationCouplingType, real> lam_i;
 608         kernelFlags = (donb_flags & ~(GMX_NONBONDED_DO_FORCE | GMX_NONBONDED_DO_SHIFTFORCE))
 609                       | GMX_NONBONDED_DO_FOREIGNLAMBDA;
 610         kernelLambda = lam_i;
 611         kernelDvdl   = dvdl_nb;
 612         gmx::ArrayRef<real> energygrp_elec =
 613                 foreignEnergyGroups_->energyGroupPairTerms[NonBondedEnergyTerms::CoulombSR];
 614         gmx::ArrayRef<real> energygrp_vdw =
 615                 foreignEnergyGroups_->energyGroupPairTerms[NonBondedEnergyTerms::LJSR];
 616
 617         for (gmx::index i = 0; i < 1 + enerd->foreignLambdaTerms.numLambdas(); i++)
 618         {
 619             std::fill(std::begin(dvdl_nb), std::end(dvdl_nb), 0);
 620             for (int j = 0; j < static_cast<int>(FreeEnergyPerturbationCouplingType::Count); j++)
 621             {
 622                 lam_i[j] = (i == 0 ? lambda[j] : fepvals->all_lambda[j][i - 1]);
 623             }
 624             foreignEnergyGroups_->clear();
 625 #pragma omp parallel for schedule(static) num_threads(nbl_fep.ssize())
 626             for (gmx::index th = 0; th < nbl_fep.ssize(); th++)
 627             {
 628                 try
 629                 {
 630                     gmx_nb_free_energy_kernel(*nbl_fep[th],
 631                                               coords,
 632                                               forceWithShiftForces,
 633                                               useSimd,
 634                                               ntype,
 635                                               rlist,
 636                                               ic,
 637                                               shiftvec,
 638                                               nbfp,
 639                                               nbfp_grid,
 640                                               chargeA,
 641                                               chargeB,
 642                                               typeA,
 643                                               typeB,
 644                                               kernelFlags,
 645                                               kernelLambda,
 646                                               kernelDvdl,
 647                                               energygrp_elec,
 648                                               energygrp_vdw,
 649                                               nrnb);
 650                 }
 651                 GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR
 652             }
 653             std::array<real, F_NRE> foreign_term = { 0 };
 654             sum_epot(*foreignEnergyGroups_, foreign_term.data());
 655             enerd->foreignLambdaTerms.accumulate(
 656                     i,
 657                     foreign_term[F_EPOT],
 658                     dvdl_nb[FreeEnergyPerturbationCouplingType::Vdw]
 659                             + dvdl_nb[FreeEnergyPerturbationCouplingType::Coul]);
 660         }
 661     }
 662     wallcycle_sub_stop(wcycle_, WallCycleSubCounter::NonbondedFep);
 663 }