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