src/gromacs/nbnxm/cuda/nbnxm_cuda_data_mgmt.cu

   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, 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 /*! \file
  37  *  \brief Define CUDA implementation of nbnxn_gpu_data_mgmt.h
  38  *
  39  *  \author Szilard Pall <pall.szilard@gmail.com>
  40  */
  41 #include "gmxpre.h"
  42
  43 #include <assert.h>
  44 #include <stdarg.h>
  45 #include <stdio.h>
  46 #include <stdlib.h>
  47
  48 // TODO We would like to move this down, but the way NbnxmGpu
  49 //      is currently declared means this has to be before gpu_types.h
  50 #include "nbnxm_cuda_types.h"
  51
  52 // TODO Remove this comment when the above order issue is resolved
  53 #include "gromacs/gpu_utils/cudautils.cuh"
  54 #include "gromacs/gpu_utils/device_stream_manager.h"
  55 #include "gromacs/gpu_utils/gpu_utils.h"
  56 #include "gromacs/gpu_utils/gpueventsynchronizer.cuh"
  57 #include "gromacs/gpu_utils/pmalloc_cuda.h"
  58 #include "gromacs/hardware/gpu_hw_info.h"
  59 #include "gromacs/math/vectypes.h"
  60 #include "gromacs/mdlib/force_flags.h"
  61 #include "gromacs/mdtypes/interaction_const.h"
  62 #include "gromacs/mdtypes/md_enums.h"
  63 #include "gromacs/nbnxm/atomdata.h"
  64 #include "gromacs/nbnxm/gpu_data_mgmt.h"
  65 #include "gromacs/nbnxm/gridset.h"
  66 #include "gromacs/nbnxm/nbnxm.h"
  67 #include "gromacs/nbnxm/nbnxm_gpu.h"
  68 #include "gromacs/nbnxm/pairlistsets.h"
  69 #include "gromacs/pbcutil/ishift.h"
  70 #include "gromacs/timing/gpu_timing.h"
  71 #include "gromacs/utility/basedefinitions.h"
  72 #include "gromacs/utility/cstringutil.h"
  73 #include "gromacs/utility/fatalerror.h"
  74 #include "gromacs/utility/real.h"
  75 #include "gromacs/utility/smalloc.h"
  76
  77 #include "nbnxm_cuda.h"
  78
  79 namespace Nbnxm
  80 {
  81
  82 /* This is a heuristically determined parameter for the Kepler
  83  * and Maxwell architectures for the minimum size of ci lists by multiplying
  84  * this constant with the # of multiprocessors on the current device.
  85  * Since the maximum number of blocks per multiprocessor is 16, the ideal
  86  * count for small systems is 32 or 48 blocks per multiprocessor. Because
  87  * there is a bit of fluctuations in the generated block counts, we use
  88  * a target of 44 instead of the ideal value of 48.
  89  */
  90 static unsigned int gpu_min_ci_balanced_factor = 44;
  91
  92 /* Fw. decl. */
  93 static void nbnxn_cuda_clear_e_fshift(NbnxmGpu* nb);
  94
  95 /* Fw. decl, */
  96 static void nbnxn_cuda_free_nbparam_table(cu_nbparam_t* nbparam);
  97
  98 /*! \brief Initialized the Ewald Coulomb correction GPU table.
  99
 100     Tabulates the Ewald Coulomb force and initializes the size/scale
 101     and the table GPU array. If called with an already allocated table,
 102     it just re-uploads the table.
 103  */
 104 static void init_ewald_coulomb_force_table(const EwaldCorrectionTables& tables,
 105                                            cu_nbparam_t*                nbp,
 106                                            const DeviceContext&         deviceContext)
 107 {
 108     if (nbp->coulomb_tab != nullptr)
 109     {
 110         nbnxn_cuda_free_nbparam_table(nbp);
 111     }
 112
 113     nbp->coulomb_tab_scale = tables.scale;
 114     initParamLookupTable(&nbp->coulomb_tab, &nbp->coulomb_tab_texobj, tables.tableF.data(),
 115                          tables.tableF.size(), deviceContext);
 116 }
 117
 118
 119 /*! Initializes the atomdata structure first time, it only gets filled at
 120     pair-search. */
 121 static void init_atomdata_first(cu_atomdata_t* ad, int ntypes, const DeviceContext& deviceContext)
 122 {
 123     ad->ntypes = ntypes;
 124     allocateDeviceBuffer(&ad->shift_vec, SHIFTS, deviceContext);
 125     ad->bShiftVecUploaded = false;
 126
 127     allocateDeviceBuffer(&ad->fshift, SHIFTS, deviceContext);
 128     allocateDeviceBuffer(&ad->e_lj, 1, deviceContext);
 129     allocateDeviceBuffer(&ad->e_el, 1, deviceContext);
 130
 131     /* initialize to nullptr poiters to data that is not allocated here and will
 132        need reallocation in nbnxn_cuda_init_atomdata */
 133     ad->xq = nullptr;
 134     ad->f  = nullptr;
 135
 136     /* size -1 indicates that the respective array hasn't been initialized yet */
 137     ad->natoms = -1;
 138     ad->nalloc = -1;
 139 }
 140
 141 /*! Selects the Ewald kernel type, analytical on SM 3.0 and later, tabulated on
 142     earlier GPUs, single or twin cut-off. */
 143 static int pick_ewald_kernel_type(const interaction_const_t& ic)
 144 {
 145     bool bTwinCut = (ic.rcoulomb != ic.rvdw);
 146     bool bUseAnalyticalEwald, bForceAnalyticalEwald, bForceTabulatedEwald;
 147     int  kernel_type;
 148
 149     /* Benchmarking/development environment variables to force the use of
 150        analytical or tabulated Ewald kernel. */
 151     bForceAnalyticalEwald = (getenv("GMX_CUDA_NB_ANA_EWALD") != nullptr);
 152     bForceTabulatedEwald  = (getenv("GMX_CUDA_NB_TAB_EWALD") != nullptr);
 153
 154     if (bForceAnalyticalEwald && bForceTabulatedEwald)
 155     {
 156         gmx_incons(
 157                 "Both analytical and tabulated Ewald CUDA non-bonded kernels "
 158                 "requested through environment variables.");
 159     }
 160
 161     /* By default use analytical Ewald. */
 162     bUseAnalyticalEwald = true;
 163     if (bForceAnalyticalEwald)
 164     {
 165         if (debug)
 166         {
 167             fprintf(debug, "Using analytical Ewald CUDA kernels\n");
 168         }
 169     }
 170     else if (bForceTabulatedEwald)
 171     {
 172         bUseAnalyticalEwald = false;
 173
 174         if (debug)
 175         {
 176             fprintf(debug, "Using tabulated Ewald CUDA kernels\n");
 177         }
 178     }
 179
 180     /* Use twin cut-off kernels if requested by bTwinCut or the env. var.
 181        forces it (use it for debugging/benchmarking only). */
 182     if (!bTwinCut && (getenv("GMX_CUDA_NB_EWALD_TWINCUT") == nullptr))
 183     {
 184         kernel_type = bUseAnalyticalEwald ? eelTypeEWALD_ANA : eelTypeEWALD_TAB;
 185     }
 186     else
 187     {
 188         kernel_type = bUseAnalyticalEwald ? eelTypeEWALD_ANA_TWIN : eelTypeEWALD_TAB_TWIN;
 189     }
 190
 191     return kernel_type;
 192 }
 193
 194 /*! Copies all parameters related to the cut-off from ic to nbp */
 195 static void set_cutoff_parameters(cu_nbparam_t* nbp, const interaction_const_t* ic, const PairlistParams& listParams)
 196 {
 197     nbp->ewald_beta        = ic->ewaldcoeff_q;
 198     nbp->sh_ewald          = ic->sh_ewald;
 199     nbp->epsfac            = ic->epsfac;
 200     nbp->two_k_rf          = 2.0 * ic->k_rf;
 201     nbp->c_rf              = ic->c_rf;
 202     nbp->rvdw_sq           = ic->rvdw * ic->rvdw;
 203     nbp->rcoulomb_sq       = ic->rcoulomb * ic->rcoulomb;
 204     nbp->rlistOuter_sq     = listParams.rlistOuter * listParams.rlistOuter;
 205     nbp->rlistInner_sq     = listParams.rlistInner * listParams.rlistInner;
 206     nbp->useDynamicPruning = listParams.useDynamicPruning;
 207
 208     nbp->sh_lj_ewald   = ic->sh_lj_ewald;
 209     nbp->ewaldcoeff_lj = ic->ewaldcoeff_lj;
 210
 211     nbp->rvdw_switch      = ic->rvdw_switch;
 212     nbp->dispersion_shift = ic->dispersion_shift;
 213     nbp->repulsion_shift  = ic->repulsion_shift;
 214     nbp->vdw_switch       = ic->vdw_switch;
 215 }
 216
 217 /*! Initializes the nonbonded parameter data structure. */
 218 static void init_nbparam(cu_nbparam_t*                   nbp,
 219                          const interaction_const_t*      ic,
 220                          const PairlistParams&           listParams,
 221                          const nbnxn_atomdata_t::Params& nbatParams,
 222                          const DeviceContext&            deviceContext)
 223 {
 224     int ntypes;
 225
 226     ntypes = nbatParams.numTypes;
 227
 228     set_cutoff_parameters(nbp, ic, listParams);
 229
 230     /* The kernel code supports LJ combination rules (geometric and LB) for
 231      * all kernel types, but we only generate useful combination rule kernels.
 232      * We currently only use LJ combination rule (geometric and LB) kernels
 233      * for plain cut-off LJ. On Maxwell the force only kernels speed up 15%
 234      * with PME and 20% with RF, the other kernels speed up about half as much.
 235      * For LJ force-switch the geometric rule would give 7% speed-up, but this
 236      * combination is rarely used. LJ force-switch with LB rule is more common,
 237      * but gives only 1% speed-up.
 238      */
 239     if (ic->vdwtype == evdwCUT)
 240     {
 241         switch (ic->vdw_modifier)
 242         {
 243             case eintmodNONE:
 244             case eintmodPOTSHIFT:
 245                 switch (nbatParams.comb_rule)
 246                 {
 247                     case ljcrNONE: nbp->vdwtype = evdwTypeCUT; break;
 248                     case ljcrGEOM: nbp->vdwtype = evdwTypeCUTCOMBGEOM; break;
 249                     case ljcrLB: nbp->vdwtype = evdwTypeCUTCOMBLB; break;
 250                     default:
 251                         gmx_incons(
 252                                 "The requested LJ combination rule is not implemented in the CUDA "
 253                                 "GPU accelerated kernels!");
 254                 }
 255                 break;
 256             case eintmodFORCESWITCH: nbp->vdwtype = evdwTypeFSWITCH; break;
 257             case eintmodPOTSWITCH: nbp->vdwtype = evdwTypePSWITCH; break;
 258             default:
 259                 gmx_incons(
 260                         "The requested VdW interaction modifier is not implemented in the CUDA GPU "
 261                         "accelerated kernels!");
 262         }
 263     }
 264     else if (ic->vdwtype == evdwPME)
 265     {
 266         if (ic->ljpme_comb_rule == ljcrGEOM)
 267         {
 268             assert(nbatParams.comb_rule == ljcrGEOM);
 269             nbp->vdwtype = evdwTypeEWALDGEOM;
 270         }
 271         else
 272         {
 273             assert(nbatParams.comb_rule == ljcrLB);
 274             nbp->vdwtype = evdwTypeEWALDLB;
 275         }
 276     }
 277     else
 278     {
 279         gmx_incons(
 280                 "The requested VdW type is not implemented in the CUDA GPU accelerated kernels!");
 281     }
 282
 283     if (ic->eeltype == eelCUT)
 284     {
 285         nbp->eeltype = eelTypeCUT;
 286     }
 287     else if (EEL_RF(ic->eeltype))
 288     {
 289         nbp->eeltype = eelTypeRF;
 290     }
 291     else if ((EEL_PME(ic->eeltype) || ic->eeltype == eelEWALD))
 292     {
 293         nbp->eeltype = pick_ewald_kernel_type(*ic);
 294     }
 295     else
 296     {
 297         /* Shouldn't happen, as this is checked when choosing Verlet-scheme */
 298         gmx_incons(
 299                 "The requested electrostatics type is not implemented in the CUDA GPU accelerated "
 300                 "kernels!");
 301     }
 302
 303     /* generate table for PME */
 304     nbp->coulomb_tab = nullptr;
 305     if (nbp->eeltype == eelTypeEWALD_TAB || nbp->eeltype == eelTypeEWALD_TAB_TWIN)
 306     {
 307         GMX_RELEASE_ASSERT(ic->coulombEwaldTables, "Need valid Coulomb Ewald correction tables");
 308         init_ewald_coulomb_force_table(*ic->coulombEwaldTables, nbp, deviceContext);
 309     }
 310
 311     /* set up LJ parameter lookup table */
 312     if (!useLjCombRule(nbp->vdwtype))
 313     {
 314         initParamLookupTable(&nbp->nbfp, &nbp->nbfp_texobj, nbatParams.nbfp.data(),
 315                              2 * ntypes * ntypes, deviceContext);
 316     }
 317
 318     /* set up LJ-PME parameter lookup table */
 319     if (ic->vdwtype == evdwPME)
 320     {
 321         initParamLookupTable(&nbp->nbfp_comb, &nbp->nbfp_comb_texobj, nbatParams.nbfp_comb.data(),
 322                              2 * ntypes, deviceContext);
 323     }
 324 }
 325
 326 /*! Re-generate the GPU Ewald force table, resets rlist, and update the
 327  *  electrostatic type switching to twin cut-off (or back) if needed. */
 328 void gpu_pme_loadbal_update_param(const nonbonded_verlet_t* nbv, const interaction_const_t* ic)
 329 {
 330     if (!nbv || !nbv->useGpu())
 331     {
 332         return;
 333     }
 334     NbnxmGpu*     nb  = nbv->gpu_nbv;
 335     cu_nbparam_t* nbp = nbv->gpu_nbv->nbparam;
 336
 337     set_cutoff_parameters(nbp, ic, nbv->pairlistSets().params());
 338
 339     nbp->eeltype = pick_ewald_kernel_type(*ic);
 340
 341     GMX_RELEASE_ASSERT(ic->coulombEwaldTables, "Need valid Coulomb Ewald correction tables");
 342     init_ewald_coulomb_force_table(*ic->coulombEwaldTables, nbp, *nb->deviceContext_);
 343 }
 344
 345 /*! Initializes the pair list data structure. */
 346 static void init_plist(cu_plist_t* pl)
 347 {
 348     /* initialize to nullptr pointers to data that is not allocated here and will
 349        need reallocation in nbnxn_gpu_init_pairlist */
 350     pl->sci   = nullptr;
 351     pl->cj4   = nullptr;
 352     pl->imask = nullptr;
 353     pl->excl  = nullptr;
 354
 355     /* size -1 indicates that the respective array hasn't been initialized yet */
 356     pl->na_c          = -1;
 357     pl->nsci          = -1;
 358     pl->sci_nalloc    = -1;
 359     pl->ncj4          = -1;
 360     pl->cj4_nalloc    = -1;
 361     pl->nimask        = -1;
 362     pl->imask_nalloc  = -1;
 363     pl->nexcl         = -1;
 364     pl->excl_nalloc   = -1;
 365     pl->haveFreshList = false;
 366 }
 367
 368 /*! Initializes the timings data structure. */
 369 static void init_timings(gmx_wallclock_gpu_nbnxn_t* t)
 370 {
 371     int i, j;
 372
 373     t->nb_h2d_t = 0.0;
 374     t->nb_d2h_t = 0.0;
 375     t->nb_c     = 0;
 376     t->pl_h2d_t = 0.0;
 377     t->pl_h2d_c = 0;
 378     for (i = 0; i < 2; i++)
 379     {
 380         for (j = 0; j < 2; j++)
 381         {
 382             t->ktime[i][j].t = 0.0;
 383             t->ktime[i][j].c = 0;
 384         }
 385     }
 386     t->pruneTime.c        = 0;
 387     t->pruneTime.t        = 0.0;
 388     t->dynamicPruneTime.c = 0;
 389     t->dynamicPruneTime.t = 0.0;
 390 }
 391
 392 /*! Initializes simulation constant data. */
 393 static void cuda_init_const(NbnxmGpu*                       nb,
 394                             const interaction_const_t*      ic,
 395                             const PairlistParams&           listParams,
 396                             const nbnxn_atomdata_t::Params& nbatParams)
 397 {
 398     init_atomdata_first(nb->atdat, nbatParams.numTypes, *nb->deviceContext_);
 399     init_nbparam(nb->nbparam, ic, listParams, nbatParams, *nb->deviceContext_);
 400
 401     /* clear energy and shift force outputs */
 402     nbnxn_cuda_clear_e_fshift(nb);
 403 }
 404
 405 NbnxmGpu* gpu_init(const gmx::DeviceStreamManager& deviceStreamManager,
 406                    const interaction_const_t*      ic,
 407                    const PairlistParams&           listParams,
 408                    const nbnxn_atomdata_t*         nbat,
 409                    bool                            bLocalAndNonlocal)
 410 {
 411     cudaError_t stat;
 412
 413     auto nb            = new NbnxmGpu();
 414     nb->deviceContext_ = &deviceStreamManager.context();
 415     snew(nb->atdat, 1);
 416     snew(nb->nbparam, 1);
 417     snew(nb->plist[InteractionLocality::Local], 1);
 418     if (bLocalAndNonlocal)
 419     {
 420         snew(nb->plist[InteractionLocality::NonLocal], 1);
 421     }
 422
 423     nb->bUseTwoStreams = bLocalAndNonlocal;
 424
 425     nb->timers = new cu_timers_t();
 426     snew(nb->timings, 1);
 427
 428     /* init nbst */
 429     pmalloc((void**)&nb->nbst.e_lj, sizeof(*nb->nbst.e_lj));
 430     pmalloc((void**)&nb->nbst.e_el, sizeof(*nb->nbst.e_el));
 431     pmalloc((void**)&nb->nbst.fshift, SHIFTS * sizeof(*nb->nbst.fshift));
 432
 433     init_plist(nb->plist[InteractionLocality::Local]);
 434
 435     /* local/non-local GPU streams */
 436     GMX_RELEASE_ASSERT(deviceStreamManager.streamIsValid(gmx::DeviceStreamType::NonBondedLocal),
 437                        "Local non-bonded stream should be initialized to use GPU for non-bonded.");
 438     nb->deviceStreams[InteractionLocality::Local] =
 439             &deviceStreamManager.stream(gmx::DeviceStreamType::NonBondedLocal);
 440     if (nb->bUseTwoStreams)
 441     {
 442         init_plist(nb->plist[InteractionLocality::NonLocal]);
 443
 444         /* Note that the device we're running on does not have to support
 445          * priorities, because we are querying the priority range which in this
 446          * case will be a single value.
 447          */
 448         GMX_RELEASE_ASSERT(deviceStreamManager.streamIsValid(gmx::DeviceStreamType::NonBondedNonLocal),
 449                            "Non-local non-bonded stream should be initialized to use GPU for "
 450                            "non-bonded with domain decomposition.");
 451         nb->deviceStreams[InteractionLocality::NonLocal] =
 452                 &deviceStreamManager.stream(gmx::DeviceStreamType::NonBondedNonLocal);
 453         ;
 454     }
 455
 456     /* init events for sychronization (timing disabled for performance reasons!) */
 457     stat = cudaEventCreateWithFlags(&nb->nonlocal_done, cudaEventDisableTiming);
 458     CU_RET_ERR(stat, "cudaEventCreate on nonlocal_done failed");
 459     stat = cudaEventCreateWithFlags(&nb->misc_ops_and_local_H2D_done, cudaEventDisableTiming);
 460     CU_RET_ERR(stat, "cudaEventCreate on misc_ops_and_local_H2D_done failed");
 461
 462     nb->xNonLocalCopyD2HDone = new GpuEventSynchronizer();
 463
 464     /* WARNING: CUDA timings are incorrect with multiple streams.
 465      *          This is the main reason why they are disabled by default.
 466      */
 467     // TODO: Consider turning on by default when we can detect nr of streams.
 468     nb->bDoTime = (getenv("GMX_ENABLE_GPU_TIMING") != nullptr);
 469
 470     if (nb->bDoTime)
 471     {
 472         init_timings(nb->timings);
 473     }
 474
 475     /* set the kernel type for the current GPU */
 476     /* pick L1 cache configuration */
 477     cuda_set_cacheconfig();
 478
 479     cuda_init_const(nb, ic, listParams, nbat->params());
 480
 481     nb->atomIndicesSize       = 0;
 482     nb->atomIndicesSize_alloc = 0;
 483     nb->ncxy_na               = 0;
 484     nb->ncxy_na_alloc         = 0;
 485     nb->ncxy_ind              = 0;
 486     nb->ncxy_ind_alloc        = 0;
 487     nb->ncell                 = 0;
 488     nb->ncell_alloc           = 0;
 489
 490     if (debug)
 491     {
 492         fprintf(debug, "Initialized CUDA data structures.\n");
 493     }
 494
 495     return nb;
 496 }
 497
 498 void gpu_init_pairlist(NbnxmGpu* nb, const NbnxnPairlistGpu* h_plist, const InteractionLocality iloc)
 499 {
 500     char                sbuf[STRLEN];
 501     bool                bDoTime      = (nb->bDoTime && !h_plist->sci.empty());
 502     const DeviceStream& deviceStream = *nb->deviceStreams[iloc];
 503     cu_plist_t*         d_plist      = nb->plist[iloc];
 504
 505     if (d_plist->na_c < 0)
 506     {
 507         d_plist->na_c = h_plist->na_ci;
 508     }
 509     else
 510     {
 511         if (d_plist->na_c != h_plist->na_ci)
 512         {
 513             sprintf(sbuf, "In cu_init_plist: the #atoms per cell has changed (from %d to %d)",
 514                     d_plist->na_c, h_plist->na_ci);
 515             gmx_incons(sbuf);
 516         }
 517     }
 518
 519     gpu_timers_t::Interaction& iTimers = nb->timers->interaction[iloc];
 520
 521     if (bDoTime)
 522     {
 523         iTimers.pl_h2d.openTimingRegion(deviceStream);
 524         iTimers.didPairlistH2D = true;
 525     }
 526
 527     const DeviceContext& deviceContext = *nb->deviceContext_;
 528
 529     reallocateDeviceBuffer(&d_plist->sci, h_plist->sci.size(), &d_plist->nsci, &d_plist->sci_nalloc,
 530                            deviceContext);
 531     copyToDeviceBuffer(&d_plist->sci, h_plist->sci.data(), 0, h_plist->sci.size(), deviceStream,
 532                        GpuApiCallBehavior::Async, bDoTime ? iTimers.pl_h2d.fetchNextEvent() : nullptr);
 533
 534     reallocateDeviceBuffer(&d_plist->cj4, h_plist->cj4.size(), &d_plist->ncj4, &d_plist->cj4_nalloc,
 535                            deviceContext);
 536     copyToDeviceBuffer(&d_plist->cj4, h_plist->cj4.data(), 0, h_plist->cj4.size(), deviceStream,
 537                        GpuApiCallBehavior::Async, bDoTime ? iTimers.pl_h2d.fetchNextEvent() : nullptr);
 538
 539     reallocateDeviceBuffer(&d_plist->imask, h_plist->cj4.size() * c_nbnxnGpuClusterpairSplit,
 540                            &d_plist->nimask, &d_plist->imask_nalloc, deviceContext);
 541
 542     reallocateDeviceBuffer(&d_plist->excl, h_plist->excl.size(), &d_plist->nexcl,
 543                            &d_plist->excl_nalloc, deviceContext);
 544     copyToDeviceBuffer(&d_plist->excl, h_plist->excl.data(), 0, h_plist->excl.size(), deviceStream,
 545                        GpuApiCallBehavior::Async, bDoTime ? iTimers.pl_h2d.fetchNextEvent() : nullptr);
 546
 547     if (bDoTime)
 548     {
 549         iTimers.pl_h2d.closeTimingRegion(deviceStream);
 550     }
 551
 552     /* the next use of thist list we be the first one, so we need to prune */
 553     d_plist->haveFreshList = true;
 554 }
 555
 556 void gpu_upload_shiftvec(NbnxmGpu* nb, const nbnxn_atomdata_t* nbatom)
 557 {
 558     cu_atomdata_t*      adat        = nb->atdat;
 559     const DeviceStream& localStream = *nb->deviceStreams[InteractionLocality::Local];
 560
 561     /* only if we have a dynamic box */
 562     if (nbatom->bDynamicBox || !adat->bShiftVecUploaded)
 563     {
 564         static_assert(sizeof(adat->shift_vec[0]) == sizeof(nbatom->shift_vec[0]),
 565                       "Sizes of host- and device-side shift vectors should be the same.");
 566         copyToDeviceBuffer(&adat->shift_vec, reinterpret_cast<const float3*>(nbatom->shift_vec.data()),
 567                            0, SHIFTS, localStream, GpuApiCallBehavior::Async, nullptr);
 568         adat->bShiftVecUploaded = true;
 569     }
 570 }
 571
 572 /*! Clears the first natoms_clear elements of the GPU nonbonded force output array. */
 573 static void nbnxn_cuda_clear_f(NbnxmGpu* nb, int natoms_clear)
 574 {
 575     cu_atomdata_t*      adat        = nb->atdat;
 576     const DeviceStream& localStream = *nb->deviceStreams[InteractionLocality::Local];
 577     clearDeviceBufferAsync(&adat->f, 0, natoms_clear, localStream);
 578 }
 579
 580 /*! Clears nonbonded shift force output array and energy outputs on the GPU. */
 581 static void nbnxn_cuda_clear_e_fshift(NbnxmGpu* nb)
 582 {
 583     cu_atomdata_t*      adat        = nb->atdat;
 584     const DeviceStream& localStream = *nb->deviceStreams[InteractionLocality::Local];
 585
 586     clearDeviceBufferAsync(&adat->fshift, 0, SHIFTS, localStream);
 587     clearDeviceBufferAsync(&adat->e_lj, 0, 1, localStream);
 588     clearDeviceBufferAsync(&adat->e_el, 0, 1, localStream);
 589 }
 590
 591 void gpu_clear_outputs(NbnxmGpu* nb, bool computeVirial)
 592 {
 593     nbnxn_cuda_clear_f(nb, nb->atdat->natoms);
 594     /* clear shift force array and energies if the outputs were
 595        used in the current step */
 596     if (computeVirial)
 597     {
 598         nbnxn_cuda_clear_e_fshift(nb);
 599     }
 600 }
 601
 602 void gpu_init_atomdata(NbnxmGpu* nb, const nbnxn_atomdata_t* nbat)
 603 {
 604     int                  nalloc, natoms;
 605     bool                 realloced;
 606     bool                 bDoTime       = nb->bDoTime;
 607     cu_timers_t*         timers        = nb->timers;
 608     cu_atomdata_t*       d_atdat       = nb->atdat;
 609     const DeviceContext& deviceContext = *nb->deviceContext_;
 610     const DeviceStream&  localStream   = *nb->deviceStreams[InteractionLocality::Local];
 611
 612     natoms    = nbat->numAtoms();
 613     realloced = false;
 614
 615     if (bDoTime)
 616     {
 617         /* time async copy */
 618         timers->atdat.openTimingRegion(localStream);
 619     }
 620
 621     /* need to reallocate if we have to copy more atoms than the amount of space
 622        available and only allocate if we haven't initialized yet, i.e d_atdat->natoms == -1 */
 623     if (natoms > d_atdat->nalloc)
 624     {
 625         nalloc = over_alloc_small(natoms);
 626
 627         /* free up first if the arrays have already been initialized */
 628         if (d_atdat->nalloc != -1)
 629         {
 630             freeDeviceBuffer(&d_atdat->f);
 631             freeDeviceBuffer(&d_atdat->xq);
 632             freeDeviceBuffer(&d_atdat->atom_types);
 633             freeDeviceBuffer(&d_atdat->lj_comb);
 634         }
 635
 636         allocateDeviceBuffer(&d_atdat->f, nalloc, deviceContext);
 637         allocateDeviceBuffer(&d_atdat->xq, nalloc, deviceContext);
 638         if (useLjCombRule(nb->nbparam->vdwtype))
 639         {
 640             allocateDeviceBuffer(&d_atdat->lj_comb, nalloc, deviceContext);
 641         }
 642         else
 643         {
 644             allocateDeviceBuffer(&d_atdat->atom_types, nalloc, deviceContext);
 645         }
 646
 647         d_atdat->nalloc = nalloc;
 648         realloced       = true;
 649     }
 650
 651     d_atdat->natoms       = natoms;
 652     d_atdat->natoms_local = nbat->natoms_local;
 653
 654     /* need to clear GPU f output if realloc happened */
 655     if (realloced)
 656     {
 657         nbnxn_cuda_clear_f(nb, nalloc);
 658     }
 659
 660     if (useLjCombRule(nb->nbparam->vdwtype))
 661     {
 662         static_assert(sizeof(d_atdat->lj_comb[0]) == sizeof(float2),
 663                       "Size of the LJ parameters element should be equal to the size of float2.");
 664         copyToDeviceBuffer(&d_atdat->lj_comb,
 665                            reinterpret_cast<const float2*>(nbat->params().lj_comb.data()), 0,
 666                            natoms, localStream, GpuApiCallBehavior::Async, nullptr);
 667     }
 668     else
 669     {
 670         static_assert(sizeof(d_atdat->atom_types[0]) == sizeof(nbat->params().type[0]),
 671                       "Sizes of host- and device-side atom types should be the same.");
 672         copyToDeviceBuffer(&d_atdat->atom_types, nbat->params().type.data(), 0, natoms, localStream,
 673                            GpuApiCallBehavior::Async, nullptr);
 674     }
 675
 676     if (bDoTime)
 677     {
 678         timers->atdat.closeTimingRegion(localStream);
 679     }
 680 }
 681
 682 static void nbnxn_cuda_free_nbparam_table(cu_nbparam_t* nbparam)
 683 {
 684     if (nbparam->eeltype == eelTypeEWALD_TAB || nbparam->eeltype == eelTypeEWALD_TAB_TWIN)
 685     {
 686         destroyParamLookupTable(&nbparam->coulomb_tab, nbparam->coulomb_tab_texobj);
 687     }
 688 }
 689
 690 void gpu_free(NbnxmGpu* nb)
 691 {
 692     cudaError_t    stat;
 693     cu_atomdata_t* atdat;
 694     cu_nbparam_t*  nbparam;
 695
 696     if (nb == nullptr)
 697     {
 698         return;
 699     }
 700
 701     atdat   = nb->atdat;
 702     nbparam = nb->nbparam;
 703
 704     nbnxn_cuda_free_nbparam_table(nbparam);
 705
 706     stat = cudaEventDestroy(nb->nonlocal_done);
 707     CU_RET_ERR(stat, "cudaEventDestroy failed on timers->nonlocal_done");
 708     stat = cudaEventDestroy(nb->misc_ops_and_local_H2D_done);
 709     CU_RET_ERR(stat, "cudaEventDestroy failed on timers->misc_ops_and_local_H2D_done");
 710
 711     delete nb->timers;
 712
 713     if (!useLjCombRule(nb->nbparam->vdwtype))
 714     {
 715         destroyParamLookupTable(&nbparam->nbfp, nbparam->nbfp_texobj);
 716     }
 717
 718     if (nbparam->vdwtype == evdwTypeEWALDGEOM || nbparam->vdwtype == evdwTypeEWALDLB)
 719     {
 720         destroyParamLookupTable(&nbparam->nbfp_comb, nbparam->nbfp_comb_texobj);
 721     }
 722
 723     freeDeviceBuffer(&atdat->shift_vec);
 724     freeDeviceBuffer(&atdat->fshift);
 725
 726     freeDeviceBuffer(&atdat->e_lj);
 727     freeDeviceBuffer(&atdat->e_el);
 728
 729     freeDeviceBuffer(&atdat->f);
 730     freeDeviceBuffer(&atdat->xq);
 731     freeDeviceBuffer(&atdat->atom_types);
 732     freeDeviceBuffer(&atdat->lj_comb);
 733
 734     /* Free plist */
 735     auto* plist = nb->plist[InteractionLocality::Local];
 736     freeDeviceBuffer(&plist->sci);
 737     freeDeviceBuffer(&plist->cj4);
 738     freeDeviceBuffer(&plist->imask);
 739     freeDeviceBuffer(&plist->excl);
 740     sfree(plist);
 741     if (nb->bUseTwoStreams)
 742     {
 743         auto* plist_nl = nb->plist[InteractionLocality::NonLocal];
 744         freeDeviceBuffer(&plist_nl->sci);
 745         freeDeviceBuffer(&plist_nl->cj4);
 746         freeDeviceBuffer(&plist_nl->imask);
 747         freeDeviceBuffer(&plist_nl->excl);
 748         sfree(plist_nl);
 749     }
 750
 751     /* Free nbst */
 752     pfree(nb->nbst.e_lj);
 753     nb->nbst.e_lj = nullptr;
 754
 755     pfree(nb->nbst.e_el);
 756     nb->nbst.e_el = nullptr;
 757
 758     pfree(nb->nbst.fshift);
 759     nb->nbst.fshift = nullptr;
 760
 761     sfree(atdat);
 762     sfree(nbparam);
 763     sfree(nb->timings);
 764     delete nb;
 765
 766     if (debug)
 767     {
 768         fprintf(debug, "Cleaned up CUDA data structures.\n");
 769     }
 770 }
 771
 772 //! This function is documented in the header file
 773 gmx_wallclock_gpu_nbnxn_t* gpu_get_timings(NbnxmGpu* nb)
 774 {
 775     return (nb != nullptr && nb->bDoTime) ? nb->timings : nullptr;
 776 }
 777
 778 void gpu_reset_timings(nonbonded_verlet_t* nbv)
 779 {
 780     if (nbv->gpu_nbv && nbv->gpu_nbv->bDoTime)
 781     {
 782         init_timings(nbv->gpu_nbv->timings);
 783     }
 784 }
 785
 786 int gpu_min_ci_balanced(NbnxmGpu* nb)
 787 {
 788     return nb != nullptr ? gpu_min_ci_balanced_factor * nb->deviceContext_->deviceInfo().prop.multiProcessorCount
 789                          : 0;
 790 }
 791
 792 gmx_bool gpu_is_kernel_ewald_analytical(const NbnxmGpu* nb)
 793 {
 794     return ((nb->nbparam->eeltype == eelTypeEWALD_ANA) || (nb->nbparam->eeltype == eelTypeEWALD_ANA_TWIN));
 795 }
 796
 797 void* gpu_get_xq(NbnxmGpu* nb)
 798 {
 799     assert(nb);
 800
 801     return static_cast<void*>(nb->atdat->xq);
 802 }
 803
 804 DeviceBuffer<gmx::RVec> gpu_get_f(NbnxmGpu* nb)
 805 {
 806     assert(nb);
 807
 808     return reinterpret_cast<DeviceBuffer<gmx::RVec>>(nb->atdat->f);
 809 }
 810
 811 DeviceBuffer<gmx::RVec> gpu_get_fshift(NbnxmGpu* nb)
 812 {
 813     assert(nb);
 814
 815     return reinterpret_cast<DeviceBuffer<gmx::RVec>>(nb->atdat->fshift);
 816 }
 817
 818 /* Initialization for X buffer operations on GPU. */
 819 /* TODO  Remove explicit pinning from host arrays from here and manage in a more natural way*/
 820 void nbnxn_gpu_init_x_to_nbat_x(const Nbnxm::GridSet& gridSet, NbnxmGpu* gpu_nbv)
 821 {
 822     const DeviceStream& deviceStream  = *gpu_nbv->deviceStreams[InteractionLocality::Local];
 823     bool                bDoTime       = gpu_nbv->bDoTime;
 824     const int           maxNumColumns = gridSet.numColumnsMax();
 825
 826     reallocateDeviceBuffer(&gpu_nbv->cxy_na, maxNumColumns * gridSet.grids().size(),
 827                            &gpu_nbv->ncxy_na, &gpu_nbv->ncxy_na_alloc, *gpu_nbv->deviceContext_);
 828     reallocateDeviceBuffer(&gpu_nbv->cxy_ind, maxNumColumns * gridSet.grids().size(),
 829                            &gpu_nbv->ncxy_ind, &gpu_nbv->ncxy_ind_alloc, *gpu_nbv->deviceContext_);
 830
 831     for (unsigned int g = 0; g < gridSet.grids().size(); g++)
 832     {
 833
 834         const Nbnxm::Grid& grid = gridSet.grids()[g];
 835
 836         const int  numColumns      = grid.numColumns();
 837         const int* atomIndices     = gridSet.atomIndices().data();
 838         const int  atomIndicesSize = gridSet.atomIndices().size();
 839         const int* cxy_na          = grid.cxy_na().data();
 840         const int* cxy_ind         = grid.cxy_ind().data();
 841
 842         reallocateDeviceBuffer(&gpu_nbv->atomIndices, atomIndicesSize, &gpu_nbv->atomIndicesSize,
 843                                &gpu_nbv->atomIndicesSize_alloc, *gpu_nbv->deviceContext_);
 844
 845         if (atomIndicesSize > 0)
 846         {
 847
 848             if (bDoTime)
 849             {
 850                 gpu_nbv->timers->xf[AtomLocality::Local].nb_h2d.openTimingRegion(deviceStream);
 851             }
 852
 853             copyToDeviceBuffer(&gpu_nbv->atomIndices, atomIndices, 0, atomIndicesSize, deviceStream,
 854                                GpuApiCallBehavior::Async, nullptr);
 855
 856             if (bDoTime)
 857             {
 858                 gpu_nbv->timers->xf[AtomLocality::Local].nb_h2d.closeTimingRegion(deviceStream);
 859             }
 860         }
 861
 862         if (numColumns > 0)
 863         {
 864             if (bDoTime)
 865             {
 866                 gpu_nbv->timers->xf[AtomLocality::Local].nb_h2d.openTimingRegion(deviceStream);
 867             }
 868
 869             int* destPtr = &gpu_nbv->cxy_na[maxNumColumns * g];
 870             copyToDeviceBuffer(&destPtr, cxy_na, 0, numColumns, deviceStream,
 871                                GpuApiCallBehavior::Async, nullptr);
 872
 873             if (bDoTime)
 874             {
 875                 gpu_nbv->timers->xf[AtomLocality::Local].nb_h2d.closeTimingRegion(deviceStream);
 876             }
 877
 878             if (bDoTime)
 879             {
 880                 gpu_nbv->timers->xf[AtomLocality::Local].nb_h2d.openTimingRegion(deviceStream);
 881             }
 882
 883             destPtr = &gpu_nbv->cxy_ind[maxNumColumns * g];
 884             copyToDeviceBuffer(&destPtr, cxy_ind, 0, numColumns, deviceStream,
 885                                GpuApiCallBehavior::Async, nullptr);
 886
 887             if (bDoTime)
 888             {
 889                 gpu_nbv->timers->xf[AtomLocality::Local].nb_h2d.closeTimingRegion(deviceStream);
 890             }
 891         }
 892     }
 893
 894     // The above data is transferred on the local stream but is a
 895     // dependency of the nonlocal stream (specifically the nonlocal X
 896     // buf ops kernel).  We therefore set a dependency to ensure
 897     // that the nonlocal stream waits on the local stream here.
 898     // This call records an event in the local stream:
 899     nbnxnInsertNonlocalGpuDependency(gpu_nbv, Nbnxm::InteractionLocality::Local);
 900     // ...and this call instructs the nonlocal stream to wait on that event:
 901     nbnxnInsertNonlocalGpuDependency(gpu_nbv, Nbnxm::InteractionLocality::NonLocal);
 902
 903     return;
 904 }
 905
 906 /* Initialization for F buffer operations on GPU. */
 907 void nbnxn_gpu_init_add_nbat_f_to_f(const int*                  cell,
 908                                     NbnxmGpu*                   gpu_nbv,
 909                                     int                         natoms_total,
 910                                     GpuEventSynchronizer* const localReductionDone)
 911 {
 912
 913     const DeviceStream& deviceStream = *gpu_nbv->deviceStreams[InteractionLocality::Local];
 914
 915     GMX_ASSERT(localReductionDone, "localReductionDone should be a valid pointer");
 916     gpu_nbv->localFReductionDone = localReductionDone;
 917
 918     if (natoms_total > 0)
 919     {
 920         reallocateDeviceBuffer(&gpu_nbv->cell, natoms_total, &gpu_nbv->ncell, &gpu_nbv->ncell_alloc,
 921                                *gpu_nbv->deviceContext_);
 922         copyToDeviceBuffer(&gpu_nbv->cell, cell, 0, natoms_total, deviceStream,
 923                            GpuApiCallBehavior::Async, nullptr);
 924     }
 925
 926     return;
 927 }
 928
 929 } // namespace Nbnxm