Sort all includes in src/gromacs
[alexxy/gromacs.git] / src / gromacs / mdlib / nbnxn_search.c
index 2f93f3036cfcb8b151579576f955cff11f58aed0..f13de8c7db485633350dd28414c1d23bb53eb237 100644 (file)
  * the research papers on the package. Check out http://www.gromacs.org.
  */
 
-#ifdef HAVE_CONFIG_H
-#include <config.h>
-#endif
+#include "gmxpre.h"
+
+#include "nbnxn_search.h"
 
+#include "config.h"
+
+#include <assert.h>
 #include <math.h>
 #include <string.h>
-#include "sysstuff.h"
-#include "smalloc.h"
-#include "macros.h"
+
+#include "gromacs/legacyheaders/gmx_omp_nthreads.h"
+#include "gromacs/legacyheaders/macros.h"
+#include "gromacs/legacyheaders/nrnb.h"
+#include "gromacs/legacyheaders/ns.h"
+#include "gromacs/legacyheaders/types/commrec.h"
 #include "gromacs/math/utilities.h"
-#include "vec.h"
-#include "pbc.h"
-#include "nbnxn_consts.h"
+#include "gromacs/math/vec.h"
+#include "gromacs/mdlib/nb_verlet.h"
+#include "gromacs/mdlib/nbnxn_atomdata.h"
+#include "gromacs/mdlib/nbnxn_consts.h"
+#include "gromacs/pbcutil/ishift.h"
+#include "gromacs/pbcutil/pbc.h"
+#include "gromacs/utility/smalloc.h"
+
 /* nbnxn_internal.h included gromacs/simd/macros.h */
-#include "nbnxn_internal.h"
-#ifdef GMX_NBNXN_SIMD
+#include "gromacs/mdlib/nbnxn_internal.h"
+#ifdef GMX_SIMD
 #include "gromacs/simd/vector_operations.h"
 #endif
-#include "nbnxn_atomdata.h"
-#include "nbnxn_search.h"
-#include "gmx_omp_nthreads.h"
-#include "nrnb.h"
-
-#include "gromacs/fileio/gmxfio.h"
 
 #ifdef NBNXN_SEARCH_BB_SIMD4
 /* Always use 4-wide SIMD for bounding box calculations */
@@ -319,9 +324,34 @@ gmx_bool nbnxn_kernel_pairlist_simple(int nb_kernel_type)
     }
 }
 
