src/gromacs/nbnxm/grid.h

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2019,2020,2021, 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 /*! \internal \file
  37  *
  38  * \brief
  39  * Declares the Grid class.
  40  *
  41  * This class provides functionality for setting up and accessing atoms
  42  * on a grid for one domain decomposition zone. This grid is used for
  43  * generating cluster pair lists for computing non-bonded pair interactions.
  44  * The grid consists of a regular array of columns along dimensions x and y.
  45  * Along z the number of cells and their boundaries vary between the columns.
  46  * Each cell can hold one or more clusters of atoms, depending on the grid
  47  * geometry, which is set by the pair-list type.
  48  *
  49  * \author Berk Hess <hess@kth.se>
  50  * \ingroup module_nbnxm
  51  */
  52
  53 #ifndef GMX_NBNXM_GRID_H
  54 #define GMX_NBNXM_GRID_H
  55
  56 #include <memory>
  57 #include <vector>
  58
  59 #include "gromacs/gpu_utils/hostallocator.h"
  60 #include "gromacs/math/vectypes.h"
  61 #include "gromacs/utility/alignedallocator.h"
  62 #include "gromacs/utility/arrayref.h"
  63 #include "gromacs/utility/range.h"
  64
  65 struct nbnxn_atomdata_t;
  66 enum class PairlistType;
  67
  68 namespace gmx
  69 {
  70 class UpdateGroupsCog;
  71 } // namespace gmx
  72
  73 namespace Nbnxm
  74 {
  75
  76 struct GridSetData;
  77 struct GridWork;
  78
  79 /*! \internal
  80  * \brief Bounding box for a nbnxm atom cluster
  81  *
  82  * \note Should be aligned in memory to enable 4-wide SIMD operations.
  83  */
  84 struct BoundingBox
  85 {
  86     /*! \internal
  87      * \brief Corner for the bounding box, padded with one element to enable 4-wide SIMD operations
  88      */
  89     struct Corner
  90     {
  91         //! Returns a corner with the minimum coordinates along each dimension
  92         static Corner min(const Corner& c1, const Corner& c2)
  93         {
  94             Corner cMin;
  95
  96             cMin.x = std::min(c1.x, c2.x);
  97             cMin.y = std::min(c1.y, c2.y);
  98             cMin.z = std::min(c1.z, c2.z);
  99             /* This value of the padding is irrelevant, as long as it
 100              * is initialized. We use min to allow auto-vectorization.
 101              */
 102             cMin.padding = std::min(c1.padding, c2.padding);
 103
 104             return cMin;
 105         }
 106
 107         //! Returns a corner with the maximum coordinates along each dimension
 108         static Corner max(const Corner& c1, const Corner& c2)
 109         {
 110             Corner cMax;
 111
 112             cMax.x       = std::max(c1.x, c2.x);
 113             cMax.y       = std::max(c1.y, c2.y);
 114             cMax.z       = std::max(c1.z, c2.z);
 115             cMax.padding = std::max(c1.padding, c2.padding);
 116
 117             return cMax;
 118         }
 119
 120         //! Returns a pointer for SIMD loading of a Corner object
 121         const float* ptr() const { return &x; }
 122
 123         //! Returns a pointer for SIMD storing of a Corner object
 124         float* ptr() { return &x; }
 125
 126         //! x coordinate
 127         float x;
 128         //! y coordinate
 129         float y;
 130         //! z coordinate
 131         float z;
 132         //! padding, unused, but should be set to avoid operations on unitialized data
 133         float padding;
 134     };
 135
 136     //! lower, along x and y and z, corner
 137     Corner lower;
 138     //! upper, along x and y and z, corner
 139     Corner upper;
 140 };
 141
 142 /*! \internal
 143  * \brief Bounding box for one dimension of a grid cell
 144  */
 145 struct BoundingBox1D
 146 {
 147     //! lower bound
 148     float lower;
 149     //! upper bound
 150     float upper;
 151 };
 152
 153 } // namespace Nbnxm
 154
 155 namespace Nbnxm
 156 {
 157
 158 /*! \internal
 159  * \brief A pair-search grid object for one domain decomposition zone
 160  *
 161  * This is a rectangular 3D grid covering a potentially non-rectangular
 162  * volume which is either the whole unit cell or the local zone or part
 163  * of a non-local zone when using domain decomposition. The grid cells
 164  * are even spaced along x/y and irregular along z. Each cell is sub-divided
 165  * into atom clusters. With a CPU geometry, each cell contains 1 or 2 clusters.
 166  * With a GPU geometry, each cell contains up to 8 clusters. The geometry is
 167  * set by the pairlist type which is the only argument of the constructor.
 168  *
 169  * When multiple grids are used, i.e. with domain decomposition, we want
 170  * to avoid the overhead of multiple coordinate arrays or extra indexing.
 171  * Therefore each grid stores a cell offset, so a contiguous cell index
 172  * can be used to index atom arrays. All methods returning atom indices
 173  * return indices which index into a full atom array.
 174  *
 175  * Note that when atom groups, instead of individual atoms, are assigned
 176  * to grid cells, individual atoms can be geometrically outside the cell
 177  * and grid that they have been assigned to (as determined by the center
 178  * or geometry of the atom group they belong to).
 179  */
 180 class Grid
 181 {
 182 public:
 183     /*! \internal
 184      * \brief The cluster and cell geometry of a grid
 185      */
 186     struct Geometry
 187     {
 188         //! Constructs the cluster/cell geometry given the type of pairlist
 189         Geometry(PairlistType pairlistType);
 190
 191         //! Is this grid simple (CPU) or hierarchical (GPU)
 192         bool isSimple;
 193         //! Number of atoms per cluster
 194         int numAtomsICluster;
 195         //! Number of atoms for list j-clusters
 196         int numAtomsJCluster;
 197         //! Number of atoms per cell
 198         int numAtomsPerCell;
 199         //! 2log of na_c
 200         int numAtomsICluster2Log;
 201     };
 202
 203     //! The physical dimensions of a grid \internal
 204     struct Dimensions
 205     {
 206         //! The lower corner of the (local) grid
 207         rvec lowerCorner;
 208         //! The upper corner of the (local) grid
 209         rvec upperCorner;
 210         //! The physical grid size: upperCorner - lowerCorner
 211         rvec gridSize;
 212         //! An estimate for the atom number density of the region targeted by the grid
 213         real atomDensity;
 214         //! The maximum distance an atom can be outside of a cell and outside of the grid
 215         real maxAtomGroupRadius;
 216         //! Size of cell along dimension x and y
 217         real cellSize[DIM - 1];
 218         //! 1/size of a cell along dimensions x and y
 219         real invCellSize[DIM - 1];
 220         //! The number of grid cells along dimensions x and y
 221         int numCells[DIM - 1];
 222     };
 223
 224     //! Constructs a grid given the type of pairlist
 225     Grid(PairlistType pairlistType, const bool& haveFep);
 226
 227     //! Returns the geometry of the grid cells
 228     const Geometry& geometry() const { return geometry_; }
 229
 230     //! Returns the dimensions of the grid
 231     const Dimensions& dimensions() const { return dimensions_; }
 232
 233     //! Returns the total number of grid columns
 234     int numColumns() const { return dimensions_.numCells[XX] * dimensions_.numCells[YY]; }
 235
 236     //! Returns the total number of grid cells
 237     int numCells() const { return numCellsTotal_; }
 238
 239     //! Returns the cell offset of (the first cell of) this grid in the list of cells combined over all grids
 240     int cellOffset() const { return cellOffset_; }
 241
 242     //! Returns the maximum number of grid cells in a column
 243     int numCellsColumnMax() const { return numCellsColumnMax_; }
 244
 245     //! Returns the start of the source atom range mapped to this grid
 246     int srcAtomBegin() const { return srcAtomBegin_; }
 247
 248     //! Returns the end of the source atom range mapped to this grid
 249     int srcAtomEnd() const { return srcAtomEnd_; }
 250
 251     //! Returns the first cell index in the grid, starting at 0 in this grid
 252     int firstCellInColumn(int columnIndex) const { return cxy_ind_[columnIndex]; }
 253
 254     //! Returns the number of cells in the column
 255     int numCellsInColumn(int columnIndex) const
 256     {
 257         return cxy_ind_[columnIndex + 1LL] - cxy_ind_[columnIndex];
 258     }
 259
 260     //! Returns the index of the first atom in the column
 261     int firstAtomInColumn(int columnIndex) const
 262     {
 263         return (cellOffset_ + cxy_ind_[columnIndex]) * geometry_.numAtomsPerCell;
 264     }
 265
 266     //! Returns the number of real atoms in the column
 267     int numAtomsInColumn(int columnIndex) const { return cxy_na_[columnIndex]; }
 268
 269     /*! \brief Returns a view of the number of non-filler, atoms for each grid column
 270      *
 271      * \todo Needs a useful name. */
 272     gmx::ArrayRef<const int> cxy_na() const { return cxy_na_; }
 273     /*! \brief Returns a view of the grid-local cell index for each grid column
 274      *
 275      * \todo Needs a useful name. */
 276     gmx::ArrayRef<const int> cxy_ind() const { return cxy_ind_; }
 277
 278     //! Returns the number of real atoms in the column
 279     int numAtomsPerCell() const { return geometry_.numAtomsPerCell; }
 280
 281     //! Returns the number of atoms in the column including padding
 282     int paddedNumAtomsInColumn(int columnIndex) const
 283     {
 284         return numCellsInColumn(columnIndex) * geometry_.numAtomsPerCell;
 285     }
 286
 287     //! Returns the end of the atom index range on the grid, including padding
 288     int atomIndexEnd() const { return (cellOffset_ + numCellsTotal_) * geometry_.numAtomsPerCell; }
 289
 290     //! Returns whether any atom in the cluster is perturbed
 291     bool clusterIsPerturbed(int clusterIndex) const { return fep_[clusterIndex] != 0U; }
 292
 293     //! Returns whether the given atom in the cluster is perturbed
 294     bool atomIsPerturbed(int clusterIndex, int atomIndexInCluster) const
 295     {
 296         return (fep_[clusterIndex] & (1 << atomIndexInCluster)) != 0U;
 297     }
 298
 299     //! Returns the free-energy perturbation bits for the cluster
 300     unsigned int fepBits(int clusterIndex) const { return fep_[clusterIndex]; }
 301
 302     //! Returns the i-bounding boxes for all clusters on the grid
 303     gmx::ArrayRef<const BoundingBox> iBoundingBoxes() const { return bb_; }
 304
 305     //! Returns the j-bounding boxes for all clusters on the grid
 306     gmx::ArrayRef<const BoundingBox> jBoundingBoxes() const { return bbj_; }
 307
 308     //! Returns the packed bounding boxes for all clusters on the grid, empty with a CPU list
 309     gmx::ArrayRef<const float> packedBoundingBoxes() const { return pbb_; }
 310
 311     //! Returns the bounding boxes along z for all cells on the grid
 312     gmx::ArrayRef<const BoundingBox1D> zBoundingBoxes() const { return bbcz_; }
 313
 314     //! Returns the flags for all clusters on the grid
 315     gmx::ArrayRef<const int> clusterFlags() const { return flags_; }
 316
 317     //! Returns the number of clusters for all cells on the grid, empty with a CPU geometry
 318     gmx::ArrayRef<const int> numClustersPerCell() const { return numClusters_; }
 319
 320     //! Returns the cluster index for an atom
 321     int atomToCluster(int atomIndex) const { return (atomIndex >> geometry_.numAtomsICluster2Log); }
 322
 323     //! Returns the total number of clusters on the grid
 324     int numClusters() const
 325     {
 326         if (geometry_.isSimple)
 327         {
 328             return numCellsTotal_;
 329         }
 330         else
 331         {
 332             return numClustersTotal_;
 333         }
 334     }
 335
 336     //! Sets the grid dimensions
 337     void setDimensions(int                ddZone,
 338                        int                numAtoms,
 339                        gmx::RVec          lowerCorner,
 340                        gmx::RVec          upperCorner,
 341                        real               atomDensity,
 342                        real               maxAtomGroupRadius,
 343                        bool               haveFep,
 344                        gmx::PinningPolicy pinningPolicy);
 345
 346     //! Sets the cell indices using indices in \p gridSetData and \p gridWork
 347     void setCellIndices(int                            ddZone,
 348                         int                            cellOffset,
 349                         GridSetData*                   gridSetData,
 350                         gmx::ArrayRef<GridWork>        gridWork,
 351                         gmx::Range<int>                atomRange,
 352                         gmx::ArrayRef<const int64_t>   atomInfo,
 353                         gmx::ArrayRef<const gmx::RVec> x,
 354                         int                            numAtomsMoved,
 355                         nbnxn_atomdata_t*              nbat);
 356
 357     //! Determine in which grid columns atoms should go, store cells and atom counts in \p cell and \p cxy_na
 358     static void calcColumnIndices(const Grid::Dimensions&        gridDims,
 359                                   const gmx::UpdateGroupsCog*    updateGroupsCog,
 360                                   gmx::Range<int>                atomRange,
 361                                   gmx::ArrayRef<const gmx::RVec> x,
 362                                   int                            dd_zone,
 363                                   const int*                     move,
 364                                   int                            thread,
 365                                   int                            nthread,
 366                                   gmx::ArrayRef<int>             cell,
 367                                   gmx::ArrayRef<int>             cxy_na);
 368
 369 private:
 370     /*! \brief Fill a pair search cell with atoms
 371      *
 372      * Potentially sorts atoms and sets the interaction flags.
 373      */
 374     void fillCell(GridSetData*                   gridSetData,
 375                   nbnxn_atomdata_t*              nbat,
 376                   int                            atomStart,
 377                   int                            atomEnd,
 378                   gmx::ArrayRef<const int64_t>   atomInfo,
 379                   gmx::ArrayRef<const gmx::RVec> x,
 380                   BoundingBox gmx_unused* bb_work_aligned);
 381
 382     //! Spatially sort the atoms within the given column range, for CPU geometry
 383     void sortColumnsCpuGeometry(GridSetData*                   gridSetData,
 384                                 int                            dd_zone,
 385                                 gmx::ArrayRef<const int64_t>   atomInfo,
 386                                 gmx::ArrayRef<const gmx::RVec> x,
 387                                 nbnxn_atomdata_t*              nbat,
 388                                 gmx::Range<int>                columnRange,
 389                                 gmx::ArrayRef<int>             sort_work);
 390
 391     //! Spatially sort the atoms within the given column range, for GPU geometry
 392     void sortColumnsGpuGeometry(GridSetData*                   gridSetData,
 393                                 int                            dd_zone,
 394                                 gmx::ArrayRef<const int64_t>   atomInfo,
 395                                 gmx::ArrayRef<const gmx::RVec> x,
 396                                 nbnxn_atomdata_t*              nbat,
 397                                 gmx::Range<int>                columnRange,
 398                                 gmx::ArrayRef<int>             sort_work);
 399
 400     /* Data members */
 401     //! The geometry of the grid clusters and cells
 402     Geometry geometry_;
 403     //! The physical dimensions of the grid
 404     Dimensions dimensions_;
 405
 406     //! The total number of cells in this grid
 407     int numCellsTotal_;
 408     //! Index in nbs->cell corresponding to cell 0
 409     int cellOffset_;
 410     //! The maximum number of cells in a column
 411     int numCellsColumnMax_;
 412
 413     //! The start of the source atom range mapped to this grid
 414     int srcAtomBegin_;
 415     //! The end of the source atom range mapped to this grid
 416     int srcAtomEnd_;
 417
 418     /* Grid data */
 419     /*! \brief The number of, non-filler, atoms for each grid column.
 420      *
 421      * \todo Needs a useful name. */
 422     gmx::HostVector<int> cxy_na_;
 423     /*! \brief The grid-local cell index for each grid column
 424      *
 425      * \todo Needs a useful name. */
 426     gmx::HostVector<int> cxy_ind_;
 427
 428     //! The number of cluster for each cell
 429     std::vector<int> numClusters_;
 430
 431     /* Bounding boxes */
 432     //! Bounding boxes in z for the cells
 433     std::vector<BoundingBox1D> bbcz_;
 434     //! 3D bounding boxes for the sub cells
 435     std::vector<BoundingBox, gmx::AlignedAllocator<BoundingBox>> bb_;
 436     //! 3D j-bounding boxes for the case where the i- and j-cluster sizes are different
 437     std::vector<BoundingBox, gmx::AlignedAllocator<BoundingBox>> bbjStorage_;
 438     //! 3D j-bounding boxes
 439     gmx::ArrayRef<BoundingBox> bbj_;
 440     //! 3D bounding boxes in packed xxxx format per cell
 441     std::vector<float, gmx::AlignedAllocator<float>> pbb_;
 442
 443     //! Tells whether we have perturbed interactions, authorative source is in GridSet (never modified)
 444     const bool& haveFep_;
 445
 446     /* Bit-flag information */
 447     //! Flags for properties of clusters in each cell
 448     std::vector<int> flags_;
 449     //! Signal bits for atoms in each cell that tell whether an atom is perturbed
 450     std::vector<unsigned int> fep_;
 451
 452     /* Statistics */
 453     //! Total number of clusters, used for printing
 454     int numClustersTotal_;
 455 };
 456
 457 } // namespace Nbnxm
 458
 459 #endif