Add nbnxm GridSet class
authorBerk Hess <hess@kth.se>
Mon, 11 Mar 2019 14:07:47 +0000 (15:07 +0100)
committerBerk Hess <hess@kth.se>
Thu, 14 Mar 2019 10:39:14 +0000 (11:39 +0100)
Collected all data that is shared between nbnxm grids to a new class
called GridSet.

TODO: Turn nbnxn_search into a class and move some functions

Change-Id: Iae0dc988d88310d88e0541656f896e552cd19260

src/gromacs/domdec/partition.cpp
src/gromacs/nbnxm/atomdata.cpp
src/gromacs/nbnxm/grid.cpp
src/gromacs/nbnxm/grid.h
src/gromacs/nbnxm/gridset.cpp [new file with mode: 0644]
src/gromacs/nbnxm/gridset.h [new file with mode: 0644]
src/gromacs/nbnxm/internal.h
src/gromacs/nbnxm/nbnxm.h
src/gromacs/nbnxm/pairlist.cpp

index 7df2f9cf2bc9e4b551c3137dd5ea0284f5c9589d..05e0d5934f8194f36da54e4eb60ddebd62ac70d1 100644 (file)
@@ -2896,7 +2896,7 @@ static void dd_sort_state(gmx_domdec_t *dd, rvec *cgcm, t_forcerec *fr, t_state
     if (fr->cutoff_scheme == ecutsVERLET)
     {
         /* The atoms are now exactly in grid order, update the grid order */
-        nbnxn_set_atomorder(fr->nbv->nbs.get());
+        fr->nbv->setLocalAtomOrder();
     }
     else
     {
@@ -3387,7 +3387,7 @@ void dd_partition_system(FILE                    *fplog,
                        fr->rlist, grid_density);
             break;
         case ecutsVERLET:
-            nbnxn_get_ncells(fr->nbv->nbs.get(), &ncells_old[XX], &ncells_old[YY]);
+            fr->nbv->getLocalNumCells(&ncells_old[XX], &ncells_old[YY]);
             break;
         default:
             gmx_incons("unimplemented");
@@ -3426,7 +3426,7 @@ void dd_partition_system(FILE                    *fplog,
                                   state_local->x,
                                   ncg_moved, bRedist ? comm->movedBuffer.data() : nullptr);
 
-                nbnxn_get_ncells(fr->nbv->nbs.get(), &ncells_new[XX], &ncells_new[YY]);
+                fr->nbv->getLocalNumCells(&ncells_new[XX], &ncells_new[YY]);
                 break;
             case ecutsGROUP:
                 fill_grid(&comm->zones, fr->ns->grid, dd->ncg_home,
index dcfd25cb244d9d0024c8785c20a76a309132d583..86a33118cca0b0e5c512c7e9f50af05c57692454 100644 (file)
@@ -719,19 +719,14 @@ static void copy_lj_to_nbat_lj_comb(gmx::ArrayRef<const real> ljparam_type,
     }
 }
 
-static int numAtomsFromGrids(const nbnxn_search &nbs)
-{
-    return nbs.grid.back().atomIndexEnd();
-}
-
 /* Sets the atom type in nbnxn_atomdata_t */
 static void nbnxn_atomdata_set_atomtypes(nbnxn_atomdata_t::Params *params,
-                                         const nbnxn_search       *nbs,
+                                         const Nbnxm::GridSet     &gridSet,
                                          const int                *type)
 {
-    params->type.resize(numAtomsFromGrids(*nbs));
+    params->type.resize(gridSet.numGridAtomsTotal());
 
-    for (const Nbnxm::Grid &grid : nbs->grid)
+    for (const Nbnxm::Grid &grid : gridSet.grids())
     {
         /* Loop over all columns and copy and fill */
         for (int i = 0; i < grid.numColumns(); i++)
@@ -739,7 +734,8 @@ static void nbnxn_atomdata_set_atomtypes(nbnxn_atomdata_t::Params *params,
             const int numAtoms   = grid.paddedNumAtomsInColumn(i);
             const int atomOffset = grid.firstAtomInColumn(i);
 
-            copy_int_to_nbat_int(nbs->a.data() + atomOffset, grid.numAtomsInColumn(i), numAtoms,
+            copy_int_to_nbat_int(gridSet.atomIndices().data() + atomOffset,
+                                 grid.numAtomsInColumn(i), numAtoms,
                                  type, params->numTypes - 1, params->type.data() + atomOffset);
         }
     }
@@ -748,13 +744,13 @@ static void nbnxn_atomdata_set_atomtypes(nbnxn_atomdata_t::Params *params,
 /* Sets the LJ combination rule parameters in nbnxn_atomdata_t */
 static void nbnxn_atomdata_set_ljcombparams(nbnxn_atomdata_t::Params *params,
                                             const int                 XFormat,
-                                            const nbnxn_search       *nbs)
+                                            const Nbnxm::GridSet     &gridSet)
 {
-    params->lj_comb.resize(numAtomsFromGrids(*nbs)*2);
+    params->lj_comb.resize(gridSet.numGridAtomsTotal()*2);
 
     if (params->comb_rule != ljcrNONE)
     {
-        for (const Nbnxm::Grid &grid : nbs->grid)
+        for (const Nbnxm::Grid &grid : gridSet.grids())
         {
             /* Loop over all columns and copy and fill */
             for (int i = 0; i < grid.numColumns(); i++)
@@ -789,16 +785,16 @@ static void nbnxn_atomdata_set_ljcombparams(nbnxn_atomdata_t::Params *params,
 }
 
 /* Sets the charges in nbnxn_atomdata_t *nbat */
-static void nbnxn_atomdata_set_charges(nbnxn_atomdata_t    *nbat,
-                                       const nbnxn_search  *nbs,
-                                       const real          *charge)
+static void nbnxn_atomdata_set_charges(nbnxn_atomdata_t     *nbat,
+                                       const Nbnxm::GridSet &gridSet,
+                                       const real           *charge)
 {
     if (nbat->XFormat != nbatXYZQ)
     {
         nbat->paramsDeprecated().q.resize(nbat->numAtoms());
     }
 
-    for (const Nbnxm::Grid &grid : nbs->grid)
+    for (const Nbnxm::Grid &grid : gridSet.grids())
     {
         /* Loop over all columns and copy and fill */
         for (int cxy = 0; cxy < grid.numColumns(); cxy++)
@@ -813,7 +809,7 @@ static void nbnxn_atomdata_set_charges(nbnxn_atomdata_t    *nbat,
                 int   i;
                 for (i = 0; i < numAtoms; i++)
                 {
-                    *q = charge[nbs->a[atomOffset + i]];
+                    *q = charge[gridSet.atomIndices()[atomOffset + i]];
                     q += STRIDE_XYZQ;
                 }
                 /* Complete the partially filled last cell with zeros */
@@ -829,7 +825,7 @@ static void nbnxn_atomdata_set_charges(nbnxn_atomdata_t    *nbat,
                 int   i;
                 for (i = 0; i < numAtoms; i++)
                 {
-                    *q = charge[nbs->a[atomOffset + i]];
+                    *q = charge[gridSet.atomIndices()[atomOffset + i]];
                     q++;
                 }
                 /* Complete the partially filled last cell with zeros */
@@ -849,8 +845,8 @@ static void nbnxn_atomdata_set_charges(nbnxn_atomdata_t    *nbat,
  * All perturbed interactions are calculated in the free energy kernel,
  * using the original charge and LJ data, not nbnxn_atomdata_t.
  */
-static void nbnxn_atomdata_mask_fep(nbnxn_atomdata_t    *nbat,
-                                    const nbnxn_search  *nbs)
+static void nbnxn_atomdata_mask_fep(nbnxn_atomdata_t     *nbat,
+                                    const Nbnxm::GridSet &gridSet)
 {
     nbnxn_atomdata_t::Params &params = nbat->paramsDeprecated();
     real                     *q;
@@ -867,7 +863,7 @@ static void nbnxn_atomdata_mask_fep(nbnxn_atomdata_t    *nbat,
         stride_q = 1;
     }
 
-    for (const Nbnxm::Grid &grid : nbs->grid)
+    for (const Nbnxm::Grid &grid : gridSet.grids())
     {
         int nsubc;
         if (grid.geometry().isSimple)
@@ -936,7 +932,7 @@ static void copy_egp_to_nbat_egps(const int *a, int na, int na_round,
 
 /* Set the energy group indices for atoms in nbnxn_atomdata_t */
 static void nbnxn_atomdata_set_energygroups(nbnxn_atomdata_t::Params *params,
-                                            const nbnxn_search       *nbs,
+                                            const Nbnxm::GridSet     &gridSet,
                                             const int                *atinfo)
 {
     if (params->nenergrp == 1)
@@ -944,9 +940,9 @@ static void nbnxn_atomdata_set_energygroups(nbnxn_atomdata_t::Params *params,
         return;
     }
 
-    params->energrp.resize(numAtomsFromGrids(*nbs));
+    params->energrp.resize(gridSet.numGridAtomsTotal());
 
-    for (const Nbnxm::Grid &grid : nbs->grid)
+    for (const Nbnxm::Grid &grid : gridSet.grids())
     {
         /* Loop over all columns and copy and fill */
         for (int i = 0; i < grid.numColumns(); i++)
@@ -954,7 +950,8 @@ static void nbnxn_atomdata_set_energygroups(nbnxn_atomdata_t::Params *params,
             const int numAtoms   = grid.paddedNumAtomsInColumn(i);
             const int atomOffset = grid.firstAtomInColumn(i);
 
-            copy_egp_to_nbat_egps(nbs->a.data() + atomOffset, grid.numAtomsInColumn(i), numAtoms,
+            copy_egp_to_nbat_egps(gridSet.atomIndices().data() + atomOffset,
+                                  grid.numAtomsInColumn(i), numAtoms,
                                   c_nbnxnCpuIClusterSize, params->neg_2log,
                                   atinfo,
                                   params->energrp.data() + grid.atomToCluster(atomOffset));
@@ -970,19 +967,21 @@ void nbnxn_atomdata_set(nbnxn_atomdata_t    *nbat,
 {
     nbnxn_atomdata_t::Params &params = nbat->paramsDeprecated();
 
-    nbnxn_atomdata_set_atomtypes(&params, nbs, mdatoms->typeA);
+    const Nbnxm::GridSet     &gridSet = nbs->gridSet();
 
-    nbnxn_atomdata_set_charges(nbat, nbs, mdatoms->chargeA);
+    nbnxn_atomdata_set_atomtypes(&params, gridSet, mdatoms->typeA);
 
-    if (nbs->bFEP)
+    nbnxn_atomdata_set_charges(nbat, gridSet, mdatoms->chargeA);
+
+    if (gridSet.haveFep())
     {
-        nbnxn_atomdata_mask_fep(nbat, nbs);
+        nbnxn_atomdata_mask_fep(nbat, gridSet);
     }
 
     /* This must be done after masking types for FEP */
-    nbnxn_atomdata_set_ljcombparams(&params, nbat->XFormat, nbs);
+    nbnxn_atomdata_set_ljcombparams(&params, nbat->XFormat, gridSet);
 
-    nbnxn_atomdata_set_energygroups(&params, nbs, atinfo);
+    nbnxn_atomdata_set_energygroups(&params, gridSet, atinfo);
 }
 
 /* Copies the shift vector array to nbnxn_atomdata_t */
@@ -1010,41 +1009,42 @@ void nbnxn_atomdata_copy_x_to_nbat_x(const nbnxn_search       *nbs,
     wallcycle_start(wcycle, ewcNB_XF_BUF_OPS);
     wallcycle_sub_start(wcycle, ewcsNB_X_BUF_OPS);
 
-    int g0 = 0, g1 = 0;
-    int nth, th;
+    const Nbnxm::GridSet &gridSet = nbs->gridSet();
 
+    int                   gridBegin = 0;
+    int                   gridEnd   = 0;
     switch (locality)
     {
         case Nbnxm::AtomLocality::All:
         case Nbnxm::AtomLocality::Count:
-            g0 = 0;
-            g1 = nbs->grid.size();
+            gridBegin = 0;
+            gridEnd   = gridSet.grids().size();
             break;
         case Nbnxm::AtomLocality::Local:
-            g0 = 0;
-            g1 = 1;
+            gridBegin = 0;
+            gridEnd   = 1;
             break;
         case Nbnxm::AtomLocality::NonLocal:
-            g0 = 1;
-            g1 = nbs->grid.size();
+            gridBegin = 1;
+            gridEnd   = gridSet.grids().size();
             break;
     }
 
     if (FillLocal)
     {
-        nbat->natoms_local = nbs->grid[0].atomIndexEnd();
+        nbat->natoms_local = gridSet.grids()[0].atomIndexEnd();
     }
 
-    nth = gmx_omp_nthreads_get(emntPairsearch);
+    const int nth = gmx_omp_nthreads_get(emntPairsearch);
 
 #pragma omp parallel for num_threads(nth) schedule(static)
-    for (th = 0; th < nth; th++)
+    for (int th = 0; th < nth; th++)
     {
         try
         {
-            for (int g = g0; g < g1; g++)
+            for (int g = gridBegin; g < gridEnd; g++)
             {
-                const Nbnxm::Grid  &grid       = nbs->grid[g];
+                const Nbnxm::Grid  &grid       = gridSet.grids()[g];
                 const int           numCellsXY = grid.numColumns();
 
                 const int           cxy0 = (numCellsXY* th      + nth - 1)/nth;
@@ -1068,7 +1068,8 @@ void nbnxn_atomdata_copy_x_to_nbat_x(const nbnxn_search       *nbs,
                          */
                         na_fill = na;
                     }
-                    copy_rvec_to_nbat_real(nbs->a.data() + ash, na, na_fill, x,
+                    copy_rvec_to_nbat_real(gridSet.atomIndices().data() + ash,
+                                           na, na_fill, x,
                                            nbat->XFormat, nbat->x().data(), ash);
                 }
             }
@@ -1166,14 +1167,14 @@ nbnxn_atomdata_reduce_reals_simd(real gmx_unused * gmx_restrict dest,
 
 /* Add part of the force array(s) from nbnxn_atomdata_t to f */
 static void
-nbnxn_atomdata_add_nbat_f_to_f_part(const nbnxn_search *nbs,
+nbnxn_atomdata_add_nbat_f_to_f_part(const Nbnxm::GridSet &gridSet,
                                     const nbnxn_atomdata_t *nbat,
                                     gmx::ArrayRef<const nbnxn_atomdata_output_t> out,
                                     int nfa,
                                     int a0, int a1,
                                     rvec *f)
 {
-    gmx::ArrayRef<const int> cell = nbs->cell;
+    gmx::ArrayRef<const int> cell = gridSet.cells();
 
     /* Loop over all columns and copy and fill */
     switch (nbat->FFormat)
@@ -1477,24 +1478,26 @@ nonbonded_verlet_t::atomdata_add_nbat_f_to_f(const Nbnxm::AtomLocality  locality
     wallcycle_start(wcycle, ewcNB_XF_BUF_OPS);
     wallcycle_sub_start(wcycle, ewcsNB_F_BUF_OPS);
 
-    int a0 = 0, na = 0;
+    const Nbnxm::GridSet &gridSet = nbs->gridSet();
 
     nbs_cycle_start(&nbs->cc[enbsCCreducef]);
 
+    int a0 = 0;
+    int na = 0;
     switch (locality)
     {
         case Nbnxm::AtomLocality::All:
         case Nbnxm::AtomLocality::Count:
             a0 = 0;
-            na = nbs->natoms_nonlocal;
+            na = gridSet.numRealAtomsTotal();
             break;
         case Nbnxm::AtomLocality::Local:
             a0 = 0;
-            na = nbs->natoms_local;
+            na = gridSet.numRealAtomsLocal();
             break;
         case Nbnxm::AtomLocality::NonLocal:
-            a0 = nbs->natoms_local;
-            na = nbs->natoms_nonlocal - nbs->natoms_local;
+            a0 = gridSet.numRealAtomsLocal();
+            na = gridSet.numRealAtomsTotal() - gridSet.numRealAtomsLocal();
             break;
     }
 
@@ -1524,7 +1527,7 @@ nonbonded_verlet_t::atomdata_add_nbat_f_to_f(const Nbnxm::AtomLocality  locality
     {
         try
         {
-            nbnxn_atomdata_add_nbat_f_to_f_part(nbs.get(), nbat.get(),
+            nbnxn_atomdata_add_nbat_f_to_f_part(gridSet, nbat.get(),
                                                 nbat->out,
                                                 1,
                                                 a0+((th+0)*na)/nth,
index 9abbe860bdb69ebc7f80838b285e014cf9c68545..d5b5e129d17cd3a1d94c865bf9c6fbfb875d541f 100644 (file)
 #include "gromacs/mdlib/updategroupscog.h"
 #include "gromacs/mdtypes/forcerec.h" // only for GET_CGINFO_*
 #include "gromacs/nbnxm/atomdata.h"
-#include "gromacs/nbnxm/nbnxm.h"
 #include "gromacs/nbnxm/nbnxm_geometry.h"
-#include "gromacs/nbnxm/pairlist.h"
-#include "gromacs/nbnxm/pairlistset.h"
 #include "gromacs/simd/simd.h"
 #include "gromacs/simd/vector_operations.h"
-#include "gromacs/utility/exceptions.h"
-#include "gromacs/utility/smalloc.h"
 
 #include "internal.h"
 
-struct gmx_domdec_zones_t;
-
 namespace Nbnxm
 {
 
@@ -106,13 +99,13 @@ static real gridAtomDensity(int        numAtoms,
     return numAtoms/(size[XX]*size[YY]*size[ZZ]);
 }
 
-void Grid::setDimensions(const nbnxn_search *nbs,
-                         const int           ddZone,
+void Grid::setDimensions(const int           ddZone,
                          const int           numAtoms,
                          const rvec          lowerCorner,
                          const rvec          upperCorner,
                          real                atomDensity,
-                         const real          maxAtomGroupRadius)
+                         const real          maxAtomGroupRadius,
+                         const bool          haveFep)
 {
     /* For the home zone we compute the density when not set (=-1) or when =0 */
     if (ddZone == 0 && atomDensity <= 0)
@@ -188,11 +181,6 @@ void Grid::setDimensions(const nbnxn_search *nbs,
     cxy_na_.resize(numColumns() + 1);
     cxy_ind_.resize(numColumns() + 2);
 
-    for (nbnxn_search_work_t &work : nbs->work)
-    {
-        work.cxy_na.resize(numColumns() + 1);
-    }
-
     /* Worst case scenario of 1 atom in each last cell */
     int maxNumCells;
     if (geometry_.numAtomsJCluster <= geometry_.numAtomsICluster)
@@ -239,7 +227,7 @@ void Grid::setDimensions(const nbnxn_search *nbs,
     }
 
     flags_.resize(maxNumCells);
-    if (nbs->bFEP)
+    if (haveFep)
     {
         fep_.resize(maxNumCells*geometry_.numAtomsPerCell/geometry_.numAtomsICluster);
     }
@@ -859,7 +847,7 @@ static void sort_cluster_on_flag(int                 numAtomsInCluster,
  *
  * Potentially sorts atoms and sets the interaction flags.
  */
-void Grid::fillCell(nbnxn_search                   *nbs,
+void Grid::fillCell(const GridSetData              &gridSetData,
                     nbnxn_atomdata_t               *nbat,
                     int                             atomStart,
                     int                             atomEnd,
@@ -867,7 +855,10 @@ void Grid::fillCell(nbnxn_search                   *nbs,
                     gmx::ArrayRef<const gmx::RVec>  x,
                     BoundingBox gmx_unused         *bb_work_aligned)
 {
-    const int numAtoms = atomEnd - atomStart;
+    const int                 numAtoms = atomEnd - atomStart;
+
+    const gmx::ArrayRef<int> &cells       = gridSetData.cells;
+    const gmx::ArrayRef<int> &atomIndices = gridSetData.atomIndices;
 
     if (geometry_.isSimple)
     {
@@ -875,11 +866,11 @@ void Grid::fillCell(nbnxn_search                   *nbs,
          * Then sort_cluster_on_flag will only set the flags and the sorting
          * will not affect the atom order.
          */
-        sort_cluster_on_flag(geometry_.numAtomsICluster, atomStart, atomEnd, atinfo, nbs->a,
+        sort_cluster_on_flag(geometry_.numAtomsICluster, atomStart, atomEnd, atinfo, atomIndices,
                              flags_.data() + atomToCluster(atomStart) - cellOffset_);
     }
 
-    if (nbs->bFEP)
+    if (gridSetData.haveFep)
     {
         /* Set the fep flag for perturbed atoms in this (sub-)cell */
 
@@ -888,7 +879,7 @@ void Grid::fillCell(nbnxn_search                   *nbs,
         fep_[cell] = 0;
         for (int at = atomStart; at < atomEnd; at++)
         {
-            if (nbs->a[at] >= 0 && GET_CGINFO_FEP(atinfo[nbs->a[at]]))
+            if (atomIndices[at] >= 0 && GET_CGINFO_FEP(atinfo[atomIndices[at]]))
             {
                 fep_[cell] |= (1 << (at - atomStart));
             }
@@ -898,10 +889,10 @@ void Grid::fillCell(nbnxn_search                   *nbs,
     /* Now we have sorted the atoms, set the cell indices */
     for (int at = atomStart; at < atomEnd; at++)
     {
-        nbs->cell[nbs->a[at]] = at;
+        cells[atomIndices[at]] = at;
     }
 
-    copy_rvec_to_nbat_real(nbs->a.data() + atomStart, numAtoms, geometry_.numAtomsICluster,
+    copy_rvec_to_nbat_real(atomIndices.data() + atomStart, numAtoms, geometry_.numAtomsICluster,
                            as_rvec_array(x.data()),
                            nbat->XFormat, nbat->x().data(), atomStart);
 
@@ -987,7 +978,7 @@ void Grid::fillCell(nbnxn_search                   *nbs,
     }
 }
 
-void Grid::sortColumnsCpuGeometry(nbnxn_search *nbs,
+void Grid::sortColumnsCpuGeometry(const GridSetData &gridSetData,
                                   int dd_zone,
                                   int atomStart, int atomEnd,
                                   const int *atinfo,
@@ -1016,7 +1007,7 @@ void Grid::sortColumnsCpuGeometry(nbnxn_search *nbs,
         /* Sort the atoms within each x,y column on z coordinate */
         sort_atoms(ZZ, FALSE, dd_zone,
                    relevantAtomsAreWithinGridBounds,
-                   nbs->a.data() + atomOffset, numAtoms, x,
+                   gridSetData.atomIndices.data() + atomOffset, numAtoms, x,
                    dimensions_.lowerCorner[ZZ],
                    1.0/dimensions_.gridSize[ZZ], numCellsZ*numAtomsPerCell,
                    sort_work);
@@ -1031,7 +1022,7 @@ void Grid::sortColumnsCpuGeometry(nbnxn_search *nbs,
             const int atomOffsetCell = atomOffset + cellZ*numAtomsPerCell;
             const int numAtomsCell   = std::min(numAtomsPerCell, numAtoms - (atomOffsetCell - atomOffset));
 
-            fillCell(nbs, nbat,
+            fillCell(gridSetData, nbat,
                      atomOffsetCell, atomOffsetCell + numAtomsCell, atinfo, x,
                      nullptr);
 
@@ -1050,13 +1041,13 @@ void Grid::sortColumnsCpuGeometry(nbnxn_search *nbs,
         /* Set the unused atom indices to -1 */
         for (int ind = numAtoms; ind < numCellsZ*numAtomsPerCell; ind++)
         {
-            nbs->a[atomOffset + ind] = -1;
+            gridSetData.atomIndices[atomOffset + ind] = -1;
         }
     }
 }
 
 /* Spatially sort the atoms within one grid column */
-void Grid::sortColumnsGpuGeometry(nbnxn_search *nbs,
+void Grid::sortColumnsGpuGeometry(const GridSetData &gridSetData,
                                   int dd_zone,
                                   int atomStart, int atomEnd,
                                   const int *atinfo,
@@ -1082,6 +1073,9 @@ void Grid::sortColumnsGpuGeometry(nbnxn_search *nbs,
     const int  subdiv_y = c_gpuNumClusterPerCellX*subdiv_x;
     const int  subdiv_z = c_gpuNumClusterPerCellY*subdiv_y;
 
+    /* Extract the atom index array that will be filled here */
+    const gmx::ArrayRef<int> &atomIndices = gridSetData.atomIndices;
+
     /* Sort the atoms within each x,y column in 3 dimensions.
      * Loop over all columns on the x/y grid.
      */
@@ -1097,7 +1091,7 @@ void Grid::sortColumnsGpuGeometry(nbnxn_search *nbs,
         /* Sort the atoms within each x,y column on z coordinate */
         sort_atoms(ZZ, FALSE, dd_zone,
                    relevantAtomsAreWithinGridBounds,
-                   nbs->a.data() + atomOffset, numAtomsInColumn, x,
+                   atomIndices.data() + atomOffset, numAtomsInColumn, x,
                    dimensions_.lowerCorner[ZZ],
                    1.0/dimensions_.gridSize[ZZ], numCellsInColumn*numAtomsPerCell,
                    sort_work);
@@ -1122,8 +1116,8 @@ void Grid::sortColumnsGpuGeometry(nbnxn_search *nbs,
                                               (numAtoms + geometry_.numAtomsICluster - 1)/ geometry_.numAtomsICluster);
 
                 /* Store the z-boundaries of the bounding box of the cell */
-                bbcz_[cell].lower = x[nbs->a[atomOffsetZ]][ZZ];
-                bbcz_[cell].upper = x[nbs->a[atomOffsetZ + numAtoms - 1]][ZZ];
+                bbcz_[cell].lower = x[atomIndices[atomOffsetZ]][ZZ];
+                bbcz_[cell].upper = x[atomIndices[atomOffsetZ + numAtoms - 1]][ZZ];
             }
 
             if (c_gpuNumClusterPerCellY > 1)
@@ -1131,7 +1125,7 @@ void Grid::sortColumnsGpuGeometry(nbnxn_search *nbs,
                 /* Sort the atoms along y */
                 sort_atoms(YY, (sub_z & 1) != 0, dd_zone,
                            relevantAtomsAreWithinGridBounds,
-                           nbs->a.data() + atomOffsetZ, numAtomsZ, x,
+                           atomIndices.data() + atomOffsetZ, numAtomsZ, x,
                            dimensions_.lowerCorner[YY] + gridY*dimensions_.cellSize[YY],
                            dimensions_.invCellSize[YY], subdiv_z,
                            sort_work);
@@ -1147,7 +1141,7 @@ void Grid::sortColumnsGpuGeometry(nbnxn_search *nbs,
                     /* Sort the atoms along x */
                     sort_atoms(XX, ((cz*c_gpuNumClusterPerCellY + sub_y) & 1) != 0, dd_zone,
                                relevantAtomsAreWithinGridBounds,
-                               nbs->a.data() + atomOffsetY, numAtomsY, x,
+                               atomIndices.data() + atomOffsetY, numAtomsY, x,
                                dimensions_.lowerCorner[XX] + gridX*dimensions_.cellSize[XX],
                                dimensions_.invCellSize[XX], subdiv_y,
                                sort_work);
@@ -1158,7 +1152,7 @@ void Grid::sortColumnsGpuGeometry(nbnxn_search *nbs,
                     const int atomOffsetX = atomOffsetY + sub_x*subdiv_x;
                     const int numAtomsX   = std::min(subdiv_x, numAtomsInColumn - (atomOffsetX - atomOffset));
 
-                    fillCell(nbs, nbat,
+                    fillCell(gridSetData, nbat,
                              atomOffsetX, atomOffsetX + numAtomsX, atinfo, x,
                              bb_work_aligned);
                 }
@@ -1168,7 +1162,7 @@ void Grid::sortColumnsGpuGeometry(nbnxn_search *nbs,
         /* Set the unused atom indices to -1 */
         for (int ind = numAtomsInColumn; ind < numCellsInColumn*numAtomsPerCell; ind++)
         {
-            nbs->a[atomOffset + ind] = -1;
+            atomIndices[atomOffset + ind] = -1;
         }
     }
 }
@@ -1285,53 +1279,35 @@ static void calc_column_indices(const Grid::Dimensions &gridDims,
 /*! \brief Resizes grid and atom data which depend on the number of cells */
 static void resizeForNumberOfCells(const int           numNbnxnAtoms,
                                    const int           numAtomsMoved,
-                                   nbnxn_search       *nbs,
+                                   GridSetData        *gridSetData,
                                    nbnxn_atomdata_t   *nbat)
 {
-    /* Note: nbs->cell was already resized before */
+    /* Note: gridSetData->cellIndices was already resized before */
 
     /* To avoid conditionals we store the moved particles at the end of a,
      * make sure we have enough space.
      */
-    nbs->a.resize(numNbnxnAtoms + numAtomsMoved);
+    gridSetData->atomIndices.resize(numNbnxnAtoms + numAtomsMoved);
 
     /* Make space in nbat for storing the atom coordinates */
     nbat->resizeCoordinateBuffer(numNbnxnAtoms);
 }
 
-void Grid::calcCellIndices(nbnxn_search                   *nbs,
-                           int                             ddZone,
-                           int                             cellOffset,
-                           const gmx::UpdateGroupsCog     *updateGroupsCog,
-                           int                             atomStart,
-                           int                             atomEnd,
-                           const int                      *atinfo,
-                           gmx::ArrayRef<const gmx::RVec>  x,
-                           int                             numAtomsMoved,
-                           const int                      *move,
-                           nbnxn_atomdata_t               *nbat)
+void Grid::setCellIndices(int                             ddZone,
+                          int                             cellOffset,
+                          GridSetData                    *gridSetData,
+                          gmx::ArrayRef<GridWork>         gridWork,
+                          int                             atomStart,
+                          int                             atomEnd,
+                          const int                      *atinfo,
+                          gmx::ArrayRef<const gmx::RVec>  x,
+                          const int                       numAtomsMoved,
+                          nbnxn_atomdata_t               *nbat)
 {
     cellOffset_ = cellOffset;
 
-    /* First compute all grid/column indices and store them in nbs->cell */
-    nbs->cell.resize(atomEnd);
-
     const int nthread = gmx_omp_nthreads_get(emntPairsearch);
 
-#pragma omp parallel for num_threads(nthread) schedule(static)
-    for (int thread = 0; thread < nthread; thread++)
-    {
-        try
-        {
-            calc_column_indices(dimensions_,
-                                updateGroupsCog,
-                                atomStart, atomEnd, x,
-                                ddZone, move, thread, nthread,
-                                nbs->cell, nbs->work[thread].cxy_na);
-        }
-        GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR;
-    }
-
     const int numAtomsPerCell = geometry_.numAtomsPerCell;
 
     /* Make the cell index as a function of x and y */
@@ -1348,10 +1324,10 @@ void Grid::calcCellIndices(nbnxn_search                   *nbs,
         {
             ncz_max = ncz;
         }
-        int cxy_na_i = nbs->work[0].cxy_na[i];
+        int cxy_na_i = gridWork[0].numAtomsPerColumn[i];
         for (int thread = 1; thread < nthread; thread++)
         {
-            cxy_na_i += nbs->work[thread].cxy_na[i];
+            cxy_na_i += gridWork[thread].numAtomsPerColumn[i];
         }
         ncz = (cxy_na_i + numAtomsPerCell - 1)/numAtomsPerCell;
         if (nbat->XFormat == nbatX8)
@@ -1365,7 +1341,8 @@ void Grid::calcCellIndices(nbnxn_search                   *nbs,
     }
     numCellsTotal_ = cxy_ind_[numColumns()] - cxy_ind_[0];
 
-    resizeForNumberOfCells(atomIndexEnd(), numAtomsMoved, nbs, nbat);
+    /* Resize grid and atom data which depend on the number of cells */
+    resizeForNumberOfCells(atomIndexEnd(), numAtomsMoved, gridSetData, nbat);
 
     if (debug)
     {
@@ -1391,9 +1368,9 @@ void Grid::calcCellIndices(nbnxn_search                   *nbs,
 
     /* Make sure the work array for sorting is large enough */
     const int worstCaseSortBufferSize = ncz_max*numAtomsPerCell*c_sortGridMaxSizeFactor;
-    if (worstCaseSortBufferSize > gmx::index(nbs->work[0].sortBuffer.size()))
+    if (worstCaseSortBufferSize > gmx::index(gridWork[0].sortBuffer.size()))
     {
-        for (nbnxn_search_work_t &work : nbs->work)
+        for (GridWork &work : gridWork)
         {
             /* Elements not in use should be -1 */
             work.sortBuffer.resize(worstCaseSortBufferSize, -1);
@@ -1403,11 +1380,13 @@ void Grid::calcCellIndices(nbnxn_search                   *nbs,
     /* Now we know the dimensions we can fill the grid.
      * This is the first, unsorted fill. We sort the columns after this.
      */
+    gmx::ArrayRef<int> cells       = gridSetData->cells;
+    gmx::ArrayRef<int> atomIndices = gridSetData->atomIndices;
     for (int i = atomStart; i < atomEnd; i++)
     {
         /* At this point nbs->cell contains the local grid x,y indices */
-        int cxy = nbs->cell[i];
-        nbs->a[firstAtomInColumn(cxy) + cxy_na_[cxy]++] = i;
+        const int cxy = cells[i];
+        atomIndices[firstAtomInColumn(cxy) + cxy_na_[cxy]++] = i;
     }
 
     if (ddZone == 0)
@@ -1419,7 +1398,7 @@ void Grid::calcCellIndices(nbnxn_search                   *nbs,
         {
             for (int i = n0; i < n1; i++)
             {
-                nbs->cell[nbs->a[i]] = i;
+                cells[atomIndices[i]] = i;
             }
         }
     }
@@ -1434,15 +1413,17 @@ void Grid::calcCellIndices(nbnxn_search                   *nbs,
             int columnEnd   = ((thread + 1)*numColumns())/nthread;
             if (geometry_.isSimple)
             {
-                sortColumnsCpuGeometry(nbs, ddZone, atomStart, atomEnd, atinfo, x, nbat,
+                sortColumnsCpuGeometry(*gridSetData, ddZone,
+                                       atomStart, atomEnd, atinfo, x, nbat,
                                        columnStart, columnEnd,
-                                       nbs->work[thread].sortBuffer);
+                                       gridWork[thread].sortBuffer);
             }
             else
             {
-                sortColumnsGpuGeometry(nbs, ddZone, atomStart, atomEnd, atinfo, x, nbat,
+                sortColumnsGpuGeometry(*gridSetData, ddZone,
+                                       atomStart, atomEnd, atinfo, x, nbat,
                                        columnStart, columnEnd,
-                                       nbs->work[thread].sortBuffer);
+                                       gridWork[thread].sortBuffer);
             }
         }
         GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR;
@@ -1479,41 +1460,31 @@ void Grid::calcCellIndices(nbnxn_search                   *nbs,
     }
 }
 
-} // namespace Nbnxm
-
-// TODO: Move the functions below to nbnxn.cpp
-
-/* Sets up a grid and puts the atoms on the grid.
- * This function only operates on one domain of the domain decompostion.
- * Note that without domain decomposition there is only one domain.
- */
-void nbnxn_put_on_grid(nonbonded_verlet_t             *nbv,
-                       const matrix                    box,
-                       int                             ddZone,
-                       const rvec                      lowerCorner,
-                       const rvec                      upperCorner,
-                       const gmx::UpdateGroupsCog     *updateGroupsCog,
-                       int                             atomStart,
-                       int                             atomEnd,
-                       real                            atomDensity,
-                       const int                      *atinfo,
-                       gmx::ArrayRef<const gmx::RVec>  x,
-                       int                             numAtomsMoved,
-                       const int                      *move)
+// TODO: Move to gridset.cpp
+void GridSet::putOnGrid(const matrix                    box,
+                        const int                       ddZone,
+                        const rvec                      lowerCorner,
+                        const rvec                      upperCorner,
+                        const gmx::UpdateGroupsCog     *updateGroupsCog,
+                        const int                       atomStart,
+                        const int                       atomEnd,
+                        real                            atomDensity,
+                        const int                      *atinfo,
+                        gmx::ArrayRef<const gmx::RVec>  x,
+                        const int                       numAtomsMoved,
+                        const int                      *move,
+                        nbnxn_atomdata_t               *nbat)
 {
-    nbnxn_search *nbs  = nbv->nbs.get();
-    Nbnxm::Grid  &grid = nbs->grid[ddZone];
-
-    nbs_cycle_start(&nbs->cc[enbsCCgrid]);
+    Nbnxm::Grid  &grid = grids_[ddZone];
 
-    int cellOffset;
+    int           cellOffset;
     if (ddZone == 0)
     {
         cellOffset = 0;
     }
     else
     {
-        const Nbnxm::Grid &previousGrid = nbs->grid[ddZone - 1];
+        const Nbnxm::Grid &previousGrid = grids_[ddZone - 1];
         cellOffset = previousGrid.atomIndexEnd()/previousGrid.geometry().numAtomsPerCell;
     }
 
@@ -1522,55 +1493,108 @@ void nbnxn_put_on_grid(nonbonded_verlet_t             *nbv,
     real      maxAtomGroupRadius;
     if (ddZone == 0)
     {
-        copy_mat(box, nbs->box);
+        copy_mat(box, box_);
 
-        nbs->natoms_local    = atomEnd - numAtomsMoved;
+        numRealAtomsLocal_ = atomEnd - numAtomsMoved;
         /* We assume that nbnxn_put_on_grid is called first
          * for the local atoms (ddZone=0).
          */
-        nbs->natoms_nonlocal = atomEnd - numAtomsMoved;
+        numRealAtomsTotal_ = atomEnd - numAtomsMoved;
 
         maxAtomGroupRadius = (updateGroupsCog ? updateGroupsCog->maxUpdateGroupRadius() : 0);
 
         if (debug)
         {
             fprintf(debug, "natoms_local = %5d atom_density = %5.1f\n",
-                    nbs->natoms_local, atomDensity);
+                    numRealAtomsLocal_, atomDensity);
         }
     }
     else
     {
-        const Nbnxm::Grid::Dimensions &dimsGrid0 = nbs->grid[0].dimensions();
+        const Nbnxm::Grid::Dimensions &dimsGrid0 = grids_[0].dimensions();
         atomDensity        = dimsGrid0.atomDensity;
         maxAtomGroupRadius = dimsGrid0.maxAtomGroupRadius;
 
-        nbs->natoms_nonlocal = std::max(nbs->natoms_nonlocal, atomEnd);
+        numRealAtomsTotal_ = std::max(numRealAtomsTotal_, atomEnd);
     }
 
     /* We always use the home zone (grid[0]) for setting the cell size,
      * since determining densities for non-local zones is difficult.
      */
-    grid.setDimensions(nbs,
-                       ddZone, n - numAtomsMoved,
+    grid.setDimensions(ddZone, n - numAtomsMoved,
                        lowerCorner, upperCorner,
                        atomDensity,
-                       maxAtomGroupRadius);
+                       maxAtomGroupRadius,
+                       haveFep_);
 
-    nbnxn_atomdata_t *nbat = nbv->nbat.get();
+    for (GridWork &work : gridWork_)
+    {
+        work.numAtomsPerColumn.resize(grid.numColumns() + 1);
+    }
 
-    grid.calcCellIndices(nbs, ddZone, cellOffset, updateGroupsCog, atomStart, atomEnd, atinfo, x, numAtomsMoved, move, nbat);
+    /* Make space for the new cell indices */
+    cells_.resize(atomEnd);
+
+    const int nthread = gmx_omp_nthreads_get(emntPairsearch);
+
+#pragma omp parallel for num_threads(nthread) schedule(static)
+    for (int thread = 0; thread < nthread; thread++)
+    {
+        try
+        {
+            calc_column_indices(grid.dimensions(),
+                                updateGroupsCog,
+                                atomStart, atomEnd, x,
+                                ddZone, move, thread, nthread,
+                                cells_, gridWork_[thread].numAtomsPerColumn);
+        }
+        GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR;
+    }
+
+    GridSetData gridSetData = getGridSetData();
+
+    /* Copy the already computed cell indices to the grid and sort, when needed */
+    grid.setCellIndices(ddZone, cellOffset, &gridSetData, gridWork_,
+                        atomStart, atomEnd, atinfo, x, numAtomsMoved, nbat);
 
     if (ddZone == 0)
     {
         nbat->natoms_local = nbat->numAtoms();
     }
-    if (ddZone == gmx::ssize(nbs->grid) - 1)
+    if (ddZone == gmx::ssize(grids_) - 1)
     {
         /* We are done setting up all grids, we can resize the force buffers */
         nbat->resizeForceBuffers();
     }
+}
 
-    nbs_cycle_stop(&nbs->cc[enbsCCgrid]);
+} // namespace Nbnxm
+
+// TODO: Move this function to a proper location after refactoring nbnxn_search
+void nbnxn_put_on_grid(nonbonded_verlet_t             *nb_verlet,
+                       const matrix                    box,
+                       int                             ddZone,
+                       const rvec                      lowerCorner,
+                       const rvec                      upperCorner,
+                       const gmx::UpdateGroupsCog     *updateGroupsCog,
+                       int                             atomStart,
+                       int                             atomEnd,
+                       real                            atomDensity,
+                       const int                      *atinfo,
+                       gmx::ArrayRef<const gmx::RVec>  x,
+                       int                             numAtomsMoved,
+                       const int                      *move)
+{
+    nbnxn_search &nbs = *nb_verlet->nbs;
+
+    nbs_cycle_start(&nbs.cc[enbsCCgrid]);
+
+    nbs.gridSet_.putOnGrid(box, ddZone, lowerCorner, upperCorner,
+                           updateGroupsCog, atomStart, atomEnd, atomDensity,
+                           atinfo, x, numAtomsMoved, move,
+                           nb_verlet->nbat.get());
+
+    nbs_cycle_stop(&nbs.cc[enbsCCgrid]);
 }
 
 /* Calls nbnxn_put_on_grid for all non-local domains */
@@ -1600,43 +1624,12 @@ void nbnxn_put_on_grid_nonlocal(nonbonded_verlet_t              *nbv,
     }
 }
 
-void nbnxn_get_ncells(const nbnxn_search *nbs, int *ncx, int *ncy)
-{
-    *ncx = nbs->grid[0].dimensions().numCells[XX];
-    *ncy = nbs->grid[0].dimensions().numCells[YY];
-}
-
 gmx::ArrayRef<const int> nbnxn_get_atomorder(const nbnxn_search *nbs)
 {
     /* Return the atom order for the home cell (index 0) */
-    const Nbnxm::Grid &grid       = nbs->grid[0];
+    const Nbnxm::Grid &grid       = nbs->gridSet().grids()[0];
 
     const int          numIndices = grid.atomIndexEnd() - grid.firstAtomInColumn(0);
 
-    return gmx::constArrayRefFromArray(nbs->a.data(), numIndices);
-}
-
-void nbnxn_set_atomorder(nbnxn_search *nbs)
-{
-    /* Set the atom order for the home cell (index 0) */
-    const Nbnxm::Grid &grid = nbs->grid[0];
-
-    int                atomIndex = 0;
-    for (int cxy = 0; cxy < grid.numColumns(); cxy++)
-    {
-        const int numAtoms  = grid.numAtomsInColumn(cxy);
-        int       cellIndex = grid.firstCellInColumn(cxy)*grid.geometry().numAtomsPerCell;
-        for (int i = 0; i < numAtoms; i++)
-        {
-            nbs->a[cellIndex]    = atomIndex;
-            nbs->cell[atomIndex] = cellIndex;
-            atomIndex++;
-            cellIndex++;
-        }
-    }
-}
-
-gmx::ArrayRef<const int> nbnxn_get_gridindices(const nbnxn_search *nbs)
-{
-    return nbs->cell;
+    return gmx::constArrayRefFromArray(nbs->gridSet().atomIndices().data(), numIndices);
 }
index 23e0571ed42c2d29830d670f5d6094cf4aba2c94..57e7ac699e41bc60cd98af0366dee115d2178338 100644 (file)
@@ -67,11 +67,6 @@ struct nbnxn_atomdata_t;
 struct nbnxn_search;
 enum class PairlistType;
 
-namespace gmx
-{
-class UpdateGroupsCog;
-}
-
 namespace Nbnxm
 {
 
@@ -202,6 +197,35 @@ static constexpr int c_numBoundingBoxBounds1D = 2;
 #endif // !DOXYGEN
 
 
+/*! \internal
+ * \brief Helper struct to pass data that is shared over all grids
+ *
+ * To enable a single coordinate and force array, a single cell range
+ * is needed which covers all grids. This helper struct contains
+ * references to the index lists mapping both ways, as well as
+ * the free-energy boolean, which is the same for all grids.
+ */
+struct GridSetData
+{
+    //! The cell indices for all atoms
+    std::vector<int> &cells;
+    //! The atom indices for all atoms stored in cell order
+    std::vector<int> &atomIndices;
+    //! Tells whether we are have perturbed non-bonded interations
+    const bool        haveFep;
+};
+
+/*! \internal
+ * \brief Working arrays for constructing a grid
+ */
+struct GridWork
+{
+    //! Number of atoms for each grid column
+    std::vector<int> numAtomsPerColumn;
+    //! Buffer for sorting integers
+    std::vector<int> sortBuffer;
+};
+
 /*! \internal
  * \brief A pair-search grid object for one domain decomposition zone
  *
@@ -407,33 +431,32 @@ class Grid
         }
 
         //! Sets the grid dimensions
-        void setDimensions(const nbnxn_search   *nbs,
-                           int                   ddZone,
+        void setDimensions(int                   ddZone,
                            int                   numAtoms,
                            const rvec            lowerCorner,
                            const rvec            upperCorner,
                            real                  atomDensity,
-                           real                  maxAtomGroupRadius);
-
-        //! Determine in which grid cells the atoms should go
-        void calcCellIndices(nbnxn_search                   *nbs,
-                             int                             ddZone,
-                             int                             cellOffset,
-                             const gmx::UpdateGroupsCog     *updateGroupsCog,
-                             int                             atomStart,
-                             int                             atomEnd,
-                             const int                      *atinfo,
-                             gmx::ArrayRef<const gmx::RVec>  x,
-                             int                             numAtomsMoved,
-                             const int                      *move,
-                             nbnxn_atomdata_t               *nbat);
+                           real                  maxAtomGroupRadius,
+                           bool                  haveFep);
+
+        //! Sets the cell indices using indices in \p gridSetData and \p gridWork
+        void setCellIndices(int                             ddZone,
+                            int                             cellOffset,
+                            GridSetData                    *gridSetData,
+                            gmx::ArrayRef<GridWork>         gridWork,
+                            int                             atomStart,
+                            int                             atomEnd,
+                            const int                      *atinfo,
+                            gmx::ArrayRef<const gmx::RVec>  x,
+                            int                             numAtomsMoved,
+                            nbnxn_atomdata_t               *nbat);
 
     private:
         /*! \brief Fill a pair search cell with atoms
          *
          * Potentially sorts atoms and sets the interaction flags.
          */
-        void fillCell(nbnxn_search                   *nbs,
+        void fillCell(const GridSetData              &gridSetData,
                       nbnxn_atomdata_t               *nbat,
                       int                             atomStart,
                       int                             atomEnd,
@@ -442,7 +465,7 @@ class Grid
                       BoundingBox gmx_unused         *bb_work_aligned);
 
         //! Spatially sort the atoms within one grid column
-        void sortColumnsCpuGeometry(nbnxn_search *nbs,
+        void sortColumnsCpuGeometry(const GridSetData &gridSetData,
                                     int dd_zone,
                                     int atomStart, int atomEnd,
                                     const int *atinfo,
@@ -452,7 +475,7 @@ class Grid
                                     gmx::ArrayRef<int> sort_work);
 
         //! Spatially sort the atoms within one grid column
-        void sortColumnsGpuGeometry(nbnxn_search *nbs,
+        void sortColumnsGpuGeometry(const GridSetData &gridSetData,
                                     int dd_zone,
                                     int atomStart, int atomEnd,
                                     const int *atinfo,
diff --git a/src/gromacs/nbnxm/gridset.cpp b/src/gromacs/nbnxm/gridset.cpp
new file mode 100644 (file)
index 0000000..e1b00b6
--- /dev/null
@@ -0,0 +1,125 @@
+/*
+ * This file is part of the GROMACS molecular simulation package.
+ *
+ * Copyright (c) 2019, by the GROMACS development team, led by
+ * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
+ * and including many others, as listed in the AUTHORS file in the
+ * top-level source directory and at http://www.gromacs.org.
+ *
+ * GROMACS is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public License
+ * as published by the Free Software Foundation; either version 2.1
+ * of the License, or (at your option) any later version.
+ *
+ * GROMACS is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with GROMACS; if not, see
+ * http://www.gnu.org/licenses, or write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
+ *
+ * If you want to redistribute modifications to GROMACS, please
+ * consider that scientific software is very special. Version
+ * control is crucial - bugs must be traceable. We will be happy to
+ * consider code for inclusion in the official distribution, but
+ * derived work must not be called official GROMACS. Details are found
+ * in the README & COPYING files - if they are missing, get the
+ * official version at http://www.gromacs.org.
+ *
+ * To help us fund GROMACS development, we humbly ask that you cite
+ * the research papers on the package. Check out http://www.gromacs.org.
+ */
+
+/*! \internal \file
+ *
+ * \brief
+ * Implements the GridSet class.
+ *
+ * \author Berk Hess <hess@kth.se>
+ * \ingroup module_nbnxm
+ */
+
+#include "gmxpre.h"
+
+#include "gridset.h"
+
+#include "gromacs/nbnxm/nbnxm.h"
+
+#include "internal.h"
+
+namespace Nbnxm
+{
+
+//! Returns the number of DD zones or 1 when \p numDDCells = nullptr
+static int numDDZones(const ivec *numDDCells)
+{
+    int  numDDZones = 1;
+    bool haveDomDec = (numDDCells != nullptr);
+    if (haveDomDec)
+    {
+        for (int d = 0; d < DIM; d++)
+        {
+            if ((*numDDCells)[d] > 1)
+            {
+                numDDZones *= 2;
+            }
+        }
+    }
+
+    return numDDZones;
+}
+
+GridSet::GridSet(const ivec         *numDDCells,
+                 const PairlistType  pairlistType,
+                 const bool          haveFep,
+                 const int           numThreads) :
+    grids_(numDDZones(numDDCells), pairlistType),
+    haveFep_(haveFep),
+    numRealAtomsLocal_(0),
+    numRealAtomsTotal_(0),
+    gridWork_(numThreads)
+{
+    clear_mat(box_);
+}
+
+void GridSet::setLocalAtomOrder()
+{
+    /* Set the atom order for the home cell (index 0) */
+    const Nbnxm::Grid &grid = grids_[0];
+
+    int                atomIndex = 0;
+    for (int cxy = 0; cxy < grid.numColumns(); cxy++)
+    {
+        const int numAtoms  = grid.numAtomsInColumn(cxy);
+        int       cellIndex = grid.firstCellInColumn(cxy)*grid.geometry().numAtomsPerCell;
+        for (int i = 0; i < numAtoms; i++)
+        {
+            atomIndices_[cellIndex] = atomIndex;
+            cells_[atomIndex]       = cellIndex;
+            atomIndex++;
+            cellIndex++;
+        }
+    }
+}
+
+} // namespace Nbnxm
+
+
+void nonbonded_verlet_t::setLocalAtomOrder()
+{
+    nbs->gridSet_.setLocalAtomOrder();
+}
+
+void nonbonded_verlet_t::getLocalNumCells(int *numCellsX,
+                                          int *numCellsY) const
+{
+    nbs->gridSet().getLocalNumCells(numCellsX, numCellsY);
+}
+
+gmx::ArrayRef<const int> nbnxn_get_gridindices(const nbnxn_search* nbs)
+{
+    return nbs->gridSet().cells();
+}
diff --git a/src/gromacs/nbnxm/gridset.h b/src/gromacs/nbnxm/gridset.h
new file mode 100644 (file)
index 0000000..bfd4320
--- /dev/null
@@ -0,0 +1,198 @@
+/*
+ * This file is part of the GROMACS molecular simulation package.
+ *
+ * Copyright (c) 2019, by the GROMACS development team, led by
+ * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
+ * and including many others, as listed in the AUTHORS file in the
+ * top-level source directory and at http://www.gromacs.org.
+ *
+ * GROMACS is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public License
+ * as published by the Free Software Foundation; either version 2.1
+ * of the License, or (at your option) any later version.
+ *
+ * GROMACS is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with GROMACS; if not, see
+ * http://www.gnu.org/licenses, or write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
+ *
+ * If you want to redistribute modifications to GROMACS, please
+ * consider that scientific software is very special. Version
+ * control is crucial - bugs must be traceable. We will be happy to
+ * consider code for inclusion in the official distribution, but
+ * derived work must not be called official GROMACS. Details are found
+ * in the README & COPYING files - if they are missing, get the
+ * official version at http://www.gromacs.org.
+ *
+ * To help us fund GROMACS development, we humbly ask that you cite
+ * the research papers on the package. Check out http://www.gromacs.org.
+ */
+
+/*! \internal \file
+ *
+ * \brief
+ * Declares the GridSet class.
+ *
+ * This class holds the grids for the local and non-local domain decomposition
+ * zones, as well as the cell and atom data that covers all grids.
+ *
+ * \author Berk Hess <hess@kth.se>
+ * \ingroup module_nbnxm
+ */
+
+#ifndef GMX_NBNXM_GRIDSET_H
+#define GMX_NBNXM_GRIDSET_H
+
+#include <memory>
+#include <vector>
+
+#include "gromacs/math/vec.h"
+#include "gromacs/math/vectypes.h"
+#include "gromacs/utility/alignedallocator.h"
+#include "gromacs/utility/arrayref.h"
+
+#include "grid.h"
+
+
+struct nbnxn_atomdata_t;
+enum class PairlistType;
+
+namespace gmx
+{
+class UpdateGroupsCog;
+}
+
+namespace Nbnxm
+{
+
+/*! \internal
+ * \brief An object holding a set of search grids for the local + non-local DD zones
+ */
+class GridSet
+{
+    public:
+        //! Constructs a grid set for 1 or multiple DD zones, when numDDCells!=nullptr
+        GridSet(const ivec   *numDDCells,
+                PairlistType  pairlistType,
+                bool          haveFep,
+                int           numThreads);
+
+        //! Puts the atoms in \p ddZone on the grid and copies the coordinates to \p nbat
+        void putOnGrid(const matrix                    box,
+                       int                             ddZone,
+                       const rvec                      lowerCorner,
+                       const rvec                      upperCorner,
+                       const gmx::UpdateGroupsCog     *updateGroupsCog,
+                       int                             atomStart,
+                       int                             atomEnd,
+                       real                            atomDensity,
+                       const int                      *atinfo,
+                       gmx::ArrayRef<const gmx::RVec>  x,
+                       int                             numAtomsMoved,
+                       const int                      *move,
+                       nbnxn_atomdata_t               *nbat);
+
+        //! Returns the number of cells along x and y for the local grid
+        void getLocalNumCells(int *numCellsX,
+                              int *numCellsY) const
+        {
+            *numCellsX = grids_[0].dimensions().numCells[XX];
+            *numCellsY = grids_[0].dimensions().numCells[YY];
+        }
+
+        //! Returns the total number of atoms in the grid set, including padding
+        int numGridAtomsTotal() const
+        {
+            return grids_.back().atomIndexEnd();
+        }
+
+        //! Returns the number of local real atoms, i.e. without padded atoms
+        int numRealAtomsLocal() const
+        {
+            return numRealAtomsLocal_;
+        }
+
+        //! Returns the number of total real atoms, i.e. without padded atoms
+        int numRealAtomsTotal() const
+        {
+            return numRealAtomsTotal_;
+        }
+
+        //! Returns the atom order on the grid for the local atoms
+        gmx::ArrayRef<const int> getLocalAtomorder() const
+        {
+            /* Return the atom order for the home cell (index 0) */
+            const int numIndices = grids_[0].atomIndexEnd() - grids_[0].firstAtomInColumn(0);
+
+            return gmx::constArrayRefFromArray(atomIndices_.data(), numIndices);
+        }
+
+        //! Sets the order of the local atoms to the order grid atom ordering
+        void setLocalAtomOrder();
+
+        //! Returns the list of grids
+        gmx::ArrayRef<const Grid> grids() const
+        {
+            return grids_;
+        }
+
+        //! Returns the grid atom indices covering all grids
+        gmx::ArrayRef<const int> cells() const
+        {
+            return cells_;
+        }
+
+        //! Returns the grid atom indices covering all grids
+        gmx::ArrayRef<const int> atomIndices() const
+        {
+            return atomIndices_;
+        }
+
+        //! Returns whether we have perturbed non-bonded interactions
+        bool haveFep() const
+        {
+            return haveFep_;
+        }
+
+        //! Returns the unit cell in \p box
+        void getBox(matrix box) const
+        {
+            copy_mat(box_, box);
+        }
+
+    private:
+        //! Returns collection of the data that covers all grids
+        const GridSetData getGridSetData()
+        {
+            GridSetData gridSetData = { cells_, atomIndices_, haveFep_ };
+
+            return gridSetData;
+        }
+
+        /* Data members */
+        //! The search grids
+        std::vector<Grid>     grids_;
+        //! The actual cell indices for all atoms, covering all grids
+        std::vector<int>      cells_;
+        //! The actual array of atom indices, covering all grids
+        std::vector<int>      atomIndices_;
+        //! Tells whether we have perturbed non-bonded interactions
+        bool                  haveFep_;
+        //! The periodic unit-cell
+        matrix                box_;
+        //! The number of local real atoms, i.e. without padded atoms, local atoms: 0 to numAtomsLocal_
+        int                   numRealAtomsLocal_;
+        //! The total number of real atoms, i.e. without padded atoms
+        int                   numRealAtomsTotal_;
+        //! Working data for constructing a single grid, one entry per thread
+        std::vector<GridWork> gridWork_;
+};
+
+} // namespace Nbnxm
+
+#endif
index 89ef3d5ce42c9f47570c39d0d191863265dd55fd..b41a42fd0f4e1e97db7ab70b95fda7333dd8b587 100644 (file)
 
 #include "gromacs/domdec/domdec.h"
 #include "gromacs/math/vectypes.h"
+#include "gromacs/nbnxm/atomdata.h"
 #include "gromacs/nbnxm/pairlist.h"
 #include "gromacs/timing/cyclecounter.h"
 #include "gromacs/utility/alignedallocator.h"
 #include "gromacs/utility/arrayref.h"
 #include "gromacs/utility/real.h"
 
-struct gmx_domdec_zones_t;
+#include "gridset.h"
 
-namespace Nbnxm
-{
-class Grid;
-}
+struct gmx_domdec_zones_t;
 
 
 /*! \brief Convenience declaration for an std::vector with aligned memory */
@@ -90,8 +88,6 @@ struct nbnxn_search_work_t
 
     gmx_cache_protect_t       cp0;          /* Buffer to avoid cache polution */
 
-    std::vector<int>          cxy_na;       /* Grid column atom counts temporary buffer */
-
     std::vector<int>          sortBuffer;   /* Temporary buffer for sorting atoms within a grid column */
 
     nbnxn_buffer_flags_t      buffer_flags; /* Flags for force buffer access */
@@ -114,33 +110,28 @@ struct nbnxn_search
      * \param[in] ePBC            The periodic boundary conditions
      * \param[in] n_dd_cells      The number of domain decomposition cells per dimension, without DD nullptr should be passed
      * \param[in] zones           The domain decomposition zone setup, without DD nullptr should be passed
-     * \param[in] nb_kernel_type  The nbnxn non-bonded kernel type
-     * \param[in] bFEP            Tells whether non-bonded interactions are perturbed
-     * \param[in] nthread_max     The maximum number of threads used in the search
+     * \param[in] haveFep         Tells whether non-bonded interactions are perturbed
+     * \param[in] maxNumThreads   The maximum number of threads used in the search
      */
-
     nbnxn_search(int                       ePBC,
                  const ivec               *n_dd_cells,
                  const gmx_domdec_zones_t *zones,
                  PairlistType              pairlistType,
-                 gmx_bool                  bFEP,
-                 int                       nthread_max);
+                 bool                      haveFep,
+                 int                       maxNumthreads);
 
-    gmx_bool                   bFEP;            /* Do we have perturbed atoms? */
-    int                        ePBC;            /* PBC type enum                              */
-    matrix                     box;             /* The periodic unit-cell                     */
+    const Nbnxm::GridSet &gridSet() const
+    {
+        return gridSet_;
+    }
 
+    int                        ePBC;            /* PBC type enum                              */
     gmx_bool                   DomDec;          /* Are we doing domain decomposition?         */
     ivec                       dd_dim;          /* Are we doing DD in x,y,z?                  */
     const gmx_domdec_zones_t  *zones;           /* The domain decomposition zones        */
 
-    std::vector<Nbnxm::Grid>   grid;            /* Array of grids, size ngrid                 */
-    std::vector<int>           cell;            /* Actual allocated cell array for all grids  */
-    std::vector<int>           a;               /* Atom index for grid, the inverse of cell   */
-
-    int                        natoms_local;    /* The local atoms run from 0 to natoms_local */
-    int                        natoms_nonlocal; /* The non-local atoms run from natoms_local
-                                                 * to natoms_nonlocal */
+    //! The local and non-local grids
+    Nbnxm::GridSet       gridSet_;
 
     gmx_bool             print_cycles;
     int                  search_count;
index dbdb9fba51b295c2a02c510295c5fbb8e3132270..3c115ff314f7f7655a74914e90df3795dc746761 100644 (file)
@@ -376,6 +376,9 @@ struct nonbonded_verlet_t
         //! Initialize the pair list sets, TODO this should be private
         void initPairlistSets(bool haveMultipleDomains);
 
+        //! Sets the order of the local atoms to the order grid atom ordering
+        void setLocalAtomOrder();
+
         //! Constructs the pairlist for the given locality
         void constructPairlist(Nbnxm::InteractionLocality  iLocality,
                                const t_blocka             *excl,
@@ -434,6 +437,10 @@ struct nonbonded_verlet_t
             return kernelSetup_;
         }
 
+        /*! \brief Returns the number of x and y cells in the local grid */
+        void getLocalNumCells(int *numCellsX,
+                              int *numCellsY) const;
+
         // TODO: Make all data members private
     public:
         //! All data related to the pair lists
@@ -501,9 +508,6 @@ void nbnxn_put_on_grid_nonlocal(nonbonded_verlet_t              *nb_verlet,
                                 const int                       *atinfo,
                                 gmx::ArrayRef<const gmx::RVec>   x);
 
-/*! \brief Returns the number of x and y cells in the local grid */
-void nbnxn_get_ncells(const nbnxn_search *nbs, int *ncx, int *ncy);
-
 /*! \brief Returns the order indices of the atoms on the pairlist search grid */
 gmx::ArrayRef<const int> nbnxn_get_atomorder(const nbnxn_search* nbs);
 
index 75dd8c854ddd414b343d6017d38adb863d7be250..cf0c0f2c8b25ff8631bc7074da1d79551c226c6a 100644 (file)
@@ -299,24 +299,19 @@ nbnxn_search_work_t::~nbnxn_search_work_t()
     free_nblist(nbl_fep.get());
 }
 
-nbnxn_search::nbnxn_search(int                       ePBC,
+nbnxn_search::nbnxn_search(const int                 ePBC,
                            const ivec               *n_dd_cells,
                            const gmx_domdec_zones_t *zones,
                            const PairlistType        pairlistType,
-                           gmx_bool                  bFEP,
-                           int                       nthread_max) :
-    bFEP(bFEP),
+                           const bool                haveFep,
+                           const int                 maxNumThreads) :
     ePBC(ePBC),
     zones(zones),
-    natoms_local(0),
-    natoms_nonlocal(0),
+    gridSet_(n_dd_cells, pairlistType, haveFep, maxNumThreads),
     search_count(0),
-    work(nthread_max)
+    work(maxNumThreads)
 {
-    // The correct value will be set during the gridding
-    clear_mat(box);
     clear_ivec(dd_dim);
-    int numGrids = 1;
     DomDec = n_dd_cells != nullptr;
     if (DomDec)
     {
@@ -325,14 +320,10 @@ nbnxn_search::nbnxn_search(int                       ePBC,
             if ((*n_dd_cells)[d] > 1)
             {
                 dd_dim[d] = 1;
-                /* Each grid matches a DD zone */
-                numGrids *= 2;
             }
         }
     }
 
-    grid.resize(numGrids, pairlistType);
-
     /* Initialize detailed nbsearch cycle counting */
     print_cycles = (getenv("GMX_NBNXN_CYCLE") != nullptr);
     nbs_cycle_clear(cc);
@@ -863,7 +854,7 @@ void nbnxn_init_pairlist_set(nbnxn_pairlist_set_t *nbl_list)
 static void print_nblist_statistics(FILE *fp, const NbnxnPairlistCpu *nbl,
                                     const nbnxn_search *nbs, real rl)
 {
-    const Grid             &grid = nbs->grid[0];
+    const Grid             &grid = nbs->gridSet().grids()[0];
     const Grid::Dimensions &dims = grid.dimensions();
 
     fprintf(fp, "nbl nci %zu ncj %d\n",
@@ -908,7 +899,7 @@ static void print_nblist_statistics(FILE *fp, const NbnxnPairlistCpu *nbl,
 static void print_nblist_statistics(FILE *fp, const NbnxnPairlistGpu *nbl,
                                     const nbnxn_search *nbs, real rl)
 {
-    const Grid             &grid = nbs->grid[0];
+    const Grid             &grid = nbs->gridSet().grids()[0];
     const Grid::Dimensions &dims = grid.dimensions();
 
     fprintf(fp, "nbl nsci %zu ncj4 %zu nsi %d excl4 %zu\n",
@@ -1475,7 +1466,7 @@ static nbnxn_sci_t *getOpenIEntry(NbnxnPairlistGpu *nbl)
  * as masks in the pair-list for simple list entry iEntry.
  */
 static void
-setExclusionsForIEntry(const nbnxn_search   *nbs,
+setExclusionsForIEntry(const Nbnxm::GridSet &gridSet,
                        NbnxnPairlistCpu     *nbl,
                        gmx_bool              diagRemoved,
                        int                   na_cj_2log,
@@ -1492,13 +1483,14 @@ setExclusionsForIEntry(const nbnxn_search   *nbs,
 
     const int                iCluster = iEntry.ci;
 
-    gmx::ArrayRef<const int> cell = nbs->cell;
+    gmx::ArrayRef<const int> cell        = gridSet.cells();
+    gmx::ArrayRef<const int> atomIndices = gridSet.atomIndices();
 
     /* Loop over the atoms in the i-cluster */
     for (int i = 0; i < nbl->na_ci; i++)
     {
         const int iIndex = iCluster*nbl->na_ci + i;
-        const int iAtom  = nbs->a[iIndex];
+        const int iAtom  = atomIndices[iIndex];
         if (iAtom >= 0)
         {
             /* Loop over the topology-based exclusions for this i-atom */
@@ -1576,18 +1568,18 @@ const int max_nrj_fep = 40;
  * LJ parameters have been zeroed in the nbnxn data structure.
  * Simultaneously make a group pair list for the perturbed pairs.
  */
-static void make_fep_list(const nbnxn_search     *nbs,
-                          const nbnxn_atomdata_t *nbat,
-                          NbnxnPairlistCpu       *nbl,
-                          gmx_bool                bDiagRemoved,
-                          nbnxn_ci_t             *nbl_ci,
-                          real gmx_unused         shx,
-                          real gmx_unused         shy,
-                          real gmx_unused         shz,
-                          real gmx_unused         rlist_fep2,
-                          const Grid             &iGrid,
-                          const Grid             &jGrid,
-                          t_nblist               *nlist)
+static void make_fep_list(gmx::ArrayRef<const int>  atomIndices,
+                          const nbnxn_atomdata_t   *nbat,
+                          NbnxnPairlistCpu         *nbl,
+                          gmx_bool                  bDiagRemoved,
+                          nbnxn_ci_t               *nbl_ci,
+                          real gmx_unused           shx,
+                          real gmx_unused           shy,
+                          real gmx_unused           shz,
+                          real gmx_unused           rlist_fep2,
+                          const Grid               &iGrid,
+                          const Grid               &jGrid,
+                          t_nblist                 *nlist)
 {
     int      ci, cj_ind_start, cj_ind_end, cja, cjr;
     int      nri_max;
@@ -1643,7 +1635,7 @@ static void make_fep_list(const nbnxn_search     *nbs,
     for (int i = 0; i < nbl->na_ci; i++)
     {
         ind_i = ci*nbl->na_ci + i;
-        ai    = nbs->a[ind_i];
+        ai    = atomIndices[ind_i];
         if (ai >= 0)
         {
             nri                  = nlist->nri;
@@ -1711,7 +1703,7 @@ static void make_fep_list(const nbnxn_search     *nbs,
                     {
                         /* Is this interaction perturbed and not excluded? */
                         ind_j = cja*nbl->na_cj + j;
-                        aj    = nbs->a[ind_j];
+                        aj    = atomIndices[ind_j];
                         if (aj >= 0 &&
                             (bFEP_i || (fep_cj & (1 << j))) &&
                             (!bDiagRemoved || ind_j >= ind_i))
@@ -1789,18 +1781,18 @@ static inline int a_mod_wj(int a)
 }
 
 /* As make_fep_list above, but for super/sub lists. */
-static void make_fep_list(const nbnxn_search     *nbs,
-                          const nbnxn_atomdata_t *nbat,
-                          NbnxnPairlistGpu       *nbl,
-                          gmx_bool                bDiagRemoved,
-                          const nbnxn_sci_t      *nbl_sci,
-                          real                    shx,
-                          real                    shy,
-                          real                    shz,
-                          real                    rlist_fep2,
-                          const Grid             &iGrid,
-                          const Grid             &jGrid,
-                          t_nblist               *nlist)
+static void make_fep_list(gmx::ArrayRef<const int>  atomIndices,
+                          const nbnxn_atomdata_t   *nbat,
+                          NbnxnPairlistGpu         *nbl,
+                          gmx_bool                  bDiagRemoved,
+                          const nbnxn_sci_t        *nbl_sci,
+                          real                      shx,
+                          real                      shy,
+                          real                      shz,
+                          real                      rlist_fep2,
+                          const Grid               &iGrid,
+                          const Grid               &jGrid,
+                          t_nblist                 *nlist)
 {
     int                nri_max;
     int                c_abs;
@@ -1844,7 +1836,7 @@ static void make_fep_list(const nbnxn_search     *nbs,
         for (int i = 0; i < nbl->na_ci; i++)
         {
             ind_i = c_abs*nbl->na_ci + i;
-            ai    = nbs->a[ind_i];
+            ai    = atomIndices[ind_i];
             if (ai >= 0)
             {
                 nri                  = nlist->nri;
@@ -1889,7 +1881,7 @@ static void make_fep_list(const nbnxn_search     *nbs,
                             {
                                 /* Is this interaction perturbed and not excluded? */
                                 ind_j = (jGrid.cellOffset()*c_gpuNumClusterPerCell + cjr)*nbl->na_cj + j;
-                                aj    = nbs->a[ind_j];
+                                aj    = atomIndices[ind_j];
                                 if (aj >= 0 &&
                                     (bFEP_i || jGrid.atomIsPerturbed(cjr, j)) &&
                                     (!bDiagRemoved || ind_j >= ind_i))
@@ -1966,7 +1958,7 @@ static void make_fep_list(const nbnxn_search     *nbs,
  * as masks in the pair-list for i-super-cluster list entry iEntry.
  */
 static void
-setExclusionsForIEntry(const nbnxn_search   *nbs,
+setExclusionsForIEntry(const Nbnxm::GridSet &gridSet,
                        NbnxnPairlistGpu     *nbl,
                        gmx_bool              diagRemoved,
                        int gmx_unused        na_cj_2log,
@@ -1993,13 +1985,14 @@ setExclusionsForIEntry(const nbnxn_search   *nbs,
 
     const int                iSuperCluster = iEntry.sci;
 
-    gmx::ArrayRef<const int> cell = nbs->cell;
+    gmx::ArrayRef<const int> atomIndices   = gridSet.atomIndices();
+    gmx::ArrayRef<const int> cell          = gridSet.cells();
 
     /* Loop over the atoms in the i super-cluster */
     for (int i = 0; i < c_superClusterSize; i++)
     {
         const int iIndex = iSuperCluster*c_superClusterSize + i;
-        const int iAtom  = nbs->a[iIndex];
+        const int iAtom  = atomIndices[iIndex];
         if (iAtom >= 0)
         {
             const int iCluster = i/c_clusterSize;
@@ -2605,7 +2598,7 @@ static void get_nsubpair_target(const nbnxn_search        *nbs,
     const int           nsubpair_target_min = 36;
     real                r_eff_sup, vol_est, nsp_est, nsp_est_nl;
 
-    const Grid         &grid = nbs->grid[0];
+    const Grid         &grid = nbs->gridSet().grids()[0];
 
     /* We don't need to balance list sizes if:
      * - We didn't request balancing.
@@ -3250,7 +3243,7 @@ static void nbnxn_make_pairlist_part(const nbnxn_search *nbs,
 {
     int               na_cj_2log;
     matrix            box;
-    real              rlist2, rl_fep2 = 0;
+    real              rl_fep2 = 0;
     float             rbb2;
     int               ci_b, ci, ci_x, ci_y, ci_xy;
     ivec              shp;
@@ -3288,11 +3281,15 @@ static void nbnxn_make_pairlist_part(const nbnxn_search *nbs,
         gridj_flag       = work->buffer_flags.flag;
     }
 
-    copy_mat(nbs->box, box);
+    const Nbnxm::GridSet &gridSet = nbs->gridSet();
+
+    gridSet.getBox(box);
+
+    const bool            haveFep = gridSet.haveFep();
 
-    rlist2 = nbl->rlist*nbl->rlist;
+    const real            rlist2  = nbl->rlist*nbl->rlist;
 
-    if (nbs->bFEP && !pairlistIsSimple(*nbl))
+    if (haveFep && !pairlistIsSimple(*nbl))
     {
         /* Determine an atom-pair list cut-off distance for FEP atom pairs.
          * We should not simply use rlist, since then we would not have
@@ -3698,16 +3695,16 @@ static void nbnxn_make_pairlist_part(const nbnxn_search *nbs,
                     }
 
                     /* Set the exclusions for this ci list */
-                    setExclusionsForIEntry(nbs,
+                    setExclusionsForIEntry(gridSet,
                                            nbl,
                                            excludeSubDiagonal,
                                            na_cj_2log,
                                            *getOpenIEntry(nbl),
                                            exclusions);
 
-                    if (nbs->bFEP)
+                    if (haveFep)
                     {
-                        make_fep_list(nbs, nbat, nbl,
+                        make_fep_list(gridSet.atomIndices(), nbat, nbl,
                                       excludeSubDiagonal,
                                       getOpenIEntry(nbl),
                                       shx, shy, shz,
@@ -3734,7 +3731,7 @@ static void nbnxn_make_pairlist_part(const nbnxn_search *nbs,
 
     nbs_cycle_stop(&work->cc[enbsCCsearch]);
 
-    checkListSizeConsistency(*nbl, nbs->bFEP);
+    checkListSizeConsistency(*nbl, haveFep);
 
     if (debug)
     {
@@ -3742,7 +3739,7 @@ static void nbnxn_make_pairlist_part(const nbnxn_search *nbs,
 
         print_nblist_statistics(debug, nbl, nbs, rlist);
 
-        if (nbs->bFEP)
+        if (haveFep)
         {
             fprintf(debug, "nbl FEP list pairs: %d\n", nbl_fep->nrj);
         }
@@ -4097,7 +4094,7 @@ nonbonded_verlet_t::PairlistSets::construct(const InteractionLocality  iLocality
             clear_pairlist(nbl_list->nblGpu[th]);
         }
 
-        if (nbs->bFEP)
+        if (nbs->gridSet().haveFep())
         {
             clear_pairlist_fep(nbl_list->nbl_fep[th]);
         }
@@ -4105,7 +4102,7 @@ nonbonded_verlet_t::PairlistSets::construct(const InteractionLocality  iLocality
 
     for (int zi = 0; zi < nzi; zi++)
     {
-        const Grid &iGrid = nbs->grid[zi];
+        const Grid &iGrid = nbs->gridSet().grids()[zi];
 
         int                 zj0;
         int                 zj1;
@@ -4125,7 +4122,7 @@ nonbonded_verlet_t::PairlistSets::construct(const InteractionLocality  iLocality
         }
         for (int zj = zj0; zj < zj1; zj++)
         {
-            const Grid &jGrid = nbs->grid[zj];
+            const Grid &jGrid = nbs->gridSet().grids()[zj];
 
             if (debug)
             {
@@ -4280,7 +4277,7 @@ nonbonded_verlet_t::PairlistSets::construct(const InteractionLocality  iLocality
         reduce_buffer_flags(nbs, nbl_list->nnbl, &nbat->buffer_flags);
     }
 
-    if (nbs->bFEP)
+    if (nbs->gridSet().haveFep())
     {
         /* Balance the free-energy lists over all the threads */
         balance_fep_lists(nbs, nbl_list);