src/gromacs/nbnxm/grid.cpp

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2012,2013,2014,2015,2016 by the GROMACS development team.
   5  * Copyright (c) 2017,2018,2019,2020, by the GROMACS development team, led by
   6  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   7  * and including many others, as listed in the AUTHORS file in the
   8  * top-level source directory and at http://www.gromacs.org.
   9  *
  10  * GROMACS is free software; you can redistribute it and/or
  11  * modify it under the terms of the GNU Lesser General Public License
  12  * as published by the Free Software Foundation; either version 2.1
  13  * of the License, or (at your option) any later version.
  14  *
  15  * GROMACS is distributed in the hope that it will be useful,
  16  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  17  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  18  * Lesser General Public License for more details.
  19  *
  20  * You should have received a copy of the GNU Lesser General Public
  21  * License along with GROMACS; if not, see
  22  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
  23  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  24  *
  25  * If you want to redistribute modifications to GROMACS, please
  26  * consider that scientific software is very special. Version
  27  * control is crucial - bugs must be traceable. We will be happy to
  28  * consider code for inclusion in the official distribution, but
  29  * derived work must not be called official GROMACS. Details are found
  30  * in the README & COPYING files - if they are missing, get the
  31  * official version at http://www.gromacs.org.
  32  *
  33  * To help us fund GROMACS development, we humbly ask that you cite
  34  * the research papers on the package. Check out http://www.gromacs.org.
  35  */
  36
  37 /*! \internal \file
  38  *
  39  * \brief
  40  * Implements the Grid class.
  41  *
  42  * \author Berk Hess <hess@kth.se>
  43  * \ingroup module_nbnxm
  44  */
  45
  46 #include "gmxpre.h"
  47
  48 #include "grid.h"
  49
  50 #include <cmath>
  51 #include <cstring>
  52
  53 #include <algorithm>
  54
  55 #include "gromacs/math/utilities.h"
  56 #include "gromacs/math/vec.h"
  57 #include "gromacs/mdlib/gmx_omp_nthreads.h"
  58 #include "gromacs/mdlib/updategroupscog.h"
  59 #include "gromacs/mdtypes/forcerec.h" // only for GET_CGINFO_*
  60 #include "gromacs/nbnxm/atomdata.h"
  61 #include "gromacs/simd/simd.h"
  62 #include "gromacs/simd/vector_operations.h"
  63
  64 #include "boundingboxes.h"
  65 #include "gridsetdata.h"
  66 #include "nbnxm_geometry.h"
  67 #include "pairlistparams.h"
  68
  69 namespace Nbnxm
  70 {
  71
  72 Grid::Geometry::Geometry(const PairlistType pairlistType) :
  73     isSimple(pairlistType != PairlistType::HierarchicalNxN),
  74     numAtomsICluster(IClusterSizePerListType[pairlistType]),
  75     numAtomsJCluster(JClusterSizePerListType[pairlistType]),
  76     numAtomsPerCell((isSimple ? 1 : c_gpuNumClusterPerCell) * numAtomsICluster),
  77     numAtomsICluster2Log(get_2log(numAtomsICluster))
  78 {
  79 }
  80
  81 Grid::Grid(const PairlistType pairlistType, const bool& haveFep) :
  82     geometry_(pairlistType),
  83     haveFep_(haveFep)
  84 {
  85 }
  86
  87 /*! \brief Returns the atom density (> 0) of a rectangular grid */
  88 static real gridAtomDensity(int numAtoms, const rvec lowerCorner, const rvec upperCorner)
  89 {
  90     rvec size;
  91
  92     if (numAtoms == 0)
  93     {
  94         /* To avoid zero density we use a minimum of 1 atom */
  95         numAtoms = 1;
  96     }
  97
  98     rvec_sub(upperCorner, lowerCorner, size);
  99
 100     return static_cast<real>(numAtoms) / (size[XX] * size[YY] * size[ZZ]);
 101 }
 102
 103 void Grid::setDimensions(const int          ddZone,
 104                          const int          numAtoms,
 105                          gmx::RVec          lowerCorner,
 106                          gmx::RVec          upperCorner,
 107                          real               atomDensity,
 108                          const real         maxAtomGroupRadius,
 109                          const bool         haveFep,
 110                          gmx::PinningPolicy pinningPolicy)
 111 {
 112     /* We allow passing lowerCorner=upperCorner, in which case we need to
 113      * create a finite sized bounding box to avoid division by zero.
 114      * We use a minimum size such that the volume fits in float with some
 115      * margin for computing and using the atom number density.
 116      */
 117     constexpr real c_minimumGridSize = 1e-10;
 118     for (int d = 0; d < DIM; d++)
 119     {
 120         GMX_ASSERT(upperCorner[d] >= lowerCorner[d],
 121                    "Upper corner should be larger than the lower corner");
 122         if (upperCorner[d] - lowerCorner[d] < c_minimumGridSize)
 123         {
 124             /* Ensure we apply a correction to the bounding box */
 125             real correction =
 126                     std::max(std::abs(lowerCorner[d]) * GMX_REAL_EPS, 0.5_real * c_minimumGridSize);
 127             lowerCorner[d] -= correction;
 128             upperCorner[d] += correction;
 129         }
 130     }
 131
 132     /* For the home zone we compute the density when not set (=-1) or when =0 */
 133     if (ddZone == 0 && atomDensity <= 0)
 134     {
 135         atomDensity = gridAtomDensity(numAtoms, lowerCorner, upperCorner);
 136     }
 137
 138     dimensions_.atomDensity        = atomDensity;
 139     dimensions_.maxAtomGroupRadius = maxAtomGroupRadius;
 140
 141     rvec size;
 142     rvec_sub(upperCorner, lowerCorner, size);
 143
 144     if (numAtoms > geometry_.numAtomsPerCell)
 145     {
 146         GMX_ASSERT(atomDensity > 0, "With one or more atoms, the density should be positive");
 147
 148         /* target cell length */
 149         real tlen_x;
 150         real tlen_y;
 151         if (geometry_.isSimple)
 152         {
 153             /* To minimize the zero interactions, we should make
 154              * the largest of the i/j cell cubic.
 155              */
 156             int numAtomsInCell = std::max(geometry_.numAtomsICluster, geometry_.numAtomsJCluster);
 157
 158             /* Approximately cubic cells */
 159             real tlen = std::cbrt(numAtomsInCell / atomDensity);
 160             tlen_x    = tlen;
 161             tlen_y    = tlen;
 162         }
 163         else
 164         {
 165             /* Approximately cubic sub cells */
 166             real tlen = std::cbrt(geometry_.numAtomsICluster / atomDensity);
 167             tlen_x    = tlen * c_gpuNumClusterPerCellX;
 168             tlen_y    = tlen * c_gpuNumClusterPerCellY;
 169         }
 170         /* We round ncx and ncy down, because we get less cell pairs
 171          * in the pairlist when the fixed cell dimensions (x,y) are
 172          * larger than the variable one (z) than the other way around.
 173          */
 174         dimensions_.numCells[XX] = std::max(1, static_cast<int>(size[XX] / tlen_x));
 175         dimensions_.numCells[YY] = std::max(1, static_cast<int>(size[YY] / tlen_y));
 176     }
 177     else
 178     {
 179         dimensions_.numCells[XX] = 1;
 180         dimensions_.numCells[YY] = 1;
 181     }
 182
 183     for (int d = 0; d < DIM - 1; d++)
 184     {
 185         dimensions_.cellSize[d]    = size[d] / dimensions_.numCells[d];
 186         dimensions_.invCellSize[d] = 1 / dimensions_.cellSize[d];
 187     }
 188
 189     if (ddZone > 0)
 190     {
 191         /* This is a non-home zone, add an extra row of cells
 192          * for particles communicated for bonded interactions.
 193          * These can be beyond the cut-off. It doesn't matter where
 194          * they end up on the grid, but for performance it's better
 195          * if they don't end up in cells that can be within cut-off range.
 196          */
 197         dimensions_.numCells[XX]++;
 198         dimensions_.numCells[YY]++;
 199     }
 200
 201     /* We need one additional cell entry for particles moved by DD */
 202     cxy_na_.resize(numColumns() + 1);
 203     cxy_ind_.resize(numColumns() + 2);
 204     changePinningPolicy(&cxy_na_, pinningPolicy);
 205     changePinningPolicy(&cxy_ind_, pinningPolicy);
 206
 207     /* Worst case scenario of 1 atom in each last cell */
 208     int maxNumCells;
 209     if (geometry_.numAtomsJCluster <= geometry_.numAtomsICluster)
 210     {
 211         maxNumCells = numAtoms / geometry_.numAtomsPerCell + numColumns();
 212     }
 213     else
 214     {
 215         maxNumCells = numAtoms / geometry_.numAtomsPerCell
 216                       + numColumns() * geometry_.numAtomsJCluster / geometry_.numAtomsICluster;
 217     }
 218
 219     if (!geometry_.isSimple)
 220     {
 221         numClusters_.resize(maxNumCells);
 222     }
 223     bbcz_.resize(maxNumCells);
 224
 225     /* This resize also zeros the contents, this avoid possible
 226      * floating exceptions in SIMD with the unused bb elements.
 227      */
 228     if (geometry_.isSimple)
 229     {
 230         bb_.resize(maxNumCells);
 231     }
 232     else
 233     {
 234 #if NBNXN_BBXXXX
 235         pbb_.resize(packedBoundingBoxesIndex(maxNumCells * c_gpuNumClusterPerCell));
 236 #else
 237         bb_.resize(maxNumCells * c_gpuNumClusterPerCell);
 238 #endif
 239     }
 240
 241     if (geometry_.numAtomsJCluster == geometry_.numAtomsICluster)
 242     {
 243         bbj_ = bb_;
 244     }
 245     else
 246     {
 247         GMX_ASSERT(geometry_.isSimple, "Only CPU lists should have different i/j cluster sizes");
 248
 249         bbjStorage_.resize(maxNumCells * geometry_.numAtomsICluster / geometry_.numAtomsJCluster);
 250         bbj_ = bbjStorage_;
 251     }
 252
 253     flags_.resize(maxNumCells);
 254     if (haveFep)
 255     {
 256         fep_.resize(maxNumCells * geometry_.numAtomsPerCell / geometry_.numAtomsICluster);
 257     }
 258
 259     copy_rvec(lowerCorner, dimensions_.lowerCorner);
 260     copy_rvec(upperCorner, dimensions_.upperCorner);
 261     copy_rvec(size, dimensions_.gridSize);
 262 }
 263
 264 /* We need to sort paricles in grid columns on z-coordinate.
 265  * As particle are very often distributed homogeneously, we use a sorting
 266  * algorithm similar to pigeonhole sort. We multiply the z-coordinate
 267  * by a factor, cast to an int and try to store in that hole. If the hole
 268  * is full, we move this or another particle. A second pass is needed to make
 269  * contiguous elements. SORT_GRID_OVERSIZE is the ratio of holes to particles.
 270  * 4 is the optimal value for homogeneous particle distribution and allows
 271  * for an O(#particles) sort up till distributions were all particles are
 272  * concentrated in 1/4 of the space. No NlogN fallback is implemented,
 273  * as it can be expensive to detect imhomogeneous particle distributions.
 274  */
 275 /*! \brief Ratio of grid cells to atoms */
 276 static constexpr int c_sortGridRatio = 4;
 277 /*! \brief Maximum ratio of holes used, in the worst case all particles
 278  * end up in the last hole and we need num. atoms extra holes at the end.
 279  */
 280 static constexpr int c_sortGridMaxSizeFactor = c_sortGridRatio + 1;
 281
 282 /*! \brief Sorts particle index a on coordinates x along dim.
 283  *
 284  * Backwards tells if we want decreasing iso increasing coordinates.
 285  * h0 is the minimum of the coordinate range.
 286  * invh is the 1/length of the sorting range.
 287  * n_per_h (>=n) is the expected average number of particles per 1/invh
 288  * sort is the sorting work array.
 289  * sort should have a size of at least n_per_h*c_sortGridRatio + n,
 290  * or easier, allocate at least n*c_sortGridMaxSizeFactor elements.
 291  */
 292 static void sort_atoms(int      dim,
 293                        gmx_bool Backwards,
 294                        int gmx_unused dd_zone,
 295                        bool gmx_unused                relevantAtomsAreWithinGridBounds,
 296                        int*                           a,
 297                        int                            n,
 298                        gmx::ArrayRef<const gmx::RVec> x,
 299                        real                           h0,
 300                        real                           invh,
 301                        int                            n_per_h,
 302                        gmx::ArrayRef<int>             sort)
 303 {
 304     if (n <= 1)
 305     {
 306         /* Nothing to do */
 307         return;
 308     }
 309
 310     GMX_ASSERT(n <= n_per_h, "We require n <= n_per_h");
 311
 312     /* Transform the inverse range height into the inverse hole height */
 313     invh *= n_per_h * c_sortGridRatio;
 314
 315     /* Set nsort to the maximum possible number of holes used.
 316      * In worst case all n elements end up in the last bin.
 317      */
 318     int nsort = n_per_h * c_sortGridRatio + n;
 319
 320     /* Determine the index range used, so we can limit it for the second pass */
 321     int zi_min = INT_MAX;
 322     int zi_max = -1;
 323
 324     /* Sort the particles using a simple index sort */
 325     for (int i = 0; i < n; i++)
 326     {
 327         /* The cast takes care of float-point rounding effects below zero.
 328          * This code assumes particles are less than 1/SORT_GRID_OVERSIZE
 329          * times the box height out of the box.
 330          */
 331         int zi = static_cast<int>((x[a[i]][dim] - h0) * invh);
 332
 333 #ifndef NDEBUG
 334         /* As we can have rounding effect, we use > iso >= here */
 335         if (relevantAtomsAreWithinGridBounds && (zi < 0 || (dd_zone == 0 && zi > n_per_h * c_sortGridRatio)))
 336         {
 337             gmx_fatal(FARGS,
 338                       "(int)((x[%d][%c]=%f - %f)*%f) = %d, not in 0 - %d*%d\n",
 339                       a[i],
 340                       'x' + dim,
 341                       x[a[i]][dim],
 342                       h0,
 343                       invh,
 344                       zi,
 345                       n_per_h,
 346                       c_sortGridRatio);
 347         }
 348 #endif
 349         if (zi < 0)
 350         {
 351             zi = 0;
 352         }
 353
 354         /* In a non-local domain, particles communcated for bonded interactions
 355          * can be far beyond the grid size, which is set by the non-bonded
 356          * cut-off distance. We sort such particles into the last cell.
 357          */
 358         if (zi > n_per_h * c_sortGridRatio)
 359         {
 360             zi = n_per_h * c_sortGridRatio;
 361         }
 362
 363         /* Ideally this particle should go in sort cell zi,
 364          * but that might already be in use,
 365          * in that case find the first empty cell higher up
 366          */
 367         if (sort[zi] < 0)
 368         {
 369             sort[zi] = a[i];
 370             zi_min   = std::min(zi_min, zi);
 371             zi_max   = std::max(zi_max, zi);
 372         }
 373         else
 374         {
 375             /* We have multiple atoms in the same sorting slot.
 376              * Sort on real z for minimal bounding box size.
 377              * There is an extra check for identical z to ensure
 378              * well-defined output order, independent of input order
 379              * to ensure binary reproducibility after restarts.
 380              */
 381             while (sort[zi] >= 0
 382                    && (x[a[i]][dim] > x[sort[zi]][dim]
 383                        || (x[a[i]][dim] == x[sort[zi]][dim] && a[i] > sort[zi])))
 384             {
 385                 zi++;
 386             }
 387
 388             if (sort[zi] >= 0)
 389             {
 390                 /* Shift all elements by one slot until we find an empty slot */
 391                 int cp  = sort[zi];
 392                 int zim = zi + 1;
 393                 while (sort[zim] >= 0)
 394                 {
 395                     int tmp   = sort[zim];
 396                     sort[zim] = cp;
 397                     cp        = tmp;
 398                     zim++;
 399                 }
 400                 sort[zim] = cp;
 401                 zi_max    = std::max(zi_max, zim);
 402             }
 403             sort[zi] = a[i];
 404             zi_max   = std::max(zi_max, zi);
 405         }
 406     }
 407
 408     int c = 0;
 409     if (!Backwards)
 410     {
 411         for (int zi = 0; zi < nsort; zi++)
 412         {
 413             if (sort[zi] >= 0)
 414             {
 415                 a[c++]   = sort[zi];
 416                 sort[zi] = -1;
 417             }
 418         }
 419     }
 420     else
 421     {
 422         for (int zi = zi_max; zi >= zi_min; zi--)
 423         {
 424             if (sort[zi] >= 0)
 425             {
 426                 a[c++]   = sort[zi];
 427                 sort[zi] = -1;
 428             }
 429         }
 430     }
 431     if (c < n)
 432     {
 433         gmx_incons("Lost particles while sorting");
 434     }
 435 }
 436
 437 #if GMX_DOUBLE
 438 //! Returns double up to one least significant float bit smaller than x
 439 static double R2F_D(const float x)
 440 {
 441     return static_cast<float>(x >= 0 ? (1 - GMX_FLOAT_EPS) * x : (1 + GMX_FLOAT_EPS) * x);
 442 }
 443 //! Returns double up to one least significant float bit larger than x
 444 static double R2F_U(const float x)
 445 {
 446     return static_cast<float>(x >= 0 ? (1 + GMX_FLOAT_EPS) * x : (1 - GMX_FLOAT_EPS) * x);
 447 }
 448 #else
 449 //! Returns x
 450 static float R2F_D(const float x)
 451 {
 452     return x;
 453 }
 454 //! Returns x
 455 static float R2F_U(const float x)
 456 {
 457     return x;
 458 }
 459 #endif
 460
 461 //! Computes the bounding box for na coordinates in order x,y,z, bb order xyz0
 462 static void calc_bounding_box(int na, int stride, const real* x, BoundingBox* bb)
 463 {
 464     int  i;
 465     real xl, xh, yl, yh, zl, zh;
 466
 467     i  = 0;
 468     xl = x[i + XX];
 469     xh = x[i + XX];
 470     yl = x[i + YY];
 471     yh = x[i + YY];
 472     zl = x[i + ZZ];
 473     zh = x[i + ZZ];
 474     i += stride;
 475     for (int j = 1; j < na; j++)
 476     {
 477         xl = std::min(xl, x[i + XX]);
 478         xh = std::max(xh, x[i + XX]);
 479         yl = std::min(yl, x[i + YY]);
 480         yh = std::max(yh, x[i + YY]);
 481         zl = std::min(zl, x[i + ZZ]);
 482         zh = std::max(zh, x[i + ZZ]);
 483         i += stride;
 484     }
 485     /* Note: possible double to float conversion here */
 486     bb->lower.x = R2F_D(xl);
 487     bb->lower.y = R2F_D(yl);
 488     bb->lower.z = R2F_D(zl);
 489     bb->upper.x = R2F_U(xh);
 490     bb->upper.y = R2F_U(yh);
 491     bb->upper.z = R2F_U(zh);
 492 }
 493
 494 /*! \brief Computes the bounding box for na packed coordinates, bb order xyz0 */
 495 static void calc_bounding_box_x_x4(int na, const real* x, BoundingBox* bb)
 496 {
 497     real xl, xh, yl, yh, zl, zh;
 498
 499     xl = x[XX * c_packX4];
 500     xh = x[XX * c_packX4];
 501     yl = x[YY * c_packX4];
 502     yh = x[YY * c_packX4];
 503     zl = x[ZZ * c_packX4];
 504     zh = x[ZZ * c_packX4];
 505     for (int j = 1; j < na; j++)
 506     {
 507         xl = std::min(xl, x[j + XX * c_packX4]);
 508         xh = std::max(xh, x[j + XX * c_packX4]);
 509         yl = std::min(yl, x[j + YY * c_packX4]);
 510         yh = std::max(yh, x[j + YY * c_packX4]);
 511         zl = std::min(zl, x[j + ZZ * c_packX4]);
 512         zh = std::max(zh, x[j + ZZ * c_packX4]);
 513     }
 514     /* Note: possible double to float conversion here */
 515     bb->lower.x = R2F_D(xl);
 516     bb->lower.y = R2F_D(yl);
 517     bb->lower.z = R2F_D(zl);
 518     bb->upper.x = R2F_U(xh);
 519     bb->upper.y = R2F_U(yh);
 520     bb->upper.z = R2F_U(zh);
 521 }
 522
 523 /*! \brief Computes the bounding box for na coordinates, bb order xyz0 */
 524 static void calc_bounding_box_x_x8(int na, const real* x, BoundingBox* bb)
 525 {
 526     real xl, xh, yl, yh, zl, zh;
 527
 528     xl = x[XX * c_packX8];
 529     xh = x[XX * c_packX8];
 530     yl = x[YY * c_packX8];
 531     yh = x[YY * c_packX8];
 532     zl = x[ZZ * c_packX8];
 533     zh = x[ZZ * c_packX8];
 534     for (int j = 1; j < na; j++)
 535     {
 536         xl = std::min(xl, x[j + XX * c_packX8]);
 537         xh = std::max(xh, x[j + XX * c_packX8]);
 538         yl = std::min(yl, x[j + YY * c_packX8]);
 539         yh = std::max(yh, x[j + YY * c_packX8]);
 540         zl = std::min(zl, x[j + ZZ * c_packX8]);
 541         zh = std::max(zh, x[j + ZZ * c_packX8]);
 542     }
 543     /* Note: possible double to float conversion here */
 544     bb->lower.x = R2F_D(xl);
 545     bb->lower.y = R2F_D(yl);
 546     bb->lower.z = R2F_D(zl);
 547     bb->upper.x = R2F_U(xh);
 548     bb->upper.y = R2F_U(yh);
 549     bb->upper.z = R2F_U(zh);
 550 }
 551
 552 /*! \brief Computes the bounding box for na packed coordinates, bb order xyz0 */
 553 gmx_unused static void calc_bounding_box_x_x4_halves(int na, const real* x, BoundingBox* bb, BoundingBox* bbj)
 554 {
 555     // TODO: During SIMDv2 transition only some archs use namespace (remove when done)
 556     using namespace gmx;
 557
 558     calc_bounding_box_x_x4(std::min(na, 2), x, bbj);
 559
 560     if (na > 2)
 561     {
 562         calc_bounding_box_x_x4(std::min(na - 2, 2), x + (c_packX4 >> 1), bbj + 1);
 563     }
 564     else
 565     {
 566         /* Set the "empty" bounding box to the same as the first one,
 567          * so we don't need to treat special cases in the rest of the code.
 568          */
 569 #if NBNXN_SEARCH_BB_SIMD4
 570         store4(bbj[1].lower.ptr(), load4(bbj[0].lower.ptr()));
 571         store4(bbj[1].upper.ptr(), load4(bbj[0].upper.ptr()));
 572 #else
 573         bbj[1] = bbj[0];
 574 #endif
 575     }
 576
 577 #if NBNXN_SEARCH_BB_SIMD4
 578     store4(bb->lower.ptr(), min(load4(bbj[0].lower.ptr()), load4(bbj[1].lower.ptr())));
 579     store4(bb->upper.ptr(), max(load4(bbj[0].upper.ptr()), load4(bbj[1].upper.ptr())));
 580 #else
 581     {
 582         bb->lower = BoundingBox::Corner::min(bbj[0].lower, bbj[1].lower);
 583         bb->upper = BoundingBox::Corner::max(bbj[0].upper, bbj[1].upper);
 584     }
 585 #endif
 586 }
 587
 588 #if NBNXN_BBXXXX
 589
 590 /*! \brief Computes the bounding box for na coordinates in order xyz, bb order xxxxyyyyzzzz */
 591 static void calc_bounding_box_xxxx(int na, int stride, const real* x, float* bb)
 592 {
 593     int  i;
 594     real xl, xh, yl, yh, zl, zh;
 595
 596     i  = 0;
 597     xl = x[i + XX];
 598     xh = x[i + XX];
 599     yl = x[i + YY];
 600     yh = x[i + YY];
 601     zl = x[i + ZZ];
 602     zh = x[i + ZZ];
 603     i += stride;
 604     for (int j = 1; j < na; j++)
 605     {
 606         xl = std::min(xl, x[i + XX]);
 607         xh = std::max(xh, x[i + XX]);
 608         yl = std::min(yl, x[i + YY]);
 609         yh = std::max(yh, x[i + YY]);
 610         zl = std::min(zl, x[i + ZZ]);
 611         zh = std::max(zh, x[i + ZZ]);
 612         i += stride;
 613     }
 614     /* Note: possible double to float conversion here */
 615     bb[0 * c_packedBoundingBoxesDimSize] = R2F_D(xl);
 616     bb[1 * c_packedBoundingBoxesDimSize] = R2F_D(yl);
 617     bb[2 * c_packedBoundingBoxesDimSize] = R2F_D(zl);
 618     bb[3 * c_packedBoundingBoxesDimSize] = R2F_U(xh);
 619     bb[4 * c_packedBoundingBoxesDimSize] = R2F_U(yh);
 620     bb[5 * c_packedBoundingBoxesDimSize] = R2F_U(zh);
 621 }
 622
 623 #endif /* NBNXN_BBXXXX */
 624
 625 #if NBNXN_SEARCH_SIMD4_FLOAT_X_BB
 626
 627 /*! \brief Computes the bounding box for na coordinates in order xyz?, bb order xyz0 */
 628 static void calc_bounding_box_simd4(int na, const float* x, BoundingBox* bb)
 629 {
 630     // TODO: During SIMDv2 transition only some archs use namespace (remove when done)
 631     using namespace gmx;
 632
 633     static_assert(sizeof(BoundingBox::Corner) == GMX_SIMD4_WIDTH * sizeof(float),
 634                   "The Corner struct should hold exactly 4 floats");
 635
 636     Simd4Float bb_0_S = load4(x);
 637     Simd4Float bb_1_S = bb_0_S;
 638
 639     for (int i = 1; i < na; i++)
 640     {
 641         Simd4Float x_S = load4(x + i * GMX_SIMD4_WIDTH);
 642         bb_0_S         = min(bb_0_S, x_S);
 643         bb_1_S         = max(bb_1_S, x_S);
 644     }
 645
 646     store4(bb->lower.ptr(), bb_0_S);
 647     store4(bb->upper.ptr(), bb_1_S);
 648 }
 649
 650 #    if NBNXN_BBXXXX
 651
 652 /*! \brief Computes the bounding box for na coordinates in order xyz?, bb order xxxxyyyyzzzz */
 653 static void calc_bounding_box_xxxx_simd4(int na, const float* x, BoundingBox* bb_work_aligned, real* bb)
 654 {
 655     calc_bounding_box_simd4(na, x, bb_work_aligned);
 656
 657     bb[0 * c_packedBoundingBoxesDimSize] = bb_work_aligned->lower.x;
 658     bb[1 * c_packedBoundingBoxesDimSize] = bb_work_aligned->lower.y;
 659     bb[2 * c_packedBoundingBoxesDimSize] = bb_work_aligned->lower.z;
 660     bb[3 * c_packedBoundingBoxesDimSize] = bb_work_aligned->upper.x;
 661     bb[4 * c_packedBoundingBoxesDimSize] = bb_work_aligned->upper.y;
 662     bb[5 * c_packedBoundingBoxesDimSize] = bb_work_aligned->upper.z;
 663 }
 664
 665 #    endif /* NBNXN_BBXXXX */
 666
 667 #endif /* NBNXN_SEARCH_SIMD4_FLOAT_X_BB */
 668
 669
 670 /*! \brief Combines pairs of consecutive bounding boxes */
 671 static void combine_bounding_box_pairs(const Grid&                      grid,
 672                                        gmx::ArrayRef<const BoundingBox> bb,
 673                                        gmx::ArrayRef<BoundingBox>       bbj)
 674 {
 675     // TODO: During SIMDv2 transition only some archs use namespace (remove when done)
 676     using namespace gmx;
 677
 678     for (int i = 0; i < grid.numColumns(); i++)
 679     {
 680         /* Starting bb in a column is expected to be 2-aligned */
 681         const int sc2 = grid.firstCellInColumn(i) >> 1;
 682         /* For odd numbers skip the last bb here */
 683         const int nc2 = (grid.numAtomsInColumn(i) + 3) >> (2 + 1);
 684         int       c2;
 685         for (c2 = sc2; c2 < sc2 + nc2; c2++)
 686         {
 687 #if NBNXN_SEARCH_BB_SIMD4
 688             Simd4Float min_S, max_S;
 689
 690             min_S = min(load4(bb[c2 * 2 + 0].lower.ptr()), load4(bb[c2 * 2 + 1].lower.ptr()));
 691             max_S = max(load4(bb[c2 * 2 + 0].upper.ptr()), load4(bb[c2 * 2 + 1].upper.ptr()));
 692             store4(bbj[c2].lower.ptr(), min_S);
 693             store4(bbj[c2].upper.ptr(), max_S);
 694 #else
 695             bbj[c2].lower = BoundingBox::Corner::min(bb[c2 * 2 + 0].lower, bb[c2 * 2 + 1].lower);
 696             bbj[c2].upper = BoundingBox::Corner::max(bb[c2 * 2 + 0].upper, bb[c2 * 2 + 1].upper);
 697 #endif
 698         }
 699         if (((grid.numAtomsInColumn(i) + 3) >> 2) & 1)
 700         {
 701             /* The bb count in this column is odd: duplicate the last bb */
 702             bbj[c2].lower = bb[c2 * 2].lower;
 703             bbj[c2].upper = bb[c2 * 2].upper;
 704         }
 705     }
 706 }
 707
 708
 709 /*! \brief Prints the average bb size, used for debug output */
 710 static void print_bbsizes_simple(FILE* fp, const Grid& grid)
 711 {
 712     dvec ba = { 0 };
 713     for (int c = 0; c < grid.numCells(); c++)
 714     {
 715         const BoundingBox& bb = grid.iBoundingBoxes()[c];
 716         ba[XX] += bb.upper.x - bb.lower.x;
 717         ba[YY] += bb.upper.y - bb.lower.y;
 718         ba[ZZ] += bb.upper.z - bb.lower.z;
 719     }
 720     dsvmul(1.0 / grid.numCells(), ba, ba);
 721
 722     const Grid::Dimensions& dims = grid.dimensions();
 723     real                    avgCellSizeZ =
 724             (dims.atomDensity > 0 ? grid.geometry().numAtomsICluster
 725                                             / (dims.atomDensity * dims.cellSize[XX] * dims.cellSize[YY])
 726                                   : 0.0);
 727
 728     fprintf(fp,
 729             "ns bb: grid %4.2f %4.2f %4.2f abs %4.2f %4.2f %4.2f rel %4.2f %4.2f %4.2f\n",
 730             dims.cellSize[XX],
 731             dims.cellSize[YY],
 732             avgCellSizeZ,
 733             ba[XX],
 734             ba[YY],
 735             ba[ZZ],
 736             ba[XX] * dims.invCellSize[XX],
 737             ba[YY] * dims.invCellSize[YY],
 738             dims.atomDensity > 0 ? ba[ZZ] / avgCellSizeZ : 0.0);
 739 }
 740
 741 /*! \brief Prints the average bb size, used for debug output */
 742 static void print_bbsizes_supersub(FILE* fp, const Grid& grid)
 743 {
 744     int  ns;
 745     dvec ba;
 746
 747     clear_dvec(ba);
 748     ns = 0;
 749     for (int c = 0; c < grid.numCells(); c++)
 750     {
 751 #if NBNXN_BBXXXX
 752         for (int s = 0; s < grid.numClustersPerCell()[c]; s += c_packedBoundingBoxesDimSize)
 753         {
 754             int  cs_w          = (c * c_gpuNumClusterPerCell + s) / c_packedBoundingBoxesDimSize;
 755             auto boundingBoxes = grid.packedBoundingBoxes().subArray(
 756                     cs_w * c_packedBoundingBoxesSize, c_packedBoundingBoxesSize);
 757             for (int i = 0; i < c_packedBoundingBoxesDimSize; i++)
 758             {
 759                 for (int d = 0; d < DIM; d++)
 760                 {
 761                     ba[d] += boundingBoxes[(DIM + d) * c_packedBoundingBoxesDimSize + i]
 762                              - boundingBoxes[(0 + d) * c_packedBoundingBoxesDimSize + i];
 763                 }
 764             }
 765         }
 766 #else
 767         for (int s = 0; s < grid.numClustersPerCell()[c]; s++)
 768         {
 769             const BoundingBox& bb = grid.iBoundingBoxes()[c * c_gpuNumClusterPerCell + s];
 770             ba[XX] += bb.upper.x - bb.lower.x;
 771             ba[YY] += bb.upper.y - bb.lower.y;
 772             ba[ZZ] += bb.upper.z - bb.lower.z;
 773         }
 774 #endif
 775         ns += grid.numClustersPerCell()[c];
 776     }
 777     dsvmul(1.0 / ns, ba, ba);
 778
 779     const Grid::Dimensions& dims = grid.dimensions();
 780     const real              avgClusterSizeZ =
 781             (dims.atomDensity > 0 ? grid.geometry().numAtomsPerCell
 782                                             / (dims.atomDensity * dims.cellSize[XX]
 783                                                * dims.cellSize[YY] * c_gpuNumClusterPerCellZ)
 784                                   : 0.0);
 785
 786     fprintf(fp,
 787             "ns bb: grid %4.2f %4.2f %4.2f abs %4.2f %4.2f %4.2f rel %4.2f %4.2f %4.2f\n",
 788             dims.cellSize[XX] / c_gpuNumClusterPerCellX,
 789             dims.cellSize[YY] / c_gpuNumClusterPerCellY,
 790             avgClusterSizeZ,
 791             ba[XX],
 792             ba[YY],
 793             ba[ZZ],
 794             ba[XX] * c_gpuNumClusterPerCellX * dims.invCellSize[XX],
 795             ba[YY] * c_gpuNumClusterPerCellY * dims.invCellSize[YY],
 796             dims.atomDensity > 0 ? ba[ZZ] / avgClusterSizeZ : 0.0);
 797 }
 798
 799 /*!\brief Set non-bonded interaction flags for the current cluster.
 800  *
 801  * Sorts atoms on LJ coefficients: !=0 first, ==0 at the end.
 802  */
 803 static void sort_cluster_on_flag(int                numAtomsInCluster,
 804                                  int                atomStart,
 805                                  int                atomEnd,
 806                                  const int*         atinfo,
 807                                  gmx::ArrayRef<int> order,
 808                                  int*               flags)
 809 {
 810     constexpr int c_maxNumAtomsInCluster = 8;
 811     int           sort1[c_maxNumAtomsInCluster];
 812     int           sort2[c_maxNumAtomsInCluster];
 813
 814     GMX_ASSERT(numAtomsInCluster <= c_maxNumAtomsInCluster,
 815                "Need to increase c_maxNumAtomsInCluster to support larger clusters");
 816
 817     *flags = 0;
 818
 819     int subc = 0;
 820     for (int s = atomStart; s < atomEnd; s += numAtomsInCluster)
 821     {
 822         /* Make lists for this (sub-)cell on atoms with and without LJ */
 823         int      n1       = 0;
 824         int      n2       = 0;
 825         gmx_bool haveQ    = FALSE;
 826         int      a_lj_max = -1;
 827         for (int a = s; a < std::min(s + numAtomsInCluster, atomEnd); a++)
 828         {
 829             haveQ = haveQ || GET_CGINFO_HAS_Q(atinfo[order[a]]);
 830
 831             if (GET_CGINFO_HAS_VDW(atinfo[order[a]]))
 832             {
 833                 sort1[n1++] = order[a];
 834                 a_lj_max    = a;
 835             }
 836             else
 837             {
 838                 sort2[n2++] = order[a];
 839             }
 840         }
 841
 842         /* If we don't have atoms with LJ, there's nothing to sort */
 843         if (n1 > 0)
 844         {
 845             *flags |= NBNXN_CI_DO_LJ(subc);
 846
 847             if (2 * n1 <= numAtomsInCluster)
 848             {
 849                 /* Only sort when strictly necessary. Ordering particles
 850                  * Ordering particles can lead to less accurate summation
 851                  * due to rounding, both for LJ and Coulomb interactions.
 852                  */
 853                 if (2 * (a_lj_max - s) >= numAtomsInCluster)
 854                 {
 855                     for (int i = 0; i < n1; i++)
 856                     {
 857                         order[atomStart + i] = sort1[i];
 858                     }
 859                     for (int j = 0; j < n2; j++)
 860                     {
 861                         order[atomStart + n1 + j] = sort2[j];
 862                     }
 863                 }
 864
 865                 *flags |= NBNXN_CI_HALF_LJ(subc);
 866             }
 867         }
 868         if (haveQ)
 869         {
 870             *flags |= NBNXN_CI_DO_COUL(subc);
 871         }
 872         subc++;
 873     }
 874 }
 875
 876 /*! \brief Fill a pair search cell with atoms.
 877  *
 878  * Potentially sorts atoms and sets the interaction flags.
 879  */
 880 void Grid::fillCell(GridSetData*                   gridSetData,
 881                     nbnxn_atomdata_t*              nbat,
 882                     int                            atomStart,
 883                     int                            atomEnd,
 884                     const int*                     atinfo,
 885                     gmx::ArrayRef<const gmx::RVec> x,
 886                     BoundingBox gmx_unused* bb_work_aligned)
 887 {
 888     const int numAtoms = atomEnd - atomStart;
 889
 890     const gmx::ArrayRef<int>& cells       = gridSetData->cells;
 891     const gmx::ArrayRef<int>& atomIndices = gridSetData->atomIndices;
 892
 893     if (geometry_.isSimple)
 894     {
 895         /* Note that non-local grids are already sorted.
 896          * Then sort_cluster_on_flag will only set the flags and the sorting
 897          * will not affect the atom order.
 898          */
 899         sort_cluster_on_flag(geometry_.numAtomsICluster,
 900                              atomStart,
 901                              atomEnd,
 902                              atinfo,
 903                              atomIndices,
 904                              flags_.data() + atomToCluster(atomStart) - cellOffset_);
 905     }
 906
 907     if (haveFep_)
 908     {
 909         /* Set the fep flag for perturbed atoms in this (sub-)cell */
 910
 911         /* The grid-local cluster/(sub-)cell index */
 912         int cell = atomToCluster(atomStart)
 913                    - cellOffset_ * (geometry_.isSimple ? 1 : c_gpuNumClusterPerCell);
 914         fep_[cell] = 0;
 915         for (int at = atomStart; at < atomEnd; at++)
 916         {
 917             if (atomIndices[at] >= 0 && GET_CGINFO_FEP(atinfo[atomIndices[at]]))
 918             {
 919                 fep_[cell] |= (1 << (at - atomStart));
 920             }
 921         }
 922     }
 923
 924     /* Now we have sorted the atoms, set the cell indices */
 925     for (int at = atomStart; at < atomEnd; at++)
 926     {
 927         cells[atomIndices[at]] = at;
 928     }
 929
 930     copy_rvec_to_nbat_real(atomIndices.data() + atomStart,
 931                            numAtoms,
 932                            geometry_.numAtomsICluster,
 933                            as_rvec_array(x.data()),
 934                            nbat->XFormat,
 935                            nbat->x().data(),
 936                            atomStart);
 937
 938     if (nbat->XFormat == nbatX4)
 939     {
 940         /* Store the bounding boxes as xyz.xyz. */
 941         size_t       offset = atomToCluster(atomStart - cellOffset_ * geometry_.numAtomsICluster);
 942         BoundingBox* bb_ptr = bb_.data() + offset;
 943
 944 #if GMX_SIMD && GMX_SIMD_REAL_WIDTH == 2
 945         if (2 * geometry_.numAtomsJCluster == geometry_.numAtomsICluster)
 946         {
 947             calc_bounding_box_x_x4_halves(numAtoms,
 948                                           nbat->x().data() + atom_to_x_index<c_packX4>(atomStart),
 949                                           bb_ptr,
 950                                           bbj_.data() + offset * 2);
 951         }
 952         else
 953 #endif
 954         {
 955             calc_bounding_box_x_x4(numAtoms, nbat->x().data() + atom_to_x_index<c_packX4>(atomStart), bb_ptr);
 956         }
 957     }
 958     else if (nbat->XFormat == nbatX8)
 959     {
 960         /* Store the bounding boxes as xyz.xyz. */
 961         size_t       offset = atomToCluster(atomStart - cellOffset_ * geometry_.numAtomsICluster);
 962         BoundingBox* bb_ptr = bb_.data() + offset;
 963
 964         calc_bounding_box_x_x8(numAtoms, nbat->x().data() + atom_to_x_index<c_packX8>(atomStart), bb_ptr);
 965     }
 966 #if NBNXN_BBXXXX
 967     else if (!geometry_.isSimple)
 968     {
 969         /* Store the bounding boxes in a format convenient
 970          * for SIMD4 calculations: xxxxyyyyzzzz...
 971          */
 972         const int clusterIndex = ((atomStart - cellOffset_ * geometry_.numAtomsPerCell)
 973                                   >> geometry_.numAtomsICluster2Log);
 974         float*    pbb_ptr      = pbb_.data() + packedBoundingBoxesIndex(clusterIndex)
 975                          + (clusterIndex & (c_packedBoundingBoxesDimSize - 1));
 976
 977 #    if NBNXN_SEARCH_SIMD4_FLOAT_X_BB
 978         if (nbat->XFormat == nbatXYZQ)
 979         {
 980             GMX_ASSERT(bb_work_aligned != nullptr, "Must have valid aligned work structure");
 981             calc_bounding_box_xxxx_simd4(
 982                     numAtoms, nbat->x().data() + atomStart * nbat->xstride, bb_work_aligned, pbb_ptr);
 983         }
 984         else
 985 #    endif
 986         {
 987             calc_bounding_box_xxxx(
 988                     numAtoms, nbat->xstride, nbat->x().data() + atomStart * nbat->xstride, pbb_ptr);
 989         }
 990         if (gmx_debug_at)
 991         {
 992             fprintf(debug,
 993                     "cell %4d bb %5.2f %5.2f %5.2f %5.2f %5.2f %5.2f\n",
 994                     atomToCluster(atomStart),
 995                     pbb_ptr[0 * c_packedBoundingBoxesDimSize],
 996                     pbb_ptr[3 * c_packedBoundingBoxesDimSize],
 997                     pbb_ptr[1 * c_packedBoundingBoxesDimSize],
 998                     pbb_ptr[4 * c_packedBoundingBoxesDimSize],
 999                     pbb_ptr[2 * c_packedBoundingBoxesDimSize],
1000                     pbb_ptr[5 * c_packedBoundingBoxesDimSize]);
1001         }
1002     }
1003 #endif
1004     else
1005     {
1006         /* Store the bounding boxes as xyz.xyz. */
1007         BoundingBox* bb_ptr =
1008                 bb_.data() + atomToCluster(atomStart - cellOffset_ * geometry_.numAtomsPerCell);
1009
1010         calc_bounding_box(numAtoms, nbat->xstride, nbat->x().data() + atomStart * nbat->xstride, bb_ptr);
1011
1012         if (gmx_debug_at)
1013         {
1014             int bbo = atomToCluster(atomStart - cellOffset_ * geometry_.numAtomsPerCell);
1015             fprintf(debug,
1016                     "cell %4d bb %5.2f %5.2f %5.2f %5.2f %5.2f %5.2f\n",
1017                     atomToCluster(atomStart),
1018                     bb_[bbo].lower.x,
1019                     bb_[bbo].lower.y,
1020                     bb_[bbo].lower.z,
1021                     bb_[bbo].upper.x,
1022                     bb_[bbo].upper.y,
1023                     bb_[bbo].upper.z);
1024         }
1025     }
1026 }
1027
1028 void Grid::sortColumnsCpuGeometry(GridSetData*                   gridSetData,
1029                                   int                            dd_zone,
1030                                   const int*                     atinfo,
1031                                   gmx::ArrayRef<const gmx::RVec> x,
1032                                   nbnxn_atomdata_t*              nbat,
1033                                   const gmx::Range<int>          columnRange,
1034                                   gmx::ArrayRef<int>             sort_work)
1035 {
1036     if (debug)
1037     {
1038         fprintf(debug,
1039                 "cell_offset %d sorting columns %d - %d\n",
1040                 cellOffset_,
1041                 *columnRange.begin(),
1042                 *columnRange.end());
1043     }
1044
1045     const bool relevantAtomsAreWithinGridBounds = (dimensions_.maxAtomGroupRadius == 0);
1046
1047     const int numAtomsPerCell = geometry_.numAtomsPerCell;
1048
1049     /* Sort the atoms within each x,y column in 3 dimensions */
1050     for (int cxy : columnRange)
1051     {
1052         const int numAtoms   = numAtomsInColumn(cxy);
1053         const int numCellsZ  = cxy_ind_[cxy + 1] - cxy_ind_[cxy];
1054         const int atomOffset = firstAtomInColumn(cxy);
1055
1056         /* Sort the atoms within each x,y column on z coordinate */
1057         sort_atoms(ZZ,
1058                    FALSE,
1059                    dd_zone,
1060                    relevantAtomsAreWithinGridBounds,
1061                    gridSetData->atomIndices.data() + atomOffset,
1062                    numAtoms,
1063                    x,
1064                    dimensions_.lowerCorner[ZZ],
1065                    1.0 / dimensions_.gridSize[ZZ],
1066                    numCellsZ * numAtomsPerCell,
1067                    sort_work);
1068
1069         /* Fill the ncz cells in this column */
1070         const int firstCell  = firstCellInColumn(cxy);
1071         int       cellFilled = firstCell;
1072         for (int cellZ = 0; cellZ < numCellsZ; cellZ++)
1073         {
1074             const int cell = firstCell + cellZ;
1075
1076             const int atomOffsetCell = atomOffset + cellZ * numAtomsPerCell;
1077             const int numAtomsCell = std::min(numAtomsPerCell, numAtoms - (atomOffsetCell - atomOffset));
1078
1079             fillCell(gridSetData, nbat, atomOffsetCell, atomOffsetCell + numAtomsCell, atinfo, x, nullptr);
1080
1081             /* This copy to bbcz is not really necessary.
1082              * But it allows to use the same grid search code
1083              * for the simple and supersub cell setups.
1084              */
1085             if (numAtomsCell > 0)
1086             {
1087                 cellFilled = cell;
1088             }
1089             bbcz_[cell].lower = bb_[cellFilled].lower.z;
1090             bbcz_[cell].upper = bb_[cellFilled].upper.z;
1091         }
1092
1093         /* Set the unused atom indices to -1 */
1094         for (int ind = numAtoms; ind < numCellsZ * numAtomsPerCell; ind++)
1095         {
1096             gridSetData->atomIndices[atomOffset + ind] = -1;
1097         }
1098     }
1099 }
1100
1101 /* Spatially sort the atoms within one grid column */
1102 void Grid::sortColumnsGpuGeometry(GridSetData*                   gridSetData,
1103                                   int                            dd_zone,
1104                                   const int*                     atinfo,
1105                                   gmx::ArrayRef<const gmx::RVec> x,
1106                                   nbnxn_atomdata_t*              nbat,
1107                                   const gmx::Range<int>          columnRange,
1108                                   gmx::ArrayRef<int>             sort_work)
1109 {
1110     BoundingBox  bb_work_array[2];
1111     BoundingBox* bb_work_aligned = reinterpret_cast<BoundingBox*>(
1112             (reinterpret_cast<std::size_t>(bb_work_array + 1)) & (~(static_cast<std::size_t>(15))));
1113
1114     if (debug)
1115     {
1116         fprintf(debug,
1117                 "cell_offset %d sorting columns %d - %d\n",
1118                 cellOffset_,
1119                 *columnRange.begin(),
1120                 *columnRange.end());
1121     }
1122
1123     const bool relevantAtomsAreWithinGridBounds = (dimensions_.maxAtomGroupRadius == 0);
1124
1125     const int numAtomsPerCell = geometry_.numAtomsPerCell;
1126
1127     const int subdiv_x = geometry_.numAtomsICluster;
1128     const int subdiv_y = c_gpuNumClusterPerCellX * subdiv_x;
1129     const int subdiv_z = c_gpuNumClusterPerCellY * subdiv_y;
1130
1131     /* Extract the atom index array that will be filled here */
1132     const gmx::ArrayRef<int>& atomIndices = gridSetData->atomIndices;
1133
1134     /* Sort the atoms within each x,y column in 3 dimensions.
1135      * Loop over all columns on the x/y grid.
1136      */
1137     for (int cxy : columnRange)
1138     {
1139         const int gridX = cxy / dimensions_.numCells[YY];
1140         const int gridY = cxy - gridX * dimensions_.numCells[YY];
1141
1142         const int numAtomsInColumn = cxy_na_[cxy];
1143         const int numCellsInColumn = cxy_ind_[cxy + 1] - cxy_ind_[cxy];
1144         const int atomOffset       = firstAtomInColumn(cxy);
1145
1146         /* Sort the atoms within each x,y column on z coordinate */
1147         sort_atoms(ZZ,
1148                    FALSE,
1149                    dd_zone,
1150                    relevantAtomsAreWithinGridBounds,
1151                    atomIndices.data() + atomOffset,
1152                    numAtomsInColumn,
1153                    x,
1154                    dimensions_.lowerCorner[ZZ],
1155                    1.0 / dimensions_.gridSize[ZZ],
1156                    numCellsInColumn * numAtomsPerCell,
1157                    sort_work);
1158
1159         /* This loop goes over the cells and clusters along z at once */
1160         for (int sub_z = 0; sub_z < numCellsInColumn * c_gpuNumClusterPerCellZ; sub_z++)
1161         {
1162             const int atomOffsetZ = atomOffset + sub_z * subdiv_z;
1163             const int numAtomsZ = std::min(subdiv_z, numAtomsInColumn - (atomOffsetZ - atomOffset));
1164             int       cz        = -1;
1165             /* We have already sorted on z */
1166
1167             if (sub_z % c_gpuNumClusterPerCellZ == 0)
1168             {
1169                 cz             = sub_z / c_gpuNumClusterPerCellZ;
1170                 const int cell = cxy_ind_[cxy] + cz;
1171
1172                 /* The number of atoms in this cell/super-cluster */
1173                 const int numAtoms =
1174                         std::min(numAtomsPerCell, numAtomsInColumn - (atomOffsetZ - atomOffset));
1175
1176                 numClusters_[cell] = std::min(
1177                         c_gpuNumClusterPerCell,
1178                         (numAtoms + geometry_.numAtomsICluster - 1) / geometry_.numAtomsICluster);
1179
1180                 /* Store the z-boundaries of the bounding box of the cell */
1181                 bbcz_[cell].lower = x[atomIndices[atomOffsetZ]][ZZ];
1182                 bbcz_[cell].upper = x[atomIndices[atomOffsetZ + numAtoms - 1]][ZZ];
1183             }
1184
1185             if (c_gpuNumClusterPerCellY > 1)
1186             {
1187                 /* Sort the atoms along y */
1188                 sort_atoms(YY,
1189                            (sub_z & 1) != 0,
1190                            dd_zone,
1191                            relevantAtomsAreWithinGridBounds,
1192                            atomIndices.data() + atomOffsetZ,
1193                            numAtomsZ,
1194                            x,
1195                            dimensions_.lowerCorner[YY] + gridY * dimensions_.cellSize[YY],
1196                            dimensions_.invCellSize[YY],
1197                            subdiv_z,
1198                            sort_work);
1199             }
1200
1201             for (int sub_y = 0; sub_y < c_gpuNumClusterPerCellY; sub_y++)
1202             {
1203                 const int atomOffsetY = atomOffsetZ + sub_y * subdiv_y;
1204                 const int numAtomsY = std::min(subdiv_y, numAtomsInColumn - (atomOffsetY - atomOffset));
1205
1206                 if (c_gpuNumClusterPerCellX > 1)
1207                 {
1208                     /* Sort the atoms along x */
1209                     sort_atoms(XX,
1210                                ((cz * c_gpuNumClusterPerCellY + sub_y) & 1) != 0,
1211                                dd_zone,
1212                                relevantAtomsAreWithinGridBounds,
1213                                atomIndices.data() + atomOffsetY,
1214                                numAtomsY,
1215                                x,
1216                                dimensions_.lowerCorner[XX] + gridX * dimensions_.cellSize[XX],
1217                                dimensions_.invCellSize[XX],
1218                                subdiv_y,
1219                                sort_work);
1220                 }
1221
1222                 for (int sub_x = 0; sub_x < c_gpuNumClusterPerCellX; sub_x++)
1223                 {
1224                     const int atomOffsetX = atomOffsetY + sub_x * subdiv_x;
1225                     const int numAtomsX =
1226                             std::min(subdiv_x, numAtomsInColumn - (atomOffsetX - atomOffset));
1227
1228                     fillCell(gridSetData, nbat, atomOffsetX, atomOffsetX + numAtomsX, atinfo, x, bb_work_aligned);
1229                 }
1230             }
1231         }
1232
1233         /* Set the unused atom indices to -1 */
1234         for (int ind = numAtomsInColumn; ind < numCellsInColumn * numAtomsPerCell; ind++)
1235         {
1236             atomIndices[atomOffset + ind] = -1;
1237         }
1238     }
1239 }
1240
1241 /*! \brief Sets the cell index in the cell array for atom \p atomIndex and increments the atom count for the grid column */
1242 static void setCellAndAtomCount(gmx::ArrayRef<int> cell, int cellIndex, gmx::ArrayRef<int> cxy_na, int atomIndex)
1243 {
1244     cell[atomIndex] = cellIndex;
1245     cxy_na[cellIndex] += 1;
1246 }
1247
1248 void Grid::calcColumnIndices(const Grid::Dimensions&        gridDims,
1249                              const gmx::UpdateGroupsCog*    updateGroupsCog,
1250                              const gmx::Range<int>          atomRange,
1251                              gmx::ArrayRef<const gmx::RVec> x,
1252                              const int                      dd_zone,
1253                              const int*                     move,
1254                              const int                      thread,
1255                              const int                      nthread,
1256                              gmx::ArrayRef<int>             cell,
1257                              gmx::ArrayRef<int>             cxy_na)
1258 {
1259     const int numColumns = gridDims.numCells[XX] * gridDims.numCells[YY];
1260
1261     /* We add one extra cell for particles which moved during DD */
1262     for (int i = 0; i < numColumns; i++)
1263     {
1264         cxy_na[i] = 0;
1265     }
1266
1267     int taskAtomStart = *atomRange.begin() + static_cast<int>((thread + 0) * atomRange.size()) / nthread;
1268     int taskAtomEnd = *atomRange.begin() + static_cast<int>((thread + 1) * atomRange.size()) / nthread;
1269
1270     if (dd_zone == 0)
1271     {
1272         /* Home zone */
1273         for (int i = taskAtomStart; i < taskAtomEnd; i++)
1274         {
1275             if (move == nullptr || move[i] >= 0)
1276             {
1277
1278                 const gmx::RVec& coord = (updateGroupsCog ? updateGroupsCog->cogForAtom(i) : x[i]);
1279
1280                 /* We need to be careful with rounding,
1281                  * particles might be a few bits outside the local zone.
1282                  * The int cast takes care of the lower bound,
1283                  * we will explicitly take care of the upper bound.
1284                  */
1285                 int cx = static_cast<int>((coord[XX] - gridDims.lowerCorner[XX])
1286                                           * gridDims.invCellSize[XX]);
1287                 int cy = static_cast<int>((coord[YY] - gridDims.lowerCorner[YY])
1288                                           * gridDims.invCellSize[YY]);
1289
1290 #ifndef NDEBUG
1291                 if (cx < 0 || cx > gridDims.numCells[XX] || cy < 0 || cy > gridDims.numCells[YY])
1292                 {
1293                     gmx_fatal(FARGS,
1294                               "grid cell cx %d cy %d out of range (max %d %d)\n"
1295                               "atom %f %f %f, grid->c0 %f %f",
1296                               cx,
1297                               cy,
1298                               gridDims.numCells[XX],
1299                               gridDims.numCells[YY],
1300                               x[i][XX],
1301                               x[i][YY],
1302                               x[i][ZZ],
1303                               gridDims.lowerCorner[XX],
1304                               gridDims.lowerCorner[YY]);
1305                 }
1306 #endif
1307                 /* Take care of potential rouding issues */
1308                 cx = std::min(cx, gridDims.numCells[XX] - 1);
1309                 cy = std::min(cy, gridDims.numCells[YY] - 1);
1310
1311                 /* For the moment cell will contain only the, grid local,
1312                  * x and y indices, not z.
1313                  */
1314                 setCellAndAtomCount(cell, cx * gridDims.numCells[YY] + cy, cxy_na, i);
1315             }
1316             else
1317             {
1318                 /* Put this moved particle after the end of the grid,
1319                  * so we can process it later without using conditionals.
1320                  */
1321                 setCellAndAtomCount(cell, numColumns, cxy_na, i);
1322             }
1323         }
1324     }
1325     else
1326     {
1327         /* Non-home zone */
1328         for (int i = taskAtomStart; i < taskAtomEnd; i++)
1329         {
1330             int cx = static_cast<int>((x[i][XX] - gridDims.lowerCorner[XX]) * gridDims.invCellSize[XX]);
1331             int cy = static_cast<int>((x[i][YY] - gridDims.lowerCorner[YY]) * gridDims.invCellSize[YY]);
1332
1333             /* For non-home zones there could be particles outside
1334              * the non-bonded cut-off range, which have been communicated
1335              * for bonded interactions only. For the result it doesn't
1336              * matter where these end up on the grid. For performance
1337              * we put them in an extra row at the border.
1338              */
1339             cx = std::max(cx, 0);
1340             cx = std::min(cx, gridDims.numCells[XX] - 1);
1341             cy = std::max(cy, 0);
1342             cy = std::min(cy, gridDims.numCells[YY] - 1);
1343
1344             /* For the moment cell will contain only the, grid local,
1345              * x and y indices, not z.
1346              */
1347             setCellAndAtomCount(cell, cx * gridDims.numCells[YY] + cy, cxy_na, i);
1348         }
1349     }
1350 }
1351
1352 /*! \brief Resizes grid and atom data which depend on the number of cells */
1353 static void resizeForNumberOfCells(const int         numNbnxnAtoms,
1354                                    const int         numAtomsMoved,
1355                                    GridSetData*      gridSetData,
1356                                    nbnxn_atomdata_t* nbat)
1357 {
1358     /* Note: gridSetData->cellIndices was already resized before */
1359
1360     /* To avoid conditionals we store the moved particles at the end of a,
1361      * make sure we have enough space.
1362      */
1363     gridSetData->atomIndices.resize(numNbnxnAtoms + numAtomsMoved);
1364
1365     /* Make space in nbat for storing the atom coordinates */
1366     nbat->resizeCoordinateBuffer(numNbnxnAtoms);
1367 }
1368
1369 void Grid::setCellIndices(int                            ddZone,
1370                           int                            cellOffset,
1371                           GridSetData*                   gridSetData,
1372                           gmx::ArrayRef<GridWork>        gridWork,
1373                           const gmx::Range<int>          atomRange,
1374                           const int*                     atinfo,
1375                           gmx::ArrayRef<const gmx::RVec> x,
1376                           const int                      numAtomsMoved,
1377                           nbnxn_atomdata_t*              nbat)
1378 {
1379     cellOffset_ = cellOffset;
1380
1381     srcAtomBegin_ = *atomRange.begin();
1382     srcAtomEnd_   = *atomRange.end();
1383
1384     const int nthread = gmx_omp_nthreads_get(emntPairsearch);
1385
1386     const int numAtomsPerCell = geometry_.numAtomsPerCell;
1387
1388     /* Make the cell index as a function of x and y */
1389     int ncz_max = 0;
1390     int ncz     = 0;
1391     cxy_ind_[0] = 0;
1392     for (int i = 0; i < numColumns() + 1; i++)
1393     {
1394         /* We set ncz_max at the beginning of the loop iso at the end
1395          * to skip i=grid->ncx*grid->numCells[YY] which are moved particles
1396          * that do not need to be ordered on the grid.
1397          */
1398         if (ncz > ncz_max)
1399         {
1400             ncz_max = ncz;
1401         }
1402         int cxy_na_i = gridWork[0].numAtomsPerColumn[i];
1403         for (int thread = 1; thread < nthread; thread++)
1404         {
1405             cxy_na_i += gridWork[thread].numAtomsPerColumn[i];
1406         }
1407         ncz = (cxy_na_i + numAtomsPerCell - 1) / numAtomsPerCell;
1408         if (nbat->XFormat == nbatX8)
1409         {
1410             /* Make the number of cell a multiple of 2 */
1411             ncz = (ncz + 1) & ~1;
1412         }
1413         cxy_ind_[i + 1] = cxy_ind_[i] + ncz;
1414         /* Clear cxy_na_, so we can reuse the array below */
1415         cxy_na_[i] = 0;
1416     }
1417     numCellsTotal_     = cxy_ind_[numColumns()] - cxy_ind_[0];
1418     numCellsColumnMax_ = ncz_max;
1419
1420     /* Resize grid and atom data which depend on the number of cells */
1421     resizeForNumberOfCells(atomIndexEnd(), numAtomsMoved, gridSetData, nbat);
1422
1423     if (debug)
1424     {
1425         fprintf(debug,
1426                 "ns na_sc %d na_c %d super-cells: %d x %d y %d z %.1f maxz %d\n",
1427                 numAtomsPerCell,
1428                 geometry_.numAtomsICluster,
1429                 numCellsTotal_,
1430                 dimensions_.numCells[XX],
1431                 dimensions_.numCells[YY],
1432                 numCellsTotal_ / (static_cast<double>(numColumns())),
1433                 ncz_max);
1434         if (gmx_debug_at)
1435         {
1436             int i = 0;
1437             for (int cy = 0; cy < dimensions_.numCells[YY]; cy++)
1438             {
1439                 for (int cx = 0; cx < dimensions_.numCells[XX]; cx++)
1440                 {
1441                     fprintf(debug, " %2d", cxy_ind_[i + 1] - cxy_ind_[i]);
1442                     i++;
1443                 }
1444                 fprintf(debug, "\n");
1445             }
1446         }
1447     }
1448
1449     /* Make sure the work array for sorting is large enough */
1450     const int worstCaseSortBufferSize = ncz_max * numAtomsPerCell * c_sortGridMaxSizeFactor;
1451     if (worstCaseSortBufferSize > gmx::index(gridWork[0].sortBuffer.size()))
1452     {
1453         for (GridWork& work : gridWork)
1454         {
1455             /* Elements not in use should be -1 */
1456             work.sortBuffer.resize(worstCaseSortBufferSize, -1);
1457         }
1458     }
1459
1460     /* Now we know the dimensions we can fill the grid.
1461      * This is the first, unsorted fill. We sort the columns after this.
1462      */
1463     gmx::ArrayRef<int> cells       = gridSetData->cells;
1464     gmx::ArrayRef<int> atomIndices = gridSetData->atomIndices;
1465     for (int i : atomRange)
1466     {
1467         /* At this point nbs->cell contains the local grid x,y indices */
1468         const int cxy                                        = cells[i];
1469         atomIndices[firstAtomInColumn(cxy) + cxy_na_[cxy]++] = i;
1470     }
1471
1472     if (ddZone == 0)
1473     {
1474         /* Set the cell indices for the moved particles */
1475         int n0 = numCellsTotal_ * numAtomsPerCell;
1476         int n1 = numCellsTotal_ * numAtomsPerCell + cxy_na_[numColumns()];
1477         for (int i = n0; i < n1; i++)
1478         {
1479             cells[atomIndices[i]] = i;
1480         }
1481     }
1482
1483     /* Sort the super-cell columns along z into the sub-cells. */
1484 #pragma omp parallel for num_threads(nthread) schedule(static)
1485     for (int thread = 0; thread < nthread; thread++)
1486     {
1487         try
1488         {
1489             gmx::Range<int> columnRange(((thread + 0) * numColumns()) / nthread,
1490                                         ((thread + 1) * numColumns()) / nthread);
1491             if (geometry_.isSimple)
1492             {
1493                 sortColumnsCpuGeometry(
1494                         gridSetData, ddZone, atinfo, x, nbat, columnRange, gridWork[thread].sortBuffer);
1495             }
1496             else
1497             {
1498                 sortColumnsGpuGeometry(
1499                         gridSetData, ddZone, atinfo, x, nbat, columnRange, gridWork[thread].sortBuffer);
1500             }
1501         }
1502         GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR
1503     }
1504
1505     if (geometry_.isSimple && nbat->XFormat == nbatX8)
1506     {
1507         combine_bounding_box_pairs(*this, bb_, bbj_);
1508     }
1509
1510     if (!geometry_.isSimple)
1511     {
1512         numClustersTotal_ = 0;
1513         for (int i = 0; i < numCellsTotal_; i++)
1514         {
1515             numClustersTotal_ += numClusters_[i];
1516         }
1517     }
1518
1519     if (debug)
1520     {
1521         if (geometry_.isSimple)
1522         {
1523             print_bbsizes_simple(debug, *this);
1524         }
1525         else
1526         {
1527             fprintf(debug,
1528                     "ns non-zero sub-cells: %d average atoms %.2f\n",
1529                     numClustersTotal_,
1530                     atomRange.size() / static_cast<double>(numClustersTotal_));
1531
1532             print_bbsizes_supersub(debug, *this);
1533         }
1534     }
1535 }
1536
1537 } // namespace Nbnxm