Move around PairSearch code
[alexxy/gromacs.git] / src / gromacs / nbnxm / pairlist.cpp
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
36 #include "gmxpre.h"
37
38 #include "pairlist.h"
39
40 #include "config.h"
41
42 #include <cassert>
43 #include <cmath>
44 #include <cstring>
45
46 #include <algorithm>
47
48 #include "gromacs/domdec/domdec_struct.h"
49 #include "gromacs/gmxlib/nrnb.h"
50 #include "gromacs/math/functions.h"
51 #include "gromacs/math/utilities.h"
52 #include "gromacs/math/vec.h"
53 #include "gromacs/mdlib/gmx_omp_nthreads.h"
54 #include "gromacs/mdlib/ns.h"
55 #include "gromacs/mdtypes/group.h"
56 #include "gromacs/mdtypes/md_enums.h"
57 #include "gromacs/nbnxm/atomdata.h"
58 #include "gromacs/nbnxm/gpu_data_mgmt.h"
59 #include "gromacs/nbnxm/nbnxm.h"
60 #include "gromacs/nbnxm/nbnxm_geometry.h"
61 #include "gromacs/nbnxm/nbnxm_simd.h"
62 #include "gromacs/nbnxm/pairlistset.h"
63 #include "gromacs/pbcutil/ishift.h"
64 #include "gromacs/pbcutil/pbc.h"
65 #include "gromacs/simd/simd.h"
66 #include "gromacs/simd/vector_operations.h"
67 #include "gromacs/topology/block.h"
68 #include "gromacs/utility/exceptions.h"
69 #include "gromacs/utility/fatalerror.h"
70 #include "gromacs/utility/gmxomp.h"
71 #include "gromacs/utility/smalloc.h"
72
73 #include "gridset.h"
74 #include "pairlistwork.h"
75 #include "pairsearch.h"
76
77 using namespace gmx;                        // TODO: Remove when this file is moved into gmx namespace
78
79 using BoundingBox   = Nbnxm::BoundingBox;   // TODO: Remove when refactoring this file
80 using BoundingBox1D = Nbnxm::BoundingBox1D; // TODO: Remove when refactoring this file
81
82 using Grid          = Nbnxm::Grid;          // TODO: Remove when refactoring this file
83
84 // Convience alias for partial Nbnxn namespace usage
85 using InteractionLocality = Nbnxm::InteractionLocality;
86
87 /* We shift the i-particles backward for PBC.
88  * This leads to more conditionals than shifting forward.
89  * We do this to get more balanced pair lists.
90  */
91 constexpr bool c_pbcShiftBackward = true;
92
93 /* Layout for the nonbonded NxN pair lists */
94 enum class NbnxnLayout
95 {
96     NoSimd4x4, // i-cluster size 4, j-cluster size 4
97     Simd4xN,   // i-cluster size 4, j-cluster size SIMD width
98     Simd2xNN,  // i-cluster size 4, j-cluster size half SIMD width
99     Gpu8x8x8   // i-cluster size 8, j-cluster size 8 + super-clustering
100 };
101
102 #if GMX_SIMD
103 /* Returns the j-cluster size */
104 template <NbnxnLayout layout>
105 static constexpr int jClusterSize()
106 {
107     static_assert(layout == NbnxnLayout::NoSimd4x4 || layout == NbnxnLayout::Simd4xN || layout == NbnxnLayout::Simd2xNN, "Currently jClusterSize only supports CPU layouts");
108
109     return layout == NbnxnLayout::Simd4xN ? GMX_SIMD_REAL_WIDTH : (layout == NbnxnLayout::Simd2xNN ? GMX_SIMD_REAL_WIDTH/2 : c_nbnxnCpuIClusterSize);
110 }
111
112 /*! \brief Returns the j-cluster index given the i-cluster index.
113  *
114  * \tparam    jClusterSize      The number of atoms in a j-cluster
115  * \tparam    jSubClusterIndex  The j-sub-cluster index (0/1), used when size(j-cluster) < size(i-cluster)
116  * \param[in] ci                The i-cluster index
117  */
118 template <int jClusterSize, int jSubClusterIndex>
119 static inline int cjFromCi(int ci)
120 {
121     static_assert(jClusterSize == c_nbnxnCpuIClusterSize/2 || jClusterSize == c_nbnxnCpuIClusterSize || jClusterSize == c_nbnxnCpuIClusterSize*2, "Only j-cluster sizes 2, 4 and 8 are currently implemented");
122
123     static_assert(jSubClusterIndex == 0 || jSubClusterIndex == 1,
124                   "Only sub-cluster indices 0 and 1 are supported");
125
126     if (jClusterSize == c_nbnxnCpuIClusterSize/2)
127     {
128         if (jSubClusterIndex == 0)
129         {
130             return ci << 1;
131         }
132         else
133         {
134             return ((ci + 1) << 1) - 1;
135         }
136     }
137     else if (jClusterSize == c_nbnxnCpuIClusterSize)
138     {
139         return ci;
140     }
141     else
142     {
143         return ci >> 1;
144     }
145 }
146
147 /*! \brief Returns the j-cluster index given the i-cluster index.
148  *
149  * \tparam    layout            The pair-list layout
150  * \tparam    jSubClusterIndex  The j-sub-cluster index (0/1), used when size(j-cluster) < size(i-cluster)
151  * \param[in] ci                The i-cluster index
152  */
153 template <NbnxnLayout layout, int jSubClusterIndex>
154 static inline int cjFromCi(int ci)
155 {
156     constexpr int clusterSize = jClusterSize<layout>();
157
158     return cjFromCi<clusterSize, jSubClusterIndex>(ci);
159 }
160
161 /* Returns the nbnxn coordinate data index given the i-cluster index */
162 template <NbnxnLayout layout>
163 static inline int xIndexFromCi(int ci)
164 {
165     constexpr int clusterSize = jClusterSize<layout>();
166
167     static_assert(clusterSize == c_nbnxnCpuIClusterSize/2 || clusterSize == c_nbnxnCpuIClusterSize || clusterSize == c_nbnxnCpuIClusterSize*2, "Only j-cluster sizes 2, 4 and 8 are currently implemented");
168
169     if (clusterSize <= c_nbnxnCpuIClusterSize)
170     {
171         /* Coordinates are stored packed in groups of 4 */
172         return ci*STRIDE_P4;
173     }
174     else
175     {
176         /* Coordinates packed in 8, i-cluster size is half the packing width */
177         return (ci >> 1)*STRIDE_P8 + (ci & 1)*(c_packX8 >> 1);
178     }
179 }
180
181 /* Returns the nbnxn coordinate data index given the j-cluster index */
182 template <NbnxnLayout layout>
183 static inline int xIndexFromCj(int cj)
184 {
185     constexpr int clusterSize = jClusterSize<layout>();
186
187     static_assert(clusterSize == c_nbnxnCpuIClusterSize/2 || clusterSize == c_nbnxnCpuIClusterSize || clusterSize == c_nbnxnCpuIClusterSize*2, "Only j-cluster sizes 2, 4 and 8 are currently implemented");
188
189     if (clusterSize == c_nbnxnCpuIClusterSize/2)
190     {
191         /* Coordinates are stored packed in groups of 4 */
192         return (cj >> 1)*STRIDE_P4 + (cj & 1)*(c_packX4 >> 1);
193     }
194     else if (clusterSize == c_nbnxnCpuIClusterSize)
195     {
196         /* Coordinates are stored packed in groups of 4 */
197         return cj*STRIDE_P4;
198     }
199     else
200     {
201         /* Coordinates are stored packed in groups of 8 */
202         return cj*STRIDE_P8;
203     }
204 }
205 #endif //GMX_SIMD
206
207
208 void nbnxn_init_pairlist_fep(t_nblist *nl)
209 {
210     nl->type        = GMX_NBLIST_INTERACTION_FREE_ENERGY;
211     nl->igeometry   = GMX_NBLIST_GEOMETRY_PARTICLE_PARTICLE;
212     /* The interaction functions are set in the free energy kernel fuction */
213     nl->ivdw        = -1;
214     nl->ivdwmod     = -1;
215     nl->ielec       = -1;
216     nl->ielecmod    = -1;
217
218     nl->maxnri      = 0;
219     nl->maxnrj      = 0;
220     nl->nri         = 0;
221     nl->nrj         = 0;
222     nl->iinr        = nullptr;
223     nl->gid         = nullptr;
224     nl->shift       = nullptr;
225     nl->jindex      = nullptr;
226     nl->jjnr        = nullptr;
227     nl->excl_fep    = nullptr;
228
229 }
230
231 static void init_buffer_flags(nbnxn_buffer_flags_t *flags,
232                               int                   natoms)
233 {
234     flags->nflag = (natoms + NBNXN_BUFFERFLAG_SIZE - 1)/NBNXN_BUFFERFLAG_SIZE;
235     if (flags->nflag > flags->flag_nalloc)
236     {
237         flags->flag_nalloc = over_alloc_large(flags->nflag);
238         srenew(flags->flag, flags->flag_nalloc);
239     }
240     for (int b = 0; b < flags->nflag; b++)
241     {
242         bitmask_clear(&(flags->flag[b]));
243     }
244 }
245
246 /* Returns the pair-list cutoff between a bounding box and a grid cell given an atom-to-atom pair-list cutoff
247  *
248  * Given a cutoff distance between atoms, this functions returns the cutoff
249  * distance2 between a bounding box of a group of atoms and a grid cell.
250  * Since atoms can be geometrically outside of the cell they have been
251  * assigned to (when atom groups instead of individual atoms are assigned
252  * to cells), this distance returned can be larger than the input.
253  */
254 static real
255 listRangeForBoundingBoxToGridCell(real                    rlist,
256                                   const Grid::Dimensions &gridDims)
257 {
258     return rlist + gridDims.maxAtomGroupRadius;
259
260 }
261 /* Returns the pair-list cutoff between a grid cells given an atom-to-atom pair-list cutoff
262  *
263  * Given a cutoff distance between atoms, this functions returns the cutoff
264  * distance2 between two grid cells.
265  * Since atoms can be geometrically outside of the cell they have been
266  * assigned to (when atom groups instead of individual atoms are assigned
267  * to cells), this distance returned can be larger than the input.
268  */
269 static real
270 listRangeForGridCellToGridCell(real                    rlist,
271                                const Grid::Dimensions &iGridDims,
272                                const Grid::Dimensions &jGridDims)
273 {
274     return rlist + iGridDims.maxAtomGroupRadius + jGridDims.maxAtomGroupRadius;
275 }
276
277 /* Determines the cell range along one dimension that
278  * the bounding box b0 - b1 sees.
279  */
280 template<int dim>
281 static void get_cell_range(real b0, real b1,
282                            const Grid::Dimensions &jGridDims,
283                            real d2, real rlist, int *cf, int *cl)
284 {
285     real listRangeBBToCell2 = gmx::square(listRangeForBoundingBoxToGridCell(rlist, jGridDims));
286     real distanceInCells    = (b0 - jGridDims.lowerCorner[dim])*jGridDims.invCellSize[dim];
287     *cf                     = std::max(static_cast<int>(distanceInCells), 0);
288
289     while (*cf > 0 &&
290            d2 + gmx::square((b0 - jGridDims.lowerCorner[dim]) - (*cf - 1 + 1)*jGridDims.cellSize[dim]) < listRangeBBToCell2)
291     {
292         (*cf)--;
293     }
294
295     *cl = std::min(static_cast<int>((b1 - jGridDims.lowerCorner[dim])*jGridDims.invCellSize[dim]), jGridDims.numCells[dim] - 1);
296     while (*cl < jGridDims.numCells[dim] - 1 &&
297            d2 + gmx::square((*cl + 1)*jGridDims.cellSize[dim] - (b1 - jGridDims.lowerCorner[dim])) < listRangeBBToCell2)
298     {
299         (*cl)++;
300     }
301 }
302
303 /* Reference code calculating the distance^2 between two bounding boxes */
304 /*
305    static float box_dist2(float bx0, float bx1, float by0,
306                        float by1, float bz0, float bz1,
307                        const BoundingBox *bb)
308    {
309     float d2;
310     float dl, dh, dm, dm0;
311
312     d2 = 0;
313
314     dl  = bx0 - bb->upper.x;
315     dh  = bb->lower.x - bx1;
316     dm  = std::max(dl, dh);
317     dm0 = std::max(dm, 0.0f);
318     d2 += dm0*dm0;
319
320     dl  = by0 - bb->upper.y;
321     dh  = bb->lower.y - by1;
322     dm  = std::max(dl, dh);
323     dm0 = std::max(dm, 0.0f);
324     d2 += dm0*dm0;
325
326     dl  = bz0 - bb->upper.z;
327     dh  = bb->lower.z - bz1;
328     dm  = std::max(dl, dh);
329     dm0 = std::max(dm, 0.0f);
330     d2 += dm0*dm0;
331
332     return d2;
333    }
334  */
335
336 #if !NBNXN_SEARCH_BB_SIMD4
337
338 /*! \brief Plain C code calculating the distance^2 between two bounding boxes in xyz0 format
339  *
340  * \param[in] bb_i  First bounding box
341  * \param[in] bb_j  Second bounding box
342  */
343 static float clusterBoundingBoxDistance2(const BoundingBox &bb_i,
344                                          const BoundingBox &bb_j)
345 {
346     float dl   = bb_i.lower.x - bb_j.upper.x;
347     float dh   = bb_j.lower.x - bb_i.upper.x;
348     float dm   = std::max(dl, dh);
349     float dm0  = std::max(dm, 0.0f);
350     float d2   = dm0*dm0;
351
352     dl         = bb_i.lower.y - bb_j.upper.y;
353     dh         = bb_j.lower.y - bb_i.upper.y;
354     dm         = std::max(dl, dh);
355     dm0        = std::max(dm, 0.0f);
356     d2        += dm0*dm0;
357
358     dl         = bb_i.lower.z - bb_j.upper.z;
359     dh         = bb_j.lower.z - bb_i.upper.z;
360     dm         = std::max(dl, dh);
361     dm0        = std::max(dm, 0.0f);
362     d2        += dm0*dm0;
363
364     return d2;
365 }
366
367 #else /* NBNXN_SEARCH_BB_SIMD4 */
368
369 /*! \brief 4-wide SIMD code calculating the distance^2 between two bounding boxes in xyz0 format
370  *
371  * \param[in] bb_i  First bounding box, should be aligned for 4-wide SIMD
372  * \param[in] bb_j  Second bounding box, should be aligned for 4-wide SIMD
373  */
374 static float clusterBoundingBoxDistance2(const BoundingBox &bb_i,
375                                          const BoundingBox &bb_j)
376 {
377     // TODO: During SIMDv2 transition only some archs use namespace (remove when done)
378     using namespace gmx;
379
380     const Simd4Float bb_i_S0 = load4(bb_i.lower.ptr());
381     const Simd4Float bb_i_S1 = load4(bb_i.upper.ptr());
382     const Simd4Float bb_j_S0 = load4(bb_j.lower.ptr());
383     const Simd4Float bb_j_S1 = load4(bb_j.upper.ptr());
384
385     const Simd4Float dl_S    = bb_i_S0 - bb_j_S1;
386     const Simd4Float dh_S    = bb_j_S0 - bb_i_S1;
387
388     const Simd4Float dm_S    = max(dl_S, dh_S);
389     const Simd4Float dm0_S   = max(dm_S, simd4SetZeroF());
390
391     return dotProduct(dm0_S, dm0_S);
392 }
393
394 /* Calculate bb bounding distances of bb_i[si,...,si+3] and store them in d2 */
395 template <int boundingBoxStart>
396 static inline void gmx_simdcall
397 clusterBoundingBoxDistance2_xxxx_simd4_inner(const float      *bb_i,
398                                              float            *d2,
399                                              const Simd4Float  xj_l,
400                                              const Simd4Float  yj_l,
401                                              const Simd4Float  zj_l,
402                                              const Simd4Float  xj_h,
403                                              const Simd4Float  yj_h,
404                                              const Simd4Float  zj_h)
405 {
406     constexpr int    stride = c_packedBoundingBoxesDimSize;
407
408     const int        shi  = boundingBoxStart*Nbnxm::c_numBoundingBoxBounds1D*DIM;
409
410     const Simd4Float zero = setZero();
411
412     const Simd4Float xi_l = load4(bb_i + shi + 0*stride);
413     const Simd4Float yi_l = load4(bb_i + shi + 1*stride);
414     const Simd4Float zi_l = load4(bb_i + shi + 2*stride);
415     const Simd4Float xi_h = load4(bb_i + shi + 3*stride);
416     const Simd4Float yi_h = load4(bb_i + shi + 4*stride);
417     const Simd4Float zi_h = load4(bb_i + shi + 5*stride);
418
419     const Simd4Float dx_0 = xi_l - xj_h;
420     const Simd4Float dy_0 = yi_l - yj_h;
421     const Simd4Float dz_0 = zi_l - zj_h;
422
423     const Simd4Float dx_1 = xj_l - xi_h;
424     const Simd4Float dy_1 = yj_l - yi_h;
425     const Simd4Float dz_1 = zj_l - zi_h;
426
427     const Simd4Float mx   = max(dx_0, dx_1);
428     const Simd4Float my   = max(dy_0, dy_1);
429     const Simd4Float mz   = max(dz_0, dz_1);
430
431     const Simd4Float m0x  = max(mx, zero);
432     const Simd4Float m0y  = max(my, zero);
433     const Simd4Float m0z  = max(mz, zero);
434
435     const Simd4Float d2x  = m0x * m0x;
436     const Simd4Float d2y  = m0y * m0y;
437     const Simd4Float d2z  = m0z * m0z;
438
439     const Simd4Float d2s  = d2x + d2y;
440     const Simd4Float d2t  = d2s + d2z;
441
442     store4(d2 + boundingBoxStart, d2t);
443 }
444
445 /* 4-wide SIMD code for nsi bb distances for bb format xxxxyyyyzzzz */
446 static void
447 clusterBoundingBoxDistance2_xxxx_simd4(const float *bb_j,
448                                        const int    nsi,
449                                        const float *bb_i,
450                                        float       *d2)
451 {
452     constexpr int    stride = c_packedBoundingBoxesDimSize;
453
454     // TODO: During SIMDv2 transition only some archs use namespace (remove when done)
455     using namespace gmx;
456
457     const Simd4Float xj_l = Simd4Float(bb_j[0*stride]);
458     const Simd4Float yj_l = Simd4Float(bb_j[1*stride]);
459     const Simd4Float zj_l = Simd4Float(bb_j[2*stride]);
460     const Simd4Float xj_h = Simd4Float(bb_j[3*stride]);
461     const Simd4Float yj_h = Simd4Float(bb_j[4*stride]);
462     const Simd4Float zj_h = Simd4Float(bb_j[5*stride]);
463
464     /* Here we "loop" over si (0,stride) from 0 to nsi with step stride.
465      * But as we know the number of iterations is 1 or 2, we unroll manually.
466      */
467     clusterBoundingBoxDistance2_xxxx_simd4_inner<0>(bb_i, d2,
468                                                     xj_l, yj_l, zj_l,
469                                                     xj_h, yj_h, zj_h);
470     if (stride < nsi)
471     {
472         clusterBoundingBoxDistance2_xxxx_simd4_inner<stride>(bb_i, d2,
473                                                              xj_l, yj_l, zj_l,
474                                                              xj_h, yj_h, zj_h);
475     }
476 }
477
478 #endif /* NBNXN_SEARCH_BB_SIMD4 */
479
480
481 /* Returns if any atom pair from two clusters is within distance sqrt(rlist2) */
482 static inline gmx_bool
483 clusterpair_in_range(const NbnxnPairlistGpuWork &work,
484                      int si,
485                      int csj, int stride, const real *x_j,
486                      real rlist2)
487 {
488 #if !GMX_SIMD4_HAVE_REAL
489
490     /* Plain C version.
491      * All coordinates are stored as xyzxyz...
492      */
493
494     const real *x_i = work.iSuperClusterData.x.data();
495
496     for (int i = 0; i < c_nbnxnGpuClusterSize; i++)
497     {
498         int i0 = (si*c_nbnxnGpuClusterSize + i)*DIM;
499         for (int j = 0; j < c_nbnxnGpuClusterSize; j++)
500         {
501             int  j0 = (csj*c_nbnxnGpuClusterSize + j)*stride;
502
503             real d2 = gmx::square(x_i[i0  ] - x_j[j0  ]) + gmx::square(x_i[i0+1] - x_j[j0+1]) + gmx::square(x_i[i0+2] - x_j[j0+2]);
504
505             if (d2 < rlist2)
506             {
507                 return TRUE;
508             }
509         }
510     }
511
512     return FALSE;
513
514 #else /* !GMX_SIMD4_HAVE_REAL */
515
516     /* 4-wide SIMD version.
517      * The coordinates x_i are stored as xxxxyyyy..., x_j is stored xyzxyz...
518      * Using 8-wide AVX(2) is not faster on Intel Sandy Bridge and Haswell.
519      */
520     static_assert(c_nbnxnGpuClusterSize == 8 || c_nbnxnGpuClusterSize == 4,
521                   "A cluster is hard-coded to 4/8 atoms.");
522
523     Simd4Real   rc2_S      = Simd4Real(rlist2);
524
525     const real *x_i        = work.iSuperClusterData.xSimd.data();
526
527     int         dim_stride = c_nbnxnGpuClusterSize*DIM;
528     Simd4Real   ix_S0      = load4(x_i + si*dim_stride + 0*GMX_SIMD4_WIDTH);
529     Simd4Real   iy_S0      = load4(x_i + si*dim_stride + 1*GMX_SIMD4_WIDTH);
530     Simd4Real   iz_S0      = load4(x_i + si*dim_stride + 2*GMX_SIMD4_WIDTH);
531
532     Simd4Real   ix_S1, iy_S1, iz_S1;
533     if (c_nbnxnGpuClusterSize == 8)
534     {
535         ix_S1      = load4(x_i + si*dim_stride + 3*GMX_SIMD4_WIDTH);
536         iy_S1      = load4(x_i + si*dim_stride + 4*GMX_SIMD4_WIDTH);
537         iz_S1      = load4(x_i + si*dim_stride + 5*GMX_SIMD4_WIDTH);
538     }
539     /* We loop from the outer to the inner particles to maximize
540      * the chance that we find a pair in range quickly and return.
541      */
542     int j0 = csj*c_nbnxnGpuClusterSize;
543     int j1 = j0 + c_nbnxnGpuClusterSize - 1;
544     while (j0 < j1)
545     {
546         Simd4Real jx0_S, jy0_S, jz0_S;
547         Simd4Real jx1_S, jy1_S, jz1_S;
548
549         Simd4Real dx_S0, dy_S0, dz_S0;
550         Simd4Real dx_S1, dy_S1, dz_S1;
551         Simd4Real dx_S2, dy_S2, dz_S2;
552         Simd4Real dx_S3, dy_S3, dz_S3;
553
554         Simd4Real rsq_S0;
555         Simd4Real rsq_S1;
556         Simd4Real rsq_S2;
557         Simd4Real rsq_S3;
558
559         Simd4Bool wco_S0;
560         Simd4Bool wco_S1;
561         Simd4Bool wco_S2;
562         Simd4Bool wco_S3;
563         Simd4Bool wco_any_S01, wco_any_S23, wco_any_S;
564
565         jx0_S = Simd4Real(x_j[j0*stride+0]);
566         jy0_S = Simd4Real(x_j[j0*stride+1]);
567         jz0_S = Simd4Real(x_j[j0*stride+2]);
568
569         jx1_S = Simd4Real(x_j[j1*stride+0]);
570         jy1_S = Simd4Real(x_j[j1*stride+1]);
571         jz1_S = Simd4Real(x_j[j1*stride+2]);
572
573         /* Calculate distance */
574         dx_S0            = ix_S0 - jx0_S;
575         dy_S0            = iy_S0 - jy0_S;
576         dz_S0            = iz_S0 - jz0_S;
577         dx_S2            = ix_S0 - jx1_S;
578         dy_S2            = iy_S0 - jy1_S;
579         dz_S2            = iz_S0 - jz1_S;
580         if (c_nbnxnGpuClusterSize == 8)
581         {
582             dx_S1            = ix_S1 - jx0_S;
583             dy_S1            = iy_S1 - jy0_S;
584             dz_S1            = iz_S1 - jz0_S;
585             dx_S3            = ix_S1 - jx1_S;
586             dy_S3            = iy_S1 - jy1_S;
587             dz_S3            = iz_S1 - jz1_S;
588         }
589
590         /* rsq = dx*dx+dy*dy+dz*dz */
591         rsq_S0           = norm2(dx_S0, dy_S0, dz_S0);
592         rsq_S2           = norm2(dx_S2, dy_S2, dz_S2);
593         if (c_nbnxnGpuClusterSize == 8)
594         {
595             rsq_S1           = norm2(dx_S1, dy_S1, dz_S1);
596             rsq_S3           = norm2(dx_S3, dy_S3, dz_S3);
597         }
598
599         wco_S0           = (rsq_S0 < rc2_S);
600         wco_S2           = (rsq_S2 < rc2_S);
601         if (c_nbnxnGpuClusterSize == 8)
602         {
603             wco_S1           = (rsq_S1 < rc2_S);
604             wco_S3           = (rsq_S3 < rc2_S);
605         }
606         if (c_nbnxnGpuClusterSize == 8)
607         {
608             wco_any_S01      = wco_S0 || wco_S1;
609             wco_any_S23      = wco_S2 || wco_S3;
610             wco_any_S        = wco_any_S01 || wco_any_S23;
611         }
612         else
613         {
614             wco_any_S = wco_S0 || wco_S2;
615         }
616
617         if (anyTrue(wco_any_S))
618         {
619             return TRUE;
620         }
621
622         j0++;
623         j1--;
624     }
625
626     return FALSE;
627
628 #endif /* !GMX_SIMD4_HAVE_REAL */
629 }
630
631 /* Returns the j-cluster index for index cjIndex in a cj list */
632 static inline int nblCj(gmx::ArrayRef<const nbnxn_cj_t> cjList,
633                         int                             cjIndex)
634 {
635     return cjList[cjIndex].cj;
636 }
637
638 /* Returns the j-cluster index for index cjIndex in a cj4 list */
639 static inline int nblCj(gmx::ArrayRef<const nbnxn_cj4_t> cj4List,
640                         int                              cjIndex)
641 {
642     return cj4List[cjIndex/c_nbnxnGpuJgroupSize].cj[cjIndex & (c_nbnxnGpuJgroupSize - 1)];
643 }
644
645 /* Returns the i-interaction mask of the j sub-cell for index cj_ind */
646 static unsigned int nbl_imask0(const NbnxnPairlistGpu *nbl, int cj_ind)
647 {
648     return nbl->cj4[cj_ind/c_nbnxnGpuJgroupSize].imei[0].imask;
649 }
650
651 /* Initializes a single NbnxnPairlistCpu data structure */
652 static void nbnxn_init_pairlist(NbnxnPairlistCpu *nbl)
653 {
654     nbl->na_ci       = c_nbnxnCpuIClusterSize;
655     nbl->na_cj       = 0;
656     nbl->ci.clear();
657     nbl->ciOuter.clear();
658     nbl->ncjInUse    = 0;
659     nbl->cj.clear();
660     nbl->cjOuter.clear();
661     nbl->nci_tot     = 0;
662
663     nbl->work        = new NbnxnPairlistCpuWork();
664 }
665
666 NbnxnPairlistGpu::NbnxnPairlistGpu(gmx::PinningPolicy pinningPolicy) :
667     na_ci(c_nbnxnGpuClusterSize),
668     na_cj(c_nbnxnGpuClusterSize),
669     na_sc(c_gpuNumClusterPerCell*c_nbnxnGpuClusterSize),
670     rlist(0),
671     sci({}, {pinningPolicy}),
672     cj4({}, {pinningPolicy}),
673     excl({}, {pinningPolicy}),
674     nci_tot(0)
675 {
676     static_assert(c_nbnxnGpuNumClusterPerSupercluster == c_gpuNumClusterPerCell,
677                   "The search code assumes that the a super-cluster matches a search grid cell");
678
679     static_assert(sizeof(cj4[0].imei[0].imask)*8 >= c_nbnxnGpuJgroupSize*c_gpuNumClusterPerCell,
680                   "The i super-cluster cluster interaction mask does not contain a sufficient number of bits");
681
682     static_assert(sizeof(excl[0])*8 >= c_nbnxnGpuJgroupSize*c_gpuNumClusterPerCell, "The GPU exclusion mask does not contain a sufficient number of bits");
683
684     // We always want a first entry without any exclusions
685     excl.resize(1);
686
687     work = new NbnxnPairlistGpuWork();
688 }
689
690 void nbnxn_init_pairlist_set(nbnxn_pairlist_set_t *nbl_list)
691 {
692     nbl_list->bSimple   =
693         (nbl_list->params.pairlistType == PairlistType::Simple4x2 ||
694          nbl_list->params.pairlistType == PairlistType::Simple4x4 ||
695          nbl_list->params.pairlistType == PairlistType::Simple4x8);
696     // Currently GPU lists are always combined
697     nbl_list->bCombined = !nbl_list->bSimple;
698
699     nbl_list->nnbl = gmx_omp_nthreads_get(emntNonbonded);
700
701     if (!nbl_list->bCombined &&
702         nbl_list->nnbl > NBNXN_BUFFERFLAG_MAX_THREADS)
703     {
704         gmx_fatal(FARGS, "%d OpenMP threads were requested. Since the non-bonded force buffer reduction is prohibitively slow with more than %d threads, we do not allow this. Use %d or less OpenMP threads.",
705                   nbl_list->nnbl, NBNXN_BUFFERFLAG_MAX_THREADS, NBNXN_BUFFERFLAG_MAX_THREADS);
706     }
707
708     if (nbl_list->bSimple)
709     {
710         snew(nbl_list->nbl, nbl_list->nnbl);
711         if (nbl_list->nnbl > 1)
712         {
713             snew(nbl_list->nbl_work, nbl_list->nnbl);
714         }
715     }
716     else
717     {
718         snew(nbl_list->nblGpu, nbl_list->nnbl);
719     }
720     nbl_list->nbl_fep.resize(nbl_list->nnbl);
721     /* Execute in order to avoid memory interleaving between threads */
722 #pragma omp parallel for num_threads(nbl_list->nnbl) schedule(static)
723     for (int i = 0; i < nbl_list->nnbl; i++)
724     {
725         try
726         {
727             /* Allocate the nblist data structure locally on each thread
728              * to optimize memory access for NUMA architectures.
729              */
730             if (nbl_list->bSimple)
731             {
732                 nbl_list->nbl[i] = new NbnxnPairlistCpu();
733
734                 nbnxn_init_pairlist(nbl_list->nbl[i]);
735                 if (nbl_list->nnbl > 1)
736                 {
737                     nbl_list->nbl_work[i] = new NbnxnPairlistCpu();
738                     nbnxn_init_pairlist(nbl_list->nbl_work[i]);
739                 }
740             }
741             else
742             {
743                 /* Only list 0 is used on the GPU, use normal allocation for i>0 */
744                 auto pinningPolicy = (i == 0 ? gmx::PinningPolicy::PinnedIfSupported : gmx::PinningPolicy::CannotBePinned);
745
746                 nbl_list->nblGpu[i] = new NbnxnPairlistGpu(pinningPolicy);
747             }
748
749             snew(nbl_list->nbl_fep[i], 1);
750             nbnxn_init_pairlist_fep(nbl_list->nbl_fep[i]);
751         }
752         GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR;
753     }
754 }
755
756 /* Print statistics of a pair list, used for debug output */
757 static void print_nblist_statistics(FILE                   *fp,
758                                     const NbnxnPairlistCpu *nbl,
759                                     const PairSearch       &pairSearch,
760                                     const real              rl)
761 {
762     const Grid             &grid = pairSearch.gridSet().grids()[0];
763     const Grid::Dimensions &dims = grid.dimensions();
764
765     fprintf(fp, "nbl nci %zu ncj %d\n",
766             nbl->ci.size(), nbl->ncjInUse);
767     const int    numAtomsJCluster = grid.geometry().numAtomsJCluster;
768     const double numAtomsPerCell  = nbl->ncjInUse/static_cast<double>(grid.numCells())*numAtomsJCluster;
769     fprintf(fp, "nbl na_cj %d rl %g ncp %d per cell %.1f atoms %.1f ratio %.2f\n",
770             nbl->na_cj, rl, nbl->ncjInUse, nbl->ncjInUse/static_cast<double>(grid.numCells()),
771             numAtomsPerCell,
772             numAtomsPerCell/(0.5*4.0/3.0*M_PI*rl*rl*rl*grid.numCells()*numAtomsJCluster/(dims.gridSize[XX]*dims.gridSize[YY]*dims.gridSize[ZZ])));
773
774     fprintf(fp, "nbl average j cell list length %.1f\n",
775             0.25*nbl->ncjInUse/std::max(static_cast<double>(nbl->ci.size()), 1.0));
776
777     int cs[SHIFTS] = { 0 };
778     int npexcl     = 0;
779     for (const nbnxn_ci_t &ciEntry : nbl->ci)
780     {
781         cs[ciEntry.shift & NBNXN_CI_SHIFT] +=
782             ciEntry.cj_ind_end - ciEntry.cj_ind_start;
783
784         int j = ciEntry.cj_ind_start;
785         while (j < ciEntry.cj_ind_end &&
786                nbl->cj[j].excl != NBNXN_INTERACTION_MASK_ALL)
787         {
788             npexcl++;
789             j++;
790         }
791     }
792     fprintf(fp, "nbl cell pairs, total: %zu excl: %d %.1f%%\n",
793             nbl->cj.size(), npexcl, 100*npexcl/std::max(static_cast<double>(nbl->cj.size()), 1.0));
794     for (int s = 0; s < SHIFTS; s++)
795     {
796         if (cs[s] > 0)
797         {
798             fprintf(fp, "nbl shift %2d ncj %3d\n", s, cs[s]);
799         }
800     }
801 }
802
803 /* Print statistics of a pair lists, used for debug output */
804 static void print_nblist_statistics(FILE                   *fp,
805                                     const NbnxnPairlistGpu *nbl,
806                                     const PairSearch       &pairSearch,
807                                     const real              rl)
808 {
809     const Grid             &grid = pairSearch.gridSet().grids()[0];
810     const Grid::Dimensions &dims = grid.dimensions();
811
812     fprintf(fp, "nbl nsci %zu ncj4 %zu nsi %d excl4 %zu\n",
813             nbl->sci.size(), nbl->cj4.size(), nbl->nci_tot, nbl->excl.size());
814     const int    numAtomsCluster = grid.geometry().numAtomsICluster;
815     const double numAtomsPerCell = nbl->nci_tot/static_cast<double>(grid.numClusters())*numAtomsCluster;
816     fprintf(fp, "nbl na_c %d rl %g ncp %d per cell %.1f atoms %.1f ratio %.2f\n",
817             nbl->na_ci, rl, nbl->nci_tot, nbl->nci_tot/static_cast<double>(grid.numClusters()),
818             numAtomsPerCell,
819             numAtomsPerCell/(0.5*4.0/3.0*M_PI*rl*rl*rl*grid.numClusters()*numAtomsCluster/(dims.gridSize[XX]*dims.gridSize[YY]*dims.gridSize[ZZ])));
820
821     double sum_nsp  = 0;
822     double sum_nsp2 = 0;
823     int    nsp_max  = 0;
824     int    c[c_gpuNumClusterPerCell + 1] = { 0 };
825     for (const nbnxn_sci_t &sci : nbl->sci)
826     {
827         int nsp = 0;
828         for (int j4 = sci.cj4_ind_start; j4 < sci.cj4_ind_end; j4++)
829         {
830             for (int j = 0; j < c_nbnxnGpuJgroupSize; j++)
831             {
832                 int b = 0;
833                 for (int si = 0; si < c_gpuNumClusterPerCell; si++)
834                 {
835                     if (nbl->cj4[j4].imei[0].imask & (1U << (j*c_gpuNumClusterPerCell + si)))
836                     {
837                         b++;
838                     }
839                 }
840                 nsp += b;
841                 c[b]++;
842             }
843         }
844         sum_nsp  += nsp;
845         sum_nsp2 += nsp*nsp;
846         nsp_max   = std::max(nsp_max, nsp);
847     }
848     if (!nbl->sci.empty())
849     {
850         sum_nsp  /= nbl->sci.size();
851         sum_nsp2 /= nbl->sci.size();
852     }
853     fprintf(fp, "nbl #cluster-pairs: av %.1f stddev %.1f max %d\n",
854             sum_nsp, std::sqrt(sum_nsp2 - sum_nsp*sum_nsp), nsp_max);
855
856     if (!nbl->cj4.empty())
857     {
858         for (int b = 0; b <= c_gpuNumClusterPerCell; b++)
859         {
860             fprintf(fp, "nbl j-list #i-subcell %d %7d %4.1f\n",
861                     b, c[b], 100.0*c[b]/size_t {nbl->cj4.size()*c_nbnxnGpuJgroupSize});
862         }
863     }
864 }
865
866 /* Returns a pointer to the exclusion mask for j-cluster-group \p cj4 and warp \p warp
867  * Generates a new exclusion entry when the j-cluster-group uses
868  * the default all-interaction mask at call time, so the returned mask
869  * can be modified when needed.
870  */
871 static nbnxn_excl_t *get_exclusion_mask(NbnxnPairlistGpu *nbl,
872                                         int               cj4,
873                                         int               warp)
874 {
875     if (nbl->cj4[cj4].imei[warp].excl_ind == 0)
876     {
877         /* No exclusions set, make a new list entry */
878         const size_t oldSize = nbl->excl.size();
879         GMX_ASSERT(oldSize >= 1, "We should always have entry [0]");
880         /* Add entry with default values: no exclusions */
881         nbl->excl.resize(oldSize + 1);
882         nbl->cj4[cj4].imei[warp].excl_ind = oldSize;
883     }
884
885     return &nbl->excl[nbl->cj4[cj4].imei[warp].excl_ind];
886 }
887
888 static void set_self_and_newton_excls_supersub(NbnxnPairlistGpu *nbl,
889                                                int cj4_ind, int sj_offset,
890                                                int i_cluster_in_cell)
891 {
892     nbnxn_excl_t *excl[c_nbnxnGpuClusterpairSplit];
893
894     /* Here we only set the set self and double pair exclusions */
895
896     /* Reserve extra elements, so the resize() in get_exclusion_mask()
897      * will not invalidate excl entries in the loop below
898      */
899     nbl->excl.reserve(nbl->excl.size() + c_nbnxnGpuClusterpairSplit);
900     for (int w = 0; w < c_nbnxnGpuClusterpairSplit; w++)
901     {
902         excl[w] = get_exclusion_mask(nbl, cj4_ind, w);
903     }
904
905     /* Only minor < major bits set */
906     for (int ej = 0; ej < nbl->na_ci; ej++)
907     {
908         int w = (ej>>2);
909         for (int ei = ej; ei < nbl->na_ci; ei++)
910         {
911             excl[w]->pair[(ej & (c_nbnxnGpuJgroupSize-1))*nbl->na_ci + ei] &=
912                 ~(1U << (sj_offset*c_gpuNumClusterPerCell + i_cluster_in_cell));
913         }
914     }
915 }
916
917 /* Returns a diagonal or off-diagonal interaction mask for plain C lists */
918 static unsigned int get_imask(gmx_bool rdiag, int ci, int cj)
919 {
920     return (rdiag && ci == cj ? NBNXN_INTERACTION_MASK_DIAG : NBNXN_INTERACTION_MASK_ALL);
921 }
922
923 /* Returns a diagonal or off-diagonal interaction mask for cj-size=2 */
924 gmx_unused static unsigned int get_imask_simd_j2(gmx_bool rdiag, int ci, int cj)
925 {
926     return (rdiag && ci*2 == cj ? NBNXN_INTERACTION_MASK_DIAG_J2_0 :
927             (rdiag && ci*2+1 == cj ? NBNXN_INTERACTION_MASK_DIAG_J2_1 :
928              NBNXN_INTERACTION_MASK_ALL));
929 }
930
931 /* Returns a diagonal or off-diagonal interaction mask for cj-size=4 */
932 gmx_unused static unsigned int get_imask_simd_j4(gmx_bool rdiag, int ci, int cj)
933 {
934     return (rdiag && ci == cj ? NBNXN_INTERACTION_MASK_DIAG : NBNXN_INTERACTION_MASK_ALL);
935 }
936
937 /* Returns a diagonal or off-diagonal interaction mask for cj-size=8 */
938 gmx_unused static unsigned int get_imask_simd_j8(gmx_bool rdiag, int ci, int cj)
939 {
940     return (rdiag && ci == cj*2 ? NBNXN_INTERACTION_MASK_DIAG_J8_0 :
941             (rdiag && ci == cj*2+1 ? NBNXN_INTERACTION_MASK_DIAG_J8_1 :
942              NBNXN_INTERACTION_MASK_ALL));
943 }
944
945 #if GMX_SIMD
946 #if GMX_SIMD_REAL_WIDTH == 2
947 #define get_imask_simd_4xn  get_imask_simd_j2
948 #endif
949 #if GMX_SIMD_REAL_WIDTH == 4
950 #define get_imask_simd_4xn  get_imask_simd_j4
951 #endif
952 #if GMX_SIMD_REAL_WIDTH == 8
953 #define get_imask_simd_4xn  get_imask_simd_j8
954 #define get_imask_simd_2xnn get_imask_simd_j4
955 #endif
956 #if GMX_SIMD_REAL_WIDTH == 16
957 #define get_imask_simd_2xnn get_imask_simd_j8
958 #endif
959 #endif
960
961 /* Plain C code for checking and adding cluster-pairs to the list.
962  *
963  * \param[in]     gridj               The j-grid
964  * \param[in,out] nbl                 The pair-list to store the cluster pairs in
965  * \param[in]     icluster            The index of the i-cluster
966  * \param[in]     jclusterFirst       The first cluster in the j-range
967  * \param[in]     jclusterLast        The last cluster in the j-range
968  * \param[in]     excludeSubDiagonal  Exclude atom pairs with i-index > j-index
969  * \param[in]     x_j                 Coordinates for the j-atom, in xyz format
970  * \param[in]     rlist2              The squared list cut-off
971  * \param[in]     rbb2                The squared cut-off for putting cluster-pairs in the list based on bounding box distance only
972  * \param[in,out] numDistanceChecks   The number of distance checks performed
973  */
974 static void
975 makeClusterListSimple(const Grid               &jGrid,
976                       NbnxnPairlistCpu *        nbl,
977                       int                       icluster,
978                       int                       jclusterFirst,
979                       int                       jclusterLast,
980                       bool                      excludeSubDiagonal,
981                       const real * gmx_restrict x_j,
982                       real                      rlist2,
983                       float                     rbb2,
984                       int * gmx_restrict        numDistanceChecks)
985 {
986     const BoundingBox * gmx_restrict bb_ci = nbl->work->iClusterData.bb.data();
987     const real * gmx_restrict        x_ci  = nbl->work->iClusterData.x.data();
988
989     gmx_bool                         InRange;
990
991     InRange = FALSE;
992     while (!InRange && jclusterFirst <= jclusterLast)
993     {
994         real d2  = clusterBoundingBoxDistance2(bb_ci[0], jGrid.jBoundingBoxes()[jclusterFirst]);
995         *numDistanceChecks += 2;
996
997         /* Check if the distance is within the distance where
998          * we use only the bounding box distance rbb,
999          * or within the cut-off and there is at least one atom pair
1000          * within the cut-off.
1001          */
1002         if (d2 < rbb2)
1003         {
1004             InRange = TRUE;
1005         }
1006         else if (d2 < rlist2)
1007         {
1008             int cjf_gl = jGrid.cellOffset() + jclusterFirst;
1009             for (int i = 0; i < c_nbnxnCpuIClusterSize && !InRange; i++)
1010             {
1011                 for (int j = 0; j < c_nbnxnCpuIClusterSize; j++)
1012                 {
1013                     InRange = InRange ||
1014                         (gmx::square(x_ci[i*STRIDE_XYZ+XX] - x_j[(cjf_gl*c_nbnxnCpuIClusterSize+j)*STRIDE_XYZ+XX]) +
1015                          gmx::square(x_ci[i*STRIDE_XYZ+YY] - x_j[(cjf_gl*c_nbnxnCpuIClusterSize+j)*STRIDE_XYZ+YY]) +
1016                          gmx::square(x_ci[i*STRIDE_XYZ+ZZ] - x_j[(cjf_gl*c_nbnxnCpuIClusterSize+j)*STRIDE_XYZ+ZZ]) < rlist2);
1017                 }
1018             }
1019             *numDistanceChecks += c_nbnxnCpuIClusterSize*c_nbnxnCpuIClusterSize;
1020         }
1021         if (!InRange)
1022         {
1023             jclusterFirst++;
1024         }
1025     }
1026     if (!InRange)
1027     {
1028         return;
1029     }
1030
1031     InRange = FALSE;
1032     while (!InRange && jclusterLast > jclusterFirst)
1033     {
1034         real d2  = clusterBoundingBoxDistance2(bb_ci[0], jGrid.jBoundingBoxes()[jclusterLast]);
1035         *numDistanceChecks += 2;
1036
1037         /* Check if the distance is within the distance where
1038          * we use only the bounding box distance rbb,
1039          * or within the cut-off and there is at least one atom pair
1040          * within the cut-off.
1041          */
1042         if (d2 < rbb2)
1043         {
1044             InRange = TRUE;
1045         }
1046         else if (d2 < rlist2)
1047         {
1048             int cjl_gl = jGrid.cellOffset() + jclusterLast;
1049             for (int i = 0; i < c_nbnxnCpuIClusterSize && !InRange; i++)
1050             {
1051                 for (int j = 0; j < c_nbnxnCpuIClusterSize; j++)
1052                 {
1053                     InRange = InRange ||
1054                         (gmx::square(x_ci[i*STRIDE_XYZ+XX] - x_j[(cjl_gl*c_nbnxnCpuIClusterSize+j)*STRIDE_XYZ+XX]) +
1055                          gmx::square(x_ci[i*STRIDE_XYZ+YY] - x_j[(cjl_gl*c_nbnxnCpuIClusterSize+j)*STRIDE_XYZ+YY]) +
1056                          gmx::square(x_ci[i*STRIDE_XYZ+ZZ] - x_j[(cjl_gl*c_nbnxnCpuIClusterSize+j)*STRIDE_XYZ+ZZ]) < rlist2);
1057                 }
1058             }
1059             *numDistanceChecks += c_nbnxnCpuIClusterSize*c_nbnxnCpuIClusterSize;
1060         }
1061         if (!InRange)
1062         {
1063             jclusterLast--;
1064         }
1065     }
1066
1067     if (jclusterFirst <= jclusterLast)
1068     {
1069         for (int jcluster = jclusterFirst; jcluster <= jclusterLast; jcluster++)
1070         {
1071             /* Store cj and the interaction mask */
1072             nbnxn_cj_t cjEntry;
1073             cjEntry.cj   = jGrid.cellOffset() + jcluster;
1074             cjEntry.excl = get_imask(excludeSubDiagonal, icluster, jcluster);
1075             nbl->cj.push_back(cjEntry);
1076         }
1077         /* Increase the closing index in the i list */
1078         nbl->ci.back().cj_ind_end = nbl->cj.size();
1079     }
1080 }
1081
1082 #ifdef GMX_NBNXN_SIMD_4XN
1083 #include "gromacs/nbnxm/pairlist_simd_4xm.h"
1084 #endif
1085 #ifdef GMX_NBNXN_SIMD_2XNN
1086 #include "gromacs/nbnxm/pairlist_simd_2xmm.h"
1087 #endif
1088
1089 /* Plain C or SIMD4 code for making a pair list of super-cell sci vs scj.
1090  * Checks bounding box distances and possibly atom pair distances.
1091  */
1092 static void make_cluster_list_supersub(const Grid         &iGrid,
1093                                        const Grid         &jGrid,
1094                                        NbnxnPairlistGpu   *nbl,
1095                                        const int           sci,
1096                                        const int           scj,
1097                                        const bool          excludeSubDiagonal,
1098                                        const int           stride,
1099                                        const real         *x,
1100                                        const real          rlist2,
1101                                        const float         rbb2,
1102                                        int                *numDistanceChecks)
1103 {
1104     NbnxnPairlistGpuWork &work   = *nbl->work;
1105
1106 #if NBNXN_BBXXXX
1107     const float          *pbb_ci = work.iSuperClusterData.bbPacked.data();
1108 #else
1109     const BoundingBox    *bb_ci  = work.iSuperClusterData.bb.data();
1110 #endif
1111
1112     assert(c_nbnxnGpuClusterSize == iGrid.geometry().numAtomsICluster);
1113     assert(c_nbnxnGpuClusterSize == jGrid.geometry().numAtomsICluster);
1114
1115     /* We generate the pairlist mainly based on bounding-box distances
1116      * and do atom pair distance based pruning on the GPU.
1117      * Only if a j-group contains a single cluster-pair, we try to prune
1118      * that pair based on atom distances on the CPU to avoid empty j-groups.
1119      */
1120 #define PRUNE_LIST_CPU_ONE 1
1121 #define PRUNE_LIST_CPU_ALL 0
1122
1123 #if PRUNE_LIST_CPU_ONE
1124     int  ci_last = -1;
1125 #endif
1126
1127     float *d2l = work.distanceBuffer.data();
1128
1129     for (int subc = 0; subc < jGrid.numClustersPerCell()[scj]; subc++)
1130     {
1131         const int    cj4_ind   = work.cj_ind/c_nbnxnGpuJgroupSize;
1132         const int    cj_offset = work.cj_ind - cj4_ind*c_nbnxnGpuJgroupSize;
1133         const int    cj        = scj*c_gpuNumClusterPerCell + subc;
1134
1135         const int    cj_gl     = jGrid.cellOffset()*c_gpuNumClusterPerCell + cj;
1136
1137         int          ci1;
1138         if (excludeSubDiagonal && sci == scj)
1139         {
1140             ci1 = subc + 1;
1141         }
1142         else
1143         {
1144             ci1 = iGrid.numClustersPerCell()[sci];
1145         }
1146
1147 #if NBNXN_BBXXXX
1148         /* Determine all ci1 bb distances in one call with SIMD4 */
1149         const int offset = packedBoundingBoxesIndex(cj) + (cj & (c_packedBoundingBoxesDimSize - 1));
1150         clusterBoundingBoxDistance2_xxxx_simd4(jGrid.packedBoundingBoxes().data() + offset,
1151                                                ci1, pbb_ci, d2l);
1152         *numDistanceChecks += c_nbnxnGpuClusterSize*2;
1153 #endif
1154
1155         int          npair = 0;
1156         unsigned int imask = 0;
1157         /* We use a fixed upper-bound instead of ci1 to help optimization */
1158         for (int ci = 0; ci < c_gpuNumClusterPerCell; ci++)
1159         {
1160             if (ci == ci1)
1161             {
1162                 break;
1163             }
1164
1165 #if !NBNXN_BBXXXX
1166             /* Determine the bb distance between ci and cj */
1167             d2l[ci]             = clusterBoundingBoxDistance2(bb_ci[ci], jGrid.jBoundingBoxes()[cj]);
1168             *numDistanceChecks += 2;
1169 #endif
1170             float d2 = d2l[ci];
1171
1172 #if PRUNE_LIST_CPU_ALL
1173             /* Check if the distance is within the distance where
1174              * we use only the bounding box distance rbb,
1175              * or within the cut-off and there is at least one atom pair
1176              * within the cut-off. This check is very costly.
1177              */
1178             *numDistanceChecks += c_nbnxnGpuClusterSize*c_nbnxnGpuClusterSize;
1179             if (d2 < rbb2 ||
1180                 (d2 < rlist2 &&
1181                  clusterpair_in_range(work, ci, cj_gl, stride, x, rlist2)))
1182 #else
1183             /* Check if the distance between the two bounding boxes
1184              * in within the pair-list cut-off.
1185              */
1186             if (d2 < rlist2)
1187 #endif
1188             {
1189                 /* Flag this i-subcell to be taken into account */
1190                 imask |= (1U << (cj_offset*c_gpuNumClusterPerCell + ci));
1191
1192 #if PRUNE_LIST_CPU_ONE
1193                 ci_last = ci;
1194 #endif
1195
1196                 npair++;
1197             }
1198         }
1199
1200 #if PRUNE_LIST_CPU_ONE
1201         /* If we only found 1 pair, check if any atoms are actually
1202          * within the cut-off, so we could get rid of it.
1203          */
1204         if (npair == 1 && d2l[ci_last] >= rbb2 &&
1205             !clusterpair_in_range(work, ci_last, cj_gl, stride, x, rlist2))
1206         {
1207             imask &= ~(1U << (cj_offset*c_gpuNumClusterPerCell + ci_last));
1208             npair--;
1209         }
1210 #endif
1211
1212         if (npair > 0)
1213         {
1214             /* We have at least one cluster pair: add a j-entry */
1215             if (static_cast<size_t>(cj4_ind) == nbl->cj4.size())
1216             {
1217                 nbl->cj4.resize(nbl->cj4.size() + 1);
1218             }
1219             nbnxn_cj4_t *cj4   = &nbl->cj4[cj4_ind];
1220
1221             cj4->cj[cj_offset] = cj_gl;
1222
1223             /* Set the exclusions for the ci==sj entry.
1224              * Here we don't bother to check if this entry is actually flagged,
1225              * as it will nearly always be in the list.
1226              */
1227             if (excludeSubDiagonal && sci == scj)
1228             {
1229                 set_self_and_newton_excls_supersub(nbl, cj4_ind, cj_offset, subc);
1230             }
1231
1232             /* Copy the cluster interaction mask to the list */
1233             for (int w = 0; w < c_nbnxnGpuClusterpairSplit; w++)
1234             {
1235                 cj4->imei[w].imask |= imask;
1236             }
1237
1238             nbl->work->cj_ind++;
1239
1240             /* Keep the count */
1241             nbl->nci_tot += npair;
1242
1243             /* Increase the closing index in i super-cell list */
1244             nbl->sci.back().cj4_ind_end =
1245                 (nbl->work->cj_ind + c_nbnxnGpuJgroupSize - 1)/c_nbnxnGpuJgroupSize;
1246         }
1247     }
1248 }
1249
1250 /* Returns how many contiguous j-clusters we have starting in the i-list */
1251 template <typename CjListType>
1252 static int numContiguousJClusters(const int                       cjIndexStart,
1253                                   const int                       cjIndexEnd,
1254                                   gmx::ArrayRef<const CjListType> cjList)
1255 {
1256     const int firstJCluster = nblCj(cjList, cjIndexStart);
1257
1258     int       numContiguous = 0;
1259
1260     while (cjIndexStart + numContiguous < cjIndexEnd &&
1261            nblCj(cjList, cjIndexStart + numContiguous) == firstJCluster + numContiguous)
1262     {
1263         numContiguous++;
1264     }
1265
1266     return numContiguous;
1267 }
1268
1269 /*! \internal
1270  * \brief Helper struct for efficient searching for excluded atoms in a j-list
1271  */
1272 struct JListRanges
1273 {
1274     /*! \brief Constructs a j-list range from \p cjList with the given index range */
1275     template <typename CjListType>
1276     JListRanges(int                             cjIndexStart,
1277                 int                             cjIndexEnd,
1278                 gmx::ArrayRef<const CjListType> cjList);
1279
1280     int cjIndexStart; //!< The start index in the j-list
1281     int cjIndexEnd;   //!< The end index in the j-list
1282     int cjFirst;      //!< The j-cluster with index cjIndexStart
1283     int cjLast;       //!< The j-cluster with index cjIndexEnd-1
1284     int numDirect;    //!< Up to cjIndexStart+numDirect the j-clusters are cjFirst + the index offset
1285 };
1286
1287 #ifndef DOXYGEN
1288 template <typename CjListType>
1289 JListRanges::JListRanges(int                             cjIndexStart,
1290                          int                             cjIndexEnd,
1291                          gmx::ArrayRef<const CjListType> cjList) :
1292     cjIndexStart(cjIndexStart),
1293     cjIndexEnd(cjIndexEnd)
1294 {
1295     GMX_ASSERT(cjIndexEnd > cjIndexStart, "JListRanges should only be called with non-empty lists");
1296
1297     cjFirst   = nblCj(cjList, cjIndexStart);
1298     cjLast    = nblCj(cjList, cjIndexEnd - 1);
1299
1300     /* Determine how many contiguous j-cells we have starting
1301      * from the first i-cell. This number can be used to directly
1302      * calculate j-cell indices for excluded atoms.
1303      */
1304     numDirect = numContiguousJClusters(cjIndexStart, cjIndexEnd, cjList);
1305 }
1306 #endif // !DOXYGEN
1307
1308 /* Return the index of \p jCluster in the given range or -1 when not present
1309  *
1310  * Note: This code is executed very often and therefore performance is
1311  *       important. It should be inlined and fully optimized.
1312  */
1313 template <typename CjListType>
1314 static inline int
1315 findJClusterInJList(int                              jCluster,
1316                     const JListRanges               &ranges,
1317                     gmx::ArrayRef<const CjListType>  cjList)
1318 {
1319     int index;
1320
1321     if (jCluster < ranges.cjFirst + ranges.numDirect)
1322     {
1323         /* We can calculate the index directly using the offset */
1324         index = ranges.cjIndexStart + jCluster - ranges.cjFirst;
1325     }
1326     else
1327     {
1328         /* Search for jCluster using bisection */
1329         index           = -1;
1330         int rangeStart  = ranges.cjIndexStart + ranges.numDirect;
1331         int rangeEnd    = ranges.cjIndexEnd;
1332         int rangeMiddle;
1333         while (index == -1 && rangeStart < rangeEnd)
1334         {
1335             rangeMiddle = (rangeStart + rangeEnd) >> 1;
1336
1337             const int clusterMiddle = nblCj(cjList, rangeMiddle);
1338
1339             if (jCluster == clusterMiddle)
1340             {
1341                 index      = rangeMiddle;
1342             }
1343             else if (jCluster < clusterMiddle)
1344             {
1345                 rangeEnd   = rangeMiddle;
1346             }
1347             else
1348             {
1349                 rangeStart = rangeMiddle + 1;
1350             }
1351         }
1352     }
1353
1354     return index;
1355 }
1356
1357 // TODO: Get rid of the two functions below by renaming sci to ci (or something better)
1358
1359 /* Return the i-entry in the list we are currently operating on */
1360 static nbnxn_ci_t *getOpenIEntry(NbnxnPairlistCpu *nbl)
1361 {
1362     return &nbl->ci.back();
1363 }
1364
1365 /* Return the i-entry in the list we are currently operating on */
1366 static nbnxn_sci_t *getOpenIEntry(NbnxnPairlistGpu *nbl)
1367 {
1368     return &nbl->sci.back();
1369 }
1370
1371 /* Set all atom-pair exclusions for a simple type list i-entry
1372  *
1373  * Set all atom-pair exclusions from the topology stored in exclusions
1374  * as masks in the pair-list for simple list entry iEntry.
1375  */
1376 static void
1377 setExclusionsForIEntry(const Nbnxm::GridSet &gridSet,
1378                        NbnxnPairlistCpu     *nbl,
1379                        gmx_bool              diagRemoved,
1380                        int                   na_cj_2log,
1381                        const nbnxn_ci_t     &iEntry,
1382                        const t_blocka       &exclusions)
1383 {
1384     if (iEntry.cj_ind_end == iEntry.cj_ind_start)
1385     {
1386         /* Empty list: no exclusions */
1387         return;
1388     }
1389
1390     const JListRanges        ranges(iEntry.cj_ind_start, iEntry.cj_ind_end, gmx::makeConstArrayRef(nbl->cj));
1391
1392     const int                iCluster = iEntry.ci;
1393
1394     gmx::ArrayRef<const int> cell        = gridSet.cells();
1395     gmx::ArrayRef<const int> atomIndices = gridSet.atomIndices();
1396
1397     /* Loop over the atoms in the i-cluster */
1398     for (int i = 0; i < nbl->na_ci; i++)
1399     {
1400         const int iIndex = iCluster*nbl->na_ci + i;
1401         const int iAtom  = atomIndices[iIndex];
1402         if (iAtom >= 0)
1403         {
1404             /* Loop over the topology-based exclusions for this i-atom */
1405             for (int exclIndex = exclusions.index[iAtom]; exclIndex < exclusions.index[iAtom + 1]; exclIndex++)
1406             {
1407                 const int jAtom = exclusions.a[exclIndex];
1408
1409                 if (jAtom == iAtom)
1410                 {
1411                     /* The self exclusion are already set, save some time */
1412                     continue;
1413                 }
1414
1415                 /* Get the index of the j-atom in the nbnxn atom data */
1416                 const int jIndex = cell[jAtom];
1417
1418                 /* Without shifts we only calculate interactions j>i
1419                  * for one-way pair-lists.
1420                  */
1421                 if (diagRemoved && jIndex <= iIndex)
1422                 {
1423                     continue;
1424                 }
1425
1426                 const int jCluster = (jIndex >> na_cj_2log);
1427
1428                 /* Could the cluster se be in our list? */
1429                 if (jCluster >= ranges.cjFirst && jCluster <= ranges.cjLast)
1430                 {
1431                     const int index =
1432                         findJClusterInJList(jCluster, ranges,
1433                                             gmx::makeConstArrayRef(nbl->cj));
1434
1435                     if (index >= 0)
1436                     {
1437                         /* We found an exclusion, clear the corresponding
1438                          * interaction bit.
1439                          */
1440                         const int innerJ     = jIndex - (jCluster << na_cj_2log);
1441
1442                         nbl->cj[index].excl &= ~(1U << ((i << na_cj_2log) + innerJ));
1443                     }
1444                 }
1445             }
1446         }
1447     }
1448 }
1449
1450 /* Add a new i-entry to the FEP list and copy the i-properties */
1451 static inline void fep_list_new_nri_copy(t_nblist *nlist)
1452 {
1453     /* Add a new i-entry */
1454     nlist->nri++;
1455
1456     assert(nlist->nri < nlist->maxnri);
1457
1458     /* Duplicate the last i-entry, except for jindex, which continues */
1459     nlist->iinr[nlist->nri]   = nlist->iinr[nlist->nri-1];
1460     nlist->shift[nlist->nri]  = nlist->shift[nlist->nri-1];
1461     nlist->gid[nlist->nri]    = nlist->gid[nlist->nri-1];
1462     nlist->jindex[nlist->nri] = nlist->nrj;
1463 }
1464
1465 /* For load balancing of the free-energy lists over threads, we set
1466  * the maximum nrj size of an i-entry to 40. This leads to good
1467  * load balancing in the worst case scenario of a single perturbed
1468  * particle on 16 threads, while not introducing significant overhead.
1469  * Note that half of the perturbed pairs will anyhow end up in very small lists,
1470  * since non perturbed i-particles will see few perturbed j-particles).
1471  */
1472 const int max_nrj_fep = 40;
1473
1474 /* Exclude the perturbed pairs from the Verlet list. This is only done to avoid
1475  * singularities for overlapping particles (0/0), since the charges and
1476  * LJ parameters have been zeroed in the nbnxn data structure.
1477  * Simultaneously make a group pair list for the perturbed pairs.
1478  */
1479 static void make_fep_list(gmx::ArrayRef<const int>  atomIndices,
1480                           const nbnxn_atomdata_t   *nbat,
1481                           NbnxnPairlistCpu         *nbl,
1482                           gmx_bool                  bDiagRemoved,
1483                           nbnxn_ci_t               *nbl_ci,
1484                           real gmx_unused           shx,
1485                           real gmx_unused           shy,
1486                           real gmx_unused           shz,
1487                           real gmx_unused           rlist_fep2,
1488                           const Grid               &iGrid,
1489                           const Grid               &jGrid,
1490                           t_nblist                 *nlist)
1491 {
1492     int      ci, cj_ind_start, cj_ind_end, cja, cjr;
1493     int      nri_max;
1494     int      gid_i = 0, gid_j, gid;
1495     int      egp_shift, egp_mask;
1496     int      gid_cj = 0;
1497     int      ind_i, ind_j, ai, aj;
1498     int      nri;
1499     gmx_bool bFEP_i, bFEP_i_all;
1500
1501     if (nbl_ci->cj_ind_end == nbl_ci->cj_ind_start)
1502     {
1503         /* Empty list */
1504         return;
1505     }
1506
1507     ci = nbl_ci->ci;
1508
1509     cj_ind_start = nbl_ci->cj_ind_start;
1510     cj_ind_end   = nbl_ci->cj_ind_end;
1511
1512     /* In worst case we have alternating energy groups
1513      * and create #atom-pair lists, which means we need the size
1514      * of a cluster pair (na_ci*na_cj) times the number of cj's.
1515      */
1516     nri_max = nbl->na_ci*nbl->na_cj*(cj_ind_end - cj_ind_start);
1517     if (nlist->nri + nri_max > nlist->maxnri)
1518     {
1519         nlist->maxnri = over_alloc_large(nlist->nri + nri_max);
1520         reallocate_nblist(nlist);
1521     }
1522
1523     const int numAtomsJCluster = jGrid.geometry().numAtomsJCluster;
1524
1525     const nbnxn_atomdata_t::Params &nbatParams = nbat->params();
1526
1527     const int ngid = nbatParams.nenergrp;
1528
1529     /* TODO: Consider adding a check in grompp and changing this to an assert */
1530     const int numBitsInEnergyGroupIdsForAtomsInJCluster = sizeof(gid_cj)*8;
1531     if (ngid*numAtomsJCluster > numBitsInEnergyGroupIdsForAtomsInJCluster)
1532     {
1533         gmx_fatal(FARGS, "The Verlet scheme with %dx%d kernels and free-energy only supports up to %zu energy groups",
1534                   iGrid.geometry().numAtomsICluster, numAtomsJCluster,
1535                   (sizeof(gid_cj)*8)/numAtomsJCluster);
1536     }
1537
1538     egp_shift = nbatParams.neg_2log;
1539     egp_mask  = (1 << egp_shift) - 1;
1540
1541     /* Loop over the atoms in the i sub-cell */
1542     bFEP_i_all = TRUE;
1543     for (int i = 0; i < nbl->na_ci; i++)
1544     {
1545         ind_i = ci*nbl->na_ci + i;
1546         ai    = atomIndices[ind_i];
1547         if (ai >= 0)
1548         {
1549             nri                  = nlist->nri;
1550             nlist->jindex[nri+1] = nlist->jindex[nri];
1551             nlist->iinr[nri]     = ai;
1552             /* The actual energy group pair index is set later */
1553             nlist->gid[nri]      = 0;
1554             nlist->shift[nri]    = nbl_ci->shift & NBNXN_CI_SHIFT;
1555
1556             bFEP_i = iGrid.atomIsPerturbed(ci - iGrid.cellOffset(), i);
1557
1558             bFEP_i_all = bFEP_i_all && bFEP_i;
1559
1560             if (nlist->nrj + (cj_ind_end - cj_ind_start)*nbl->na_cj > nlist->maxnrj)
1561             {
1562                 nlist->maxnrj = over_alloc_small(nlist->nrj + (cj_ind_end - cj_ind_start)*nbl->na_cj);
1563                 srenew(nlist->jjnr,     nlist->maxnrj);
1564                 srenew(nlist->excl_fep, nlist->maxnrj);
1565             }
1566
1567             if (ngid > 1)
1568             {
1569                 gid_i = (nbatParams.energrp[ci] >> (egp_shift*i)) & egp_mask;
1570             }
1571
1572             for (int cj_ind = cj_ind_start; cj_ind < cj_ind_end; cj_ind++)
1573             {
1574                 unsigned int fep_cj;
1575
1576                 cja = nbl->cj[cj_ind].cj;
1577
1578                 if (numAtomsJCluster == jGrid.geometry().numAtomsICluster)
1579                 {
1580                     cjr    = cja - jGrid.cellOffset();
1581                     fep_cj = jGrid.fepBits(cjr);
1582                     if (ngid > 1)
1583                     {
1584                         gid_cj = nbatParams.energrp[cja];
1585                     }
1586                 }
1587                 else if (2*numAtomsJCluster == jGrid.geometry().numAtomsICluster)
1588                 {
1589                     cjr    = cja - jGrid.cellOffset()*2;
1590                     /* Extract half of the ci fep/energrp mask */
1591                     fep_cj = (jGrid.fepBits(cjr >> 1) >> ((cjr & 1)*numAtomsJCluster)) & ((1 << numAtomsJCluster) - 1);
1592                     if (ngid > 1)
1593                     {
1594                         gid_cj = nbatParams.energrp[cja >> 1] >> ((cja & 1)*numAtomsJCluster*egp_shift) & ((1 << (numAtomsJCluster*egp_shift)) - 1);
1595                     }
1596                 }
1597                 else
1598                 {
1599                     cjr    = cja - (jGrid.cellOffset() >> 1);
1600                     /* Combine two ci fep masks/energrp */
1601                     fep_cj = jGrid.fepBits(cjr*2) + (jGrid.fepBits(cjr*2 + 1) << jGrid.geometry().numAtomsICluster);
1602                     if (ngid > 1)
1603                     {
1604                         gid_cj = nbatParams.energrp[cja*2] + (nbatParams.energrp[cja*2+1] << (jGrid.geometry().numAtomsICluster*egp_shift));
1605                     }
1606                 }
1607
1608                 if (bFEP_i || fep_cj != 0)
1609                 {
1610                     for (int j = 0; j < nbl->na_cj; j++)
1611                     {
1612                         /* Is this interaction perturbed and not excluded? */
1613                         ind_j = cja*nbl->na_cj + j;
1614                         aj    = atomIndices[ind_j];
1615                         if (aj >= 0 &&
1616                             (bFEP_i || (fep_cj & (1 << j))) &&
1617                             (!bDiagRemoved || ind_j >= ind_i))
1618                         {
1619                             if (ngid > 1)
1620                             {
1621                                 gid_j = (gid_cj >> (j*egp_shift)) & egp_mask;
1622                                 gid   = GID(gid_i, gid_j, ngid);
1623
1624                                 if (nlist->nrj > nlist->jindex[nri] &&
1625                                     nlist->gid[nri] != gid)
1626                                 {
1627                                     /* Energy group pair changed: new list */
1628                                     fep_list_new_nri_copy(nlist);
1629                                     nri = nlist->nri;
1630                                 }
1631                                 nlist->gid[nri] = gid;
1632                             }
1633
1634                             if (nlist->nrj - nlist->jindex[nri] >= max_nrj_fep)
1635                             {
1636                                 fep_list_new_nri_copy(nlist);
1637                                 nri = nlist->nri;
1638                             }
1639
1640                             /* Add it to the FEP list */
1641                             nlist->jjnr[nlist->nrj]     = aj;
1642                             nlist->excl_fep[nlist->nrj] = (nbl->cj[cj_ind].excl >> (i*nbl->na_cj + j)) & 1;
1643                             nlist->nrj++;
1644
1645                             /* Exclude it from the normal list.
1646                              * Note that the charge has been set to zero,
1647                              * but we need to avoid 0/0, as perturbed atoms
1648                              * can be on top of each other.
1649                              */
1650                             nbl->cj[cj_ind].excl &= ~(1U << (i*nbl->na_cj + j));
1651                         }
1652                     }
1653                 }
1654             }
1655
1656             if (nlist->nrj > nlist->jindex[nri])
1657             {
1658                 /* Actually add this new, non-empty, list */
1659                 nlist->nri++;
1660                 nlist->jindex[nlist->nri] = nlist->nrj;
1661             }
1662         }
1663     }
1664
1665     if (bFEP_i_all)
1666     {
1667         /* All interactions are perturbed, we can skip this entry */
1668         nbl_ci->cj_ind_end = cj_ind_start;
1669         nbl->ncjInUse     -= cj_ind_end - cj_ind_start;
1670     }
1671 }
1672
1673 /* Return the index of atom a within a cluster */
1674 static inline int cj_mod_cj4(int cj)
1675 {
1676     return cj & (c_nbnxnGpuJgroupSize - 1);
1677 }
1678
1679 /* Convert a j-cluster to a cj4 group */
1680 static inline int cj_to_cj4(int cj)
1681 {
1682     return cj/c_nbnxnGpuJgroupSize;
1683 }
1684
1685 /* Return the index of an j-atom within a warp */
1686 static inline int a_mod_wj(int a)
1687 {
1688     return a & (c_nbnxnGpuClusterSize/c_nbnxnGpuClusterpairSplit - 1);
1689 }
1690
1691 /* As make_fep_list above, but for super/sub lists. */
1692 static void make_fep_list(gmx::ArrayRef<const int>  atomIndices,
1693                           const nbnxn_atomdata_t   *nbat,
1694                           NbnxnPairlistGpu         *nbl,
1695                           gmx_bool                  bDiagRemoved,
1696                           const nbnxn_sci_t        *nbl_sci,
1697                           real                      shx,
1698                           real                      shy,
1699                           real                      shz,
1700                           real                      rlist_fep2,
1701                           const Grid               &iGrid,
1702                           const Grid               &jGrid,
1703                           t_nblist                 *nlist)
1704 {
1705     int                nri_max;
1706     int                c_abs;
1707     int                ind_i, ind_j, ai, aj;
1708     int                nri;
1709     gmx_bool           bFEP_i;
1710     real               xi, yi, zi;
1711     const nbnxn_cj4_t *cj4;
1712
1713     const int          numJClusterGroups = nbl_sci->numJClusterGroups();
1714     if (numJClusterGroups == 0)
1715     {
1716         /* Empty list */
1717         return;
1718     }
1719
1720     const int sci           = nbl_sci->sci;
1721
1722     const int cj4_ind_start = nbl_sci->cj4_ind_start;
1723     const int cj4_ind_end   = nbl_sci->cj4_ind_end;
1724
1725     /* Here we process one super-cell, max #atoms na_sc, versus a list
1726      * cj4 entries, each with max c_nbnxnGpuJgroupSize cj's, each
1727      * of size na_cj atoms.
1728      * On the GPU we don't support energy groups (yet).
1729      * So for each of the na_sc i-atoms, we need max one FEP list
1730      * for each max_nrj_fep j-atoms.
1731      */
1732     nri_max = nbl->na_sc*nbl->na_cj*(1 + (numJClusterGroups*c_nbnxnGpuJgroupSize)/max_nrj_fep);
1733     if (nlist->nri + nri_max > nlist->maxnri)
1734     {
1735         nlist->maxnri = over_alloc_large(nlist->nri + nri_max);
1736         reallocate_nblist(nlist);
1737     }
1738
1739     /* Loop over the atoms in the i super-cluster */
1740     for (int c = 0; c < c_gpuNumClusterPerCell; c++)
1741     {
1742         c_abs = sci*c_gpuNumClusterPerCell + c;
1743
1744         for (int i = 0; i < nbl->na_ci; i++)
1745         {
1746             ind_i = c_abs*nbl->na_ci + i;
1747             ai    = atomIndices[ind_i];
1748             if (ai >= 0)
1749             {
1750                 nri                  = nlist->nri;
1751                 nlist->jindex[nri+1] = nlist->jindex[nri];
1752                 nlist->iinr[nri]     = ai;
1753                 /* With GPUs, energy groups are not supported */
1754                 nlist->gid[nri]      = 0;
1755                 nlist->shift[nri]    = nbl_sci->shift & NBNXN_CI_SHIFT;
1756
1757                 bFEP_i = iGrid.atomIsPerturbed(c_abs - iGrid.cellOffset()*c_gpuNumClusterPerCell, i);
1758
1759                 xi = nbat->x()[ind_i*nbat->xstride+XX] + shx;
1760                 yi = nbat->x()[ind_i*nbat->xstride+YY] + shy;
1761                 zi = nbat->x()[ind_i*nbat->xstride+ZZ] + shz;
1762
1763                 const int nrjMax = nlist->nrj + numJClusterGroups*c_nbnxnGpuJgroupSize*nbl->na_cj;
1764                 if (nrjMax > nlist->maxnrj)
1765                 {
1766                     nlist->maxnrj = over_alloc_small(nrjMax);
1767                     srenew(nlist->jjnr,     nlist->maxnrj);
1768                     srenew(nlist->excl_fep, nlist->maxnrj);
1769                 }
1770
1771                 for (int cj4_ind = cj4_ind_start; cj4_ind < cj4_ind_end; cj4_ind++)
1772                 {
1773                     cj4 = &nbl->cj4[cj4_ind];
1774
1775                     for (int gcj = 0; gcj < c_nbnxnGpuJgroupSize; gcj++)
1776                     {
1777                         if ((cj4->imei[0].imask & (1U << (gcj*c_gpuNumClusterPerCell + c))) == 0)
1778                         {
1779                             /* Skip this ci for this cj */
1780                             continue;
1781                         }
1782
1783                         const int cjr =
1784                             cj4->cj[gcj] - jGrid.cellOffset()*c_gpuNumClusterPerCell;
1785
1786                         if (bFEP_i || jGrid.clusterIsPerturbed(cjr))
1787                         {
1788                             for (int j = 0; j < nbl->na_cj; j++)
1789                             {
1790                                 /* Is this interaction perturbed and not excluded? */
1791                                 ind_j = (jGrid.cellOffset()*c_gpuNumClusterPerCell + cjr)*nbl->na_cj + j;
1792                                 aj    = atomIndices[ind_j];
1793                                 if (aj >= 0 &&
1794                                     (bFEP_i || jGrid.atomIsPerturbed(cjr, j)) &&
1795                                     (!bDiagRemoved || ind_j >= ind_i))
1796                                 {
1797                                     int           excl_pair;
1798                                     unsigned int  excl_bit;
1799                                     real          dx, dy, dz;
1800
1801                                     const int     jHalf = j/(c_nbnxnGpuClusterSize/c_nbnxnGpuClusterpairSplit);
1802                                     nbnxn_excl_t *excl  =
1803                                         get_exclusion_mask(nbl, cj4_ind, jHalf);
1804
1805                                     excl_pair = a_mod_wj(j)*nbl->na_ci + i;
1806                                     excl_bit  = (1U << (gcj*c_gpuNumClusterPerCell + c));
1807
1808                                     dx = nbat->x()[ind_j*nbat->xstride+XX] - xi;
1809                                     dy = nbat->x()[ind_j*nbat->xstride+YY] - yi;
1810                                     dz = nbat->x()[ind_j*nbat->xstride+ZZ] - zi;
1811
1812                                     /* The unpruned GPU list has more than 2/3
1813                                      * of the atom pairs beyond rlist. Using
1814                                      * this list will cause a lot of overhead
1815                                      * in the CPU FEP kernels, especially
1816                                      * relative to the fast GPU kernels.
1817                                      * So we prune the FEP list here.
1818                                      */
1819                                     if (dx*dx + dy*dy + dz*dz < rlist_fep2)
1820                                     {
1821                                         if (nlist->nrj - nlist->jindex[nri] >= max_nrj_fep)
1822                                         {
1823                                             fep_list_new_nri_copy(nlist);
1824                                             nri = nlist->nri;
1825                                         }
1826
1827                                         /* Add it to the FEP list */
1828                                         nlist->jjnr[nlist->nrj]     = aj;
1829                                         nlist->excl_fep[nlist->nrj] = (excl->pair[excl_pair] & excl_bit) ? 1 : 0;
1830                                         nlist->nrj++;
1831                                     }
1832
1833                                     /* Exclude it from the normal list.
1834                                      * Note that the charge and LJ parameters have
1835                                      * been set to zero, but we need to avoid 0/0,
1836                                      * as perturbed atoms can be on top of each other.
1837                                      */
1838                                     excl->pair[excl_pair] &= ~excl_bit;
1839                                 }
1840                             }
1841
1842                             /* Note that we could mask out this pair in imask
1843                              * if all i- and/or all j-particles are perturbed.
1844                              * But since the perturbed pairs on the CPU will
1845                              * take an order of magnitude more time, the GPU
1846                              * will finish before the CPU and there is no gain.
1847                              */
1848                         }
1849                     }
1850                 }
1851
1852                 if (nlist->nrj > nlist->jindex[nri])
1853                 {
1854                     /* Actually add this new, non-empty, list */
1855                     nlist->nri++;
1856                     nlist->jindex[nlist->nri] = nlist->nrj;
1857                 }
1858             }
1859         }
1860     }
1861 }
1862
1863 /* Set all atom-pair exclusions for a GPU type list i-entry
1864  *
1865  * Sets all atom-pair exclusions from the topology stored in exclusions
1866  * as masks in the pair-list for i-super-cluster list entry iEntry.
1867  */
1868 static void
1869 setExclusionsForIEntry(const Nbnxm::GridSet &gridSet,
1870                        NbnxnPairlistGpu     *nbl,
1871                        gmx_bool              diagRemoved,
1872                        int gmx_unused        na_cj_2log,
1873                        const nbnxn_sci_t    &iEntry,
1874                        const t_blocka       &exclusions)
1875 {
1876     if (iEntry.numJClusterGroups() == 0)
1877     {
1878         /* Empty list */
1879         return;
1880     }
1881
1882     /* Set the search ranges using start and end j-cluster indices.
1883      * Note that here we can not use cj4_ind_end, since the last cj4
1884      * can be only partially filled, so we use cj_ind.
1885      */
1886     const JListRanges ranges(iEntry.cj4_ind_start*c_nbnxnGpuJgroupSize,
1887                              nbl->work->cj_ind,
1888                              gmx::makeConstArrayRef(nbl->cj4));
1889
1890     GMX_ASSERT(nbl->na_ci == c_nbnxnGpuClusterSize, "na_ci should match the GPU cluster size");
1891     constexpr int            c_clusterSize      = c_nbnxnGpuClusterSize;
1892     constexpr int            c_superClusterSize = c_nbnxnGpuNumClusterPerSupercluster*c_nbnxnGpuClusterSize;
1893
1894     const int                iSuperCluster = iEntry.sci;
1895
1896     gmx::ArrayRef<const int> atomIndices   = gridSet.atomIndices();
1897     gmx::ArrayRef<const int> cell          = gridSet.cells();
1898
1899     /* Loop over the atoms in the i super-cluster */
1900     for (int i = 0; i < c_superClusterSize; i++)
1901     {
1902         const int iIndex = iSuperCluster*c_superClusterSize + i;
1903         const int iAtom  = atomIndices[iIndex];
1904         if (iAtom >= 0)
1905         {
1906             const int iCluster = i/c_clusterSize;
1907
1908             /* Loop over the topology-based exclusions for this i-atom */
1909             for (int exclIndex = exclusions.index[iAtom]; exclIndex < exclusions.index[iAtom + 1]; exclIndex++)
1910             {
1911                 const int jAtom = exclusions.a[exclIndex];
1912
1913                 if (jAtom == iAtom)
1914                 {
1915                     /* The self exclusions are already set, save some time */
1916                     continue;
1917                 }
1918
1919                 /* Get the index of the j-atom in the nbnxn atom data */
1920                 const int jIndex = cell[jAtom];
1921
1922                 /* Without shifts we only calculate interactions j>i
1923                  * for one-way pair-lists.
1924                  */
1925                 /* NOTE: We would like to use iIndex on the right hand side,
1926                  * but that makes this routine 25% slower with gcc6/7.
1927                  * Even using c_superClusterSize makes it slower.
1928                  * Either of these changes triggers peeling of the exclIndex
1929                  * loop, which apparently leads to far less efficient code.
1930                  */
1931                 if (diagRemoved && jIndex <= iSuperCluster*nbl->na_sc + i)
1932                 {
1933                     continue;
1934                 }
1935
1936                 const int jCluster = jIndex/c_clusterSize;
1937
1938                 /* Check whether the cluster is in our list? */
1939                 if (jCluster >= ranges.cjFirst && jCluster <= ranges.cjLast)
1940                 {
1941                     const int index =
1942                         findJClusterInJList(jCluster, ranges,
1943                                             gmx::makeConstArrayRef(nbl->cj4));
1944
1945                     if (index >= 0)
1946                     {
1947                         /* We found an exclusion, clear the corresponding
1948                          * interaction bit.
1949                          */
1950                         const unsigned int pairMask = (1U << (cj_mod_cj4(index)*c_gpuNumClusterPerCell + iCluster));
1951                         /* Check if the i-cluster interacts with the j-cluster */
1952                         if (nbl_imask0(nbl, index) & pairMask)
1953                         {
1954                             const int innerI = (i      & (c_clusterSize - 1));
1955                             const int innerJ = (jIndex & (c_clusterSize - 1));
1956
1957                             /* Determine which j-half (CUDA warp) we are in */
1958                             const int     jHalf = innerJ/(c_clusterSize/c_nbnxnGpuClusterpairSplit);
1959
1960                             nbnxn_excl_t *interactionMask =
1961                                 get_exclusion_mask(nbl, cj_to_cj4(index), jHalf);
1962
1963                             interactionMask->pair[a_mod_wj(innerJ)*c_clusterSize + innerI] &= ~pairMask;
1964                         }
1965                     }
1966                 }
1967             }
1968         }
1969     }
1970 }
1971
1972 /* Make a new ci entry at the back of nbl->ci */
1973 static void addNewIEntry(NbnxnPairlistCpu *nbl, int ci, int shift, int flags)
1974 {
1975     nbnxn_ci_t ciEntry;
1976     ciEntry.ci            = ci;
1977     ciEntry.shift         = shift;
1978     /* Store the interaction flags along with the shift */
1979     ciEntry.shift        |= flags;
1980     ciEntry.cj_ind_start  = nbl->cj.size();
1981     ciEntry.cj_ind_end    = nbl->cj.size();
1982     nbl->ci.push_back(ciEntry);
1983 }
1984
1985 /* Make a new sci entry at index nbl->nsci */
1986 static void addNewIEntry(NbnxnPairlistGpu *nbl, int sci, int shift, int gmx_unused flags)
1987 {
1988     nbnxn_sci_t sciEntry;
1989     sciEntry.sci           = sci;
1990     sciEntry.shift         = shift;
1991     sciEntry.cj4_ind_start = nbl->cj4.size();
1992     sciEntry.cj4_ind_end   = nbl->cj4.size();
1993
1994     nbl->sci.push_back(sciEntry);
1995 }
1996
1997 /* Sort the simple j-list cj on exclusions.
1998  * Entries with exclusions will all be sorted to the beginning of the list.
1999  */
2000 static void sort_cj_excl(nbnxn_cj_t *cj, int ncj,
2001                          NbnxnPairlistCpuWork *work)
2002 {
2003     work->cj.resize(ncj);
2004
2005     /* Make a list of the j-cells involving exclusions */
2006     int jnew = 0;
2007     for (int j = 0; j < ncj; j++)
2008     {
2009         if (cj[j].excl != NBNXN_INTERACTION_MASK_ALL)
2010         {
2011             work->cj[jnew++] = cj[j];
2012         }
2013     }
2014     /* Check if there are exclusions at all or not just the first entry */
2015     if (!((jnew == 0) ||
2016           (jnew == 1 && cj[0].excl != NBNXN_INTERACTION_MASK_ALL)))
2017     {
2018         for (int j = 0; j < ncj; j++)
2019         {
2020             if (cj[j].excl == NBNXN_INTERACTION_MASK_ALL)
2021             {
2022                 work->cj[jnew++] = cj[j];
2023             }
2024         }
2025         for (int j = 0; j < ncj; j++)
2026         {
2027             cj[j] = work->cj[j];
2028         }
2029     }
2030 }
2031
2032 /* Close this simple list i entry */
2033 static void closeIEntry(NbnxnPairlistCpu    *nbl,
2034                         int gmx_unused       sp_max_av,
2035                         gmx_bool gmx_unused  progBal,
2036                         float gmx_unused     nsp_tot_est,
2037                         int gmx_unused       thread,
2038                         int gmx_unused       nthread)
2039 {
2040     nbnxn_ci_t &ciEntry = nbl->ci.back();
2041
2042     /* All content of the new ci entry have already been filled correctly,
2043      * we only need to sort and increase counts or remove the entry when empty.
2044      */
2045     const int jlen = ciEntry.cj_ind_end - ciEntry.cj_ind_start;
2046     if (jlen > 0)
2047     {
2048         sort_cj_excl(nbl->cj.data() + ciEntry.cj_ind_start, jlen, nbl->work);
2049
2050         /* The counts below are used for non-bonded pair/flop counts
2051          * and should therefore match the available kernel setups.
2052          */
2053         if (!(ciEntry.shift & NBNXN_CI_DO_COUL(0)))
2054         {
2055             nbl->work->ncj_noq += jlen;
2056         }
2057         else if ((ciEntry.shift & NBNXN_CI_HALF_LJ(0)) ||
2058                  !(ciEntry.shift & NBNXN_CI_DO_LJ(0)))
2059         {
2060             nbl->work->ncj_hlj += jlen;
2061         }
2062     }
2063     else
2064     {
2065         /* Entry is empty: remove it  */
2066         nbl->ci.pop_back();
2067     }
2068 }
2069
2070 /* Split sci entry for load balancing on the GPU.
2071  * Splitting ensures we have enough lists to fully utilize the whole GPU.
2072  * With progBal we generate progressively smaller lists, which improves
2073  * load balancing. As we only know the current count on our own thread,
2074  * we will need to estimate the current total amount of i-entries.
2075  * As the lists get concatenated later, this estimate depends
2076  * both on nthread and our own thread index.
2077  */
2078 static void split_sci_entry(NbnxnPairlistGpu *nbl,
2079                             int nsp_target_av,
2080                             gmx_bool progBal, float nsp_tot_est,
2081                             int thread, int nthread)
2082 {
2083     int nsp_max;
2084
2085     if (progBal)
2086     {
2087         float nsp_est;
2088
2089         /* Estimate the total numbers of ci's of the nblist combined
2090          * over all threads using the target number of ci's.
2091          */
2092         nsp_est = (nsp_tot_est*thread)/nthread + nbl->nci_tot;
2093
2094         /* The first ci blocks should be larger, to avoid overhead.
2095          * The last ci blocks should be smaller, to improve load balancing.
2096          * The factor 3/2 makes the first block 3/2 times the target average
2097          * and ensures that the total number of blocks end up equal to
2098          * that of equally sized blocks of size nsp_target_av.
2099          */
2100         nsp_max = static_cast<int>(nsp_target_av*(nsp_tot_est*1.5/(nsp_est + nsp_tot_est)));
2101     }
2102     else
2103     {
2104         nsp_max = nsp_target_av;
2105     }
2106
2107     const int cj4_start = nbl->sci.back().cj4_ind_start;
2108     const int cj4_end   = nbl->sci.back().cj4_ind_end;
2109     const int j4len     = cj4_end - cj4_start;
2110
2111     if (j4len > 1 && j4len*c_gpuNumClusterPerCell*c_nbnxnGpuJgroupSize > nsp_max)
2112     {
2113         /* Modify the last ci entry and process the cj4's again */
2114
2115         int nsp        = 0;
2116         int nsp_sci    = 0;
2117         int nsp_cj4_e  = 0;
2118         int nsp_cj4    = 0;
2119         for (int cj4 = cj4_start; cj4 < cj4_end; cj4++)
2120         {
2121             int nsp_cj4_p = nsp_cj4;
2122             /* Count the number of cluster pairs in this cj4 group */
2123             nsp_cj4   = 0;
2124             for (int p = 0; p < c_gpuNumClusterPerCell*c_nbnxnGpuJgroupSize; p++)
2125             {
2126                 nsp_cj4 += (nbl->cj4[cj4].imei[0].imask >> p) & 1;
2127             }
2128
2129             /* If adding the current cj4 with nsp_cj4 pairs get us further
2130              * away from our target nsp_max, split the list before this cj4.
2131              */
2132             if (nsp > 0 && nsp_max - nsp < nsp + nsp_cj4 - nsp_max)
2133             {
2134                 /* Split the list at cj4 */
2135                 nbl->sci.back().cj4_ind_end = cj4;
2136                 /* Create a new sci entry */
2137                 nbnxn_sci_t sciNew;
2138                 sciNew.sci           = nbl->sci.back().sci;
2139                 sciNew.shift         = nbl->sci.back().shift;
2140                 sciNew.cj4_ind_start = cj4;
2141                 nbl->sci.push_back(sciNew);
2142
2143                 nsp_sci              = nsp;
2144                 nsp_cj4_e            = nsp_cj4_p;
2145                 nsp                  = 0;
2146             }
2147             nsp += nsp_cj4;
2148         }
2149
2150         /* Put the remaining cj4's in the last sci entry */
2151         nbl->sci.back().cj4_ind_end = cj4_end;
2152
2153         /* Possibly balance out the last two sci's
2154          * by moving the last cj4 of the second last sci.
2155          */
2156         if (nsp_sci - nsp_cj4_e >= nsp + nsp_cj4_e)
2157         {
2158             GMX_ASSERT(nbl->sci.size() >= 2, "We expect at least two elements");
2159             nbl->sci[nbl->sci.size() - 2].cj4_ind_end--;
2160             nbl->sci[nbl->sci.size() - 1].cj4_ind_start--;
2161         }
2162     }
2163 }
2164
2165 /* Clost this super/sub list i entry */
2166 static void closeIEntry(NbnxnPairlistGpu *nbl,
2167                         int nsp_max_av,
2168                         gmx_bool progBal, float nsp_tot_est,
2169                         int thread, int nthread)
2170 {
2171     nbnxn_sci_t &sciEntry = *getOpenIEntry(nbl);
2172
2173     /* All content of the new ci entry have already been filled correctly,
2174      * we only need to, potentially, split or remove the entry when empty.
2175      */
2176     int j4len = sciEntry.numJClusterGroups();
2177     if (j4len > 0)
2178     {
2179         /* We can only have complete blocks of 4 j-entries in a list,
2180          * so round the count up before closing.
2181          */
2182         int ncj4          = (nbl->work->cj_ind + c_nbnxnGpuJgroupSize - 1)/c_nbnxnGpuJgroupSize;
2183         nbl->work->cj_ind = ncj4*c_nbnxnGpuJgroupSize;
2184
2185         if (nsp_max_av > 0)
2186         {
2187             /* Measure the size of the new entry and potentially split it */
2188             split_sci_entry(nbl, nsp_max_av, progBal, nsp_tot_est,
2189                             thread, nthread);
2190         }
2191     }
2192     else
2193     {
2194         /* Entry is empty: remove it  */
2195         nbl->sci.pop_back();
2196     }
2197 }
2198
2199 /* Syncs the working array before adding another grid pair to the GPU list */
2200 static void sync_work(NbnxnPairlistCpu gmx_unused *nbl)
2201 {
2202 }
2203
2204 /* Syncs the working array before adding another grid pair to the GPU list */
2205 static void sync_work(NbnxnPairlistGpu *nbl)
2206 {
2207     nbl->work->cj_ind   = nbl->cj4.size()*c_nbnxnGpuJgroupSize;
2208 }
2209
2210 /* Clears an NbnxnPairlistCpu data structure */
2211 static void clear_pairlist(NbnxnPairlistCpu *nbl)
2212 {
2213     nbl->ci.clear();
2214     nbl->cj.clear();
2215     nbl->ncjInUse      = 0;
2216     nbl->nci_tot       = 0;
2217     nbl->ciOuter.clear();
2218     nbl->cjOuter.clear();
2219
2220     nbl->work->ncj_noq = 0;
2221     nbl->work->ncj_hlj = 0;
2222 }
2223
2224 /* Clears an NbnxnPairlistGpu data structure */
2225 static void clear_pairlist(NbnxnPairlistGpu *nbl)
2226 {
2227     nbl->sci.clear();
2228     nbl->cj4.clear();
2229     nbl->excl.resize(1);
2230     nbl->nci_tot = 0;
2231 }
2232
2233 /* Clears a group scheme pair list */
2234 static void clear_pairlist_fep(t_nblist *nl)
2235 {
2236     nl->nri = 0;
2237     nl->nrj = 0;
2238     if (nl->jindex == nullptr)
2239     {
2240         snew(nl->jindex, 1);
2241     }
2242     nl->jindex[0] = 0;
2243 }
2244
2245 /* Sets a simple list i-cell bounding box, including PBC shift */
2246 static inline void set_icell_bb_simple(gmx::ArrayRef<const BoundingBox> bb,
2247                                        int ci,
2248                                        real shx, real shy, real shz,
2249                                        BoundingBox *bb_ci)
2250 {
2251     bb_ci->lower.x = bb[ci].lower.x + shx;
2252     bb_ci->lower.y = bb[ci].lower.y + shy;
2253     bb_ci->lower.z = bb[ci].lower.z + shz;
2254     bb_ci->upper.x = bb[ci].upper.x + shx;
2255     bb_ci->upper.y = bb[ci].upper.y + shy;
2256     bb_ci->upper.z = bb[ci].upper.z + shz;
2257 }
2258
2259 /* Sets a simple list i-cell bounding box, including PBC shift */
2260 static inline void set_icell_bb(const Grid &iGrid,
2261                                 int ci,
2262                                 real shx, real shy, real shz,
2263                                 NbnxnPairlistCpuWork *work)
2264 {
2265     set_icell_bb_simple(iGrid.iBoundingBoxes(), ci, shx, shy, shz,
2266                         &work->iClusterData.bb[0]);
2267 }
2268
2269 #if NBNXN_BBXXXX
2270 /* Sets a super-cell and sub cell bounding boxes, including PBC shift */
2271 static void set_icell_bbxxxx_supersub(gmx::ArrayRef<const float> bb,
2272                                       int ci,
2273                                       real shx, real shy, real shz,
2274                                       float *bb_ci)
2275 {
2276     constexpr int cellBBStride = packedBoundingBoxesIndex(c_gpuNumClusterPerCell);
2277     constexpr int pbbStride    = c_packedBoundingBoxesDimSize;
2278     const int     ia           = ci*cellBBStride;
2279     for (int m = 0; m < cellBBStride; m += c_packedBoundingBoxesSize)
2280     {
2281         for (int i = 0; i < pbbStride; i++)
2282         {
2283             bb_ci[m + 0*pbbStride + i] = bb[ia + m + 0*pbbStride + i] + shx;
2284             bb_ci[m + 1*pbbStride + i] = bb[ia + m + 1*pbbStride + i] + shy;
2285             bb_ci[m + 2*pbbStride + i] = bb[ia + m + 2*pbbStride + i] + shz;
2286             bb_ci[m + 3*pbbStride + i] = bb[ia + m + 3*pbbStride + i] + shx;
2287             bb_ci[m + 4*pbbStride + i] = bb[ia + m + 4*pbbStride + i] + shy;
2288             bb_ci[m + 5*pbbStride + i] = bb[ia + m + 5*pbbStride + i] + shz;
2289         }
2290     }
2291 }
2292 #endif
2293
2294 /* Sets a super-cell and sub cell bounding boxes, including PBC shift */
2295 gmx_unused static void set_icell_bb_supersub(gmx::ArrayRef<const BoundingBox> bb,
2296                                              int ci,
2297                                              real shx, real shy, real shz,
2298                                              BoundingBox *bb_ci)
2299 {
2300     for (int i = 0; i < c_gpuNumClusterPerCell; i++)
2301     {
2302         set_icell_bb_simple(bb, ci*c_gpuNumClusterPerCell+i,
2303                             shx, shy, shz,
2304                             &bb_ci[i]);
2305     }
2306 }
2307
2308 /* Sets a super-cell and sub cell bounding boxes, including PBC shift */
2309 gmx_unused static void set_icell_bb(const Grid &iGrid,
2310                                     int ci,
2311                                     real shx, real shy, real shz,
2312                                     NbnxnPairlistGpuWork *work)
2313 {
2314 #if NBNXN_BBXXXX
2315     set_icell_bbxxxx_supersub(iGrid.packedBoundingBoxes(), ci, shx, shy, shz,
2316                               work->iSuperClusterData.bbPacked.data());
2317 #else
2318     set_icell_bb_supersub(iGrid.iBoundingBoxes(), ci, shx, shy, shz,
2319                           work->iSuperClusterData.bb.data());
2320 #endif
2321 }
2322
2323 /* Copies PBC shifted i-cell atom coordinates x,y,z to working array */
2324 static void icell_set_x_simple(int ci,
2325                                real shx, real shy, real shz,
2326                                int stride, const real *x,
2327                                NbnxnPairlistCpuWork::IClusterData *iClusterData)
2328 {
2329     const int ia = ci*c_nbnxnCpuIClusterSize;
2330
2331     for (int i = 0; i < c_nbnxnCpuIClusterSize; i++)
2332     {
2333         iClusterData->x[i*STRIDE_XYZ+XX] = x[(ia+i)*stride+XX] + shx;
2334         iClusterData->x[i*STRIDE_XYZ+YY] = x[(ia+i)*stride+YY] + shy;
2335         iClusterData->x[i*STRIDE_XYZ+ZZ] = x[(ia+i)*stride+ZZ] + shz;
2336     }
2337 }
2338
2339 static void icell_set_x(int ci,
2340                         real shx, real shy, real shz,
2341                         int stride, const real *x,
2342                         const Nbnxm::KernelType kernelType,
2343                         NbnxnPairlistCpuWork *work)
2344 {
2345     switch (kernelType)
2346     {
2347 #if GMX_SIMD
2348 #ifdef GMX_NBNXN_SIMD_4XN
2349         case Nbnxm::KernelType::Cpu4xN_Simd_4xN:
2350             icell_set_x_simd_4xn(ci, shx, shy, shz, stride, x, work);
2351             break;
2352 #endif
2353 #ifdef GMX_NBNXN_SIMD_2XNN
2354         case Nbnxm::KernelType::Cpu4xN_Simd_2xNN:
2355             icell_set_x_simd_2xnn(ci, shx, shy, shz, stride, x, work);
2356             break;
2357 #endif
2358 #endif
2359         case Nbnxm::KernelType::Cpu4x4_PlainC:
2360             icell_set_x_simple(ci, shx, shy, shz, stride, x, &work->iClusterData);
2361             break;
2362         default:
2363             GMX_ASSERT(false, "Unhandled case");
2364             break;
2365     }
2366 }
2367
2368 /* Copies PBC shifted super-cell atom coordinates x,y,z to working array */
2369 static void icell_set_x(int ci,
2370                         real shx, real shy, real shz,
2371                         int stride, const real *x,
2372                         Nbnxm::KernelType gmx_unused kernelType,
2373                         NbnxnPairlistGpuWork *work)
2374 {
2375 #if !GMX_SIMD4_HAVE_REAL
2376
2377     real * x_ci = work->iSuperClusterData.x.data();
2378
2379     int    ia = ci*c_gpuNumClusterPerCell*c_nbnxnGpuClusterSize;
2380     for (int i = 0; i < c_gpuNumClusterPerCell*c_nbnxnGpuClusterSize; i++)
2381     {
2382         x_ci[i*DIM + XX] = x[(ia+i)*stride + XX] + shx;
2383         x_ci[i*DIM + YY] = x[(ia+i)*stride + YY] + shy;
2384         x_ci[i*DIM + ZZ] = x[(ia+i)*stride + ZZ] + shz;
2385     }
2386
2387 #else /* !GMX_SIMD4_HAVE_REAL */
2388
2389     real * x_ci = work->iSuperClusterData.xSimd.data();
2390
2391     for (int si = 0; si < c_gpuNumClusterPerCell; si++)
2392     {
2393         for (int i = 0; i < c_nbnxnGpuClusterSize; i += GMX_SIMD4_WIDTH)
2394         {
2395             int io = si*c_nbnxnGpuClusterSize + i;
2396             int ia = ci*c_gpuNumClusterPerCell*c_nbnxnGpuClusterSize + io;
2397             for (int j = 0; j < GMX_SIMD4_WIDTH; j++)
2398             {
2399                 x_ci[io*DIM + j + XX*GMX_SIMD4_WIDTH] = x[(ia + j)*stride + XX] + shx;
2400                 x_ci[io*DIM + j + YY*GMX_SIMD4_WIDTH] = x[(ia + j)*stride + YY] + shy;
2401                 x_ci[io*DIM + j + ZZ*GMX_SIMD4_WIDTH] = x[(ia + j)*stride + ZZ] + shz;
2402             }
2403         }
2404     }
2405
2406 #endif /* !GMX_SIMD4_HAVE_REAL */
2407 }
2408
2409 static real minimum_subgrid_size_xy(const Grid &grid)
2410 {
2411     const Grid::Dimensions &dims = grid.dimensions();
2412
2413     if (grid.geometry().isSimple)
2414     {
2415         return std::min(dims.cellSize[XX], dims.cellSize[YY]);
2416     }
2417     else
2418     {
2419         return std::min(dims.cellSize[XX]/c_gpuNumClusterPerCellX,
2420                         dims.cellSize[YY]/c_gpuNumClusterPerCellY);
2421     }
2422 }
2423
2424 static real effective_buffer_1x1_vs_MxN(const Grid &iGrid,
2425                                         const Grid &jGrid)
2426 {
2427     const real eff_1x1_buffer_fac_overest = 0.1;
2428
2429     /* Determine an atom-pair list cut-off buffer size for atom pairs,
2430      * to be added to rlist (including buffer) used for MxN.
2431      * This is for converting an MxN list to a 1x1 list. This means we can't
2432      * use the normal buffer estimate, as we have an MxN list in which
2433      * some atom pairs beyond rlist are missing. We want to capture
2434      * the beneficial effect of buffering by extra pairs just outside rlist,
2435      * while removing the useless pairs that are further away from rlist.
2436      * (Also the buffer could have been set manually not using the estimate.)
2437      * This buffer size is an overestimate.
2438      * We add 10% of the smallest grid sub-cell dimensions.
2439      * Note that the z-size differs per cell and we don't use this,
2440      * so we overestimate.
2441      * With PME, the 10% value gives a buffer that is somewhat larger
2442      * than the effective buffer with a tolerance of 0.005 kJ/mol/ps.
2443      * Smaller tolerances or using RF lead to a smaller effective buffer,
2444      * so 10% gives a safe overestimate.
2445      */
2446     return eff_1x1_buffer_fac_overest*(minimum_subgrid_size_xy(iGrid) +
2447                                        minimum_subgrid_size_xy(jGrid));
2448 }
2449
2450 /* Estimates the interaction volume^2 for non-local interactions */
2451 static real nonlocal_vol2(const struct gmx_domdec_zones_t *zones, const rvec ls, real r)
2452 {
2453     real cl, ca, za;
2454     real vold_est;
2455     real vol2_est_tot;
2456
2457     vol2_est_tot = 0;
2458
2459     /* Here we simply add up the volumes of 1, 2 or 3 1D decomposition
2460      * not home interaction volume^2. As these volumes are not additive,
2461      * this is an overestimate, but it would only be significant in the limit
2462      * of small cells, where we anyhow need to split the lists into
2463      * as small parts as possible.
2464      */
2465
2466     for (int z = 0; z < zones->n; z++)
2467     {
2468         if (zones->shift[z][XX] + zones->shift[z][YY] + zones->shift[z][ZZ] == 1)
2469         {
2470             cl = 0;
2471             ca = 1;
2472             za = 1;
2473             for (int d = 0; d < DIM; d++)
2474             {
2475                 if (zones->shift[z][d] == 0)
2476                 {
2477                     cl += 0.5*ls[d];
2478                     ca *= ls[d];
2479                     za *= zones->size[z].x1[d] - zones->size[z].x0[d];
2480                 }
2481             }
2482
2483             /* 4 octants of a sphere */
2484             vold_est  = 0.25*M_PI*r*r*r*r;
2485             /* 4 quarter pie slices on the edges */
2486             vold_est += 4*cl*M_PI/6.0*r*r*r;
2487             /* One rectangular volume on a face */
2488             vold_est += ca*0.5*r*r;
2489
2490             vol2_est_tot += vold_est*za;
2491         }
2492     }
2493
2494     return vol2_est_tot;
2495 }
2496
2497 /* Estimates the average size of a full j-list for super/sub setup */
2498 static void get_nsubpair_target(const PairSearch          &pairSearch,
2499                                 const InteractionLocality  iloc,
2500                                 const real                 rlist,
2501                                 const int                  min_ci_balanced,
2502                                 int                       *nsubpair_target,
2503                                 float                     *nsubpair_tot_est)
2504 {
2505     /* The target value of 36 seems to be the optimum for Kepler.
2506      * Maxwell is less sensitive to the exact value.
2507      */
2508     const int           nsubpair_target_min = 36;
2509     real                r_eff_sup, vol_est, nsp_est, nsp_est_nl;
2510
2511     const Grid         &grid = pairSearch.gridSet().grids()[0];
2512
2513     /* We don't need to balance list sizes if:
2514      * - We didn't request balancing.
2515      * - The number of grid cells >= the number of lists requested,
2516      *   since we will always generate at least #cells lists.
2517      * - We don't have any cells, since then there won't be any lists.
2518      */
2519     if (min_ci_balanced <= 0 || grid.numCells() >= min_ci_balanced || grid.numCells() == 0)
2520     {
2521         /* nsubpair_target==0 signals no balancing */
2522         *nsubpair_target  = 0;
2523         *nsubpair_tot_est = 0;
2524
2525         return;
2526     }
2527
2528     gmx::RVec               ls;
2529     const int               numAtomsCluster = grid.geometry().numAtomsICluster;
2530     const Grid::Dimensions &dims            = grid.dimensions();
2531
2532     ls[XX] = dims.cellSize[XX]/c_gpuNumClusterPerCellX;
2533     ls[YY] = dims.cellSize[YY]/c_gpuNumClusterPerCellY;
2534     ls[ZZ] = numAtomsCluster/(dims.atomDensity*ls[XX]*ls[YY]);
2535
2536     /* The formulas below are a heuristic estimate of the average nsj per si*/
2537     r_eff_sup = rlist + nbnxn_get_rlist_effective_inc(numAtomsCluster, ls);
2538
2539     if (!pairSearch.domainSetup().haveDomDec ||
2540         pairSearch.domainSetup().zones->n == 1)
2541     {
2542         nsp_est_nl = 0;
2543     }
2544     else
2545     {
2546         nsp_est_nl =
2547             gmx::square(dims.atomDensity/numAtomsCluster)*
2548             nonlocal_vol2(pairSearch.domainSetup().zones, ls, r_eff_sup);
2549     }
2550
2551     if (iloc == InteractionLocality::Local)
2552     {
2553         /* Sub-cell interacts with itself */
2554         vol_est  = ls[XX]*ls[YY]*ls[ZZ];
2555         /* 6/2 rectangular volume on the faces */
2556         vol_est += (ls[XX]*ls[YY] + ls[XX]*ls[ZZ] + ls[YY]*ls[ZZ])*r_eff_sup;
2557         /* 12/2 quarter pie slices on the edges */
2558         vol_est += 2*(ls[XX] + ls[YY] + ls[ZZ])*0.25*M_PI*gmx::square(r_eff_sup);
2559         /* 4 octants of a sphere */
2560         vol_est += 0.5*4.0/3.0*M_PI*gmx::power3(r_eff_sup);
2561
2562         /* Estimate the number of cluster pairs as the local number of
2563          * clusters times the volume they interact with times the density.
2564          */
2565         nsp_est = grid.numClusters()*vol_est*dims.atomDensity/numAtomsCluster;
2566
2567         /* Subtract the non-local pair count */
2568         nsp_est -= nsp_est_nl;
2569
2570         /* For small cut-offs nsp_est will be an underesimate.
2571          * With DD nsp_est_nl is an overestimate so nsp_est can get negative.
2572          * So to avoid too small or negative nsp_est we set a minimum of
2573          * all cells interacting with all 3^3 direct neighbors (3^3-1)/2+1=14.
2574          * This might be a slight overestimate for small non-periodic groups of
2575          * atoms as will occur for a local domain with DD, but for small
2576          * groups of atoms we'll anyhow be limited by nsubpair_target_min,
2577          * so this overestimation will not matter.
2578          */
2579         nsp_est = std::max(nsp_est, grid.numClusters()*14._real);
2580
2581         if (debug)
2582         {
2583             fprintf(debug, "nsp_est local %5.1f non-local %5.1f\n",
2584                     nsp_est, nsp_est_nl);
2585         }
2586     }
2587     else
2588     {
2589         nsp_est = nsp_est_nl;
2590     }
2591
2592     /* Thus the (average) maximum j-list size should be as follows.
2593      * Since there is overhead, we shouldn't make the lists too small
2594      * (and we can't chop up j-groups) so we use a minimum target size of 36.
2595      */
2596     *nsubpair_target  = std::max(nsubpair_target_min,
2597                                  roundToInt(nsp_est/min_ci_balanced));
2598     *nsubpair_tot_est = static_cast<int>(nsp_est);
2599
2600     if (debug)
2601     {
2602         fprintf(debug, "nbl nsp estimate %.1f, nsubpair_target %d\n",
2603                 nsp_est, *nsubpair_target);
2604     }
2605 }
2606
2607 /* Debug list print function */
2608 static void print_nblist_ci_cj(FILE *fp, const NbnxnPairlistCpu *nbl)
2609 {
2610     for (const nbnxn_ci_t &ciEntry : nbl->ci)
2611     {
2612         fprintf(fp, "ci %4d  shift %2d  ncj %3d\n",
2613                 ciEntry.ci, ciEntry.shift,
2614                 ciEntry.cj_ind_end - ciEntry.cj_ind_start);
2615
2616         for (int j = ciEntry.cj_ind_start; j < ciEntry.cj_ind_end; j++)
2617         {
2618             fprintf(fp, "  cj %5d  imask %x\n",
2619                     nbl->cj[j].cj,
2620                     nbl->cj[j].excl);
2621         }
2622     }
2623 }
2624
2625 /* Debug list print function */
2626 static void print_nblist_sci_cj(FILE *fp, const NbnxnPairlistGpu *nbl)
2627 {
2628     for (const nbnxn_sci_t &sci : nbl->sci)
2629     {
2630         fprintf(fp, "ci %4d  shift %2d  ncj4 %2d\n",
2631                 sci.sci, sci.shift,
2632                 sci.numJClusterGroups());
2633
2634         int ncp = 0;
2635         for (int j4 = sci.cj4_ind_start; j4 < sci.cj4_ind_end; j4++)
2636         {
2637             for (int j = 0; j < c_nbnxnGpuJgroupSize; j++)
2638             {
2639                 fprintf(fp, "  sj %5d  imask %x\n",
2640                         nbl->cj4[j4].cj[j],
2641                         nbl->cj4[j4].imei[0].imask);
2642                 for (int si = 0; si < c_gpuNumClusterPerCell; si++)
2643                 {
2644                     if (nbl->cj4[j4].imei[0].imask & (1U << (j*c_gpuNumClusterPerCell + si)))
2645                     {
2646                         ncp++;
2647                     }
2648                 }
2649             }
2650         }
2651         fprintf(fp, "ci %4d  shift %2d  ncj4 %2d ncp %3d\n",
2652                 sci.sci, sci.shift,
2653                 sci.numJClusterGroups(),
2654                 ncp);
2655     }
2656 }
2657
2658 /* Combine pair lists *nbl generated on multiple threads nblc */
2659 static void combine_nblists(int nnbl, NbnxnPairlistGpu **nbl,
2660                             NbnxnPairlistGpu *nblc)
2661 {
2662     int nsci  = nblc->sci.size();
2663     int ncj4  = nblc->cj4.size();
2664     int nexcl = nblc->excl.size();
2665     for (int i = 0; i < nnbl; i++)
2666     {
2667         nsci  += nbl[i]->sci.size();
2668         ncj4  += nbl[i]->cj4.size();
2669         nexcl += nbl[i]->excl.size();
2670     }
2671
2672     /* Resize with the final, combined size, so we can fill in parallel */
2673     /* NOTE: For better performance we should use default initialization */
2674     nblc->sci.resize(nsci);
2675     nblc->cj4.resize(ncj4);
2676     nblc->excl.resize(nexcl);
2677
2678     /* Each thread should copy its own data to the combined arrays,
2679      * as otherwise data will go back and forth between different caches.
2680      */
2681 #if GMX_OPENMP && !(defined __clang_analyzer__)
2682     int nthreads = gmx_omp_nthreads_get(emntPairsearch);
2683 #endif
2684
2685 #pragma omp parallel for num_threads(nthreads) schedule(static)
2686     for (int n = 0; n < nnbl; n++)
2687     {
2688         try
2689         {
2690             /* Determine the offset in the combined data for our thread.
2691              * Note that the original sizes in nblc are lost.
2692              */
2693             int sci_offset  = nsci;
2694             int cj4_offset  = ncj4;
2695             int excl_offset = nexcl;
2696
2697             for (int i = n; i < nnbl; i++)
2698             {
2699                 sci_offset  -= nbl[i]->sci.size();
2700                 cj4_offset  -= nbl[i]->cj4.size();
2701                 excl_offset -= nbl[i]->excl.size();
2702             }
2703
2704             const NbnxnPairlistGpu &nbli = *nbl[n];
2705
2706             for (size_t i = 0; i < nbli.sci.size(); i++)
2707             {
2708                 nblc->sci[sci_offset + i]                = nbli.sci[i];
2709                 nblc->sci[sci_offset + i].cj4_ind_start += cj4_offset;
2710                 nblc->sci[sci_offset + i].cj4_ind_end   += cj4_offset;
2711             }
2712
2713             for (size_t j4 = 0; j4 < nbli.cj4.size(); j4++)
2714             {
2715                 nblc->cj4[cj4_offset + j4]                   = nbli.cj4[j4];
2716                 nblc->cj4[cj4_offset + j4].imei[0].excl_ind += excl_offset;
2717                 nblc->cj4[cj4_offset + j4].imei[1].excl_ind += excl_offset;
2718             }
2719
2720             for (size_t j4 = 0; j4 < nbli.excl.size(); j4++)
2721             {
2722                 nblc->excl[excl_offset + j4] = nbli.excl[j4];
2723             }
2724         }
2725         GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR;
2726     }
2727
2728     for (int n = 0; n < nnbl; n++)
2729     {
2730         nblc->nci_tot += nbl[n]->nci_tot;
2731     }
2732 }
2733
2734 static void balance_fep_lists(gmx::ArrayRef<PairsearchWork>       work,
2735                               nbnxn_pairlist_set_t               *nbl_lists)
2736 {
2737     int       nnbl;
2738     int       nri_tot, nrj_tot, nrj_target;
2739     int       th_dest;
2740     t_nblist *nbld;
2741
2742     nnbl = nbl_lists->nnbl;
2743
2744     if (nnbl == 1)
2745     {
2746         /* Nothing to balance */
2747         return;
2748     }
2749
2750     /* Count the total i-lists and pairs */
2751     nri_tot = 0;
2752     nrj_tot = 0;
2753     for (int th = 0; th < nnbl; th++)
2754     {
2755         nri_tot += nbl_lists->nbl_fep[th]->nri;
2756         nrj_tot += nbl_lists->nbl_fep[th]->nrj;
2757     }
2758
2759     nrj_target = (nrj_tot + nnbl - 1)/nnbl;
2760
2761     assert(gmx_omp_nthreads_get(emntNonbonded) == nnbl);
2762
2763 #pragma omp parallel for schedule(static) num_threads(nnbl)
2764     for (int th = 0; th < nnbl; th++)
2765     {
2766         try
2767         {
2768             t_nblist *nbl = work[th].nbl_fep.get();
2769
2770             /* Note that here we allocate for the total size, instead of
2771              * a per-thread esimate (which is hard to obtain).
2772              */
2773             if (nri_tot > nbl->maxnri)
2774             {
2775                 nbl->maxnri = over_alloc_large(nri_tot);
2776                 reallocate_nblist(nbl);
2777             }
2778             if (nri_tot > nbl->maxnri || nrj_tot > nbl->maxnrj)
2779             {
2780                 nbl->maxnrj = over_alloc_small(nrj_tot);
2781                 srenew(nbl->jjnr, nbl->maxnrj);
2782                 srenew(nbl->excl_fep, nbl->maxnrj);
2783             }
2784
2785             clear_pairlist_fep(nbl);
2786         }
2787         GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR;
2788     }
2789
2790     /* Loop over the source lists and assign and copy i-entries */
2791     th_dest = 0;
2792     nbld    = work[th_dest].nbl_fep.get();
2793     for (int th = 0; th < nnbl; th++)
2794     {
2795         t_nblist *nbls;
2796
2797         nbls = nbl_lists->nbl_fep[th];
2798
2799         for (int i = 0; i < nbls->nri; i++)
2800         {
2801             int nrj;
2802
2803             /* The number of pairs in this i-entry */
2804             nrj = nbls->jindex[i+1] - nbls->jindex[i];
2805
2806             /* Decide if list th_dest is too large and we should procede
2807              * to the next destination list.
2808              */
2809             if (th_dest+1 < nnbl && nbld->nrj > 0 &&
2810                 nbld->nrj + nrj - nrj_target > nrj_target - nbld->nrj)
2811             {
2812                 th_dest++;
2813                 nbld = work[th_dest].nbl_fep.get();
2814             }
2815
2816             nbld->iinr[nbld->nri]  = nbls->iinr[i];
2817             nbld->gid[nbld->nri]   = nbls->gid[i];
2818             nbld->shift[nbld->nri] = nbls->shift[i];
2819
2820             for (int j = nbls->jindex[i]; j < nbls->jindex[i+1]; j++)
2821             {
2822                 nbld->jjnr[nbld->nrj]     = nbls->jjnr[j];
2823                 nbld->excl_fep[nbld->nrj] = nbls->excl_fep[j];
2824                 nbld->nrj++;
2825             }
2826             nbld->nri++;
2827             nbld->jindex[nbld->nri] = nbld->nrj;
2828         }
2829     }
2830
2831     /* Swap the list pointers */
2832     for (int th = 0; th < nnbl; th++)
2833     {
2834         t_nblist *nbl_tmp      = work[th].nbl_fep.release();
2835         work[th].nbl_fep.reset(nbl_lists->nbl_fep[th]);
2836         nbl_lists->nbl_fep[th] = nbl_tmp;
2837
2838         if (debug)
2839         {
2840             fprintf(debug, "nbl_fep[%d] nri %4d nrj %4d\n",
2841                     th,
2842                     nbl_lists->nbl_fep[th]->nri,
2843                     nbl_lists->nbl_fep[th]->nrj);
2844         }
2845     }
2846 }
2847
2848 /* Returns the next ci to be processes by our thread */
2849 static gmx_bool next_ci(const Grid &grid,
2850                         int nth, int ci_block,
2851                         int *ci_x, int *ci_y,
2852                         int *ci_b, int *ci)
2853 {
2854     (*ci_b)++;
2855     (*ci)++;
2856
2857     if (*ci_b == ci_block)
2858     {
2859         /* Jump to the next block assigned to this task */
2860         *ci   += (nth - 1)*ci_block;
2861         *ci_b  = 0;
2862     }
2863
2864     if (*ci >= grid.numCells())
2865     {
2866         return FALSE;
2867     }
2868
2869     while (*ci >= grid.firstCellInColumn(*ci_x*grid.dimensions().numCells[YY] + *ci_y + 1))
2870     {
2871         *ci_y += 1;
2872         if (*ci_y == grid.dimensions().numCells[YY])
2873         {
2874             *ci_x += 1;
2875             *ci_y  = 0;
2876         }
2877     }
2878
2879     return TRUE;
2880 }
2881
2882 /* Returns the distance^2 for which we put cell pairs in the list
2883  * without checking atom pair distances. This is usually < rlist^2.
2884  */
2885 static float boundingbox_only_distance2(const Grid::Dimensions &iGridDims,
2886                                         const Grid::Dimensions &jGridDims,
2887                                         real                    rlist,
2888                                         gmx_bool                simple)
2889 {
2890     /* If the distance between two sub-cell bounding boxes is less
2891      * than this distance, do not check the distance between
2892      * all particle pairs in the sub-cell, since then it is likely
2893      * that the box pair has atom pairs within the cut-off.
2894      * We use the nblist cut-off minus 0.5 times the average x/y diagonal
2895      * spacing of the sub-cells. Around 40% of the checked pairs are pruned.
2896      * Using more than 0.5 gains at most 0.5%.
2897      * If forces are calculated more than twice, the performance gain
2898      * in the force calculation outweighs the cost of checking.
2899      * Note that with subcell lists, the atom-pair distance check
2900      * is only performed when only 1 out of 8 sub-cells in within range,
2901      * this is because the GPU is much faster than the cpu.
2902      */
2903     real bbx, bby;
2904     real rbb2;
2905
2906     bbx = 0.5*(iGridDims.cellSize[XX] + jGridDims.cellSize[XX]);
2907     bby = 0.5*(iGridDims.cellSize[YY] + jGridDims.cellSize[YY]);
2908     if (!simple)
2909     {
2910         bbx /= c_gpuNumClusterPerCellX;
2911         bby /= c_gpuNumClusterPerCellY;
2912     }
2913
2914     rbb2 = std::max(0.0, rlist - 0.5*std::sqrt(bbx*bbx + bby*bby));
2915     rbb2 = rbb2 * rbb2;
2916
2917 #if !GMX_DOUBLE
2918     return rbb2;
2919 #else
2920     return (float)((1+GMX_FLOAT_EPS)*rbb2);
2921 #endif
2922 }
2923
2924 static int get_ci_block_size(const Grid &iGrid,
2925                              gmx_bool bDomDec, int nth)
2926 {
2927     const int ci_block_enum      = 5;
2928     const int ci_block_denom     = 11;
2929     const int ci_block_min_atoms = 16;
2930     int       ci_block;
2931
2932     /* Here we decide how to distribute the blocks over the threads.
2933      * We use prime numbers to try to avoid that the grid size becomes
2934      * a multiple of the number of threads, which would lead to some
2935      * threads getting "inner" pairs and others getting boundary pairs,
2936      * which in turns will lead to load imbalance between threads.
2937      * Set the block size as 5/11/ntask times the average number of cells
2938      * in a y,z slab. This should ensure a quite uniform distribution
2939      * of the grid parts of the different thread along all three grid
2940      * zone boundaries with 3D domain decomposition. At the same time
2941      * the blocks will not become too small.
2942      */
2943     ci_block = (iGrid.numCells()*ci_block_enum)/(ci_block_denom*iGrid.dimensions().numCells[XX]*nth);
2944
2945     const int numAtomsPerCell = iGrid.geometry().numAtomsPerCell;
2946
2947     /* Ensure the blocks are not too small: avoids cache invalidation */
2948     if (ci_block*numAtomsPerCell < ci_block_min_atoms)
2949     {
2950         ci_block = (ci_block_min_atoms + numAtomsPerCell - 1)/numAtomsPerCell;
2951     }
2952
2953     /* Without domain decomposition
2954      * or with less than 3 blocks per task, divide in nth blocks.
2955      */
2956     if (!bDomDec || nth*3*ci_block > iGrid.numCells())
2957     {
2958         ci_block = (iGrid.numCells() + nth - 1)/nth;
2959     }
2960
2961     if (ci_block > 1 && (nth - 1)*ci_block >= iGrid.numCells())
2962     {
2963         /* Some threads have no work. Although reducing the block size
2964          * does not decrease the block count on the first few threads,
2965          * with GPUs better mixing of "upper" cells that have more empty
2966          * clusters results in a somewhat lower max load over all threads.
2967          * Without GPUs the regime of so few atoms per thread is less
2968          * performance relevant, but with 8-wide SIMD the same reasoning
2969          * applies, since the pair list uses 4 i-atom "sub-clusters".
2970          */
2971         ci_block--;
2972     }
2973
2974     return ci_block;
2975 }
2976
2977 /* Returns the number of bits to right-shift a cluster index to obtain
2978  * the corresponding force buffer flag index.
2979  */
2980 static int getBufferFlagShift(int numAtomsPerCluster)
2981 {
2982     int bufferFlagShift = 0;
2983     while ((numAtomsPerCluster << bufferFlagShift) < NBNXN_BUFFERFLAG_SIZE)
2984     {
2985         bufferFlagShift++;
2986     }
2987
2988     return bufferFlagShift;
2989 }
2990
2991 static bool pairlistIsSimple(const NbnxnPairlistCpu gmx_unused &pairlist)
2992 {
2993     return true;
2994 }
2995
2996 static bool pairlistIsSimple(const NbnxnPairlistGpu gmx_unused &pairlist)
2997 {
2998     return false;
2999 }
3000
3001 static void makeClusterListWrapper(NbnxnPairlistCpu       *nbl,
3002                                    const Grid gmx_unused  &iGrid,
3003                                    const int               ci,
3004                                    const Grid             &jGrid,
3005                                    const int               firstCell,
3006                                    const int               lastCell,
3007                                    const bool              excludeSubDiagonal,
3008                                    const nbnxn_atomdata_t *nbat,
3009                                    const real              rlist2,
3010                                    const real              rbb2,
3011                                    const Nbnxm::KernelType kernelType,
3012                                    int                    *numDistanceChecks)
3013 {
3014     switch (kernelType)
3015     {
3016         case Nbnxm::KernelType::Cpu4x4_PlainC:
3017             makeClusterListSimple(jGrid,
3018                                   nbl, ci, firstCell, lastCell,
3019                                   excludeSubDiagonal,
3020                                   nbat->x().data(),
3021                                   rlist2, rbb2,
3022                                   numDistanceChecks);
3023             break;
3024 #ifdef GMX_NBNXN_SIMD_4XN
3025         case Nbnxm::KernelType::Cpu4xN_Simd_4xN:
3026             makeClusterListSimd4xn(jGrid,
3027                                    nbl, ci, firstCell, lastCell,
3028                                    excludeSubDiagonal,
3029                                    nbat->x().data(),
3030                                    rlist2, rbb2,
3031                                    numDistanceChecks);
3032             break;
3033 #endif
3034 #ifdef GMX_NBNXN_SIMD_2XNN
3035         case Nbnxm::KernelType::Cpu4xN_Simd_2xNN:
3036             makeClusterListSimd2xnn(jGrid,
3037                                     nbl, ci, firstCell, lastCell,
3038                                     excludeSubDiagonal,
3039                                     nbat->x().data(),
3040                                     rlist2, rbb2,
3041                                     numDistanceChecks);
3042             break;
3043 #endif
3044         default:
3045             GMX_ASSERT(false, "Unhandled kernel type");
3046     }
3047 }
3048
3049 static void makeClusterListWrapper(NbnxnPairlistGpu             *nbl,
3050                                    const Grid &gmx_unused        iGrid,
3051                                    const int                     ci,
3052                                    const Grid                   &jGrid,
3053                                    const int                     firstCell,
3054                                    const int                     lastCell,
3055                                    const bool                    excludeSubDiagonal,
3056                                    const nbnxn_atomdata_t       *nbat,
3057                                    const real                    rlist2,
3058                                    const real                    rbb2,
3059                                    Nbnxm::KernelType gmx_unused  kernelType,
3060                                    int                          *numDistanceChecks)
3061 {
3062     for (int cj = firstCell; cj <= lastCell; cj++)
3063     {
3064         make_cluster_list_supersub(iGrid, jGrid,
3065                                    nbl, ci, cj,
3066                                    excludeSubDiagonal,
3067                                    nbat->xstride, nbat->x().data(),
3068                                    rlist2, rbb2,
3069                                    numDistanceChecks);
3070     }
3071 }
3072
3073 static int getNumSimpleJClustersInList(const NbnxnPairlistCpu &nbl)
3074 {
3075     return nbl.cj.size();
3076 }
3077
3078 static int getNumSimpleJClustersInList(const gmx_unused NbnxnPairlistGpu &nbl)
3079 {
3080     return 0;
3081 }
3082
3083 static void incrementNumSimpleJClustersInList(NbnxnPairlistCpu *nbl,
3084                                               int               ncj_old_j)
3085 {
3086     nbl->ncjInUse += nbl->cj.size() - ncj_old_j;
3087 }
3088
3089 static void incrementNumSimpleJClustersInList(NbnxnPairlistGpu gmx_unused *nbl,
3090                                               int              gmx_unused  ncj_old_j)
3091 {
3092 }
3093
3094 static void checkListSizeConsistency(const NbnxnPairlistCpu &nbl,
3095                                      const bool              haveFreeEnergy)
3096 {
3097     GMX_RELEASE_ASSERT(static_cast<size_t>(nbl.ncjInUse) == nbl.cj.size() || haveFreeEnergy,
3098                        "Without free-energy all cj pair-list entries should be in use. "
3099                        "Note that subsequent code does not make use of the equality, "
3100                        "this check is only here to catch bugs");
3101 }
3102
3103 static void checkListSizeConsistency(const NbnxnPairlistGpu gmx_unused &nbl,
3104                                      bool gmx_unused                    haveFreeEnergy)
3105 {
3106     /* We currently can not check consistency here */
3107 }
3108
3109 /* Set the buffer flags for newly added entries in the list */
3110 static void setBufferFlags(const NbnxnPairlistCpu &nbl,
3111                            const int               ncj_old_j,
3112                            const int               gridj_flag_shift,
3113                            gmx_bitmask_t          *gridj_flag,
3114                            const int               th)
3115 {
3116     if (gmx::ssize(nbl.cj) > ncj_old_j)
3117     {
3118         int cbFirst = nbl.cj[ncj_old_j].cj >> gridj_flag_shift;
3119         int cbLast  = nbl.cj.back().cj >> gridj_flag_shift;
3120         for (int cb = cbFirst; cb <= cbLast; cb++)
3121         {
3122             bitmask_init_bit(&gridj_flag[cb], th);
3123         }
3124     }
3125 }
3126
3127 static void setBufferFlags(const NbnxnPairlistGpu gmx_unused &nbl,
3128                            int gmx_unused                     ncj_old_j,
3129                            int gmx_unused                     gridj_flag_shift,
3130                            gmx_bitmask_t gmx_unused          *gridj_flag,
3131                            int gmx_unused                     th)
3132 {
3133     GMX_ASSERT(false, "This function should never be called");
3134 }
3135
3136 /* Generates the part of pair-list nbl assigned to our thread */
3137 template <typename T>
3138 static void nbnxn_make_pairlist_part(const PairSearch &pairSearch,
3139                                      const Grid &iGrid,
3140                                      const Grid &jGrid,
3141                                      PairsearchWork *work,
3142                                      const nbnxn_atomdata_t *nbat,
3143                                      const t_blocka &exclusions,
3144                                      real rlist,
3145                                      const Nbnxm::KernelType kernelType,
3146                                      int ci_block,
3147                                      gmx_bool bFBufferFlag,
3148                                      int nsubpair_max,
3149                                      gmx_bool progBal,
3150                                      float nsubpair_tot_est,
3151                                      int th, int nth,
3152                                      T *nbl,
3153                                      t_nblist *nbl_fep)
3154 {
3155     int               na_cj_2log;
3156     matrix            box;
3157     real              rl_fep2 = 0;
3158     float             rbb2;
3159     int               ci_b, ci, ci_x, ci_y, ci_xy;
3160     ivec              shp;
3161     real              bx0, bx1, by0, by1, bz0, bz1;
3162     real              bz1_frac;
3163     real              d2cx, d2z, d2z_cx, d2z_cy, d2zx, d2zxy, d2xy;
3164     int               cxf, cxl, cyf, cyf_x, cyl;
3165     int               numDistanceChecks;
3166     int               gridi_flag_shift = 0, gridj_flag_shift = 0;
3167     gmx_bitmask_t    *gridj_flag       = nullptr;
3168     int               ncj_old_i, ncj_old_j;
3169
3170     if (jGrid.geometry().isSimple != pairlistIsSimple(*nbl) ||
3171         iGrid.geometry().isSimple != pairlistIsSimple(*nbl))
3172     {
3173         gmx_incons("Grid incompatible with pair-list");
3174     }
3175
3176     sync_work(nbl);
3177     GMX_ASSERT(nbl->na_ci == jGrid.geometry().numAtomsICluster,
3178                "The cluster sizes in the list and grid should match");
3179     nbl->na_cj = Nbnxm::JClusterSizePerKernelType[kernelType];
3180     na_cj_2log = get_2log(nbl->na_cj);
3181
3182     nbl->rlist  = rlist;
3183
3184     if (bFBufferFlag)
3185     {
3186         /* Determine conversion of clusters to flag blocks */
3187         gridi_flag_shift = getBufferFlagShift(nbl->na_ci);
3188         gridj_flag_shift = getBufferFlagShift(nbl->na_cj);
3189
3190         gridj_flag       = work->buffer_flags.flag;
3191     }
3192
3193     const Nbnxm::GridSet &gridSet = pairSearch.gridSet();
3194
3195     gridSet.getBox(box);
3196
3197     const bool            haveFep = gridSet.haveFep();
3198
3199     const real            rlist2  = nbl->rlist*nbl->rlist;
3200
3201     if (haveFep && !pairlistIsSimple(*nbl))
3202     {
3203         /* Determine an atom-pair list cut-off distance for FEP atom pairs.
3204          * We should not simply use rlist, since then we would not have
3205          * the small, effective buffering of the NxN lists.
3206          * The buffer is on overestimate, but the resulting cost for pairs
3207          * beyond rlist is neglible compared to the FEP pairs within rlist.
3208          */
3209         rl_fep2 = nbl->rlist + effective_buffer_1x1_vs_MxN(iGrid, jGrid);
3210
3211         if (debug)
3212         {
3213             fprintf(debug, "nbl_fep atom-pair rlist %f\n", rl_fep2);
3214         }
3215         rl_fep2 = rl_fep2*rl_fep2;
3216     }
3217
3218     const Grid::Dimensions &iGridDims = iGrid.dimensions();
3219     const Grid::Dimensions &jGridDims = jGrid.dimensions();
3220
3221     rbb2 = boundingbox_only_distance2(iGridDims, jGridDims, nbl->rlist, pairlistIsSimple(*nbl));
3222
3223     if (debug)
3224     {
3225         fprintf(debug, "nbl bounding box only distance %f\n", std::sqrt(rbb2));
3226     }
3227
3228     const bool isIntraGridList = (&iGrid == &jGrid);
3229
3230     /* Set the shift range */
3231     for (int d = 0; d < DIM; d++)
3232     {
3233         /* Check if we need periodicity shifts.
3234          * Without PBC or with domain decomposition we don't need them.
3235          */
3236         if (d >= ePBC2npbcdim(pairSearch.domainSetup().ePBC) ||
3237             pairSearch.domainSetup().haveDomDecPerDim[d])
3238         {
3239             shp[d] = 0;
3240         }
3241         else
3242         {
3243             const real listRangeCellToCell =
3244                 listRangeForGridCellToGridCell(rlist, iGrid.dimensions(), jGrid.dimensions());
3245             if (d == XX &&
3246                 box[XX][XX] - fabs(box[YY][XX]) - fabs(box[ZZ][XX]) < listRangeCellToCell)
3247             {
3248                 shp[d] = 2;
3249             }
3250             else
3251             {
3252                 shp[d] = 1;
3253             }
3254         }
3255     }
3256     const bool bSimple = pairlistIsSimple(*nbl);
3257     gmx::ArrayRef<const BoundingBox> bb_i;
3258 #if NBNXN_BBXXXX
3259     gmx::ArrayRef<const float>       pbb_i;
3260     if (bSimple)
3261     {
3262         bb_i  = iGrid.iBoundingBoxes();
3263     }
3264     else
3265     {
3266         pbb_i = iGrid.packedBoundingBoxes();
3267     }
3268 #else
3269     /* We use the normal bounding box format for both grid types */
3270     bb_i  = iGrid.iBoundingBoxes();
3271 #endif
3272     gmx::ArrayRef<const BoundingBox1D> bbcz_i  = iGrid.zBoundingBoxes();
3273     gmx::ArrayRef<const int>           flags_i = iGrid.clusterFlags();
3274     gmx::ArrayRef<const BoundingBox1D> bbcz_j  = jGrid.zBoundingBoxes();
3275     int                                cell0_i = iGrid.cellOffset();
3276
3277     if (debug)
3278     {
3279         fprintf(debug, "nbl nc_i %d col.av. %.1f ci_block %d\n",
3280                 iGrid.numCells(), iGrid.numCells()/static_cast<double>(iGrid.numColumns()), ci_block);
3281     }
3282
3283     numDistanceChecks = 0;
3284
3285     const real listRangeBBToJCell2 = gmx::square(listRangeForBoundingBoxToGridCell(rlist, jGrid.dimensions()));
3286
3287     /* Initially ci_b and ci to 1 before where we want them to start,
3288      * as they will both be incremented in next_ci.
3289      */
3290     ci_b = -1;
3291     ci   = th*ci_block - 1;
3292     ci_x = 0;
3293     ci_y = 0;
3294     while (next_ci(iGrid, nth, ci_block, &ci_x, &ci_y, &ci_b, &ci))
3295     {
3296         if (bSimple && flags_i[ci] == 0)
3297         {
3298             continue;
3299         }
3300
3301         ncj_old_i = getNumSimpleJClustersInList(*nbl);
3302
3303         d2cx = 0;
3304         if (!isIntraGridList && shp[XX] == 0)
3305         {
3306             if (bSimple)
3307             {
3308                 bx1 = bb_i[ci].upper.x;
3309             }
3310             else
3311             {
3312                 bx1 = iGridDims.lowerCorner[XX] + (ci_x+1)*iGridDims.cellSize[XX];
3313             }
3314             if (bx1 < jGridDims.lowerCorner[XX])
3315             {
3316                 d2cx = gmx::square(jGridDims.lowerCorner[XX] - bx1);
3317
3318                 if (d2cx >= listRangeBBToJCell2)
3319                 {
3320                     continue;
3321                 }
3322             }
3323         }
3324
3325         ci_xy = ci_x*iGridDims.numCells[YY] + ci_y;
3326
3327         /* Loop over shift vectors in three dimensions */
3328         for (int tz = -shp[ZZ]; tz <= shp[ZZ]; tz++)
3329         {
3330             const real shz = tz*box[ZZ][ZZ];
3331
3332             bz0 = bbcz_i[ci].lower + shz;
3333             bz1 = bbcz_i[ci].upper + shz;
3334
3335             if (tz == 0)
3336             {
3337                 d2z = 0;
3338             }
3339             else if (tz < 0)
3340             {
3341                 d2z = gmx::square(bz1);
3342             }
3343             else
3344             {
3345                 d2z = gmx::square(bz0 - box[ZZ][ZZ]);
3346             }
3347
3348             d2z_cx = d2z + d2cx;
3349
3350             if (d2z_cx >= rlist2)
3351             {
3352                 continue;
3353             }
3354
3355             bz1_frac = bz1/iGrid.numCellsInColumn(ci_xy);
3356             if (bz1_frac < 0)
3357             {
3358                 bz1_frac = 0;
3359             }
3360             /* The check with bz1_frac close to or larger than 1 comes later */
3361
3362             for (int ty = -shp[YY]; ty <= shp[YY]; ty++)
3363             {
3364                 const real shy = ty*box[YY][YY] + tz*box[ZZ][YY];
3365
3366                 if (bSimple)
3367                 {
3368                     by0 = bb_i[ci].lower.y + shy;
3369                     by1 = bb_i[ci].upper.y + shy;
3370                 }
3371                 else
3372                 {
3373                     by0 = iGridDims.lowerCorner[YY] + (ci_y    )*iGridDims.cellSize[YY] + shy;
3374                     by1 = iGridDims.lowerCorner[YY] + (ci_y + 1)*iGridDims.cellSize[YY] + shy;
3375                 }
3376
3377                 get_cell_range<YY>(by0, by1,
3378                                    jGridDims,
3379                                    d2z_cx, rlist,
3380                                    &cyf, &cyl);
3381
3382                 if (cyf > cyl)
3383                 {
3384                     continue;
3385                 }
3386
3387                 d2z_cy = d2z;
3388                 if (by1 < jGridDims.lowerCorner[YY])
3389                 {
3390                     d2z_cy += gmx::square(jGridDims.lowerCorner[YY] - by1);
3391                 }
3392                 else if (by0 > jGridDims.upperCorner[YY])
3393                 {
3394                     d2z_cy += gmx::square(by0 - jGridDims.upperCorner[YY]);
3395                 }
3396
3397                 for (int tx = -shp[XX]; tx <= shp[XX]; tx++)
3398                 {
3399                     const int  shift              = XYZ2IS(tx, ty, tz);
3400
3401                     const bool excludeSubDiagonal = (isIntraGridList && shift == CENTRAL);
3402
3403                     if (c_pbcShiftBackward && isIntraGridList && shift > CENTRAL)
3404                     {
3405                         continue;
3406                     }
3407
3408                     const real shx = tx*box[XX][XX] + ty*box[YY][XX] + tz*box[ZZ][XX];
3409
3410                     if (bSimple)
3411                     {
3412                         bx0 = bb_i[ci].lower.x + shx;
3413                         bx1 = bb_i[ci].upper.x + shx;
3414                     }
3415                     else
3416                     {
3417                         bx0 = iGridDims.lowerCorner[XX] + (ci_x  )*iGridDims.cellSize[XX] + shx;
3418                         bx1 = iGridDims.lowerCorner[XX] + (ci_x+1)*iGridDims.cellSize[XX] + shx;
3419                     }
3420
3421                     get_cell_range<XX>(bx0, bx1,
3422                                        jGridDims,
3423                                        d2z_cy, rlist,
3424                                        &cxf, &cxl);
3425
3426                     if (cxf > cxl)
3427                     {
3428                         continue;
3429                     }
3430
3431                     addNewIEntry(nbl, cell0_i+ci, shift, flags_i[ci]);
3432
3433                     if ((!c_pbcShiftBackward || excludeSubDiagonal) &&
3434                         cxf < ci_x)
3435                     {
3436                         /* Leave the pairs with i > j.
3437                          * x is the major index, so skip half of it.
3438                          */
3439                         cxf = ci_x;
3440                     }
3441
3442                     set_icell_bb(iGrid, ci, shx, shy, shz,
3443                                  nbl->work);
3444
3445                     icell_set_x(cell0_i+ci, shx, shy, shz,
3446                                 nbat->xstride, nbat->x().data(),
3447                                 kernelType,
3448                                 nbl->work);
3449
3450                     for (int cx = cxf; cx <= cxl; cx++)
3451                     {
3452                         d2zx = d2z;
3453                         if (jGridDims.lowerCorner[XX] + cx*jGridDims.cellSize[XX] > bx1)
3454                         {
3455                             d2zx += gmx::square(jGridDims.lowerCorner[XX] + cx*jGridDims.cellSize[XX] - bx1);
3456                         }
3457                         else if (jGridDims.lowerCorner[XX] + (cx+1)*jGridDims.cellSize[XX] < bx0)
3458                         {
3459                             d2zx += gmx::square(jGridDims.lowerCorner[XX] + (cx+1)*jGridDims.cellSize[XX] - bx0);
3460                         }
3461
3462                         if (isIntraGridList &&
3463                             cx == 0 &&
3464                             (!c_pbcShiftBackward || shift == CENTRAL) &&
3465                             cyf < ci_y)
3466                         {
3467                             /* Leave the pairs with i > j.
3468                              * Skip half of y when i and j have the same x.
3469                              */
3470                             cyf_x = ci_y;
3471                         }
3472                         else
3473                         {
3474                             cyf_x = cyf;
3475                         }
3476
3477                         for (int cy = cyf_x; cy <= cyl; cy++)
3478                         {
3479                             const int columnStart = jGrid.firstCellInColumn(cx*jGridDims.numCells[YY] + cy);
3480                             const int columnEnd   = jGrid.firstCellInColumn(cx*jGridDims.numCells[YY] + cy + 1);
3481
3482                             d2zxy = d2zx;
3483                             if (jGridDims.lowerCorner[YY] + cy*jGridDims.cellSize[YY] > by1)
3484                             {
3485                                 d2zxy += gmx::square(jGridDims.lowerCorner[YY] + cy*jGridDims.cellSize[YY] - by1);
3486                             }
3487                             else if (jGridDims.lowerCorner[YY] + (cy + 1)*jGridDims.cellSize[YY] < by0)
3488                             {
3489                                 d2zxy += gmx::square(jGridDims.lowerCorner[YY] + (cy + 1)*jGridDims.cellSize[YY] - by0);
3490                             }
3491                             if (columnStart < columnEnd && d2zxy < listRangeBBToJCell2)
3492                             {
3493                                 /* To improve efficiency in the common case
3494                                  * of a homogeneous particle distribution,
3495                                  * we estimate the index of the middle cell
3496                                  * in range (midCell). We search down and up
3497                                  * starting from this index.
3498                                  *
3499                                  * Note that the bbcz_j array contains bounds
3500                                  * for i-clusters, thus for clusters of 4 atoms.
3501                                  * For the common case where the j-cluster size
3502                                  * is 8, we could step with a stride of 2,
3503                                  * but we do not do this because it would
3504                                  * complicate this code even more.
3505                                  */
3506                                 int midCell = columnStart + static_cast<int>(bz1_frac*(columnEnd - columnStart));
3507                                 if (midCell >= columnEnd)
3508                                 {
3509                                     midCell = columnEnd - 1;
3510                                 }
3511
3512                                 d2xy = d2zxy - d2z;
3513
3514                                 /* Find the lowest cell that can possibly
3515                                  * be within range.
3516                                  * Check if we hit the bottom of the grid,
3517                                  * if the j-cell is below the i-cell and if so,
3518                                  * if it is within range.
3519                                  */
3520                                 int downTestCell = midCell;
3521                                 while (downTestCell >= columnStart &&
3522                                        (bbcz_j[downTestCell].upper >= bz0 ||
3523                                         d2xy + gmx::square(bbcz_j[downTestCell].upper - bz0) < rlist2))
3524                                 {
3525                                     downTestCell--;
3526                                 }
3527                                 int firstCell = downTestCell + 1;
3528
3529                                 /* Find the highest cell that can possibly
3530                                  * be within range.
3531                                  * Check if we hit the top of the grid,
3532                                  * if the j-cell is above the i-cell and if so,
3533                                  * if it is within range.
3534                                  */
3535                                 int upTestCell = midCell + 1;
3536                                 while (upTestCell < columnEnd &&
3537                                        (bbcz_j[upTestCell].lower <= bz1 ||
3538                                         d2xy + gmx::square(bbcz_j[upTestCell].lower - bz1) < rlist2))
3539                                 {
3540                                     upTestCell++;
3541                                 }
3542                                 int lastCell = upTestCell - 1;
3543
3544 #define NBNXN_REFCODE 0
3545 #if NBNXN_REFCODE
3546                                 {
3547                                     /* Simple reference code, for debugging,
3548                                      * overrides the more complex code above.
3549                                      */
3550                                     firstCell = columnEnd;
3551                                     lastCell  = -1;
3552                                     for (int k = columnStart; k < columnEnd; k++)
3553                                     {
3554                                         if (d2xy + gmx::square(bbcz_j[k*NNBSBB_D + 1] - bz0) < rlist2 &&
3555                                             k < firstCell)
3556                                         {
3557                                             firstCell = k;
3558                                         }
3559                                         if (d2xy + gmx::square(bbcz_j[k*NNBSBB_D] - bz1) < rlist2 &&
3560                                             k > lastCell)
3561                                         {
3562                                             lastCell = k;
3563                                         }
3564                                     }
3565                                 }
3566 #endif
3567
3568                                 if (isIntraGridList)
3569                                 {
3570                                     /* We want each atom/cell pair only once,
3571                                      * only use cj >= ci.
3572                                      */
3573                                     if (!c_pbcShiftBackward || shift == CENTRAL)
3574                                     {
3575                                         firstCell = std::max(firstCell, ci);
3576                                     }
3577                                 }
3578
3579                                 if (firstCell <= lastCell)
3580                                 {
3581                                     GMX_ASSERT(firstCell >= columnStart && lastCell < columnEnd, "The range should reside within the current grid column");
3582
3583                                     /* For f buffer flags with simple lists */
3584                                     ncj_old_j = getNumSimpleJClustersInList(*nbl);
3585
3586                                     makeClusterListWrapper(nbl,
3587                                                            iGrid, ci,
3588                                                            jGrid, firstCell, lastCell,
3589                                                            excludeSubDiagonal,
3590                                                            nbat,
3591                                                            rlist2, rbb2,
3592                                                            kernelType,
3593                                                            &numDistanceChecks);
3594
3595                                     if (bFBufferFlag)
3596                                     {
3597                                         setBufferFlags(*nbl, ncj_old_j, gridj_flag_shift,
3598                                                        gridj_flag, th);
3599                                     }
3600
3601                                     incrementNumSimpleJClustersInList(nbl, ncj_old_j);
3602                                 }
3603                             }
3604                         }
3605                     }
3606
3607                     /* Set the exclusions for this ci list */
3608                     setExclusionsForIEntry(gridSet,
3609                                            nbl,
3610                                            excludeSubDiagonal,
3611                                            na_cj_2log,
3612                                            *getOpenIEntry(nbl),
3613                                            exclusions);
3614
3615                     if (haveFep)
3616                     {
3617                         make_fep_list(gridSet.atomIndices(), nbat, nbl,
3618                                       excludeSubDiagonal,
3619                                       getOpenIEntry(nbl),
3620                                       shx, shy, shz,
3621                                       rl_fep2,
3622                                       iGrid, jGrid, nbl_fep);
3623                     }
3624
3625                     /* Close this ci list */
3626                     closeIEntry(nbl,
3627                                 nsubpair_max,
3628                                 progBal, nsubpair_tot_est,
3629                                 th, nth);
3630                 }
3631             }
3632         }
3633
3634         if (bFBufferFlag && getNumSimpleJClustersInList(*nbl) > ncj_old_i)
3635         {
3636             bitmask_init_bit(&(work->buffer_flags.flag[(iGrid.cellOffset() + ci) >> gridi_flag_shift]), th);
3637         }
3638     }
3639
3640     work->ndistc = numDistanceChecks;
3641
3642     checkListSizeConsistency(*nbl, haveFep);
3643
3644     if (debug)
3645     {
3646         fprintf(debug, "number of distance checks %d\n", numDistanceChecks);
3647
3648         print_nblist_statistics(debug, nbl, pairSearch, rlist);
3649
3650         if (haveFep)
3651         {
3652             fprintf(debug, "nbl FEP list pairs: %d\n", nbl_fep->nrj);
3653         }
3654     }
3655 }
3656
3657 static void reduce_buffer_flags(const PairSearch           &pairSearch,
3658                                 int                         nsrc,
3659                                 const nbnxn_buffer_flags_t *dest)
3660 {
3661     for (int s = 0; s < nsrc; s++)
3662     {
3663         gmx_bitmask_t * flag = pairSearch.work()[s].buffer_flags.flag;
3664
3665         for (int b = 0; b < dest->nflag; b++)
3666         {
3667             bitmask_union(&(dest->flag[b]), flag[b]);
3668         }
3669     }
3670 }
3671
3672 static void print_reduction_cost(const nbnxn_buffer_flags_t *flags, int nout)
3673 {
3674     int           nelem, nkeep, ncopy, nred, out;
3675     gmx_bitmask_t mask_0;
3676
3677     nelem = 0;
3678     nkeep = 0;
3679     ncopy = 0;
3680     nred  = 0;
3681     bitmask_init_bit(&mask_0, 0);
3682     for (int b = 0; b < flags->nflag; b++)
3683     {
3684         if (bitmask_is_equal(flags->flag[b], mask_0))
3685         {
3686             /* Only flag 0 is set, no copy of reduction required */
3687             nelem++;
3688             nkeep++;
3689         }
3690         else if (!bitmask_is_zero(flags->flag[b]))
3691         {
3692             int c = 0;
3693             for (out = 0; out < nout; out++)
3694             {
3695                 if (bitmask_is_set(flags->flag[b], out))
3696                 {
3697                     c++;
3698                 }
3699             }
3700             nelem += c;
3701             if (c == 1)
3702             {
3703                 ncopy++;
3704             }
3705             else
3706             {
3707                 nred += c;
3708             }
3709         }
3710     }
3711
3712     fprintf(debug, "nbnxn reduction: #flag %d #list %d elem %4.2f, keep %4.2f copy %4.2f red %4.2f\n",
3713             flags->nflag, nout,
3714             nelem/static_cast<double>(flags->nflag),
3715             nkeep/static_cast<double>(flags->nflag),
3716             ncopy/static_cast<double>(flags->nflag),
3717             nred/static_cast<double>(flags->nflag));
3718 }
3719
3720 /* Copies the list entries from src to dest when cjStart <= *cjGlobal < cjEnd.
3721  * *cjGlobal is updated with the cj count in src.
3722  * When setFlags==true, flag bit t is set in flag for all i and j clusters.
3723  */
3724 template<bool setFlags>
3725 static void copySelectedListRange(const nbnxn_ci_t * gmx_restrict srcCi,
3726                                   const NbnxnPairlistCpu * gmx_restrict src,
3727                                   NbnxnPairlistCpu * gmx_restrict dest,
3728                                   gmx_bitmask_t *flag,
3729                                   int iFlagShift, int jFlagShift, int t)
3730 {
3731     const int ncj = srcCi->cj_ind_end - srcCi->cj_ind_start;
3732
3733     dest->ci.push_back(*srcCi);
3734     dest->ci.back().cj_ind_start = dest->cj.size();
3735     dest->ci.back().cj_ind_end   = dest->cj.size() + ncj;
3736
3737     if (setFlags)
3738     {
3739         bitmask_init_bit(&flag[srcCi->ci >> iFlagShift], t);
3740     }
3741
3742     for (int j = srcCi->cj_ind_start; j < srcCi->cj_ind_end; j++)
3743     {
3744         dest->cj.push_back(src->cj[j]);
3745
3746         if (setFlags)
3747         {
3748             /* NOTE: This is relatively expensive, since this
3749              * operation is done for all elements in the list,
3750              * whereas at list generation this is done only
3751              * once for each flag entry.
3752              */
3753             bitmask_init_bit(&flag[src->cj[j].cj >> jFlagShift], t);
3754         }
3755     }
3756 }
3757
3758 /* This routine re-balances the pairlists such that all are nearly equally
3759  * sized. Only whole i-entries are moved between lists. These are moved
3760  * between the ends of the lists, such that the buffer reduction cost should
3761  * not change significantly.
3762  * Note that all original reduction flags are currently kept. This can lead
3763  * to reduction of parts of the force buffer that could be avoided. But since
3764  * the original lists are quite balanced, this will only give minor overhead.
3765  */
3766 static void rebalanceSimpleLists(int                                  numLists,
3767                                  NbnxnPairlistCpu * const * const     srcSet,
3768                                  NbnxnPairlistCpu                   **destSet,
3769                                  gmx::ArrayRef<PairsearchWork>        searchWork)
3770 {
3771     int ncjTotal = 0;
3772     for (int s = 0; s < numLists; s++)
3773     {
3774         ncjTotal += srcSet[s]->ncjInUse;
3775     }
3776     int ncjTarget = (ncjTotal + numLists - 1)/numLists;
3777
3778 #pragma omp parallel num_threads(numLists)
3779     {
3780         int t       = gmx_omp_get_thread_num();
3781
3782         int cjStart = ncjTarget* t;
3783         int cjEnd   = ncjTarget*(t + 1);
3784
3785         /* The destination pair-list for task/thread t */
3786         NbnxnPairlistCpu *dest = destSet[t];
3787
3788         clear_pairlist(dest);
3789         dest->na_cj   = srcSet[0]->na_cj;
3790
3791         /* Note that the flags in the work struct (still) contain flags
3792          * for all entries that are present in srcSet->nbl[t].
3793          */
3794         gmx_bitmask_t *flag       = searchWork[t].buffer_flags.flag;
3795
3796         int            iFlagShift = getBufferFlagShift(dest->na_ci);
3797         int            jFlagShift = getBufferFlagShift(dest->na_cj);
3798
3799         int            cjGlobal   = 0;
3800         for (int s = 0; s < numLists && cjGlobal < cjEnd; s++)
3801         {
3802             const NbnxnPairlistCpu *src = srcSet[s];
3803
3804             if (cjGlobal + src->ncjInUse > cjStart)
3805             {
3806                 for (gmx::index i = 0; i < gmx::ssize(src->ci) && cjGlobal < cjEnd; i++)
3807                 {
3808                     const nbnxn_ci_t *srcCi = &src->ci[i];
3809                     int               ncj   = srcCi->cj_ind_end - srcCi->cj_ind_start;
3810                     if (cjGlobal >= cjStart)
3811                     {
3812                         /* If the source list is not our own, we need to set
3813                          * extra flags (the template bool parameter).
3814                          */
3815                         if (s != t)
3816                         {
3817                             copySelectedListRange
3818                             <true>
3819                                 (srcCi, src, dest,
3820                                 flag, iFlagShift, jFlagShift, t);
3821                         }
3822                         else
3823                         {
3824                             copySelectedListRange
3825                             <false>
3826                                 (srcCi, src,
3827                                 dest, flag, iFlagShift, jFlagShift, t);
3828                         }
3829                     }
3830                     cjGlobal += ncj;
3831                 }
3832             }
3833             else
3834             {
3835                 cjGlobal += src->ncjInUse;
3836             }
3837         }
3838
3839         dest->ncjInUse = dest->cj.size();
3840     }
3841
3842 #ifndef NDEBUG
3843     int ncjTotalNew = 0;
3844     for (int s = 0; s < numLists; s++)
3845     {
3846         ncjTotalNew += destSet[s]->ncjInUse;
3847     }
3848     GMX_RELEASE_ASSERT(ncjTotalNew == ncjTotal, "The total size of the lists before and after rebalancing should match");
3849 #endif
3850 }
3851
3852 /* Returns if the pairlists are so imbalanced that it is worth rebalancing. */
3853 static bool checkRebalanceSimpleLists(const nbnxn_pairlist_set_t *listSet)
3854 {
3855     int numLists = listSet->nnbl;
3856     int ncjMax   = 0;
3857     int ncjTotal = 0;
3858     for (int s = 0; s < numLists; s++)
3859     {
3860         ncjMax    = std::max(ncjMax, listSet->nbl[s]->ncjInUse);
3861         ncjTotal += listSet->nbl[s]->ncjInUse;
3862     }
3863     if (debug)
3864     {
3865         fprintf(debug, "Pair-list ncjMax %d ncjTotal %d\n", ncjMax, ncjTotal);
3866     }
3867     /* The rebalancing adds 3% extra time to the search. Heuristically we
3868      * determined that under common conditions the non-bonded kernel balance
3869      * improvement will outweigh this when the imbalance is more than 3%.
3870      * But this will, obviously, depend on search vs kernel time and nstlist.
3871      */
3872     const real rebalanceTolerance = 1.03;
3873
3874     return numLists*ncjMax > ncjTotal*rebalanceTolerance;
3875 }
3876
3877 /* Perform a count (linear) sort to sort the smaller lists to the end.
3878  * This avoids load imbalance on the GPU, as large lists will be
3879  * scheduled and executed first and the smaller lists later.
3880  * Load balancing between multi-processors only happens at the end
3881  * and there smaller lists lead to more effective load balancing.
3882  * The sorting is done on the cj4 count, not on the actual pair counts.
3883  * Not only does this make the sort faster, but it also results in
3884  * better load balancing than using a list sorted on exact load.
3885  * This function swaps the pointer in the pair list to avoid a copy operation.
3886  */
3887 static void sort_sci(NbnxnPairlistGpu *nbl)
3888 {
3889     if (nbl->cj4.size() <= nbl->sci.size())
3890     {
3891         /* nsci = 0 or all sci have size 1, sorting won't change the order */
3892         return;
3893     }
3894
3895     NbnxnPairlistGpuWork &work = *nbl->work;
3896
3897     /* We will distinguish differences up to double the average */
3898     const int m = (2*nbl->cj4.size())/nbl->sci.size();
3899
3900     /* Resize work.sci_sort so we can sort into it */
3901     work.sci_sort.resize(nbl->sci.size());
3902
3903     std::vector<int> &sort = work.sortBuffer;
3904     /* Set up m + 1 entries in sort, initialized at 0 */
3905     sort.clear();
3906     sort.resize(m + 1, 0);
3907     /* Count the entries of each size */
3908     for (const nbnxn_sci_t &sci : nbl->sci)
3909     {
3910         int i = std::min(m, sci.numJClusterGroups());
3911         sort[i]++;
3912     }
3913     /* Calculate the offset for each count */
3914     int s0  = sort[m];
3915     sort[m] = 0;
3916     for (int i = m - 1; i >= 0; i--)
3917     {
3918         int s1  = sort[i];
3919         sort[i] = sort[i + 1] + s0;
3920         s0      = s1;
3921     }
3922
3923     /* Sort entries directly into place */
3924     gmx::ArrayRef<nbnxn_sci_t> sci_sort = work.sci_sort;
3925     for (const nbnxn_sci_t &sci : nbl->sci)
3926     {
3927         int i = std::min(m, sci.numJClusterGroups());
3928         sci_sort[sort[i]++] = sci;
3929     }
3930
3931     /* Swap the sci pointers so we use the new, sorted list */
3932     std::swap(nbl->sci, work.sci_sort);
3933 }
3934
3935 void
3936 nonbonded_verlet_t::PairlistSets::construct(const InteractionLocality  iLocality,
3937                                             PairSearch                *pairSearch,
3938                                             nbnxn_atomdata_t          *nbat,
3939                                             const t_blocka            *excl,
3940                                             const Nbnxm::KernelType    kernelType,
3941                                             const int64_t              step,
3942                                             t_nrnb                    *nrnb)
3943 {
3944     nbnxn_pairlist_set_t *nbl_list = &pairlistSet(iLocality);
3945
3946     const real            rlist    = nbl_list->params.rlistOuter;
3947
3948     int                nsubpair_target;
3949     float              nsubpair_tot_est;
3950     int                nnbl;
3951     int                ci_block;
3952     gmx_bool           CombineNBLists;
3953     gmx_bool           progBal;
3954     int                np_tot, np_noq, np_hlj, nap;
3955
3956     nnbl            = nbl_list->nnbl;
3957     CombineNBLists  = nbl_list->bCombined;
3958
3959     if (debug)
3960     {
3961         fprintf(debug, "ns making %d nblists\n", nnbl);
3962     }
3963
3964     nbat->bUseBufferFlags = (nbat->out.size() > 1);
3965     /* We should re-init the flags before making the first list */
3966     if (nbat->bUseBufferFlags && iLocality == InteractionLocality::Local)
3967     {
3968         init_buffer_flags(&nbat->buffer_flags, nbat->numAtoms());
3969     }
3970
3971     int nzi;
3972     if (iLocality == InteractionLocality::Local)
3973     {
3974         /* Only zone (grid) 0 vs 0 */
3975         nzi = 1;
3976     }
3977     else
3978     {
3979         nzi = pairSearch->domainSetup().zones->nizone;
3980     }
3981
3982     if (!nbl_list->bSimple && minimumIlistCountForGpuBalancing_ > 0)
3983     {
3984         get_nsubpair_target(*pairSearch, iLocality, rlist, minimumIlistCountForGpuBalancing_,
3985                             &nsubpair_target, &nsubpair_tot_est);
3986     }
3987     else
3988     {
3989         nsubpair_target  = 0;
3990         nsubpair_tot_est = 0;
3991     }
3992
3993     /* Clear all pair-lists */
3994     for (int th = 0; th < nnbl; th++)
3995     {
3996         if (nbl_list->bSimple)
3997         {
3998             clear_pairlist(nbl_list->nbl[th]);
3999         }
4000         else
4001         {
4002             clear_pairlist(nbl_list->nblGpu[th]);
4003         }
4004
4005         if (pairSearch->gridSet().haveFep())
4006         {
4007             clear_pairlist_fep(nbl_list->nbl_fep[th]);
4008         }
4009     }
4010
4011     const gmx_domdec_zones_t *ddZones = pairSearch->domainSetup().zones;
4012
4013     for (int zi = 0; zi < nzi; zi++)
4014     {
4015         const Grid &iGrid = pairSearch->gridSet().grids()[zi];
4016
4017         int                 zj0;
4018         int                 zj1;
4019         if (iLocality == InteractionLocality::Local)
4020         {
4021             zj0 = 0;
4022             zj1 = 1;
4023         }
4024         else
4025         {
4026             zj0 = ddZones->izone[zi].j0;
4027             zj1 = ddZones->izone[zi].j1;
4028             if (zi == 0)
4029             {
4030                 zj0++;
4031             }
4032         }
4033         for (int zj = zj0; zj < zj1; zj++)
4034         {
4035             const Grid &jGrid = pairSearch->gridSet().grids()[zj];
4036
4037             if (debug)
4038             {
4039                 fprintf(debug, "ns search grid %d vs %d\n", zi, zj);
4040             }
4041
4042             pairSearch->cycleCounting_.start(PairSearch::enbsCCsearch);
4043
4044             ci_block = get_ci_block_size(iGrid, pairSearch->domainSetup().haveDomDec, nnbl);
4045
4046             /* With GPU: generate progressively smaller lists for
4047              * load balancing for local only or non-local with 2 zones.
4048              */
4049             progBal = (iLocality == InteractionLocality::Local || ddZones->n <= 2);
4050
4051 #pragma omp parallel for num_threads(nnbl) schedule(static)
4052             for (int th = 0; th < nnbl; th++)
4053             {
4054                 try
4055                 {
4056                     /* Re-init the thread-local work flag data before making
4057                      * the first list (not an elegant conditional).
4058                      */
4059                     if (nbat->bUseBufferFlags && ((zi == 0 && zj == 0)))
4060                     {
4061                         init_buffer_flags(&pairSearch->work()[th].buffer_flags, nbat->numAtoms());
4062                     }
4063
4064                     if (CombineNBLists && th > 0)
4065                     {
4066                         GMX_ASSERT(!nbl_list->bSimple, "Can only combine GPU lists");
4067
4068                         clear_pairlist(nbl_list->nblGpu[th]);
4069                     }
4070
4071                     auto &searchWork = pairSearch->work()[th];
4072
4073                     searchWork.cycleCounter.start();
4074
4075                     /* Divide the i super cell equally over the nblists */
4076                     if (nbl_list->bSimple)
4077                     {
4078                         nbnxn_make_pairlist_part(*pairSearch, iGrid, jGrid,
4079                                                  &searchWork, nbat, *excl,
4080                                                  rlist,
4081                                                  kernelType,
4082                                                  ci_block,
4083                                                  nbat->bUseBufferFlags,
4084                                                  nsubpair_target,
4085                                                  progBal, nsubpair_tot_est,
4086                                                  th, nnbl,
4087                                                  nbl_list->nbl[th],
4088                                                  nbl_list->nbl_fep[th]);
4089                     }
4090                     else
4091                     {
4092                         nbnxn_make_pairlist_part(*pairSearch, iGrid, jGrid,
4093                                                  &searchWork, nbat, *excl,
4094                                                  rlist,
4095                                                  kernelType,
4096                                                  ci_block,
4097                                                  nbat->bUseBufferFlags,
4098                                                  nsubpair_target,
4099                                                  progBal, nsubpair_tot_est,
4100                                                  th, nnbl,
4101                                                  nbl_list->nblGpu[th],
4102                                                  nbl_list->nbl_fep[th]);
4103                     }
4104
4105                     searchWork.cycleCounter.stop();
4106                 }
4107                 GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR;
4108             }
4109             pairSearch->cycleCounting_.stop(PairSearch::enbsCCsearch);
4110
4111             np_tot = 0;
4112             np_noq = 0;
4113             np_hlj = 0;
4114             for (int th = 0; th < nnbl; th++)
4115             {
4116                 inc_nrnb(nrnb, eNR_NBNXN_DIST2, pairSearch->work()[th].ndistc);
4117
4118                 if (nbl_list->bSimple)
4119                 {
4120                     NbnxnPairlistCpu *nbl = nbl_list->nbl[th];
4121                     np_tot += nbl->cj.size();
4122                     np_noq += nbl->work->ncj_noq;
4123                     np_hlj += nbl->work->ncj_hlj;
4124                 }
4125                 else
4126                 {
4127                     NbnxnPairlistGpu *nbl = nbl_list->nblGpu[th];
4128                     /* This count ignores potential subsequent pair pruning */
4129                     np_tot += nbl->nci_tot;
4130                 }
4131             }
4132             if (nbl_list->bSimple)
4133             {
4134                 nap               = nbl_list->nbl[0]->na_ci*nbl_list->nbl[0]->na_cj;
4135             }
4136             else
4137             {
4138                 nap               = gmx::square(nbl_list->nblGpu[0]->na_ci);
4139             }
4140             nbl_list->natpair_ljq = (np_tot - np_noq)*nap - np_hlj*nap/2;
4141             nbl_list->natpair_lj  = np_noq*nap;
4142             nbl_list->natpair_q   = np_hlj*nap/2;
4143
4144             if (CombineNBLists && nnbl > 1)
4145             {
4146                 GMX_ASSERT(!nbl_list->bSimple, "Can only combine GPU lists");
4147                 NbnxnPairlistGpu **nbl = nbl_list->nblGpu;
4148
4149                 pairSearch->cycleCounting_.start(PairSearch::enbsCCcombine);
4150
4151                 combine_nblists(nnbl-1, nbl+1, nbl[0]);
4152
4153                 pairSearch->cycleCounting_.stop(PairSearch::enbsCCcombine);
4154             }
4155         }
4156     }
4157
4158     if (nbl_list->bSimple)
4159     {
4160         if (nnbl > 1 && checkRebalanceSimpleLists(nbl_list))
4161         {
4162             rebalanceSimpleLists(nbl_list->nnbl, nbl_list->nbl, nbl_list->nbl_work, pairSearch->work());
4163
4164             /* Swap the pointer of the sets of pair lists */
4165             NbnxnPairlistCpu **tmp = nbl_list->nbl;
4166             nbl_list->nbl          = nbl_list->nbl_work;
4167             nbl_list->nbl_work     = tmp;
4168         }
4169     }
4170     else
4171     {
4172         /* Sort the entries on size, large ones first */
4173         if (CombineNBLists || nnbl == 1)
4174         {
4175             sort_sci(nbl_list->nblGpu[0]);
4176         }
4177         else
4178         {
4179 #pragma omp parallel for num_threads(nnbl) schedule(static)
4180             for (int th = 0; th < nnbl; th++)
4181             {
4182                 try
4183                 {
4184                     sort_sci(nbl_list->nblGpu[th]);
4185                 }
4186                 GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR;
4187             }
4188         }
4189     }
4190
4191     if (nbat->bUseBufferFlags)
4192     {
4193         reduce_buffer_flags(*pairSearch, nbl_list->nnbl, &nbat->buffer_flags);
4194     }
4195
4196     if (pairSearch->gridSet().haveFep())
4197     {
4198         /* Balance the free-energy lists over all the threads */
4199         balance_fep_lists(pairSearch->work(), nbl_list);
4200     }
4201
4202     if (nbl_list->bSimple)
4203     {
4204         /* This is a fresh list, so not pruned, stored using ci.
4205          * ciOuter is invalid at this point.
4206          */
4207         GMX_ASSERT(nbl_list->nbl[0]->ciOuter.empty(), "ciOuter is invalid so it should be empty");
4208     }
4209
4210     if (iLocality == Nbnxm::InteractionLocality::Local)
4211     {
4212         outerListCreationStep_ = step;
4213     }
4214     else
4215     {
4216         GMX_RELEASE_ASSERT(outerListCreationStep_ == step,
4217                            "Outer list should be created at the same step as the inner list");
4218     }
4219
4220     /* Special performance logging stuff (env.var. GMX_NBNXN_CYCLE) */
4221     if (iLocality == InteractionLocality::Local)
4222     {
4223         pairSearch->cycleCounting_.searchCount_++;
4224     }
4225     if (pairSearch->cycleCounting_.recordCycles_ &&
4226         (!pairSearch->domainSetup().haveDomDec || iLocality == InteractionLocality::NonLocal) &&
4227         pairSearch->cycleCounting_.searchCount_ % 100 == 0)
4228     {
4229         pairSearch->cycleCounting_.printCycles(stderr, pairSearch->work());
4230     }
4231
4232     /* If we have more than one list, they either got rebalancing (CPU)
4233      * or combined (GPU), so we should dump the final result to debug.
4234      */
4235     if (debug && nbl_list->nnbl > 1)
4236     {
4237         if (nbl_list->bSimple)
4238         {
4239             for (int t = 0; t < nbl_list->nnbl; t++)
4240             {
4241                 print_nblist_statistics(debug, nbl_list->nbl[t], *pairSearch, rlist);
4242             }
4243         }
4244         else
4245         {
4246             print_nblist_statistics(debug, nbl_list->nblGpu[0], *pairSearch, rlist);
4247         }
4248     }
4249
4250     if (debug)
4251     {
4252         if (gmx_debug_at)
4253         {
4254             if (nbl_list->bSimple)
4255             {
4256                 for (int t = 0; t < nbl_list->nnbl; t++)
4257                 {
4258                     print_nblist_ci_cj(debug, nbl_list->nbl[t]);
4259                 }
4260             }
4261             else
4262             {
4263                 print_nblist_sci_cj(debug, nbl_list->nblGpu[0]);
4264             }
4265         }
4266
4267         if (nbat->bUseBufferFlags)
4268         {
4269             print_reduction_cost(&nbat->buffer_flags, nbl_list->nnbl);
4270         }
4271     }
4272
4273     if (params_.useDynamicPruning && nbl_list->bSimple)
4274     {
4275         nbnxnPrepareListForDynamicPruning(nbl_list);
4276     }
4277 }
4278
4279 void
4280 nonbonded_verlet_t::constructPairlist(const Nbnxm::InteractionLocality  iLocality,
4281                                       const t_blocka                   *excl,
4282                                       int64_t                           step,
4283                                       t_nrnb                           *nrnb)
4284 {
4285     pairlistSets_->construct(iLocality, pairSearch_.get(), nbat.get(), excl,
4286                              kernelSetup_.kernelType,
4287                              step, nrnb);
4288
4289     if (useGpu())
4290     {
4291         /* Launch the transfer of the pairlist to the GPU.
4292          *
4293          * NOTE: The launch overhead is currently not timed separately
4294          */
4295         Nbnxm::gpu_init_pairlist(gpu_nbv,
4296                                  pairlistSets().pairlistSet(iLocality).nblGpu[0],
4297                                  iLocality);
4298     }
4299 }
4300
4301 void nbnxnPrepareListForDynamicPruning(nbnxn_pairlist_set_t *listSet)
4302 {
4303     GMX_RELEASE_ASSERT(listSet->bSimple, "Should only be called for simple lists");
4304
4305     /* TODO: Restructure the lists so we have actual outer and inner
4306      *       list objects so we can set a single pointer instead of
4307      *       swapping several pointers.
4308      */
4309
4310     for (int i = 0; i < listSet->nnbl; i++)
4311     {
4312         NbnxnPairlistCpu &list = *listSet->nbl[i];
4313
4314         /* The search produced a list in ci/cj.
4315          * Swap the list pointers so we get the outer list is ciOuter,cjOuter
4316          * and we can prune that to get an inner list in ci/cj.
4317          */
4318         GMX_RELEASE_ASSERT(list.ciOuter.empty() && list.cjOuter.empty(),
4319                            "The outer lists should be empty before preparation");
4320
4321         std::swap(list.ci, list.ciOuter);
4322         std::swap(list.cj, list.cjOuter);
4323     }
4324 }