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