+/* Initializes a single nbnxn_pairlist_t data structure */
+static void nbnxn_init_pairlist_fep(t_nblist *nl)
+{
+    nl->type        = GMX_NBLIST_INTERACTION_FREE_ENERGY;
+    nl->igeometry   = GMX_NBLIST_GEOMETRY_PARTICLE_PARTICLE;
+    /* The interaction functions are set in the free energy kernel fuction */
+    nl->ivdw        = -1;
+    nl->ivdwmod     = -1;
+    nl->ielec       = -1;
+    nl->ielecmod    = -1;
+
+    nl->maxnri      = 0;
+    nl->maxnrj      = 0;
+    nl->nri         = 0;
+    nl->nrj         = 0;
+    nl->iinr        = NULL;
+    nl->gid         = NULL;
+    nl->shift       = NULL;
+    nl->jindex      = NULL;
+    nl->jjnr        = NULL;
+    nl->excl_fep    = NULL;
+
+}
+
 void nbnxn_init_search(nbnxn_search_t    * nbs_ptr,
                        ivec               *n_dd_cells,
                        gmx_domdec_zones_t *zones,
+                       gmx_bool            bFEP,
                        int                 nthread_max)
 {
     nbnxn_search_t nbs;
@@ -330,6 +360,8 @@ void nbnxn_init_search(nbnxn_search_t    * nbs_ptr,
     snew(nbs, 1);
     *nbs_ptr = nbs;
 
+    nbs->bFEP   = bFEP;
+
     nbs->DomDec = (n_dd_cells != NULL);
 
     clear_ivec(nbs->dd_dim);
@@ -369,6 +401,9 @@ void nbnxn_init_search(nbnxn_search_t    * nbs_ptr,
         nbs->work[t].cxy_na_nalloc    = 0;
         nbs->work[t].sort_work        = NULL;
         nbs->work[t].sort_work_nalloc = 0;
+
+        snew(nbs->work[t].nbl_fep, 1);
+        nbnxn_init_pairlist_fep(nbs->work[t].nbl_fep);
     }
 
     /* Initialize detailed nbsearch cycle counting */
@@ -385,6 +420,12 @@ static real grid_atom_density(int n, rvec corner0, rvec corner1)
 {
     rvec size;
 
+    if (n == 0)
+    {
+        /* To avoid zero density we use a minimum of 1 atom */
+        n = 1;
+    }
+
     rvec_sub(corner1, corner0, size);
 
     return n/(size[XX]*size[YY]*size[ZZ]);
@@ -405,6 +446,8 @@ static int set_grid_size_xy(const nbnxn_search_t nbs,
 
     if (n > grid->na_sc)
     {
+        assert(atom_density > 0);
+
         /* target cell length */
         if (grid->bSimple)
         {
@@ -521,6 +564,10 @@ static int set_grid_size_xy(const nbnxn_search_t nbs,
         }
 
         srenew(grid->flags, grid->nc_nalloc);
+        if (nbs->bFEP)
+        {
+            srenew(grid->fep, grid->nc_nalloc*grid->na_sc/grid->na_c);
+        }
     }
 
     copy_rvec(corner0, grid->c0);
@@ -1052,7 +1099,7 @@ void sort_on_lj(int na_c,
     int      subc, s, a, n1, n2, a_lj_max, i, j;
     int      sort1[NBNXN_NA_SC_MAX/GPU_NSUBCELL];
     int      sort2[NBNXN_NA_SC_MAX/GPU_NSUBCELL];
-    gmx_bool haveQ;
+    gmx_bool haveQ, bFEP;
 
     *flags = 0;
 
@@ -1079,7 +1126,7 @@ void sort_on_lj(int na_c,
             }
         }
 
-        /* If we don't have atom with LJ, there's nothing to sort */
+        /* If we don't have atoms with LJ, there's nothing to sort */
         if (n1 > 0)
         {
             *flags |= NBNXN_CI_DO_LJ(subc);
@@ -1140,6 +1187,23 @@ void fill_cell(const nbnxn_search_t nbs,
                    grid->flags+(a0>>grid->na_c_2log)-grid->cell0);
     }
 
+    if (nbs->bFEP)
+    {
+        /* Set the fep flag for perturbed atoms in this (sub-)cell */
+        int c, at;
+
+        /* The grid-local cluster/(sub-)cell index */
+        c            = (a0 >> grid->na_c_2log) - grid->cell0*(grid->bSimple ? 1 : GPU_NSUBCELL);
+        grid->fep[c] = 0;
+        for (at = a0; at < a1; at++)
+        {
+            if (nbs->a[at] >= 0 && GET_CGINFO_FEP(atinfo[nbs->a[at]]))
+            {
+                grid->fep[c] |= (1 << (at - a0));
+            }
+        }
+    }
+
     /* Now we have sorted the atoms, set the cell indices */
     for (a = a0; a < a1; a++)
     {
@@ -1319,14 +1383,13 @@ static void sort_columns_supersub(const nbnxn_search_t nbs,
                                   int cxy_start, int cxy_end,
                                   int *sort_work)
 {
-    int  cxy;
-    int  cx, cy, cz = -1, c = -1, ncz;
-    int  na, ash, na_c, ind, a;
-    int  subdiv_z, sub_z, na_z, ash_z;
-    int  subdiv_y, sub_y, na_y, ash_y;
-    int  subdiv_x, sub_x, na_x, ash_x;
+    int        cxy;
+    int        cx, cy, cz = -1, c = -1, ncz;
+    int        na, ash, na_c, ind, a;
+    int        subdiv_z, sub_z, na_z, ash_z;
+    int        subdiv_y, sub_y, na_y, ash_y;
+    int        subdiv_x, sub_x, na_x, ash_x;
 
-    /* cppcheck-suppress unassignedVariable */
     nbnxn_bb_t bb_work_array[2], *bb_work_aligned;
 
     bb_work_aligned = (nbnxn_bb_t *)(((size_t)(bb_work_array+1)) & (~((size_t)15)));
@@ -1642,8 +1705,8 @@ static void calc_cell_indices(const nbnxn_search_t nbs,
     }
 
     /* Sort the super-cell columns along z into the sub-cells. */
-#pragma omp parallel for num_threads(nbs->nthread_max) schedule(static)
-    for (thread = 0; thread < nbs->nthread_max; thread++)
+#pragma omp parallel for num_threads(nthread) schedule(static)
+    for (thread = 0; thread < nthread; thread++)
     {
         if (grid->bSimple)
         {
@@ -1759,7 +1822,8 @@ void nbnxn_put_on_grid(nbnxn_search_t nbs,
         nbs->ePBC = ePBC;
         copy_mat(box, nbs->box);
 
-        if (atom_density >= 0)
+        /* Avoid zero density */
+        if (atom_density > 0)
         {
             grid->atom_density = atom_density;
         }
@@ -1775,12 +1839,21 @@ void nbnxn_put_on_grid(nbnxn_search_t nbs,
          * for the local atoms (dd_zone=0).
          */
         nbs->natoms_nonlocal = a1 - nmoved;
+
+        if (debug)
+        {
+            fprintf(debug, "natoms_local = %5d atom_density = %5.1f\n",
+                    nbs->natoms_local, grid->atom_density);
+        }
     }
     else
     {
         nbs->natoms_nonlocal = max(nbs->natoms_nonlocal, a1);
     }
 
+    /* We always use the home zone (grid[0]) for setting the cell size,
+     * since determining densities for non-local zones is difficult.
+     */
     nc_max_grid = set_grid_size_xy(nbs, grid,
                                    dd_zone, n-nmoved, corner0, corner1,
                                    nbs->grid[0].atom_density);
@@ -2323,7 +2396,7 @@ static int nbl_cj(const nbnxn_pairlist_t *nbl, int cj_ind)
 }
 
 /* Returns the i-interaction mask of the j sub-cell for index cj_ind */
-static unsigned nbl_imask0(const nbnxn_pairlist_t *nbl, int cj_ind)
+static unsigned int nbl_imask0(const nbnxn_pairlist_t *nbl, int cj_ind)
 {
     return nbl->cj4[cj_ind >> NBNXN_GPU_JGROUP_SIZE_2LOG].imei[0].imask;
 }
@@ -2506,6 +2579,7 @@ void nbnxn_init_pairlist_set(nbnxn_pairlist_set_t *nbl_list,
     }
 
     snew(nbl_list->nbl, nbl_list->nnbl);
+    snew(nbl_list->nbl_fep, nbl_list->nnbl);
     /* Execute in order to avoid memory interleaving between threads */
 #pragma omp parallel for num_threads(nbl_list->nnbl) schedule(static)
     for (i = 0; i < nbl_list->nnbl; i++)
@@ -2524,6 +2598,9 @@ void nbnxn_init_pairlist_set(nbnxn_pairlist_set_t *nbl_list,
         {
             nbnxn_init_pairlist(nbl_list->nbl[i], nbl_list->bSimple, NULL, NULL);
         }
+
+        snew(nbl_list->nbl_fep[i], 1);
+        nbnxn_init_pairlist_fep(nbl_list->nbl_fep[i]);
     }
 }
 
@@ -2650,7 +2727,7 @@ static void low_get_nbl_exclusions(nbnxn_pairlist_t *nbl, int cj4,
 }
 
 /* Returns a pointer to the exclusion mask for cj4-unit cj4, warp warp,
- * allocates extra memory, if necessary.
+ * generates a new element and allocates extra memory, if necessary.
  */
 static void get_nbl_exclusions_1(nbnxn_pairlist_t *nbl, int cj4,
                                  int warp, nbnxn_excl_t **excl)
@@ -2664,7 +2741,7 @@ static void get_nbl_exclusions_1(nbnxn_pairlist_t *nbl, int cj4,
 }
 
 /* Returns pointers to the exclusion mask for cj4-unit cj4 for both warps,
- * allocates extra memory, if necessary.
+ * generates a new element and allocates extra memory, if necessary.
  */
 static void get_nbl_exclusions_2(nbnxn_pairlist_t *nbl, int cj4,
                                  nbnxn_excl_t **excl_w0,
@@ -2865,10 +2942,10 @@ static void make_cluster_list_simple(const nbnxn_grid_t *gridj,
 }
 
 #ifdef GMX_NBNXN_SIMD_4XN
-#include "nbnxn_search_simd_4xn.h"
+#include "gromacs/mdlib/nbnxn_search_simd_4xn.h"
 #endif
 #ifdef GMX_NBNXN_SIMD_2XNN
-#include "nbnxn_search_simd_2xnn.h"
+#include "gromacs/mdlib/nbnxn_search_simd_2xnn.h"
 #endif
 
 /* Plain C or SIMD4 code for making a pair list of super-cell sci vs scj.
@@ -2887,7 +2964,7 @@ static void make_cluster_list_supersub(const nbnxn_grid_t *gridi,
     int               npair;
     int               cjo, ci1, ci, cj, cj_gl;
     int               cj4_ind, cj_offset;
-    unsigned          imask;
+    unsigned int      imask;
     nbnxn_cj4_t      *cj4;
 #ifdef NBNXN_BBXXXX
     const float      *pbb_ci;
@@ -3193,6 +3270,408 @@ static void set_ci_top_excls(const nbnxn_search_t nbs,
     }
 }
 
+/* Add a new i-entry to the FEP list and copy the i-properties */
+static gmx_inline void fep_list_new_nri_copy(t_nblist *nlist)
+{
+    /* Add a new i-entry */
+    nlist->nri++;
+
+    assert(nlist->nri < nlist->maxnri);
+
+    /* Duplicate the last i-entry, except for jindex, which continues */
+    nlist->iinr[nlist->nri]   = nlist->iinr[nlist->nri-1];
+    nlist->shift[nlist->nri]  = nlist->shift[nlist->nri-1];
+    nlist->gid[nlist->nri]    = nlist->gid[nlist->nri-1];
+    nlist->jindex[nlist->nri] = nlist->nrj;
+}
+
+/* For load balancing of the free-energy lists over threads, we set
+ * the maximum nrj size of an i-entry to 40. This leads to good
+ * load balancing in the worst case scenario of a single perturbed
+ * particle on 16 threads, while not introducing significant overhead.
+ * Note that half of the perturbed pairs will anyhow end up in very small lists,
+ * since non perturbed i-particles will see few perturbed j-particles).
+ */
+const int max_nrj_fep = 40;
+
+/* Exclude the perturbed pairs from the Verlet list. This is only done to avoid
+ * singularities for overlapping particles (0/0), since the charges and
+ * 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_t    nbs,
+                          const nbnxn_atomdata_t *nbat,
+                          nbnxn_pairlist_t       *nbl,
+                          gmx_bool                bDiagRemoved,
+                          nbnxn_ci_t             *nbl_ci,
+                          const nbnxn_grid_t     *gridi,
+                          const nbnxn_grid_t     *gridj,
+                          t_nblist               *nlist)
+{
+    int      ci, cj_ind_start, cj_ind_end, cj_ind, cja, cjr;
+    int      nri_max;
+    int      ngid, gid_i = 0, gid_j, gid;
+    int      egp_shift, egp_mask;
+    int      gid_cj = 0;
+    int      i, j, ind_i, ind_j, ai, aj;
+    int      nri;
+    gmx_bool bFEP_i, bFEP_i_all;
+
+    if (nbl_ci->cj_ind_end == nbl_ci->cj_ind_start)
+    {
+        /* Empty list */
+        return;
+    }
+
+    ci = nbl_ci->ci;
+
+    cj_ind_start = nbl_ci->cj_ind_start;
+    cj_ind_end   = nbl_ci->cj_ind_end;
+
+    /* In worst case we have alternating energy groups
+     * and create #atom-pair lists, which means we need the size
+     * of a cluster pair (na_ci*na_cj) times the number of cj's.
+     */
+    nri_max = nbl->na_ci*nbl->na_cj*(cj_ind_end - cj_ind_start);
+    if (nlist->nri + nri_max > nlist->maxnri)
+    {
+        nlist->maxnri = over_alloc_large(nlist->nri + nri_max);
+        reallocate_nblist(nlist);
+    }
+
+    ngid = nbat->nenergrp;
+
+    if (ngid*gridj->na_cj > sizeof(gid_cj)*8)
+    {
+        gmx_fatal(FARGS, "The Verlet scheme with %dx%d kernels and free-energy only supports up to %d energy groups",
+                  gridi->na_c, gridj->na_cj, (sizeof(gid_cj)*8)/gridj->na_cj);
+    }
+
+    egp_shift = nbat->neg_2log;
+    egp_mask  = (1<<nbat->neg_2log) - 1;
+
+    /* Loop over the atoms in the i sub-cell */
+    bFEP_i_all = TRUE;
+    for (i = 0; i < nbl->na_ci; i++)
+    {
+        ind_i = ci*nbl->na_ci + i;
+        ai    = nbs->a[ind_i];
+        if (ai >= 0)
+        {
+            nri                  = nlist->nri;
+            nlist->jindex[nri+1] = nlist->jindex[nri];
+            nlist->iinr[nri]     = ai;
+            /* The actual energy group pair index is set later */
+            nlist->gid[nri]      = 0;
+            nlist->shift[nri]    = nbl_ci->shift & NBNXN_CI_SHIFT;
+
+            bFEP_i = gridi->fep[ci - gridi->cell0] & (1 << i);
+
+            bFEP_i_all = bFEP_i_all && bFEP_i;
+
+            if ((nlist->nrj + cj_ind_end - cj_ind_start)*nbl->na_cj > nlist->maxnrj)
+            {
+                nlist->maxnrj = over_alloc_small((nlist->nrj + cj_ind_end - cj_ind_start)*nbl->na_cj);
+                srenew(nlist->jjnr,     nlist->maxnrj);
+                srenew(nlist->excl_fep, nlist->maxnrj);
+            }
+
+            if (ngid > 1)
+            {
+                gid_i = (nbat->energrp[ci] >> (egp_shift*i)) & egp_mask;
+            }
+
+            for (cj_ind = cj_ind_start; cj_ind < cj_ind_end; cj_ind++)
+            {
+                unsigned int fep_cj;
+
+                cja = nbl->cj[cj_ind].cj;
+
+                if (gridj->na_cj == gridj->na_c)
+                {
+                    cjr    = cja - gridj->cell0;
+                    fep_cj = gridj->fep[cjr];
+                    if (ngid > 1)
+                    {
+                        gid_cj = nbat->energrp[cja];
+                    }
+                }
+                else if (2*gridj->na_cj == gridj->na_c)
+                {
+                    cjr    = cja - gridj->cell0*2;
+                    /* Extract half of the ci fep/energrp mask */
+                    fep_cj = (gridj->fep[cjr>>1] >> ((cjr&1)*gridj->na_cj)) & ((1<<gridj->na_cj) - 1);
+                    if (ngid > 1)
+                    {
+                        gid_cj = nbat->energrp[cja>>1] >> ((cja&1)*gridj->na_cj*egp_shift) & ((1<<(gridj->na_cj*egp_shift)) - 1);
+                    }
+                }
+                else
+                {
+                    cjr    = cja - (gridj->cell0>>1);
+                    /* Combine two ci fep masks/energrp */
+                    fep_cj = gridj->fep[cjr*2] + (gridj->fep[cjr*2+1] << gridj->na_c);
+                    if (ngid > 1)
+                    {
+                        gid_cj = nbat->energrp[cja*2] + (nbat->energrp[cja*2+1] << (gridj->na_c*egp_shift));
+                    }
+                }
+
+                if (bFEP_i || fep_cj != 0)
+                {
+                    for (j = 0; j < nbl->na_cj; j++)
+                    {
+                        /* Is this interaction perturbed and not excluded? */
+                        ind_j = cja*nbl->na_cj + j;
+                        aj    = nbs->a[ind_j];
+                        if (aj >= 0 &&
+                            (bFEP_i || (fep_cj & (1 << j))) &&
+                            (!bDiagRemoved || ind_j >= ind_i))
+                        {
+                            if (ngid > 1)
+                            {
+                                gid_j = (gid_cj >> (j*egp_shift)) & egp_mask;
+                                gid   = GID(gid_i, gid_j, ngid);
+
+                                if (nlist->nrj > nlist->jindex[nri] &&
+                                    nlist->gid[nri] != gid)
+                                {
+                                    /* Energy group pair changed: new list */
+                                    fep_list_new_nri_copy(nlist);
+                                    nri = nlist->nri;
+                                }
+                                nlist->gid[nri] = gid;
+                            }
+
+                            if (nlist->nrj - nlist->jindex[nri] >= max_nrj_fep)
+                            {
+                                fep_list_new_nri_copy(nlist);
+                                nri = nlist->nri;
+                            }
+
+                            /* Add it to the FEP list */
+                            nlist->jjnr[nlist->nrj]     = aj;
+                            nlist->excl_fep[nlist->nrj] = (nbl->cj[cj_ind].excl >> (i*nbl->na_cj + j)) & 1;
+                            nlist->nrj++;
+
+                            /* Exclude it from the normal list.
+                             * Note that the charge has been set to zero,
+                             * but we need to avoid 0/0, as perturbed atoms
+                             * can be on top of each other.
+                             */
+                            nbl->cj[cj_ind].excl &= ~(1U << (i*nbl->na_cj + j));
+                        }
+                    }
+                }
+            }
+
+            if (nlist->nrj > nlist->jindex[nri])
+            {
+                /* Actually add this new, non-empty, list */
+                nlist->nri++;
+                nlist->jindex[nlist->nri] = nlist->nrj;
+            }
+        }
+    }
+
+    if (bFEP_i_all)
+    {
+        /* All interactions are perturbed, we can skip this entry */
+        nbl_ci->cj_ind_end = cj_ind_start;
+    }
+}
+
+/* Return the index of atom a within a cluster */
+static gmx_inline int cj_mod_cj4(int cj)
+{
+    return cj & (NBNXN_GPU_JGROUP_SIZE - 1);
+}
+
+/* Convert a j-cluster to a cj4 group */
+static gmx_inline int cj_to_cj4(int cj)
+{
+    return cj >> NBNXN_GPU_JGROUP_SIZE_2LOG;
+}
+
+/* Return the index of an j-atom within a warp */
+static gmx_inline int a_mod_wj(int a)
+{
+    return a & (NBNXN_GPU_CLUSTER_SIZE/2 - 1);
+}
+
+/* As make_fep_list above, but for super/sub lists. */
+static void make_fep_list_supersub(const nbnxn_search_t    nbs,
+                                   const nbnxn_atomdata_t *nbat,
+                                   nbnxn_pairlist_t       *nbl,
+                                   gmx_bool                bDiagRemoved,
+                                   const nbnxn_sci_t      *nbl_sci,
+                                   real                    shx,
+                                   real                    shy,
+                                   real                    shz,
+                                   real                    rlist_fep2,
+                                   const nbnxn_grid_t     *gridi,
+                                   const nbnxn_grid_t     *gridj,
+                                   t_nblist               *nlist)
+{
+    int                sci, cj4_ind_start, cj4_ind_end, cj4_ind, gcj, cjr;
+    int                nri_max;
+    int                c, c_abs;
+    int                i, j, ind_i, ind_j, ai, aj;
+    int                nri;
+    gmx_bool           bFEP_i;
+    real               xi, yi, zi;
+    const nbnxn_cj4_t *cj4;
+
+    if (nbl_sci->cj4_ind_end == nbl_sci->cj4_ind_start)
+    {
+        /* Empty list */
+        return;
+    }
+
+    sci = nbl_sci->sci;
+
+    cj4_ind_start = nbl_sci->cj4_ind_start;
+    cj4_ind_end   = nbl_sci->cj4_ind_end;
+
+    /* Here we process one super-cell, max #atoms na_sc, versus a list
+     * cj4 entries, each with max NBNXN_GPU_JGROUP_SIZE cj's, each
+     * of size na_cj atoms.
+     * On the GPU we don't support energy groups (yet).
+     * So for each of the na_sc i-atoms, we need max one FEP list
+     * for each max_nrj_fep j-atoms.
+     */
+    nri_max = nbl->na_sc*nbl->na_cj*(1 + ((cj4_ind_end - cj4_ind_start)*NBNXN_GPU_JGROUP_SIZE)/max_nrj_fep);
+    if (nlist->nri + nri_max > nlist->maxnri)
+    {
+        nlist->maxnri = over_alloc_large(nlist->nri + nri_max);
+        reallocate_nblist(nlist);
+    }
+
+    /* Loop over the atoms in the i super-cluster */
+    for (c = 0; c < GPU_NSUBCELL; c++)
+    {
+        c_abs = sci*GPU_NSUBCELL + c;
+
+        for (i = 0; i < nbl->na_ci; i++)
+        {
+            ind_i = c_abs*nbl->na_ci + i;
+            ai    = nbs->a[ind_i];
+            if (ai >= 0)
+            {
+                nri                  = nlist->nri;
+                nlist->jindex[nri+1] = nlist->jindex[nri];
+                nlist->iinr[nri]     = ai;
+                /* With GPUs, energy groups are not supported */
+                nlist->gid[nri]      = 0;
+                nlist->shift[nri]    = nbl_sci->shift & NBNXN_CI_SHIFT;
+
+                bFEP_i = (gridi->fep[c_abs - gridi->cell0] & (1 << i));
+
+                xi = nbat->x[ind_i*nbat->xstride+XX] + shx;
+                yi = nbat->x[ind_i*nbat->xstride+YY] + shy;
+                zi = nbat->x[ind_i*nbat->xstride+ZZ] + shz;
+
+                if ((nlist->nrj + cj4_ind_end - cj4_ind_start)*NBNXN_GPU_JGROUP_SIZE*nbl->na_cj > nlist->maxnrj)
+                {
+                    nlist->maxnrj = over_alloc_small((nlist->nrj + cj4_ind_end - cj4_ind_start)*NBNXN_GPU_JGROUP_SIZE*nbl->na_cj);
+                    srenew(nlist->jjnr,     nlist->maxnrj);
+                    srenew(nlist->excl_fep, nlist->maxnrj);
+                }
+
+                for (cj4_ind = cj4_ind_start; cj4_ind < cj4_ind_end; cj4_ind++)
+                {
+                    cj4 = &nbl->cj4[cj4_ind];
+
+                    for (gcj = 0; gcj < NBNXN_GPU_JGROUP_SIZE; gcj++)
+                    {
+                        unsigned int fep_cj;
+
+                        if ((cj4->imei[0].imask & (1U << (gcj*GPU_NSUBCELL + c))) == 0)
+                        {
+                            /* Skip this ci for this cj */
+                            continue;
+                        }
+
+                        cjr = cj4->cj[gcj] - gridj->cell0*GPU_NSUBCELL;
+
+                        fep_cj = gridj->fep[cjr];
+
+                        if (bFEP_i || fep_cj != 0)
+                        {
+                            for (j = 0; j < nbl->na_cj; j++)
+                            {
+                                /* Is this interaction perturbed and not excluded? */
+                                ind_j = (gridj->cell0*GPU_NSUBCELL + cjr)*nbl->na_cj + j;
+                                aj    = nbs->a[ind_j];
+                                if (aj >= 0 &&
+                                    (bFEP_i || (fep_cj & (1 << j))) &&
+                                    (!bDiagRemoved || ind_j >= ind_i))
+                                {
+                                    nbnxn_excl_t *excl;
+                                    int           excl_pair;
+                                    unsigned int  excl_bit;
+                                    real          dx, dy, dz;
+
+                                    get_nbl_exclusions_1(nbl, cj4_ind, j>>2, &excl);
+
+                                    excl_pair = a_mod_wj(j)*nbl->na_ci + i;
+                                    excl_bit  = (1U << (gcj*GPU_NSUBCELL + c));
+
+                                    dx = nbat->x[ind_j*nbat->xstride+XX] - xi;
+                                    dy = nbat->x[ind_j*nbat->xstride+YY] - yi;
+                                    dz = nbat->x[ind_j*nbat->xstride+ZZ] - zi;
+
+                                    /* The unpruned GPU list has more than 2/3
+                                     * of the atom pairs beyond rlist. Using
+                                     * this list will cause a lot of overhead
+                                     * in the CPU FEP kernels, especially
+                                     * relative to the fast GPU kernels.
+                                     * So we prune the FEP list here.
+                                     */
+                                    if (dx*dx + dy*dy + dz*dz < rlist_fep2)
+                                    {
+                                        if (nlist->nrj - nlist->jindex[nri] >= max_nrj_fep)
+                                        {
+                                            fep_list_new_nri_copy(nlist);
+                                            nri = nlist->nri;
+                                        }
+
+                                        /* Add it to the FEP list */
+                                        nlist->jjnr[nlist->nrj]     = aj;
+                                        nlist->excl_fep[nlist->nrj] = (excl->pair[excl_pair] & excl_bit) ? 1 : 0;
+                                        nlist->nrj++;
+                                    }
+
+                                    /* Exclude it from the normal list.
+                                     * Note that the charge and LJ parameters have
+                                     * been set to zero, but we need to avoid 0/0,
+                                     * as perturbed atoms can be on top of each other.
+                                     */
+                                    excl->pair[excl_pair] &= ~excl_bit;
+                                }
+                            }
+
+                            /* Note that we could mask out this pair in imask
+                             * if all i- and/or all j-particles are perturbed.
+                             * But since the perturbed pairs on the CPU will
+                             * take an order of magnitude more time, the GPU
+                             * will finish before the CPU and there is no gain.
+                             */
+                        }
+                    }
+                }
+
+                if (nlist->nrj > nlist->jindex[nri])
+                {
+                    /* Actually add this new, non-empty, list */
+                    nlist->nri++;
+                    nlist->jindex[nlist->nri] = nlist->nrj;
+                }
+            }
+        }
+    }
+}
+
 /* Set all atom-pair exclusions from the topology stored in excl
  * as masks in the pair-list for i-super-cell entry nbl_sci
  */
@@ -3316,26 +3795,15 @@ static void set_sci_top_excls(const nbnxn_search_t nbs,
                         inner_i = i  - si*na_c;
                         inner_e = ge - se*na_c;
 
-/* Macro for getting the index of atom a within a cluster */
-#define AMODCJ4(a)  ((a) & (NBNXN_GPU_JGROUP_SIZE - 1))
-/* Macro for converting an atom number to a cluster number */
-#define A2CJ4(a)    ((a) >> NBNXN_GPU_JGROUP_SIZE_2LOG)
-/* Macro for getting the index of an i-atom within a warp */
-#define AMODWI(a)   ((a) & (NBNXN_GPU_CLUSTER_SIZE/2 - 1))
-
-                        if (nbl_imask0(nbl, found) & (1U << (AMODCJ4(found)*GPU_NSUBCELL + si)))
+                        if (nbl_imask0(nbl, found) & (1U << (cj_mod_cj4(found)*GPU_NSUBCELL + si)))
                         {
                             w       = (inner_e >> 2);
 
-                            get_nbl_exclusions_1(nbl, A2CJ4(found), w, &nbl_excl);
+                            get_nbl_exclusions_1(nbl, cj_to_cj4(found), w, &nbl_excl);
 
-                            nbl_excl->pair[AMODWI(inner_e)*nbl->na_ci+inner_i] &=
-                                ~(1U << (AMODCJ4(found)*GPU_NSUBCELL + si));
+                            nbl_excl->pair[a_mod_wj(inner_e)*nbl->na_ci+inner_i] &=
+                                ~(1U << (cj_mod_cj4(found)*GPU_NSUBCELL + si));
                         }
-
-#undef AMODCJ4
-#undef A2CJ4
-#undef AMODWI
                     }
                 }
             }
@@ -3615,6 +4083,18 @@ static void clear_pairlist(nbnxn_pairlist_t *nbl)
     nbl->work->ncj_hlj = 0;
 }
 
+/* Clears a group scheme pair list */
+static void clear_pairlist_fep(t_nblist *nl)
+{
+    nl->nri = 0;
+    nl->nrj = 0;
+    if (nl->jindex == NULL)
+    {
+        snew(nl->jindex, 1);
+    }
+    nl->jindex[0] = 0;
+}
+
 /* Sets a simple list i-cell bounding box, including PBC shift */
 static gmx_inline void set_icell_bb_simple(const nbnxn_bb_t *bb, int ci,
                                            real shx, real shy, real shz,
@@ -3693,7 +4173,7 @@ static void icell_set_x_supersub(int ci,
                                  int stride, const real *x,
                                  nbnxn_list_work_t *work)
 {
-    int  ia, i;
+    int   ia, i;
     real *x_ci;
 
     x_ci = work->x_ci;
@@ -3715,7 +4195,7 @@ static void icell_set_x_supersub_simd4(int ci,
                                        int stride, const real *x,
                                        nbnxn_list_work_t *work)
 {
-    int  si, io, ia, i, j;
+    int   si, io, ia, i, j;
     real *x_ci;
 
     x_ci = work->x_ci;
@@ -3737,6 +4217,44 @@ static void icell_set_x_supersub_simd4(int ci,
 }
 #endif
 
+static real minimum_subgrid_size_xy(const nbnxn_grid_t *grid)
+{
+    if (grid->bSimple)
+    {
+        return min(grid->sx, grid->sy);
+    }
+    else
+    {
+        return min(grid->sx/GPU_NSUBCELL_X, grid->sy/GPU_NSUBCELL_Y);
+    }
+}
+
+static real effective_buffer_1x1_vs_MxN(const nbnxn_grid_t *gridi,
+                                        const nbnxn_grid_t *gridj)
+{
+    const real eff_1x1_buffer_fac_overest = 0.1;
+
+    /* Determine an atom-pair list cut-off buffer size for atom pairs,
+     * to be added to rlist (including buffer) used for MxN.
+     * This is for converting an MxN list to a 1x1 list. This means we can't
+     * use the normal buffer estimate, as we have an MxN list in which
+     * some atom pairs beyond rlist are missing. We want to capture
+     * the beneficial effect of buffering by extra pairs just outside rlist,
+     * while removing the useless pairs that are further away from rlist.
+     * (Also the buffer could have been set manually not using the estimate.)
+     * This buffer size is an overestimate.
+     * We add 10% of the smallest grid sub-cell dimensions.
+     * Note that the z-size differs per cell and we don't use this,
+     * so we overestimate.
+     * With PME, the 10% value gives a buffer that is somewhat larger
+     * than the effective buffer with a tolerance of 0.005 kJ/mol/ps.
+     * Smaller tolerances or using RF lead to a smaller effective buffer,
+     * so 10% gives a safe overestimate.
+     */
+    return eff_1x1_buffer_fac_overest*(minimum_subgrid_size_xy(gridi) +
+                                       minimum_subgrid_size_xy(gridj));
+}
+
 /* Clusters at the cut-off only increase rlist by 60% of their size */
 static real nbnxn_rlist_inc_outside_fac = 0.6;
 
@@ -3814,9 +4332,9 @@ static int get_nsubpair_max(const nbnxn_search_t nbs,
                             int                  min_ci_balanced)
 {
     const nbnxn_grid_t *grid;
-    rvec ls;
-    real xy_diag2, r_eff_sup, vol_est, nsp_est, nsp_est_nl;
-    int  nsubpair_max;
+    rvec                ls;
+    real                xy_diag2, r_eff_sup, vol_est, nsp_est, nsp_est_nl;
+    int                 nsubpair_max;
 
     grid = &nbs->grid[0];
 
@@ -3998,11 +4516,11 @@ static void combine_nblists(int nnbl, nbnxn_pairlist_t **nbl,
 #pragma omp parallel for num_threads(gmx_omp_nthreads_get(emntPairsearch)) schedule(static)
     for (n = 0; n < nnbl; n++)
     {
-        int sci_offset;
-        int cj4_offset;
-        int ci_offset;
-        int excl_offset;
-        int i, j4;
+        int                     sci_offset;
+        int                     cj4_offset;
+        int                     ci_offset;
+        int                     excl_offset;
+        int                     i, j4;
         const nbnxn_pairlist_t *nbli;
 
         /* Determine the offset in the combined data for our thread */
@@ -4050,6 +4568,121 @@ static void combine_nblists(int nnbl, nbnxn_pairlist_t **nbl,
     }
 }
 
+static void balance_fep_lists(const nbnxn_search_t  nbs,
+                              nbnxn_pairlist_set_t *nbl_lists)
+{
+    int       nnbl, th;
+    int       nri_tot, nrj_tot, nrj_target;
+    int       th_dest;
+    t_nblist *nbld;
+
+    nnbl = nbl_lists->nnbl;
+
+    if (nnbl == 1)
+    {
+        /* Nothing to balance */
+        return;
+    }
+
+    /* Count the total i-lists and pairs */
+    nri_tot = 0;
+    nrj_tot = 0;
+    for (th = 0; th < nnbl; th++)
+    {
+        nri_tot += nbl_lists->nbl_fep[th]->nri;
+        nrj_tot += nbl_lists->nbl_fep[th]->nrj;
+    }
+
+    nrj_target = (nrj_tot + nnbl - 1)/nnbl;
+
+    assert(gmx_omp_nthreads_get(emntNonbonded) == nnbl);
+
+#pragma omp parallel for schedule(static) num_threads(nnbl)
+    for (th = 0; th < nnbl; th++)
+    {
+        t_nblist *nbl;
+
+        nbl = nbs->work[th].nbl_fep;
+
+        /* Note that here we allocate for the total size, instead of
+         * a per-thread esimate (which is hard to obtain).
+         */
+        if (nri_tot > nbl->maxnri)
+        {
+            nbl->maxnri = over_alloc_large(nri_tot);
+            reallocate_nblist(nbl);
+        }
+        if (nri_tot > nbl->maxnri || nrj_tot > nbl->maxnrj)
+        {
+            nbl->maxnrj = over_alloc_small(nrj_tot);
+            srenew(nbl->jjnr, nbl->maxnrj);
+            srenew(nbl->excl_fep, nbl->maxnrj);
+        }
+
+        clear_pairlist_fep(nbl);
+    }
+
+    /* Loop over the source lists and assign and copy i-entries */
+    th_dest = 0;
+    nbld    = nbs->work[th_dest].nbl_fep;
+    for (th = 0; th < nnbl; th++)
+    {
+        t_nblist *nbls;
+        int       i, j;
+
+        nbls = nbl_lists->nbl_fep[th];
+
+        for (i = 0; i < nbls->nri; i++)
+        {
+            int nrj;
+
+            /* The number of pairs in this i-entry */
+            nrj = nbls->jindex[i+1] - nbls->jindex[i];
+
+            /* Decide if list th_dest is too large and we should procede
+             * to the next destination list.
+             */
+            if (th_dest+1 < nnbl && nbld->nrj > 0 &&
+                nbld->nrj + nrj - nrj_target > nrj_target - nbld->nrj)
+            {
+                th_dest++;
+                nbld = nbs->work[th_dest].nbl_fep;
+            }
+
+            nbld->iinr[nbld->nri]  = nbls->iinr[i];
+            nbld->gid[nbld->nri]   = nbls->gid[i];
+            nbld->shift[nbld->nri] = nbls->shift[i];
+
+            for (j = nbls->jindex[i]; j < nbls->jindex[i+1]; j++)
+            {
+                nbld->jjnr[nbld->nrj]     = nbls->jjnr[j];
+                nbld->excl_fep[nbld->nrj] = nbls->excl_fep[j];
+                nbld->nrj++;
+            }
+            nbld->nri++;
+            nbld->jindex[nbld->nri] = nbld->nrj;
+        }
+    }
+
+    /* Swap the list pointers */
+    for (th = 0; th < nnbl; th++)
+    {
+        t_nblist *nbl_tmp;
+
+        nbl_tmp                = nbl_lists->nbl_fep[th];
+        nbl_lists->nbl_fep[th] = nbs->work[th].nbl_fep;
+        nbs->work[th].nbl_fep  = nbl_tmp;
+
+        if (debug)
+        {
+            fprintf(debug, "nbl_fep[%d] nri %4d nrj %4d\n",
+                    th,
+                    nbl_lists->nbl_fep[th]->nri,
+                    nbl_lists->nbl_fep[th]->nrj);
+        }
+    }
+}
+
 /* Returns the next ci to be processes by our thread */
 static gmx_bool next_ci(const nbnxn_grid_t *grid,
                         int conv,
@@ -4132,7 +4765,7 @@ static int get_ci_block_size(const nbnxn_grid_t *gridi,
     const int ci_block_enum      = 5;
     const int ci_block_denom     = 11;
     const int ci_block_min_atoms = 16;
-    int ci_block;
+    int       ci_block;
 
     /* Here we decide how to distribute the blocks over the threads.
      * We use prime numbers to try to avoid that the grid size becomes
@@ -4179,37 +4812,38 @@ static void nbnxn_make_pairlist_part(const nbnxn_search_t nbs,
                                      gmx_bool progBal,
                                      int min_ci_balanced,
                                      int th, int nth,
-                                     nbnxn_pairlist_t *nbl)
+                                     nbnxn_pairlist_t *nbl,
+                                     t_nblist *nbl_fep)
 {
-    int  na_cj_2log;
-    matrix box;
-    real rl2;
-    float rbb2;
-    int  d;
-    int  ci_b, ci, ci_x, ci_y, ci_xy, cj;
-    ivec shp;
-    int  tx, ty, tz;
-    int  shift;
-    gmx_bool bMakeList;
-    real shx, shy, shz;
-    int  conv_i, cell0_i;
+    int               na_cj_2log;
+    matrix            box;
+    real              rl2, rl_fep2 = 0;
+    float             rbb2;
+    int               d;
+    int               ci_b, ci, ci_x, ci_y, ci_xy, cj;
+    ivec              shp;
+    int               tx, ty, tz;
+    int               shift;
+    gmx_bool          bMakeList;
+    real              shx, shy, shz;
+    int               conv_i, cell0_i;
     const nbnxn_bb_t *bb_i = NULL;
 #ifdef NBNXN_BBXXXX
-    const float *pbb_i = NULL;
+    const float      *pbb_i = NULL;
 #endif
-    const float *bbcz_i, *bbcz_j;
-    const int *flags_i;
-    real bx0, bx1, by0, by1, bz0, bz1;
-    real bz1_frac;
-    real d2cx, d2z, d2z_cx, d2z_cy, d2zx, d2zxy, d2xy;
-    int  cxf, cxl, cyf, cyf_x, cyl;
-    int  cx, cy;
-    int  c0, c1, cs, cf, cl;
-    int  ndistc;
-    int  ncpcheck;
-    int  gridi_flag_shift = 0, gridj_flag_shift = 0;
-    unsigned *gridj_flag  = NULL;
-    int  ncj_old_i, ncj_old_j;
+    const float      *bbcz_i, *bbcz_j;
+    const int        *flags_i;
+    real              bx0, bx1, by0, by1, bz0, bz1;
+    real              bz1_frac;
+    real              d2cx, d2z, d2z_cx, d2z_cy, d2zx, d2zxy, d2xy;
+    int               cxf, cxl, cyf, cyf_x, cyl;
+    int               cx, cy;
+    int               c0, c1, cs, cf, cl;
+    int               ndistc;
+    int               ncpcheck;
+    int               gridi_flag_shift = 0, gridj_flag_shift = 0;
+    unsigned int     *gridj_flag       = NULL;
+    int               ncj_old_i, ncj_old_j;
 
     nbs_cycle_start(&work->cc[enbsCCsearch]);
 
@@ -4247,6 +4881,23 @@ static void nbnxn_make_pairlist_part(const nbnxn_search_t nbs,
 
     rl2 = nbl->rlist*nbl->rlist;
 
+    if (nbs->bFEP && !nbl->bSimple)
+    {
+        /* Determine an atom-pair list cut-off distance for FEP atom pairs.
+         * We should not simply use rlist, since then we would not have
+         * the small, effective buffering of the NxN lists.
+         * The buffer is on overestimate, but the resulting cost for pairs
+         * beyond rlist is neglible compared to the FEP pairs within rlist.
+         */
+        rl_fep2 = nbl->rlist + effective_buffer_1x1_vs_MxN(gridi, gridj);
+
+        if (debug)
+        {
+            fprintf(debug, "nbl_fep atom-pair rlist %f\n", rl_fep2);
+        }
+        rl_fep2 = rl_fep2*rl_fep2;
+    }
+
     rbb2 = boundingbox_only_distance2(gridi, gridj, nbl->rlist, nbl->bSimple);
 
     if (debug)
@@ -4715,6 +5366,14 @@ static void nbnxn_make_pairlist_part(const nbnxn_search_t nbs,
                                          na_cj_2log,
                                          &(nbl->ci[nbl->nci]),
                                          excl);
+
+                        if (nbs->bFEP)
+                        {
+                            make_fep_list(nbs, nbat, nbl,
+                                          shift == CENTRAL && gridi == gridj,
+                                          &(nbl->ci[nbl->nci]),
+                                          gridi, gridj, nbl_fep);
+                        }
                     }
                     else
                     {
@@ -4724,6 +5383,16 @@ static void nbnxn_make_pairlist_part(const nbnxn_search_t nbs,
                                           gridj->na_c_2log,
                                           &(nbl->sci[nbl->nsci]),
                                           excl);
+
+                        if (nbs->bFEP)
+                        {
+                            make_fep_list_supersub(nbs, nbat, nbl,
+                                                   shift == CENTRAL && gridi == gridj,
+                                                   &(nbl->sci[nbl->nsci]),
+                                                   shx, shy, shz,
+                                                   rl_fep2,
+                                                   gridi, gridj, nbl_fep);
+                        }
                     }
 
                     /* Close this ci list */
@@ -4767,6 +5436,10 @@ static void nbnxn_make_pairlist_part(const nbnxn_search_t nbs,
             print_nblist_statistics_supersub(debug, nbl, nbs, rlist);
         }
 
+        if (nbs->bFEP)
+        {
+            fprintf(debug, "nbl FEP list pairs: %d\n", nbl_fep->nrj);
+        }
     }
 }
 
@@ -4774,8 +5447,8 @@ static void reduce_buffer_flags(const nbnxn_search_t        nbs,
                                 int                         nsrc,
                                 const nbnxn_buffer_flags_t *dest)
 {
-    int s, b;
-    const unsigned *flag;
+    int                 s, b;
+    const unsigned int *flag;
 
     for (s = 0; s < nsrc; s++)
     {
@@ -4920,17 +5593,17 @@ void nbnxn_make_pairlist(const nbnxn_search_t  nbs,
                          int                   nb_kernel_type,
                          t_nrnb               *nrnb)
 {
-    nbnxn_grid_t *gridi, *gridj;
-    gmx_bool bGPUCPU;
-    int nzi, zi, zj0, zj1, zj;
-    int nsubpair_max;
-    int th;
-    int nnbl;
+    nbnxn_grid_t      *gridi, *gridj;
+    gmx_bool           bGPUCPU;
+    int                nzi, zi, zj0, zj1, zj;
+    int                nsubpair_max;
+    int                th;
+    int                nnbl;
     nbnxn_pairlist_t **nbl;
-    int ci_block;
-    gmx_bool CombineNBLists;
-    gmx_bool progBal;
-    int np_tot, np_noq, np_hlj, nap;
+    int                ci_block;
+    gmx_bool           CombineNBLists;
+    gmx_bool           progBal;
+    int                np_tot, np_noq, np_hlj, nap;
 
     /* Check if we are running hybrid GPU + CPU nbnxn mode */
     bGPUCPU = (!nbs->grid[0].bSimple && nbl_list->bSimple);
@@ -5004,6 +5677,11 @@ void nbnxn_make_pairlist(const nbnxn_search_t  nbs,
     for (th = 0; th < nnbl; th++)
     {
         clear_pairlist(nbl[th]);
+
+        if (nbs->bFEP)
+        {
+            clear_pairlist_fep(nbl_list->nbl_fep[th]);
+        }
     }
 
     for (zi = 0; zi < nzi; zi++)
@@ -5072,7 +5750,8 @@ void nbnxn_make_pairlist(const nbnxn_search_t  nbs,
                                          nsubpair_max,
                                          progBal, min_ci_balanced,
                                          th, nnbl,
-                                         nbl[th]);
+                                         nbl[th],
+                                         nbl_list->nbl_fep[th]);
             }
             nbs_cycle_stop(&nbs->cc[enbsCCsearch]);
 
@@ -5133,6 +5812,12 @@ void nbnxn_make_pairlist(const nbnxn_search_t  nbs,
         reduce_buffer_flags(nbs, nnbl, &nbat->buffer_flags);
     }
 
+    if (nbs->bFEP)
+    {
+        /* Balance the free-energy lists over all the threads */
+        balance_fep_lists(nbs, nbl_list);
+    }
+
     /* Special performance logging stuff (env.var. GMX_NBNXN_CYCLE) */
     if (LOCAL_I(iloc))
     {