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