c711136b470a9577da0c7e8ee460dc73ec399f51
[alexxy/gromacs.git] / 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,2017,2018,2019, by the GROMACS development team, led by
5  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
6  * and including many others, as listed in the AUTHORS file in the
7  * top-level source directory and at http://www.gromacs.org.
8  *
9  * GROMACS is free software; you can redistribute it and/or
10  * modify it under the terms of the GNU Lesser General Public License
11  * as published by the Free Software Foundation; either version 2.1
12  * of the License, or (at your option) any later version.
13  *
14  * GROMACS is distributed in the hope that it will be useful,
15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
17  * Lesser General Public License for more details.
18  *
19  * You should have received a copy of the GNU Lesser General Public
20  * License along with GROMACS; if not, see
21  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
22  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
23  *
24  * If you want to redistribute modifications to GROMACS, please
25  * consider that scientific software is very special. Version
26  * control is crucial - bugs must be traceable. We will be happy to
27  * consider code for inclusion in the official distribution, but
28  * derived work must not be called official GROMACS. Details are found
29  * in the README & COPYING files - if they are missing, get the
30  * official version at http://www.gromacs.org.
31  *
32  * To help us fund GROMACS development, we humbly ask that you cite
33  * the research papers on the package. Check out http://www.gromacs.org.
34  */
35 /*! \file
36  *  \brief Define CUDA implementation of nbnxn_gpu_data_mgmt.h
37  *
38  *  \author Szilard Pall <pall.szilard@gmail.com>
39  */
40 #include "gmxpre.h"
41
42 #include <assert.h>
43 #include <stdarg.h>
44 #include <stdio.h>
45 #include <stdlib.h>
46
47 // TODO We would like to move this down, but the way gmx_nbnxn_gpu_t
48 //      is currently declared means this has to be before gpu_types.h
49 #include "nbnxm_cuda_types.h"
50
51 // TODO Remove this comment when the above order issue is resolved
52 #include "gromacs/gpu_utils/cudautils.cuh"
53 #include "gromacs/gpu_utils/gpu_utils.h"
54 #include "gromacs/gpu_utils/pmalloc_cuda.h"
55 #include "gromacs/hardware/gpu_hw_info.h"
56 #include "gromacs/math/vectypes.h"
57 #include "gromacs/mdlib/force_flags.h"
58 #include "gromacs/mdtypes/interaction_const.h"
59 #include "gromacs/mdtypes/md_enums.h"
60 #include "gromacs/nbnxm/gpu_data_mgmt.h"
61 #include "gromacs/nbnxm/nbnxm.h"
62 #include "gromacs/pbcutil/ishift.h"
63 #include "gromacs/timing/gpu_timing.h"
64 #include "gromacs/utility/basedefinitions.h"
65 #include "gromacs/utility/cstringutil.h"
66 #include "gromacs/utility/fatalerror.h"
67 #include "gromacs/utility/real.h"
68 #include "gromacs/utility/smalloc.h"
69
70 #include "nbnxm_cuda.h"
71
72 namespace Nbnxm
73 {
74
75 /* This is a heuristically determined parameter for the Kepler
76  * and Maxwell architectures for the minimum size of ci lists by multiplying
77  * this constant with the # of multiprocessors on the current device.
78  * Since the maximum number of blocks per multiprocessor is 16, the ideal
79  * count for small systems is 32 or 48 blocks per multiprocessor. Because
80  * there is a bit of fluctuations in the generated block counts, we use
81  * a target of 44 instead of the ideal value of 48.
82  */
83 static unsigned int gpu_min_ci_balanced_factor = 44;
84
85 /* Fw. decl. */
86 static void nbnxn_cuda_clear_e_fshift(gmx_nbnxn_cuda_t *nb);
87
88 /* Fw. decl, */
89 static void nbnxn_cuda_free_nbparam_table(cu_nbparam_t            *nbparam);
90
91 /*! \brief Return whether combination rules are used.
92  *
93  * \param[in]   pointer to nonbonded paramter struct
94  * \return      true if combination rules are used in this run, false otherwise
95  */
96 static inline bool useLjCombRule(const cu_nbparam_t  *nbparam)
97 {
98     return (nbparam->vdwtype == evdwCuCUTCOMBGEOM ||
99             nbparam->vdwtype == evdwCuCUTCOMBLB);
100 }
101
102 /*! \brief Initialized the Ewald Coulomb correction GPU table.
103
104     Tabulates the Ewald Coulomb force and initializes the size/scale
105     and the table GPU array. If called with an already allocated table,
106     it just re-uploads the table.
107  */
108 static void init_ewald_coulomb_force_table(const interaction_const_t *ic,
109                                            cu_nbparam_t              *nbp)
110 {
111     if (nbp->coulomb_tab != nullptr)
112     {
113         nbnxn_cuda_free_nbparam_table(nbp);
114     }
115
116     nbp->coulomb_tab_scale = ic->tabq_scale;
117     initParamLookupTable(nbp->coulomb_tab, nbp->coulomb_tab_texobj,
118                          ic->tabq_coul_F, ic->tabq_size);
119 }
120
121
122 /*! Initializes the atomdata structure first time, it only gets filled at
123     pair-search. */
124 static void init_atomdata_first(cu_atomdata_t *ad, int ntypes)
125 {
126     cudaError_t stat;
127
128     ad->ntypes  = ntypes;
129     stat        = cudaMalloc((void**)&ad->shift_vec, SHIFTS*sizeof(*ad->shift_vec));
130     CU_RET_ERR(stat, "cudaMalloc failed on ad->shift_vec");
131     ad->bShiftVecUploaded = false;
132
133     stat = cudaMalloc((void**)&ad->fshift, SHIFTS*sizeof(*ad->fshift));
134     CU_RET_ERR(stat, "cudaMalloc failed on ad->fshift");
135
136     stat = cudaMalloc((void**)&ad->e_lj, sizeof(*ad->e_lj));
137     CU_RET_ERR(stat, "cudaMalloc failed on ad->e_lj");
138     stat = cudaMalloc((void**)&ad->e_el, sizeof(*ad->e_el));
139     CU_RET_ERR(stat, "cudaMalloc failed on ad->e_el");
140
141     /* initialize to nullptr poiters to data that is not allocated here and will
142        need reallocation in nbnxn_cuda_init_atomdata */
143     ad->xq = nullptr;
144     ad->f  = nullptr;
145
146     /* size -1 indicates that the respective array hasn't been initialized yet */
147     ad->natoms = -1;
148     ad->nalloc = -1;
149 }
150
151 /*! Selects the Ewald kernel type, analytical on SM 3.0 and later, tabulated on
152     earlier GPUs, single or twin cut-off. */
153 static int pick_ewald_kernel_type(bool                     bTwinCut)
154 {
155     bool bUseAnalyticalEwald, bForceAnalyticalEwald, bForceTabulatedEwald;
156     int  kernel_type;
157
158     /* Benchmarking/development environment variables to force the use of
159        analytical or tabulated Ewald kernel. */
160     bForceAnalyticalEwald = (getenv("GMX_CUDA_NB_ANA_EWALD") != nullptr);
161     bForceTabulatedEwald  = (getenv("GMX_CUDA_NB_TAB_EWALD") != nullptr);
162
163     if (bForceAnalyticalEwald && bForceTabulatedEwald)
164     {
165         gmx_incons("Both analytical and tabulated Ewald CUDA non-bonded kernels "
166                    "requested through environment variables.");
167     }
168
169     /* By default use analytical Ewald. */
170     bUseAnalyticalEwald = true;
171     if (bForceAnalyticalEwald)
172     {
173         if (debug)
174         {
175             fprintf(debug, "Using analytical Ewald CUDA kernels\n");
176         }
177     }
178     else if (bForceTabulatedEwald)
179     {
180         bUseAnalyticalEwald = false;
181
182         if (debug)
183         {
184             fprintf(debug, "Using tabulated Ewald CUDA kernels\n");
185         }
186     }
187
188     /* Use twin cut-off kernels if requested by bTwinCut or the env. var.
189        forces it (use it for debugging/benchmarking only). */
190     if (!bTwinCut && (getenv("GMX_CUDA_NB_EWALD_TWINCUT") == nullptr))
191     {
192         kernel_type = bUseAnalyticalEwald ? eelCuEWALD_ANA : eelCuEWALD_TAB;
193     }
194     else
195     {
196         kernel_type = bUseAnalyticalEwald ? eelCuEWALD_ANA_TWIN : eelCuEWALD_TAB_TWIN;
197     }
198
199     return kernel_type;
200 }
201
202 /*! Copies all parameters related to the cut-off from ic to nbp */
203 static void set_cutoff_parameters(cu_nbparam_t              *nbp,
204                                   const interaction_const_t *ic,
205                                   const NbnxnListParameters *listParams)
206 {
207     nbp->ewald_beta        = ic->ewaldcoeff_q;
208     nbp->sh_ewald          = ic->sh_ewald;
209     nbp->epsfac            = ic->epsfac;
210     nbp->two_k_rf          = 2.0 * ic->k_rf;
211     nbp->c_rf              = ic->c_rf;
212     nbp->rvdw_sq           = ic->rvdw * ic->rvdw;
213     nbp->rcoulomb_sq       = ic->rcoulomb * ic->rcoulomb;
214     nbp->rlistOuter_sq     = listParams->rlistOuter * listParams->rlistOuter;
215     nbp->rlistInner_sq     = listParams->rlistInner * listParams->rlistInner;
216     nbp->useDynamicPruning = listParams->useDynamicPruning;
217
218     nbp->sh_lj_ewald       = ic->sh_lj_ewald;
219     nbp->ewaldcoeff_lj     = ic->ewaldcoeff_lj;
220
221     nbp->rvdw_switch       = ic->rvdw_switch;
222     nbp->dispersion_shift  = ic->dispersion_shift;
223     nbp->repulsion_shift   = ic->repulsion_shift;
224     nbp->vdw_switch        = ic->vdw_switch;
225 }
226
227 /*! Initializes the nonbonded parameter data structure. */
228 static void init_nbparam(cu_nbparam_t                   *nbp,
229                          const interaction_const_t      *ic,
230                          const NbnxnListParameters      *listParams,
231                          const nbnxn_atomdata_t::Params &nbatParams)
232 {
233     int         ntypes;
234
235     ntypes  = nbatParams.numTypes;
236
237     set_cutoff_parameters(nbp, ic, listParams);
238
239     /* The kernel code supports LJ combination rules (geometric and LB) for
240      * all kernel types, but we only generate useful combination rule kernels.
241      * We currently only use LJ combination rule (geometric and LB) kernels
242      * for plain cut-off LJ. On Maxwell the force only kernels speed up 15%
243      * with PME and 20% with RF, the other kernels speed up about half as much.
244      * For LJ force-switch the geometric rule would give 7% speed-up, but this
245      * combination is rarely used. LJ force-switch with LB rule is more common,
246      * but gives only 1% speed-up.
247      */
248     if (ic->vdwtype == evdwCUT)
249     {
250         switch (ic->vdw_modifier)
251         {
252             case eintmodNONE:
253             case eintmodPOTSHIFT:
254                 switch (nbatParams.comb_rule)
255                 {
256                     case ljcrNONE:
257                         nbp->vdwtype = evdwCuCUT;
258                         break;
259                     case ljcrGEOM:
260                         nbp->vdwtype = evdwCuCUTCOMBGEOM;
261                         break;
262                     case ljcrLB:
263                         nbp->vdwtype = evdwCuCUTCOMBLB;
264                         break;
265                     default:
266                         gmx_incons("The requested LJ combination rule is not implemented in the CUDA GPU accelerated kernels!");
267                 }
268                 break;
269             case eintmodFORCESWITCH:
270                 nbp->vdwtype = evdwCuFSWITCH;
271                 break;
272             case eintmodPOTSWITCH:
273                 nbp->vdwtype = evdwCuPSWITCH;
274                 break;
275             default:
276                 gmx_incons("The requested VdW interaction modifier is not implemented in the CUDA GPU accelerated kernels!");
277         }
278     }
279     else if (ic->vdwtype == evdwPME)
280     {
281         if (ic->ljpme_comb_rule == ljcrGEOM)
282         {
283             assert(nbatParams.comb_rule == ljcrGEOM);
284             nbp->vdwtype = evdwCuEWALDGEOM;
285         }
286         else
287         {
288             assert(nbatParams.comb_rule == ljcrLB);
289             nbp->vdwtype = evdwCuEWALDLB;
290         }
291     }
292     else
293     {
294         gmx_incons("The requested VdW type is not implemented in the CUDA GPU accelerated kernels!");
295     }
296
297     if (ic->eeltype == eelCUT)
298     {
299         nbp->eeltype = eelCuCUT;
300     }
301     else if (EEL_RF(ic->eeltype))
302     {
303         nbp->eeltype = eelCuRF;
304     }
305     else if ((EEL_PME(ic->eeltype) || ic->eeltype == eelEWALD))
306     {
307         /* Initially rcoulomb == rvdw, so it's surely not twin cut-off. */
308         nbp->eeltype = pick_ewald_kernel_type(false);
309     }
310     else
311     {
312         /* Shouldn't happen, as this is checked when choosing Verlet-scheme */
313         gmx_incons("The requested electrostatics type is not implemented in the CUDA GPU accelerated kernels!");
314     }
315
316     /* generate table for PME */
317     nbp->coulomb_tab = nullptr;
318     if (nbp->eeltype == eelCuEWALD_TAB || nbp->eeltype == eelCuEWALD_TAB_TWIN)
319     {
320         init_ewald_coulomb_force_table(ic, nbp);
321     }
322
323     /* set up LJ parameter lookup table */
324     if (!useLjCombRule(nbp))
325     {
326         initParamLookupTable(nbp->nbfp, nbp->nbfp_texobj,
327                              nbatParams.nbfp.data(), 2*ntypes*ntypes);
328     }
329
330     /* set up LJ-PME parameter lookup table */
331     if (ic->vdwtype == evdwPME)
332     {
333         initParamLookupTable(nbp->nbfp_comb, nbp->nbfp_comb_texobj,
334                              nbatParams.nbfp_comb.data(), 2*ntypes);
335     }
336 }
337
338 /*! Re-generate the GPU Ewald force table, resets rlist, and update the
339  *  electrostatic type switching to twin cut-off (or back) if needed. */
340 void gpu_pme_loadbal_update_param(const nonbonded_verlet_t    *nbv,
341                                   const interaction_const_t   *ic,
342                                   const NbnxnListParameters   *listParams)
343 {
344     if (!nbv || !nbv->useGpu())
345     {
346         return;
347     }
348     gmx_nbnxn_cuda_t *nb    = nbv->gpu_nbv;
349     cu_nbparam_t     *nbp   = nb->nbparam;
350
351     set_cutoff_parameters(nbp, ic, listParams);
352
353     nbp->eeltype        = pick_ewald_kernel_type(ic->rcoulomb != ic->rvdw);
354
355     init_ewald_coulomb_force_table(ic, nb->nbparam);
356 }
357
358 /*! Initializes the pair list data structure. */
359 static void init_plist(cu_plist_t *pl)
360 {
361     /* initialize to nullptr pointers to data that is not allocated here and will
362        need reallocation in nbnxn_gpu_init_pairlist */
363     pl->sci      = nullptr;
364     pl->cj4      = nullptr;
365     pl->imask    = nullptr;
366     pl->excl     = nullptr;
367
368     /* size -1 indicates that the respective array hasn't been initialized yet */
369     pl->na_c           = -1;
370     pl->nsci           = -1;
371     pl->sci_nalloc     = -1;
372     pl->ncj4           = -1;
373     pl->cj4_nalloc     = -1;
374     pl->nimask         = -1;
375     pl->imask_nalloc   = -1;
376     pl->nexcl          = -1;
377     pl->excl_nalloc    = -1;
378     pl->haveFreshList  = false;
379 }
380
381 /*! Initializes the timings data structure. */
382 static void init_timings(gmx_wallclock_gpu_nbnxn_t *t)
383 {
384     int i, j;
385
386     t->nb_h2d_t = 0.0;
387     t->nb_d2h_t = 0.0;
388     t->nb_c     = 0;
389     t->pl_h2d_t = 0.0;
390     t->pl_h2d_c = 0;
391     for (i = 0; i < 2; i++)
392     {
393         for (j = 0; j < 2; j++)
394         {
395             t->ktime[i][j].t = 0.0;
396             t->ktime[i][j].c = 0;
397         }
398     }
399     t->pruneTime.c        = 0;
400     t->pruneTime.t        = 0.0;
401     t->dynamicPruneTime.c = 0;
402     t->dynamicPruneTime.t = 0.0;
403 }
404
405 /*! Initializes simulation constant data. */
406 static void cuda_init_const(gmx_nbnxn_cuda_t               *nb,
407                             const interaction_const_t      *ic,
408                             const NbnxnListParameters      *listParams,
409                             const nbnxn_atomdata_t::Params &nbatParams)
410 {
411     init_atomdata_first(nb->atdat, nbatParams.numTypes);
412     init_nbparam(nb->nbparam, ic, listParams, nbatParams);
413
414     /* clear energy and shift force outputs */
415     nbnxn_cuda_clear_e_fshift(nb);
416 }
417
418 void gpu_init(gmx_nbnxn_cuda_t         **p_nb,
419               const gmx_device_info_t   *deviceInfo,
420               const interaction_const_t *ic,
421               const NbnxnListParameters *listParams,
422               const nbnxn_atomdata_t    *nbat,
423               int                        /*rank*/,
424               gmx_bool                   bLocalAndNonlocal)
425 {
426     cudaError_t       stat;
427     gmx_nbnxn_cuda_t *nb;
428
429     if (p_nb == nullptr)
430     {
431         return;
432     }
433
434     snew(nb, 1);
435     snew(nb->atdat, 1);
436     snew(nb->nbparam, 1);
437     snew(nb->plist[InteractionLocality::Local], 1);
438     if (bLocalAndNonlocal)
439     {
440         snew(nb->plist[InteractionLocality::NonLocal], 1);
441     }
442
443     nb->bUseTwoStreams = bLocalAndNonlocal;
444
445     nb->timers = new cu_timers_t();
446     snew(nb->timings, 1);
447
448     /* init nbst */
449     pmalloc((void**)&nb->nbst.e_lj, sizeof(*nb->nbst.e_lj));
450     pmalloc((void**)&nb->nbst.e_el, sizeof(*nb->nbst.e_el));
451     pmalloc((void**)&nb->nbst.fshift, SHIFTS * sizeof(*nb->nbst.fshift));
452
453     init_plist(nb->plist[InteractionLocality::Local]);
454
455     /* set device info, just point it to the right GPU among the detected ones */
456     nb->dev_info = deviceInfo;
457
458     /* local/non-local GPU streams */
459     stat = cudaStreamCreate(&nb->stream[InteractionLocality::Local]);
460     CU_RET_ERR(stat, "cudaStreamCreate on stream[InterationLocality::Local] failed");
461     if (nb->bUseTwoStreams)
462     {
463         init_plist(nb->plist[InteractionLocality::NonLocal]);
464
465         /* Note that the device we're running on does not have to support
466          * priorities, because we are querying the priority range which in this
467          * case will be a single value.
468          */
469         int highest_priority;
470         stat = cudaDeviceGetStreamPriorityRange(nullptr, &highest_priority);
471         CU_RET_ERR(stat, "cudaDeviceGetStreamPriorityRange failed");
472
473         stat = cudaStreamCreateWithPriority(&nb->stream[InteractionLocality::NonLocal],
474                                             cudaStreamDefault,
475                                             highest_priority);
476         CU_RET_ERR(stat, "cudaStreamCreateWithPriority on stream[InteractionLocality::NonLocal] failed");
477     }
478
479     /* init events for sychronization (timing disabled for performance reasons!) */
480     stat = cudaEventCreateWithFlags(&nb->nonlocal_done, cudaEventDisableTiming);
481     CU_RET_ERR(stat, "cudaEventCreate on nonlocal_done failed");
482     stat = cudaEventCreateWithFlags(&nb->misc_ops_and_local_H2D_done, cudaEventDisableTiming);
483     CU_RET_ERR(stat, "cudaEventCreate on misc_ops_and_local_H2D_done failed");
484
485     /* WARNING: CUDA timings are incorrect with multiple streams.
486      *          This is the main reason why they are disabled by default.
487      */
488     // TODO: Consider turning on by default when we can detect nr of streams.
489     nb->bDoTime = (getenv("GMX_ENABLE_GPU_TIMING") != nullptr);
490
491     if (nb->bDoTime)
492     {
493         init_timings(nb->timings);
494     }
495
496     /* set the kernel type for the current GPU */
497     /* pick L1 cache configuration */
498     cuda_set_cacheconfig();
499
500     cuda_init_const(nb, ic, listParams, nbat->params());
501
502     *p_nb = nb;
503
504     if (debug)
505     {
506         fprintf(debug, "Initialized CUDA data structures.\n");
507     }
508 }
509
510 void gpu_init_pairlist(gmx_nbnxn_cuda_t          *nb,
511                        const NbnxnPairlistGpu    *h_plist,
512                        const InteractionLocality  iloc)
513 {
514     char          sbuf[STRLEN];
515     bool          bDoTime    =  (nb->bDoTime && !h_plist->sci.empty());
516     cudaStream_t  stream     = nb->stream[iloc];
517     cu_plist_t   *d_plist    = nb->plist[iloc];
518
519     if (d_plist->na_c < 0)
520     {
521         d_plist->na_c = h_plist->na_ci;
522     }
523     else
524     {
525         if (d_plist->na_c != h_plist->na_ci)
526         {
527             sprintf(sbuf, "In cu_init_plist: the #atoms per cell has changed (from %d to %d)",
528                     d_plist->na_c, h_plist->na_ci);
529             gmx_incons(sbuf);
530         }
531     }
532
533     gpu_timers_t::Interaction &iTimers = nb->timers->interaction[iloc];
534
535     if (bDoTime)
536     {
537         iTimers.pl_h2d.openTimingRegion(stream);
538         iTimers.didPairlistH2D = true;
539     }
540
541     Context context = nullptr;
542
543     reallocateDeviceBuffer(&d_plist->sci, h_plist->sci.size(),
544                            &d_plist->nsci, &d_plist->sci_nalloc, context);
545     copyToDeviceBuffer(&d_plist->sci, h_plist->sci.data(), 0, h_plist->sci.size(),
546                        stream, GpuApiCallBehavior::Async,
547                        bDoTime ? iTimers.pl_h2d.fetchNextEvent() : nullptr);
548
549     reallocateDeviceBuffer(&d_plist->cj4, h_plist->cj4.size(),
550                            &d_plist->ncj4, &d_plist->cj4_nalloc, context);
551     copyToDeviceBuffer(&d_plist->cj4, h_plist->cj4.data(), 0, h_plist->cj4.size(),
552                        stream, GpuApiCallBehavior::Async,
553                        bDoTime ? iTimers.pl_h2d.fetchNextEvent() : nullptr);
554
555     reallocateDeviceBuffer(&d_plist->imask, h_plist->cj4.size()*c_nbnxnGpuClusterpairSplit,
556                            &d_plist->nimask, &d_plist->imask_nalloc, context);
557
558     reallocateDeviceBuffer(&d_plist->excl, h_plist->excl.size(),
559                            &d_plist->nexcl, &d_plist->excl_nalloc, context);
560     copyToDeviceBuffer(&d_plist->excl, h_plist->excl.data(), 0, h_plist->excl.size(),
561                        stream, GpuApiCallBehavior::Async,
562                        bDoTime ? iTimers.pl_h2d.fetchNextEvent() : nullptr);
563
564     if (bDoTime)
565     {
566         iTimers.pl_h2d.closeTimingRegion(stream);
567     }
568
569     /* the next use of thist list we be the first one, so we need to prune */
570     d_plist->haveFreshList = true;
571 }
572
573 void gpu_upload_shiftvec(gmx_nbnxn_cuda_t       *nb,
574                          const nbnxn_atomdata_t *nbatom)
575 {
576     cu_atomdata_t *adat  = nb->atdat;
577     cudaStream_t   ls    = nb->stream[InteractionLocality::Local];
578
579     /* only if we have a dynamic box */
580     if (nbatom->bDynamicBox || !adat->bShiftVecUploaded)
581     {
582         cu_copy_H2D_async(adat->shift_vec, nbatom->shift_vec.data(),
583                           SHIFTS * sizeof(*adat->shift_vec), ls);
584         adat->bShiftVecUploaded = true;
585     }
586 }
587
588 /*! Clears the first natoms_clear elements of the GPU nonbonded force output array. */
589 static void nbnxn_cuda_clear_f(gmx_nbnxn_cuda_t *nb, int natoms_clear)
590 {
591     cudaError_t    stat;
592     cu_atomdata_t *adat  = nb->atdat;
593     cudaStream_t   ls    = nb->stream[InteractionLocality::Local];
594
595     stat = cudaMemsetAsync(adat->f, 0, natoms_clear * sizeof(*adat->f), ls);
596     CU_RET_ERR(stat, "cudaMemsetAsync on f falied");
597 }
598
599 /*! Clears nonbonded shift force output array and energy outputs on the GPU. */
600 static void nbnxn_cuda_clear_e_fshift(gmx_nbnxn_cuda_t *nb)
601 {
602     cudaError_t    stat;
603     cu_atomdata_t *adat  = nb->atdat;
604     cudaStream_t   ls    = nb->stream[InteractionLocality::Local];
605
606     stat = cudaMemsetAsync(adat->fshift, 0, SHIFTS * sizeof(*adat->fshift), ls);
607     CU_RET_ERR(stat, "cudaMemsetAsync on fshift falied");
608     stat = cudaMemsetAsync(adat->e_lj, 0, sizeof(*adat->e_lj), ls);
609     CU_RET_ERR(stat, "cudaMemsetAsync on e_lj falied");
610     stat = cudaMemsetAsync(adat->e_el, 0, sizeof(*adat->e_el), ls);
611     CU_RET_ERR(stat, "cudaMemsetAsync on e_el falied");
612 }
613
614 void gpu_clear_outputs(gmx_nbnxn_cuda_t *nb, int flags)
615 {
616     nbnxn_cuda_clear_f(nb, nb->atdat->natoms);
617     /* clear shift force array and energies if the outputs were
618        used in the current step */
619     if (flags & GMX_FORCE_VIRIAL)
620     {
621         nbnxn_cuda_clear_e_fshift(nb);
622     }
623 }
624
625 void gpu_init_atomdata(gmx_nbnxn_cuda_t       *nb,
626                        const nbnxn_atomdata_t *nbat)
627 {
628     cudaError_t    stat;
629     int            nalloc, natoms;
630     bool           realloced;
631     bool           bDoTime   = nb->bDoTime;
632     cu_timers_t   *timers    = nb->timers;
633     cu_atomdata_t *d_atdat   = nb->atdat;
634     cudaStream_t   ls        = nb->stream[InteractionLocality::Local];
635
636     natoms    = nbat->numAtoms();
637     realloced = false;
638
639     if (bDoTime)
640     {
641         /* time async copy */
642         timers->atdat.openTimingRegion(ls);
643     }
644
645     /* need to reallocate if we have to copy more atoms than the amount of space
646        available and only allocate if we haven't initialized yet, i.e d_atdat->natoms == -1 */
647     if (natoms > d_atdat->nalloc)
648     {
649         nalloc = over_alloc_small(natoms);
650
651         /* free up first if the arrays have already been initialized */
652         if (d_atdat->nalloc != -1)
653         {
654             freeDeviceBuffer(&d_atdat->f);
655             freeDeviceBuffer(&d_atdat->xq);
656             freeDeviceBuffer(&d_atdat->atom_types);
657             freeDeviceBuffer(&d_atdat->lj_comb);
658         }
659
660         stat = cudaMalloc((void **)&d_atdat->f, nalloc*sizeof(*d_atdat->f));
661         CU_RET_ERR(stat, "cudaMalloc failed on d_atdat->f");
662         stat = cudaMalloc((void **)&d_atdat->xq, nalloc*sizeof(*d_atdat->xq));
663         CU_RET_ERR(stat, "cudaMalloc failed on d_atdat->xq");
664         if (useLjCombRule(nb->nbparam))
665         {
666             stat = cudaMalloc((void **)&d_atdat->lj_comb, nalloc*sizeof(*d_atdat->lj_comb));
667             CU_RET_ERR(stat, "cudaMalloc failed on d_atdat->lj_comb");
668         }
669         else
670         {
671             stat = cudaMalloc((void **)&d_atdat->atom_types, nalloc*sizeof(*d_atdat->atom_types));
672             CU_RET_ERR(stat, "cudaMalloc failed on d_atdat->atom_types");
673         }
674
675         d_atdat->nalloc = nalloc;
676         realloced       = true;
677     }
678
679     d_atdat->natoms       = natoms;
680     d_atdat->natoms_local = nbat->natoms_local;
681
682     /* need to clear GPU f output if realloc happened */
683     if (realloced)
684     {
685         nbnxn_cuda_clear_f(nb, nalloc);
686     }
687
688     if (useLjCombRule(nb->nbparam))
689     {
690         cu_copy_H2D_async(d_atdat->lj_comb, nbat->params().lj_comb.data(),
691                           natoms*sizeof(*d_atdat->lj_comb), ls);
692     }
693     else
694     {
695         cu_copy_H2D_async(d_atdat->atom_types, nbat->params().type.data(),
696                           natoms*sizeof(*d_atdat->atom_types), ls);
697     }
698
699     if (bDoTime)
700     {
701         timers->atdat.closeTimingRegion(ls);
702     }
703 }
704
705 static void nbnxn_cuda_free_nbparam_table(cu_nbparam_t            *nbparam)
706 {
707     if (nbparam->eeltype == eelCuEWALD_TAB || nbparam->eeltype == eelCuEWALD_TAB_TWIN)
708     {
709         destroyParamLookupTable(nbparam->coulomb_tab, nbparam->coulomb_tab_texobj);
710     }
711 }
712
713 void gpu_free(gmx_nbnxn_cuda_t *nb)
714 {
715     cudaError_t      stat;
716     cu_atomdata_t   *atdat;
717     cu_nbparam_t    *nbparam;
718
719     if (nb == nullptr)
720     {
721         return;
722     }
723
724     atdat       = nb->atdat;
725     nbparam     = nb->nbparam;
726
727     nbnxn_cuda_free_nbparam_table(nbparam);
728
729     stat = cudaEventDestroy(nb->nonlocal_done);
730     CU_RET_ERR(stat, "cudaEventDestroy failed on timers->nonlocal_done");
731     stat = cudaEventDestroy(nb->misc_ops_and_local_H2D_done);
732     CU_RET_ERR(stat, "cudaEventDestroy failed on timers->misc_ops_and_local_H2D_done");
733
734     delete nb->timers;
735     if (nb->bDoTime)
736     {
737         /* The non-local counters/stream (second in the array) are needed only with DD. */
738         for (int i = 0; i <= (nb->bUseTwoStreams ? 1 : 0); i++)
739         {
740             stat = cudaStreamDestroy(nb->stream[i]);
741             CU_RET_ERR(stat, "cudaStreamDestroy failed on stream");
742         }
743     }
744
745     if (!useLjCombRule(nb->nbparam))
746     {
747         destroyParamLookupTable(nbparam->nbfp, nbparam->nbfp_texobj);
748
749     }
750
751     if (nbparam->vdwtype == evdwCuEWALDGEOM || nbparam->vdwtype == evdwCuEWALDLB)
752     {
753         destroyParamLookupTable(nbparam->nbfp_comb, nbparam->nbfp_comb_texobj);
754     }
755
756     stat = cudaFree(atdat->shift_vec);
757     CU_RET_ERR(stat, "cudaFree failed on atdat->shift_vec");
758     stat = cudaFree(atdat->fshift);
759     CU_RET_ERR(stat, "cudaFree failed on atdat->fshift");
760
761     stat = cudaFree(atdat->e_lj);
762     CU_RET_ERR(stat, "cudaFree failed on atdat->e_lj");
763     stat = cudaFree(atdat->e_el);
764     CU_RET_ERR(stat, "cudaFree failed on atdat->e_el");
765
766     freeDeviceBuffer(&atdat->f);
767     freeDeviceBuffer(&atdat->xq);
768     freeDeviceBuffer(&atdat->atom_types);
769     freeDeviceBuffer(&atdat->lj_comb);
770
771     /* Free plist */
772     auto *plist = nb->plist[InteractionLocality::Local];
773     freeDeviceBuffer(&plist->sci);
774     freeDeviceBuffer(&plist->cj4);
775     freeDeviceBuffer(&plist->imask);
776     freeDeviceBuffer(&plist->excl);
777     sfree(plist);
778     if (nb->bUseTwoStreams)
779     {
780         auto *plist_nl = nb->plist[InteractionLocality::NonLocal];
781         freeDeviceBuffer(&plist_nl->sci);
782         freeDeviceBuffer(&plist_nl->cj4);
783         freeDeviceBuffer(&plist_nl->imask);
784         freeDeviceBuffer(&plist_nl->excl);
785         sfree(plist_nl);
786     }
787
788     /* Free nbst */
789     pfree(nb->nbst.e_lj);
790     nb->nbst.e_lj = nullptr;
791
792     pfree(nb->nbst.e_el);
793     nb->nbst.e_el = nullptr;
794
795     pfree(nb->nbst.fshift);
796     nb->nbst.fshift = nullptr;
797
798     sfree(atdat);
799     sfree(nbparam);
800     sfree(nb->timings);
801     sfree(nb);
802
803     if (debug)
804     {
805         fprintf(debug, "Cleaned up CUDA data structures.\n");
806     }
807 }
808
809 //! This function is documented in the header file
810 gmx_wallclock_gpu_nbnxn_t *gpu_get_timings(gmx_nbnxn_cuda_t *nb)
811 {
812     return (nb != nullptr && nb->bDoTime) ? nb->timings : nullptr;
813 }
814
815 void gpu_reset_timings(nonbonded_verlet_t* nbv)
816 {
817     if (nbv->gpu_nbv && nbv->gpu_nbv->bDoTime)
818     {
819         init_timings(nbv->gpu_nbv->timings);
820     }
821 }
822
823 int gpu_min_ci_balanced(gmx_nbnxn_cuda_t *nb)
824 {
825     return nb != nullptr ?
826            gpu_min_ci_balanced_factor*nb->dev_info->prop.multiProcessorCount : 0;
827
828 }
829
830 gmx_bool gpu_is_kernel_ewald_analytical(const gmx_nbnxn_cuda_t *nb)
831 {
832     return ((nb->nbparam->eeltype == eelCuEWALD_ANA) ||
833             (nb->nbparam->eeltype == eelCuEWALD_ANA_TWIN));
834 }
835
836 void *gpu_get_command_stream(gmx_nbnxn_gpu_t           *nb,
837                              const InteractionLocality  iloc)
838 {
839     assert(nb);
840
841     return static_cast<void *>(&nb->stream[iloc]);
842 }
843
844 void *gpu_get_xq(gmx_nbnxn_gpu_t *nb)
845 {
846     assert(nb);
847
848     return static_cast<void *>(nb->atdat->xq);
849 }
850
851 void *gpu_get_f(gmx_nbnxn_gpu_t *nb)
852 {
853     assert(nb);
854
855     return static_cast<void *>(nb->atdat->f);
856 }
857
858 rvec *gpu_get_fshift(gmx_nbnxn_gpu_t *nb)
859 {
860     assert(nb);
861
862     return reinterpret_cast<rvec *>(nb->atdat->fshift);
863 }
864
865 } // namespace Nbnxm