src/gromacs/mdlib/pme.c

   1 /* -*- mode: c; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4; c-file-style: "stroustrup"; -*-
   2  *
   3  *
   4  *                This source code is part of
   5  *
   6  *                 G   R   O   M   A   C   S
   7  *
   8  *          GROningen MAchine for Chemical Simulations
   9  *
  10  *                        VERSION 3.2.0
  11  * Written by David van der Spoel, Erik Lindahl, Berk Hess, and others.
  12  * Copyright (c) 1991-2000, University of Groningen, The Netherlands.
  13  * Copyright (c) 2001-2004, The GROMACS development team,
  14  * check out http://www.gromacs.org for more information.
  15
  16  * This program is free software; you can redistribute it and/or
  17  * modify it under the terms of the GNU General Public License
  18  * as published by the Free Software Foundation; either version 2
  19  * of the License, or (at your option) any later version.
  20  *
  21  * If you want to redistribute modifications, please consider that
  22  * scientific software is very special. Version control is crucial -
  23  * bugs must be traceable. We will be happy to consider code for
  24  * inclusion in the official distribution, but derived work must not
  25  * be called official GROMACS. Details are found in the README & COPYING
  26  * files - if they are missing, get the official version at www.gromacs.org.
  27  *
  28  * To help us fund GROMACS development, we humbly ask that you cite
  29  * the papers on the package - you can find them in the top README file.
  30  *
  31  * For more info, check our website at http://www.gromacs.org
  32  *
  33  * And Hey:
  34  * GROwing Monsters And Cloning Shrimps
  35  */
  36 /* IMPORTANT FOR DEVELOPERS:
  37  *
  38  * Triclinic pme stuff isn't entirely trivial, and we've experienced
  39  * some bugs during development (many of them due to me). To avoid
  40  * this in the future, please check the following things if you make
  41  * changes in this file:
  42  *
  43  * 1. You should obtain identical (at least to the PME precision)
  44  *    energies, forces, and virial for
  45  *    a rectangular box and a triclinic one where the z (or y) axis is
  46  *    tilted a whole box side. For instance you could use these boxes:
  47  *
  48  *    rectangular       triclinic
  49  *     2  0  0           2  0  0
  50  *     0  2  0           0  2  0
  51  *     0  0  6           2  2  6
  52  *
  53  * 2. You should check the energy conservation in a triclinic box.
  54  *
  55  * It might seem an overkill, but better safe than sorry.
  56  * /Erik 001109
  57  */
  58
  59 #ifdef HAVE_CONFIG_H
  60 #include <config.h>
  61 #endif
  62
  63 #include "gromacs/fft/parallel_3dfft.h"
  64 #include "gromacs/utility/gmxmpi.h"
  65
  66 #include <stdio.h>
  67 #include <string.h>
  68 #include <math.h>
  69 #include <assert.h>
  70 #include "typedefs.h"
  71 #include "txtdump.h"
  72 #include "vec.h"
  73 #include "gmxcomplex.h"
  74 #include "smalloc.h"
  75 #include "futil.h"
  76 #include "coulomb.h"
  77 #include "gmx_fatal.h"
  78 #include "pme.h"
  79 #include "network.h"
  80 #include "physics.h"
  81 #include "nrnb.h"
  82 #include "gmx_wallcycle.h"
  83 #include "pdbio.h"
  84 #include "gmx_cyclecounter.h"
  85 #include "gmx_omp.h"
  86 #include "macros.h"
  87
  88
  89 /* Include the SIMD macro file and then check for support */
  90 #include "gmx_simd_macros.h"
  91 #if defined GMX_HAVE_SIMD_MACROS && defined GMX_SIMD_HAVE_EXP
  92 /* Turn on SIMD intrinsics for PME solve */
  93 #define PME_SIMD
  94 #endif
  95
  96 /* SIMD spread+gather only in single precision with SSE2 or higher available.
  97  * We might want to switch to use gmx_simd_macros.h, but this is somewhat
  98  * complicated, as we use unaligned and/or 4-wide only loads.
  99  */
 100 #if defined(GMX_X86_SSE2) && !defined(GMX_DOUBLE)
 101 #define PME_SSE_SPREAD_GATHER
 102 #include <emmintrin.h>
 103 /* Some old AMD processors could have problems with unaligned loads+stores */
 104 #ifndef GMX_FAHCORE
 105 #define PME_SSE_UNALIGNED
 106 #endif
 107 #endif
 108
 109 #define DFT_TOL 1e-7
 110 /* #define PRT_FORCE */
 111 /* conditions for on the fly time-measurement */
 112 /* #define TAKETIME (step > 1 && timesteps < 10) */
 113 #define TAKETIME FALSE
 114
 115 /* #define PME_TIME_THREADS */
 116
 117 #ifdef GMX_DOUBLE
 118 #define mpi_type MPI_DOUBLE
 119 #else
 120 #define mpi_type MPI_FLOAT
 121 #endif
 122
 123 /* GMX_CACHE_SEP should be a multiple of 16 to preserve alignment */
 124 #define GMX_CACHE_SEP 64
 125
 126 /* We only define a maximum to be able to use local arrays without allocation.
 127  * An order larger than 12 should never be needed, even for test cases.
 128  * If needed it can be changed here.
 129  */
 130 #define PME_ORDER_MAX 12
 131
 132 /* Internal datastructures */
 133 typedef struct {
 134     int send_index0;
 135     int send_nindex;
 136     int recv_index0;
 137     int recv_nindex;
 138     int recv_size;   /* Receive buffer width, used with OpenMP */
 139 } pme_grid_comm_t;
 140
 141 typedef struct {
 142 #ifdef GMX_MPI
 143     MPI_Comm         mpi_comm;
 144 #endif
 145     int              nnodes, nodeid;
 146     int             *s2g0;
 147     int             *s2g1;
 148     int              noverlap_nodes;
 149     int             *send_id, *recv_id;
 150     int              send_size; /* Send buffer width, used with OpenMP */
 151     pme_grid_comm_t *comm_data;
 152     real            *sendbuf;
 153     real            *recvbuf;
 154 } pme_overlap_t;
 155
 156 typedef struct {
 157     int *n;      /* Cumulative counts of the number of particles per thread */
 158     int  nalloc; /* Allocation size of i */
 159     int *i;      /* Particle indices ordered on thread index (n) */
 160 } thread_plist_t;
 161
 162 typedef struct {
 163     int      *thread_one;
 164     int       n;
 165     int      *ind;
 166     splinevec theta;
 167     real     *ptr_theta_z;
 168     splinevec dtheta;
 169     real     *ptr_dtheta_z;
 170 } splinedata_t;
 171
 172 typedef struct {
 173     int      dimind;        /* The index of the dimension, 0=x, 1=y */
 174     int      nslab;
 175     int      nodeid;
 176 #ifdef GMX_MPI
 177     MPI_Comm mpi_comm;
 178 #endif
 179
 180     int     *node_dest;     /* The nodes to send x and q to with DD */
 181     int     *node_src;      /* The nodes to receive x and q from with DD */
 182     int     *buf_index;     /* Index for commnode into the buffers */
 183
 184     int      maxshift;
 185
 186     int      npd;
 187     int      pd_nalloc;
 188     int     *pd;
 189     int     *count;         /* The number of atoms to send to each node */
 190     int    **count_thread;
 191     int     *rcount;        /* The number of atoms to receive */
 192
 193     int      n;
 194     int      nalloc;
 195     rvec    *x;
 196     real    *q;
 197     rvec    *f;
 198     gmx_bool bSpread;       /* These coordinates are used for spreading */
 199     int      pme_order;
 200     ivec    *idx;
 201     rvec    *fractx;            /* Fractional coordinate relative to the
 202                                  * lower cell boundary
 203                                  */
 204     int             nthread;
 205     int            *thread_idx; /* Which thread should spread which charge */
 206     thread_plist_t *thread_plist;
 207     splinedata_t   *spline;
 208 } pme_atomcomm_t;
 209
 210 #define FLBS  3
 211 #define FLBSZ 4
 212
 213 typedef struct {
 214     ivec  ci;     /* The spatial location of this grid         */
 215     ivec  n;      /* The used size of *grid, including order-1 */
 216     ivec  offset; /* The grid offset from the full node grid   */
 217     int   order;  /* PME spreading order                       */
 218     ivec  s;      /* The allocated size of *grid, s >= n       */
 219     real *grid;   /* The grid local thread, size n             */
 220 } pmegrid_t;
 221
 222 typedef struct {
 223     pmegrid_t  grid;         /* The full node grid (non thread-local)            */
 224     int        nthread;      /* The number of threads operating on this grid     */
 225     ivec       nc;           /* The local spatial decomposition over the threads */
 226     pmegrid_t *grid_th;      /* Array of grids for each thread                   */
 227     real      *grid_all;     /* Allocated array for the grids in *grid_th        */
 228     int      **g2t;          /* The grid to thread index                         */
 229     ivec       nthread_comm; /* The number of threads to communicate with        */
 230 } pmegrids_t;
 231
 232
 233 typedef struct {
 234 #ifdef PME_SSE_SPREAD_GATHER
 235     /* Masks for SSE aligned spreading and gathering */
 236     __m128 mask_SSE0[6], mask_SSE1[6];
 237 #else
 238     int    dummy; /* C89 requires that struct has at least one member */
 239 #endif
 240 } pme_spline_work_t;
 241
 242 typedef struct {
 243     /* work data for solve_pme */
 244     int      nalloc;
 245     real *   mhx;
 246     real *   mhy;
 247     real *   mhz;
 248     real *   m2;
 249     real *   denom;
 250     real *   tmp1_alloc;
 251     real *   tmp1;
 252     real *   eterm;
 253     real *   m2inv;
 254
 255     real     energy;
 256     matrix   vir;
 257 } pme_work_t;
 258
 259 typedef struct gmx_pme {
 260     int           ndecompdim; /* The number of decomposition dimensions */
 261     int           nodeid;     /* Our nodeid in mpi->mpi_comm */
 262     int           nodeid_major;
 263     int           nodeid_minor;
 264     int           nnodes;    /* The number of nodes doing PME */
 265     int           nnodes_major;
 266     int           nnodes_minor;
 267
 268     MPI_Comm      mpi_comm;
 269     MPI_Comm      mpi_comm_d[2]; /* Indexed on dimension, 0=x, 1=y */
 270 #ifdef GMX_MPI
 271     MPI_Datatype  rvec_mpi;      /* the pme vector's MPI type */
 272 #endif
 273
 274     gmx_bool   bUseThreads;   /* Does any of the PME ranks have nthread>1 ?  */
 275     int        nthread;       /* The number of threads doing PME on our rank */
 276
 277     gmx_bool   bPPnode;       /* Node also does particle-particle forces */
 278     gmx_bool   bFEP;          /* Compute Free energy contribution */
 279     int        nkx, nky, nkz; /* Grid dimensions */
 280     gmx_bool   bP3M;          /* Do P3M: optimize the influence function */
 281     int        pme_order;
 282     real       epsilon_r;
 283
 284     pmegrids_t pmegridA;  /* Grids on which we do spreading/interpolation, includes overlap */
 285     pmegrids_t pmegridB;
 286     /* The PME charge spreading grid sizes/strides, includes pme_order-1 */
 287     int        pmegrid_nx, pmegrid_ny, pmegrid_nz;
 288     /* pmegrid_nz might be larger than strictly necessary to ensure
 289      * memory alignment, pmegrid_nz_base gives the real base size.
 290      */
 291     int     pmegrid_nz_base;
 292     /* The local PME grid starting indices */
 293     int     pmegrid_start_ix, pmegrid_start_iy, pmegrid_start_iz;
 294
 295     /* Work data for spreading and gathering */
 296     pme_spline_work_t    *spline_work;
 297
 298     real                 *fftgridA; /* Grids for FFT. With 1D FFT decomposition this can be a pointer */
 299     real                 *fftgridB; /* inside the interpolation grid, but separate for 2D PME decomp. */
 300     int                   fftgrid_nx, fftgrid_ny, fftgrid_nz;
 301
 302     t_complex            *cfftgridA;  /* Grids for complex FFT data */
 303     t_complex            *cfftgridB;
 304     int                   cfftgrid_nx, cfftgrid_ny, cfftgrid_nz;
 305
 306     gmx_parallel_3dfft_t  pfft_setupA;
 307     gmx_parallel_3dfft_t  pfft_setupB;
 308
 309     int                  *nnx, *nny, *nnz;
 310     real                 *fshx, *fshy, *fshz;
 311
 312     pme_atomcomm_t        atc[2]; /* Indexed on decomposition index */
 313     matrix                recipbox;
 314     splinevec             bsp_mod;
 315
 316     pme_overlap_t         overlap[2]; /* Indexed on dimension, 0=x, 1=y */
 317
 318     pme_atomcomm_t        atc_energy; /* Only for gmx_pme_calc_energy */
 319
 320     rvec                 *bufv;       /* Communication buffer */
 321     real                 *bufr;       /* Communication buffer */
 322     int                   buf_nalloc; /* The communication buffer size */
 323
 324     /* thread local work data for solve_pme */
 325     pme_work_t *work;
 326
 327     /* Work data for PME_redist */
 328     gmx_bool redist_init;
 329     int *    scounts;
 330     int *    rcounts;
 331     int *    sdispls;
 332     int *    rdispls;
 333     int *    sidx;
 334     int *    idxa;
 335     real *   redist_buf;
 336     int      redist_buf_nalloc;
 337
 338     /* Work data for sum_qgrid */
 339     real *   sum_qgrid_tmp;
 340     real *   sum_qgrid_dd_tmp;
 341 } t_gmx_pme;
 342
 343
 344 static void calc_interpolation_idx(gmx_pme_t pme, pme_atomcomm_t *atc,
 345                                    int start, int end, int thread)
 346 {
 347     int             i;
 348     int            *idxptr, tix, tiy, tiz;
 349     real           *xptr, *fptr, tx, ty, tz;
 350     real            rxx, ryx, ryy, rzx, rzy, rzz;
 351     int             nx, ny, nz;
 352     int             start_ix, start_iy, start_iz;
 353     int            *g2tx, *g2ty, *g2tz;
 354     gmx_bool        bThreads;
 355     int            *thread_idx = NULL;
 356     thread_plist_t *tpl        = NULL;
 357     int            *tpl_n      = NULL;
 358     int             thread_i;
 359
 360     nx  = pme->nkx;
 361     ny  = pme->nky;
 362     nz  = pme->nkz;
 363
 364     start_ix = pme->pmegrid_start_ix;
 365     start_iy = pme->pmegrid_start_iy;
 366     start_iz = pme->pmegrid_start_iz;
 367
 368     rxx = pme->recipbox[XX][XX];
 369     ryx = pme->recipbox[YY][XX];
 370     ryy = pme->recipbox[YY][YY];
 371     rzx = pme->recipbox[ZZ][XX];
 372     rzy = pme->recipbox[ZZ][YY];
 373     rzz = pme->recipbox[ZZ][ZZ];
 374
 375     g2tx = pme->pmegridA.g2t[XX];
 376     g2ty = pme->pmegridA.g2t[YY];
 377     g2tz = pme->pmegridA.g2t[ZZ];
 378
 379     bThreads = (atc->nthread > 1);
 380     if (bThreads)
 381     {
 382         thread_idx = atc->thread_idx;
 383
 384         tpl   = &atc->thread_plist[thread];
 385         tpl_n = tpl->n;
 386         for (i = 0; i < atc->nthread; i++)
 387         {
 388             tpl_n[i] = 0;
 389         }
 390     }
 391
 392     for (i = start; i < end; i++)
 393     {
 394         xptr   = atc->x[i];
 395         idxptr = atc->idx[i];
 396         fptr   = atc->fractx[i];
 397
 398         /* Fractional coordinates along box vectors, add 2.0 to make 100% sure we are positive for triclinic boxes */
 399         tx = nx * ( xptr[XX] * rxx + xptr[YY] * ryx + xptr[ZZ] * rzx + 2.0 );
 400         ty = ny * (                  xptr[YY] * ryy + xptr[ZZ] * rzy + 2.0 );
 401         tz = nz * (                                   xptr[ZZ] * rzz + 2.0 );
 402
 403         tix = (int)(tx);
 404         tiy = (int)(ty);
 405         tiz = (int)(tz);
 406
 407         /* Because decomposition only occurs in x and y,
 408          * we never have a fraction correction in z.
 409          */
 410         fptr[XX] = tx - tix + pme->fshx[tix];
 411         fptr[YY] = ty - tiy + pme->fshy[tiy];
 412         fptr[ZZ] = tz - tiz;
 413
 414         idxptr[XX] = pme->nnx[tix];
 415         idxptr[YY] = pme->nny[tiy];
 416         idxptr[ZZ] = pme->nnz[tiz];
 417
 418 #ifdef DEBUG
 419         range_check(idxptr[XX], 0, pme->pmegrid_nx);
 420         range_check(idxptr[YY], 0, pme->pmegrid_ny);
 421         range_check(idxptr[ZZ], 0, pme->pmegrid_nz);
 422 #endif
 423
 424         if (bThreads)
 425         {
 426             thread_i      = g2tx[idxptr[XX]] + g2ty[idxptr[YY]] + g2tz[idxptr[ZZ]];
 427             thread_idx[i] = thread_i;
 428             tpl_n[thread_i]++;
 429         }
 430     }
 431
 432     if (bThreads)
 433     {
 434         /* Make a list of particle indices sorted on thread */
 435
 436         /* Get the cumulative count */
 437         for (i = 1; i < atc->nthread; i++)
 438         {
 439             tpl_n[i] += tpl_n[i-1];
 440         }
 441         /* The current implementation distributes particles equally
 442          * over the threads, so we could actually allocate for that
 443          * in pme_realloc_atomcomm_things.
 444          */
 445         if (tpl_n[atc->nthread-1] > tpl->nalloc)
 446         {
 447             tpl->nalloc = over_alloc_large(tpl_n[atc->nthread-1]);
 448             srenew(tpl->i, tpl->nalloc);
 449         }
 450         /* Set tpl_n to the cumulative start */
 451         for (i = atc->nthread-1; i >= 1; i--)
 452         {
 453             tpl_n[i] = tpl_n[i-1];
 454         }
 455         tpl_n[0] = 0;
 456
 457         /* Fill our thread local array with indices sorted on thread */
 458         for (i = start; i < end; i++)
 459         {
 460             tpl->i[tpl_n[atc->thread_idx[i]]++] = i;
 461         }
 462         /* Now tpl_n contains the cummulative count again */
 463     }
 464 }
 465
 466 static void make_thread_local_ind(pme_atomcomm_t *atc,
 467                                   int thread, splinedata_t *spline)
 468 {
 469     int             n, t, i, start, end;
 470     thread_plist_t *tpl;
 471
 472     /* Combine the indices made by each thread into one index */
 473
 474     n     = 0;
 475     start = 0;
 476     for (t = 0; t < atc->nthread; t++)
 477     {
 478         tpl = &atc->thread_plist[t];
 479         /* Copy our part (start - end) from the list of thread t */
 480         if (thread > 0)
 481         {
 482             start = tpl->n[thread-1];
 483         }
 484         end = tpl->n[thread];
 485         for (i = start; i < end; i++)
 486         {
 487             spline->ind[n++] = tpl->i[i];
 488         }
 489     }
 490
 491     spline->n = n;
 492 }
 493
 494
 495 static void pme_calc_pidx(int start, int end,
 496                           matrix recipbox, rvec x[],
 497                           pme_atomcomm_t *atc, int *count)
 498 {
 499     int   nslab, i;
 500     int   si;
 501     real *xptr, s;
 502     real  rxx, ryx, rzx, ryy, rzy;
 503     int  *pd;
 504
 505     /* Calculate PME task index (pidx) for each grid index.
 506      * Here we always assign equally sized slabs to each node
 507      * for load balancing reasons (the PME grid spacing is not used).
 508      */
 509
 510     nslab = atc->nslab;
 511     pd    = atc->pd;
 512
 513     /* Reset the count */
 514     for (i = 0; i < nslab; i++)
 515     {
 516         count[i] = 0;
 517     }
 518
 519     if (atc->dimind == 0)
 520     {
 521         rxx = recipbox[XX][XX];
 522         ryx = recipbox[YY][XX];
 523         rzx = recipbox[ZZ][XX];
 524         /* Calculate the node index in x-dimension */
 525         for (i = start; i < end; i++)
 526         {
 527             xptr   = x[i];
 528             /* Fractional coordinates along box vectors */
 529             s     = nslab*(xptr[XX]*rxx + xptr[YY]*ryx + xptr[ZZ]*rzx);
 530             si    = (int)(s + 2*nslab) % nslab;
 531             pd[i] = si;
 532             count[si]++;
 533         }
 534     }
 535     else
 536     {
 537         ryy = recipbox[YY][YY];
 538         rzy = recipbox[ZZ][YY];
 539         /* Calculate the node index in y-dimension */
 540         for (i = start; i < end; i++)
 541         {
 542             xptr   = x[i];
 543             /* Fractional coordinates along box vectors */
 544             s     = nslab*(xptr[YY]*ryy + xptr[ZZ]*rzy);
 545             si    = (int)(s + 2*nslab) % nslab;
 546             pd[i] = si;
 547             count[si]++;
 548         }
 549     }
 550 }
 551
 552 static void pme_calc_pidx_wrapper(int natoms, matrix recipbox, rvec x[],
 553                                   pme_atomcomm_t *atc)
 554 {
 555     int nthread, thread, slab;
 556
 557     nthread = atc->nthread;
 558
 559 #pragma omp parallel for num_threads(nthread) schedule(static)
 560     for (thread = 0; thread < nthread; thread++)
 561     {
 562         pme_calc_pidx(natoms* thread   /nthread,
 563                       natoms*(thread+1)/nthread,
 564                       recipbox, x, atc, atc->count_thread[thread]);
 565     }
 566     /* Non-parallel reduction, since nslab is small */
 567
 568     for (thread = 1; thread < nthread; thread++)
 569     {
 570         for (slab = 0; slab < atc->nslab; slab++)
 571         {
 572             atc->count_thread[0][slab] += atc->count_thread[thread][slab];
 573         }
 574     }
 575 }
 576
 577 static void realloc_splinevec(splinevec th, real **ptr_z, int nalloc)
 578 {
 579     const int padding = 4;
 580     int       i;
 581
 582     srenew(th[XX], nalloc);
 583     srenew(th[YY], nalloc);
 584     /* In z we add padding, this is only required for the aligned SSE code */
 585     srenew(*ptr_z, nalloc+2*padding);
 586     th[ZZ] = *ptr_z + padding;
 587
 588     for (i = 0; i < padding; i++)
 589     {
 590         (*ptr_z)[               i] = 0;
 591         (*ptr_z)[padding+nalloc+i] = 0;
 592     }
 593 }
 594
 595 static void pme_realloc_splinedata(splinedata_t *spline, pme_atomcomm_t *atc)
 596 {
 597     int i, d;
 598
 599     srenew(spline->ind, atc->nalloc);
 600     /* Initialize the index to identity so it works without threads */
 601     for (i = 0; i < atc->nalloc; i++)
 602     {
 603         spline->ind[i] = i;
 604     }
 605
 606     realloc_splinevec(spline->theta, &spline->ptr_theta_z,
 607                       atc->pme_order*atc->nalloc);
 608     realloc_splinevec(spline->dtheta, &spline->ptr_dtheta_z,
 609                       atc->pme_order*atc->nalloc);
 610 }
 611
 612 static void pme_realloc_atomcomm_things(pme_atomcomm_t *atc)
 613 {
 614     int nalloc_old, i, j, nalloc_tpl;
 615
 616     /* We have to avoid a NULL pointer for atc->x to avoid
 617      * possible fatal errors in MPI routines.
 618      */
 619     if (atc->n > atc->nalloc || atc->nalloc == 0)
 620     {
 621         nalloc_old  = atc->nalloc;
 622         atc->nalloc = over_alloc_dd(max(atc->n, 1));
 623
 624         if (atc->nslab > 1)
 625         {
 626             srenew(atc->x, atc->nalloc);
 627             srenew(atc->q, atc->nalloc);
 628             srenew(atc->f, atc->nalloc);
 629             for (i = nalloc_old; i < atc->nalloc; i++)
 630             {
 631                 clear_rvec(atc->f[i]);
 632             }
 633         }
 634         if (atc->bSpread)
 635         {
 636             srenew(atc->fractx, atc->nalloc);
 637             srenew(atc->idx, atc->nalloc);
 638
 639             if (atc->nthread > 1)
 640             {
 641                 srenew(atc->thread_idx, atc->nalloc);
 642             }
 643
 644             for (i = 0; i < atc->nthread; i++)
 645             {
 646                 pme_realloc_splinedata(&atc->spline[i], atc);
 647             }
 648         }
 649     }
 650 }
 651
 652 static void pmeredist_pd(gmx_pme_t pme, gmx_bool forw,
 653                          int n, gmx_bool bXF, rvec *x_f, real *charge,
 654                          pme_atomcomm_t *atc)
 655 /* Redistribute particle data for PME calculation */
 656 /* domain decomposition by x coordinate           */
 657 {
 658     int *idxa;
 659     int  i, ii;
 660
 661     if (FALSE == pme->redist_init)
 662     {
 663         snew(pme->scounts, atc->nslab);
 664         snew(pme->rcounts, atc->nslab);
 665         snew(pme->sdispls, atc->nslab);
 666         snew(pme->rdispls, atc->nslab);
 667         snew(pme->sidx, atc->nslab);
 668         pme->redist_init = TRUE;
 669     }
 670     if (n > pme->redist_buf_nalloc)
 671     {
 672         pme->redist_buf_nalloc = over_alloc_dd(n);
 673         srenew(pme->redist_buf, pme->redist_buf_nalloc*DIM);
 674     }
 675
 676     pme->idxa = atc->pd;
 677
 678 #ifdef GMX_MPI
 679     if (forw && bXF)
 680     {
 681         /* forward, redistribution from pp to pme */
 682
 683         /* Calculate send counts and exchange them with other nodes */
 684         for (i = 0; (i < atc->nslab); i++)
 685         {
 686             pme->scounts[i] = 0;
 687         }
 688         for (i = 0; (i < n); i++)
 689         {
 690             pme->scounts[pme->idxa[i]]++;
 691         }
 692         MPI_Alltoall( pme->scounts, 1, MPI_INT, pme->rcounts, 1, MPI_INT, atc->mpi_comm);
 693
 694         /* Calculate send and receive displacements and index into send
 695            buffer */
 696         pme->sdispls[0] = 0;
 697         pme->rdispls[0] = 0;
 698         pme->sidx[0]    = 0;
 699         for (i = 1; i < atc->nslab; i++)
 700         {
 701             pme->sdispls[i] = pme->sdispls[i-1]+pme->scounts[i-1];
 702             pme->rdispls[i] = pme->rdispls[i-1]+pme->rcounts[i-1];
 703             pme->sidx[i]    = pme->sdispls[i];
 704         }
 705         /* Total # of particles to be received */
 706         atc->n = pme->rdispls[atc->nslab-1] + pme->rcounts[atc->nslab-1];
 707
 708         pme_realloc_atomcomm_things(atc);
 709
 710         /* Copy particle coordinates into send buffer and exchange*/
 711         for (i = 0; (i < n); i++)
 712         {
 713             ii = DIM*pme->sidx[pme->idxa[i]];
 714             pme->sidx[pme->idxa[i]]++;
 715             pme->redist_buf[ii+XX] = x_f[i][XX];
 716             pme->redist_buf[ii+YY] = x_f[i][YY];
 717             pme->redist_buf[ii+ZZ] = x_f[i][ZZ];
 718         }
 719         MPI_Alltoallv(pme->redist_buf, pme->scounts, pme->sdispls,
 720                       pme->rvec_mpi, atc->x, pme->rcounts, pme->rdispls,
 721                       pme->rvec_mpi, atc->mpi_comm);
 722     }
 723     if (forw)
 724     {
 725         /* Copy charge into send buffer and exchange*/
 726         for (i = 0; i < atc->nslab; i++)
 727         {
 728             pme->sidx[i] = pme->sdispls[i];
 729         }
 730         for (i = 0; (i < n); i++)
 731         {
 732             ii = pme->sidx[pme->idxa[i]];
 733             pme->sidx[pme->idxa[i]]++;
 734             pme->redist_buf[ii] = charge[i];
 735         }
 736         MPI_Alltoallv(pme->redist_buf, pme->scounts, pme->sdispls, mpi_type,
 737                       atc->q, pme->rcounts, pme->rdispls, mpi_type,
 738                       atc->mpi_comm);
 739     }
 740     else   /* backward, redistribution from pme to pp */
 741     {
 742         MPI_Alltoallv(atc->f, pme->rcounts, pme->rdispls, pme->rvec_mpi,
 743                       pme->redist_buf, pme->scounts, pme->sdispls,
 744                       pme->rvec_mpi, atc->mpi_comm);
 745
 746         /* Copy data from receive buffer */
 747         for (i = 0; i < atc->nslab; i++)
 748         {
 749             pme->sidx[i] = pme->sdispls[i];
 750         }
 751         for (i = 0; (i < n); i++)
 752         {
 753             ii          = DIM*pme->sidx[pme->idxa[i]];
 754             x_f[i][XX] += pme->redist_buf[ii+XX];
 755             x_f[i][YY] += pme->redist_buf[ii+YY];
 756             x_f[i][ZZ] += pme->redist_buf[ii+ZZ];
 757             pme->sidx[pme->idxa[i]]++;
 758         }
 759     }
 760 #endif
 761 }
 762
 763 static void pme_dd_sendrecv(pme_atomcomm_t *atc,
 764                             gmx_bool bBackward, int shift,
 765                             void *buf_s, int nbyte_s,
 766                             void *buf_r, int nbyte_r)
 767 {
 768 #ifdef GMX_MPI
 769     int        dest, src;
 770     MPI_Status stat;
 771
 772     if (bBackward == FALSE)
 773     {
 774         dest = atc->node_dest[shift];
 775         src  = atc->node_src[shift];
 776     }
 777     else
 778     {
 779         dest = atc->node_src[shift];
 780         src  = atc->node_dest[shift];
 781     }
 782
 783     if (nbyte_s > 0 && nbyte_r > 0)
 784     {
 785         MPI_Sendrecv(buf_s, nbyte_s, MPI_BYTE,
 786                      dest, shift,
 787                      buf_r, nbyte_r, MPI_BYTE,
 788                      src, shift,
 789                      atc->mpi_comm, &stat);
 790     }
 791     else if (nbyte_s > 0)
 792     {
 793         MPI_Send(buf_s, nbyte_s, MPI_BYTE,
 794                  dest, shift,
 795                  atc->mpi_comm);
 796     }
 797     else if (nbyte_r > 0)
 798     {
 799         MPI_Recv(buf_r, nbyte_r, MPI_BYTE,
 800                  src, shift,
 801                  atc->mpi_comm, &stat);
 802     }
 803 #endif
 804 }
 805
 806 static void dd_pmeredist_x_q(gmx_pme_t pme,
 807                              int n, gmx_bool bX, rvec *x, real *charge,
 808                              pme_atomcomm_t *atc)
 809 {
 810     int *commnode, *buf_index;
 811     int  nnodes_comm, i, nsend, local_pos, buf_pos, node, scount, rcount;
 812
 813     commnode  = atc->node_dest;
 814     buf_index = atc->buf_index;
 815
 816     nnodes_comm = min(2*atc->maxshift, atc->nslab-1);
 817
 818     nsend = 0;
 819     for (i = 0; i < nnodes_comm; i++)
 820     {
 821         buf_index[commnode[i]] = nsend;
 822         nsend                 += atc->count[commnode[i]];
 823     }
 824     if (bX)
 825     {
 826         if (atc->count[atc->nodeid] + nsend != n)
 827         {
 828             gmx_fatal(FARGS, "%d particles communicated to PME node %d are more than 2/3 times the cut-off out of the domain decomposition cell of their charge group in dimension %c.\n"
 829                       "This usually means that your system is not well equilibrated.",
 830                       n - (atc->count[atc->nodeid] + nsend),
 831                       pme->nodeid, 'x'+atc->dimind);
 832         }
 833
 834         if (nsend > pme->buf_nalloc)
 835         {
 836             pme->buf_nalloc = over_alloc_dd(nsend);
 837             srenew(pme->bufv, pme->buf_nalloc);
 838             srenew(pme->bufr, pme->buf_nalloc);
 839         }
 840
 841         atc->n = atc->count[atc->nodeid];
 842         for (i = 0; i < nnodes_comm; i++)
 843         {
 844             scount = atc->count[commnode[i]];
 845             /* Communicate the count */
 846             if (debug)
 847             {
 848                 fprintf(debug, "dimind %d PME node %d send to node %d: %d\n",
 849                         atc->dimind, atc->nodeid, commnode[i], scount);
 850             }
 851             pme_dd_sendrecv(atc, FALSE, i,
 852                             &scount, sizeof(int),
 853                             &atc->rcount[i], sizeof(int));
 854             atc->n += atc->rcount[i];
 855         }
 856
 857         pme_realloc_atomcomm_things(atc);
 858     }
 859
 860     local_pos = 0;
 861     for (i = 0; i < n; i++)
 862     {
 863         node = atc->pd[i];
 864         if (node == atc->nodeid)
 865         {
 866             /* Copy direct to the receive buffer */
 867             if (bX)
 868             {
 869                 copy_rvec(x[i], atc->x[local_pos]);
 870             }
 871             atc->q[local_pos] = charge[i];
 872             local_pos++;
 873         }
 874         else
 875         {
 876             /* Copy to the send buffer */
 877             if (bX)
 878             {
 879                 copy_rvec(x[i], pme->bufv[buf_index[node]]);
 880             }
 881             pme->bufr[buf_index[node]] = charge[i];
 882             buf_index[node]++;
 883         }
 884     }
 885
 886     buf_pos = 0;
 887     for (i = 0; i < nnodes_comm; i++)
 888     {
 889         scount = atc->count[commnode[i]];
 890         rcount = atc->rcount[i];
 891         if (scount > 0 || rcount > 0)
 892         {
 893             if (bX)
 894             {
 895                 /* Communicate the coordinates */
 896                 pme_dd_sendrecv(atc, FALSE, i,
 897                                 pme->bufv[buf_pos], scount*sizeof(rvec),
 898                                 atc->x[local_pos], rcount*sizeof(rvec));
 899             }
 900             /* Communicate the charges */
 901             pme_dd_sendrecv(atc, FALSE, i,
 902                             pme->bufr+buf_pos, scount*sizeof(real),
 903                             atc->q+local_pos, rcount*sizeof(real));
 904             buf_pos   += scount;
 905             local_pos += atc->rcount[i];
 906         }
 907     }
 908 }
 909
 910 static void dd_pmeredist_f(gmx_pme_t pme, pme_atomcomm_t *atc,
 911                            int n, rvec *f,
 912                            gmx_bool bAddF)
 913 {
 914     int *commnode, *buf_index;
 915     int  nnodes_comm, local_pos, buf_pos, i, scount, rcount, node;
 916
 917     commnode  = atc->node_dest;
 918     buf_index = atc->buf_index;
 919
 920     nnodes_comm = min(2*atc->maxshift, atc->nslab-1);
 921
 922     local_pos = atc->count[atc->nodeid];
 923     buf_pos   = 0;
 924     for (i = 0; i < nnodes_comm; i++)
 925     {
 926         scount = atc->rcount[i];
 927         rcount = atc->count[commnode[i]];
 928         if (scount > 0 || rcount > 0)
 929         {
 930             /* Communicate the forces */
 931             pme_dd_sendrecv(atc, TRUE, i,
 932                             atc->f[local_pos], scount*sizeof(rvec),
 933                             pme->bufv[buf_pos], rcount*sizeof(rvec));
 934             local_pos += scount;
 935         }
 936         buf_index[commnode[i]] = buf_pos;
 937         buf_pos               += rcount;
 938     }
 939
 940     local_pos = 0;
 941     if (bAddF)
 942     {
 943         for (i = 0; i < n; i++)
 944         {
 945             node = atc->pd[i];
 946             if (node == atc->nodeid)
 947             {
 948                 /* Add from the local force array */
 949                 rvec_inc(f[i], atc->f[local_pos]);
 950                 local_pos++;
 951             }
 952             else
 953             {
 954                 /* Add from the receive buffer */
 955                 rvec_inc(f[i], pme->bufv[buf_index[node]]);
 956                 buf_index[node]++;
 957             }
 958         }
 959     }
 960     else
 961     {
 962         for (i = 0; i < n; i++)
 963         {
 964             node = atc->pd[i];
 965             if (node == atc->nodeid)
 966             {
 967                 /* Copy from the local force array */
 968                 copy_rvec(atc->f[local_pos], f[i]);
 969                 local_pos++;
 970             }
 971             else
 972             {
 973                 /* Copy from the receive buffer */
 974                 copy_rvec(pme->bufv[buf_index[node]], f[i]);
 975                 buf_index[node]++;
 976             }
 977         }
 978     }
 979 }
 980
 981 #ifdef GMX_MPI
 982 static void
 983 gmx_sum_qgrid_dd(gmx_pme_t pme, real *grid, int direction)
 984 {
 985     pme_overlap_t *overlap;
 986     int            send_index0, send_nindex;
 987     int            recv_index0, recv_nindex;
 988     MPI_Status     stat;
 989     int            i, j, k, ix, iy, iz, icnt;
 990     int            ipulse, send_id, recv_id, datasize;
 991     real          *p;
 992     real          *sendptr, *recvptr;
 993
 994     /* Start with minor-rank communication. This is a bit of a pain since it is not contiguous */
 995     overlap = &pme->overlap[1];
 996
 997     for (ipulse = 0; ipulse < overlap->noverlap_nodes; ipulse++)
 998     {
 999         /* Since we have already (un)wrapped the overlap in the z-dimension,
1000          * we only have to communicate 0 to nkz (not pmegrid_nz).
1001          */
1002         if (direction == GMX_SUM_QGRID_FORWARD)
1003         {
1004             send_id       = overlap->send_id[ipulse];
1005             recv_id       = overlap->recv_id[ipulse];
1006             send_index0   = overlap->comm_data[ipulse].send_index0;
1007             send_nindex   = overlap->comm_data[ipulse].send_nindex;
1008             recv_index0   = overlap->comm_data[ipulse].recv_index0;
1009             recv_nindex   = overlap->comm_data[ipulse].recv_nindex;
1010         }
1011         else
1012         {
1013             send_id       = overlap->recv_id[ipulse];
1014             recv_id       = overlap->send_id[ipulse];
1015             send_index0   = overlap->comm_data[ipulse].recv_index0;
1016             send_nindex   = overlap->comm_data[ipulse].recv_nindex;
1017             recv_index0   = overlap->comm_data[ipulse].send_index0;
1018             recv_nindex   = overlap->comm_data[ipulse].send_nindex;
1019         }
1020
1021         /* Copy data to contiguous send buffer */
1022         if (debug)
1023         {
1024             fprintf(debug, "PME send node %d %d -> %d grid start %d Communicating %d to %d\n",
1025                     pme->nodeid, overlap->nodeid, send_id,
1026                     pme->pmegrid_start_iy,
1027                     send_index0-pme->pmegrid_start_iy,
1028                     send_index0-pme->pmegrid_start_iy+send_nindex);
1029         }
1030         icnt = 0;
1031         for (i = 0; i < pme->pmegrid_nx; i++)
1032         {
1033             ix = i;
1034             for (j = 0; j < send_nindex; j++)
1035             {
1036                 iy = j + send_index0 - pme->pmegrid_start_iy;
1037                 for (k = 0; k < pme->nkz; k++)
1038                 {
1039                     iz = k;
1040                     overlap->sendbuf[icnt++] = grid[ix*(pme->pmegrid_ny*pme->pmegrid_nz)+iy*(pme->pmegrid_nz)+iz];
1041                 }
1042             }
1043         }
1044
1045         datasize      = pme->pmegrid_nx * pme->nkz;
1046
1047         MPI_Sendrecv(overlap->sendbuf, send_nindex*datasize, GMX_MPI_REAL,
1048                      send_id, ipulse,
1049                      overlap->recvbuf, recv_nindex*datasize, GMX_MPI_REAL,
1050                      recv_id, ipulse,
1051                      overlap->mpi_comm, &stat);
1052
1053         /* Get data from contiguous recv buffer */
1054         if (debug)
1055         {
1056             fprintf(debug, "PME recv node %d %d <- %d grid start %d Communicating %d to %d\n",
1057                     pme->nodeid, overlap->nodeid, recv_id,
1058                     pme->pmegrid_start_iy,
1059                     recv_index0-pme->pmegrid_start_iy,
1060                     recv_index0-pme->pmegrid_start_iy+recv_nindex);
1061         }
1062         icnt = 0;
1063         for (i = 0; i < pme->pmegrid_nx; i++)
1064         {
1065             ix = i;
1066             for (j = 0; j < recv_nindex; j++)
1067             {
1068                 iy = j + recv_index0 - pme->pmegrid_start_iy;
1069                 for (k = 0; k < pme->nkz; k++)
1070                 {
1071                     iz = k;
1072                     if (direction == GMX_SUM_QGRID_FORWARD)
1073                     {
1074                         grid[ix*(pme->pmegrid_ny*pme->pmegrid_nz)+iy*(pme->pmegrid_nz)+iz] += overlap->recvbuf[icnt++];
1075                     }
1076                     else
1077                     {
1078                         grid[ix*(pme->pmegrid_ny*pme->pmegrid_nz)+iy*(pme->pmegrid_nz)+iz]  = overlap->recvbuf[icnt++];
1079                     }
1080                 }
1081             }
1082         }
1083     }
1084
1085     /* Major dimension is easier, no copying required,
1086      * but we might have to sum to separate array.
1087      * Since we don't copy, we have to communicate up to pmegrid_nz,
1088      * not nkz as for the minor direction.
1089      */
1090     overlap = &pme->overlap[0];
1091
1092     for (ipulse = 0; ipulse < overlap->noverlap_nodes; ipulse++)
1093     {
1094         if (direction == GMX_SUM_QGRID_FORWARD)
1095         {
1096             send_id       = overlap->send_id[ipulse];
1097             recv_id       = overlap->recv_id[ipulse];
1098             send_index0   = overlap->comm_data[ipulse].send_index0;
1099             send_nindex   = overlap->comm_data[ipulse].send_nindex;
1100             recv_index0   = overlap->comm_data[ipulse].recv_index0;
1101             recv_nindex   = overlap->comm_data[ipulse].recv_nindex;
1102             recvptr       = overlap->recvbuf;
1103         }
1104         else
1105         {
1106             send_id       = overlap->recv_id[ipulse];
1107             recv_id       = overlap->send_id[ipulse];
1108             send_index0   = overlap->comm_data[ipulse].recv_index0;
1109             send_nindex   = overlap->comm_data[ipulse].recv_nindex;
1110             recv_index0   = overlap->comm_data[ipulse].send_index0;
1111             recv_nindex   = overlap->comm_data[ipulse].send_nindex;
1112             recvptr       = grid + (recv_index0-pme->pmegrid_start_ix)*(pme->pmegrid_ny*pme->pmegrid_nz);
1113         }
1114
1115         sendptr       = grid + (send_index0-pme->pmegrid_start_ix)*(pme->pmegrid_ny*pme->pmegrid_nz);
1116         datasize      = pme->pmegrid_ny * pme->pmegrid_nz;
1117
1118         if (debug)
1119         {
1120             fprintf(debug, "PME send node %d %d -> %d grid start %d Communicating %d to %d\n",
1121                     pme->nodeid, overlap->nodeid, send_id,
1122                     pme->pmegrid_start_ix,
1123                     send_index0-pme->pmegrid_start_ix,
1124                     send_index0-pme->pmegrid_start_ix+send_nindex);
1125             fprintf(debug, "PME recv node %d %d <- %d grid start %d Communicating %d to %d\n",
1126                     pme->nodeid, overlap->nodeid, recv_id,
1127                     pme->pmegrid_start_ix,
1128                     recv_index0-pme->pmegrid_start_ix,
1129                     recv_index0-pme->pmegrid_start_ix+recv_nindex);
1130         }
1131
1132         MPI_Sendrecv(sendptr, send_nindex*datasize, GMX_MPI_REAL,
1133                      send_id, ipulse,
1134                      recvptr, recv_nindex*datasize, GMX_MPI_REAL,
1135                      recv_id, ipulse,
1136                      overlap->mpi_comm, &stat);
1137
1138         /* ADD data from contiguous recv buffer */
1139         if (direction == GMX_SUM_QGRID_FORWARD)
1140         {
1141             p = grid + (recv_index0-pme->pmegrid_start_ix)*(pme->pmegrid_ny*pme->pmegrid_nz);
1142             for (i = 0; i < recv_nindex*datasize; i++)
1143             {
1144                 p[i] += overlap->recvbuf[i];
1145             }
1146         }
1147     }
1148 }
1149 #endif
1150
1151
1152 static int
1153 copy_pmegrid_to_fftgrid(gmx_pme_t pme, real *pmegrid, real *fftgrid)
1154 {
1155     ivec    local_fft_ndata, local_fft_offset, local_fft_size;
1156     ivec    local_pme_size;
1157     int     i, ix, iy, iz;
1158     int     pmeidx, fftidx;
1159
1160     /* Dimensions should be identical for A/B grid, so we just use A here */
1161     gmx_parallel_3dfft_real_limits(pme->pfft_setupA,
1162                                    local_fft_ndata,
1163                                    local_fft_offset,
1164                                    local_fft_size);
1165
1166     local_pme_size[0] = pme->pmegrid_nx;
1167     local_pme_size[1] = pme->pmegrid_ny;
1168     local_pme_size[2] = pme->pmegrid_nz;
1169
1170     /* The fftgrid is always 'justified' to the lower-left corner of the PME grid,
1171        the offset is identical, and the PME grid always has more data (due to overlap)
1172      */
1173     {
1174 #ifdef DEBUG_PME
1175         FILE *fp, *fp2;
1176         char  fn[STRLEN], format[STRLEN];
1177         real  val;
1178         sprintf(fn, "pmegrid%d.pdb", pme->nodeid);
1179         fp = ffopen(fn, "w");
1180         sprintf(fn, "pmegrid%d.txt", pme->nodeid);
1181         fp2 = ffopen(fn, "w");
1182         sprintf(format, "%s%s\n", pdbformat, "%6.2f%6.2f");
1183 #endif
1184
1185         for (ix = 0; ix < local_fft_ndata[XX]; ix++)
1186         {
1187             for (iy = 0; iy < local_fft_ndata[YY]; iy++)
1188             {
1189                 for (iz = 0; iz < local_fft_ndata[ZZ]; iz++)
1190                 {
1191                     pmeidx          = ix*(local_pme_size[YY]*local_pme_size[ZZ])+iy*(local_pme_size[ZZ])+iz;
1192                     fftidx          = ix*(local_fft_size[YY]*local_fft_size[ZZ])+iy*(local_fft_size[ZZ])+iz;
1193                     fftgrid[fftidx] = pmegrid[pmeidx];
1194 #ifdef DEBUG_PME
1195                     val = 100*pmegrid[pmeidx];
1196                     if (pmegrid[pmeidx] != 0)
1197                     {
1198                         fprintf(fp, format, "ATOM", pmeidx, "CA", "GLY", ' ', pmeidx, ' ',
1199                                 5.0*ix, 5.0*iy, 5.0*iz, 1.0, val);
1200                     }
1201                     if (pmegrid[pmeidx] != 0)
1202                     {
1203                         fprintf(fp2, "%-12s  %5d  %5d  %5d  %12.5e\n",
1204                                 "qgrid",
1205                                 pme->pmegrid_start_ix + ix,
1206                                 pme->pmegrid_start_iy + iy,
1207                                 pme->pmegrid_start_iz + iz,
1208                                 pmegrid[pmeidx]);
1209                     }
1210 #endif
1211                 }
1212             }
1213         }
1214 #ifdef DEBUG_PME
1215         ffclose(fp);
1216         ffclose(fp2);
1217 #endif
1218     }
1219     return 0;
1220 }
1221
1222
1223 static gmx_cycles_t omp_cyc_start()
1224 {
1225     return gmx_cycles_read();
1226 }
1227
1228 static gmx_cycles_t omp_cyc_end(gmx_cycles_t c)
1229 {
1230     return gmx_cycles_read() - c;
1231 }
1232
1233
1234 static int
1235 copy_fftgrid_to_pmegrid(gmx_pme_t pme, const real *fftgrid, real *pmegrid,
1236                         int nthread, int thread)
1237 {
1238     ivec          local_fft_ndata, local_fft_offset, local_fft_size;
1239     ivec          local_pme_size;
1240     int           ixy0, ixy1, ixy, ix, iy, iz;
1241     int           pmeidx, fftidx;
1242 #ifdef PME_TIME_THREADS
1243     gmx_cycles_t  c1;
1244     static double cs1 = 0;
1245     static int    cnt = 0;
1246 #endif
1247
1248 #ifdef PME_TIME_THREADS
1249     c1 = omp_cyc_start();
1250 #endif
1251     /* Dimensions should be identical for A/B grid, so we just use A here */
1252     gmx_parallel_3dfft_real_limits(pme->pfft_setupA,
1253                                    local_fft_ndata,
1254                                    local_fft_offset,
1255                                    local_fft_size);
1256
1257     local_pme_size[0] = pme->pmegrid_nx;
1258     local_pme_size[1] = pme->pmegrid_ny;
1259     local_pme_size[2] = pme->pmegrid_nz;
1260
1261     /* The fftgrid is always 'justified' to the lower-left corner of the PME grid,
1262        the offset is identical, and the PME grid always has more data (due to overlap)
1263      */
1264     ixy0 = ((thread  )*local_fft_ndata[XX]*local_fft_ndata[YY])/nthread;
1265     ixy1 = ((thread+1)*local_fft_ndata[XX]*local_fft_ndata[YY])/nthread;
1266
1267     for (ixy = ixy0; ixy < ixy1; ixy++)
1268     {
1269         ix = ixy/local_fft_ndata[YY];
1270         iy = ixy - ix*local_fft_ndata[YY];
1271
1272         pmeidx = (ix*local_pme_size[YY] + iy)*local_pme_size[ZZ];
1273         fftidx = (ix*local_fft_size[YY] + iy)*local_fft_size[ZZ];
1274         for (iz = 0; iz < local_fft_ndata[ZZ]; iz++)
1275         {
1276             pmegrid[pmeidx+iz] = fftgrid[fftidx+iz];
1277         }
1278     }
1279
1280 #ifdef PME_TIME_THREADS
1281     c1   = omp_cyc_end(c1);
1282     cs1 += (double)c1;
1283     cnt++;
1284     if (cnt % 20 == 0)
1285     {
1286         printf("copy %.2f\n", cs1*1e-9);
1287     }
1288 #endif
1289
1290     return 0;
1291 }
1292
1293
1294 static void
1295 wrap_periodic_pmegrid(gmx_pme_t pme, real *pmegrid)
1296 {
1297     int     nx, ny, nz, pnx, pny, pnz, ny_x, overlap, ix, iy, iz;
1298
1299     nx = pme->nkx;
1300     ny = pme->nky;
1301     nz = pme->nkz;
1302
1303     pnx = pme->pmegrid_nx;
1304     pny = pme->pmegrid_ny;
1305     pnz = pme->pmegrid_nz;
1306
1307     overlap = pme->pme_order - 1;
1308
1309     /* Add periodic overlap in z */
1310     for (ix = 0; ix < pme->pmegrid_nx; ix++)
1311     {
1312         for (iy = 0; iy < pme->pmegrid_ny; iy++)
1313         {
1314             for (iz = 0; iz < overlap; iz++)
1315             {
1316                 pmegrid[(ix*pny+iy)*pnz+iz] +=
1317                     pmegrid[(ix*pny+iy)*pnz+nz+iz];
1318             }
1319         }
1320     }
1321
1322     if (pme->nnodes_minor == 1)
1323     {
1324         for (ix = 0; ix < pme->pmegrid_nx; ix++)
1325         {
1326             for (iy = 0; iy < overlap; iy++)
1327             {
1328                 for (iz = 0; iz < nz; iz++)
1329                 {
1330                     pmegrid[(ix*pny+iy)*pnz+iz] +=
1331                         pmegrid[(ix*pny+ny+iy)*pnz+iz];
1332                 }
1333             }
1334         }
1335     }
1336
1337     if (pme->nnodes_major == 1)
1338     {
1339         ny_x = (pme->nnodes_minor == 1 ? ny : pme->pmegrid_ny);
1340
1341         for (ix = 0; ix < overlap; ix++)
1342         {
1343             for (iy = 0; iy < ny_x; iy++)
1344             {
1345                 for (iz = 0; iz < nz; iz++)
1346                 {
1347                     pmegrid[(ix*pny+iy)*pnz+iz] +=
1348                         pmegrid[((nx+ix)*pny+iy)*pnz+iz];
1349                 }
1350             }
1351         }
1352     }
1353 }
1354
1355
1356 static void
1357 unwrap_periodic_pmegrid(gmx_pme_t pme, real *pmegrid)
1358 {
1359     int     nx, ny, nz, pnx, pny, pnz, ny_x, overlap, ix;
1360
1361     nx = pme->nkx;
1362     ny = pme->nky;
1363     nz = pme->nkz;
1364
1365     pnx = pme->pmegrid_nx;
1366     pny = pme->pmegrid_ny;
1367     pnz = pme->pmegrid_nz;
1368
1369     overlap = pme->pme_order - 1;
1370
1371     if (pme->nnodes_major == 1)
1372     {
1373         ny_x = (pme->nnodes_minor == 1 ? ny : pme->pmegrid_ny);
1374
1375         for (ix = 0; ix < overlap; ix++)
1376         {
1377             int iy, iz;
1378
1379             for (iy = 0; iy < ny_x; iy++)
1380             {
1381                 for (iz = 0; iz < nz; iz++)
1382                 {
1383                     pmegrid[((nx+ix)*pny+iy)*pnz+iz] =
1384                         pmegrid[(ix*pny+iy)*pnz+iz];
1385                 }
1386             }
1387         }
1388     }
1389
1390     if (pme->nnodes_minor == 1)
1391     {
1392 #pragma omp parallel for num_threads(pme->nthread) schedule(static)
1393         for (ix = 0; ix < pme->pmegrid_nx; ix++)
1394         {
1395             int iy, iz;
1396
1397             for (iy = 0; iy < overlap; iy++)
1398             {
1399                 for (iz = 0; iz < nz; iz++)
1400                 {
1401                     pmegrid[(ix*pny+ny+iy)*pnz+iz] =
1402                         pmegrid[(ix*pny+iy)*pnz+iz];
1403                 }
1404             }
1405         }
1406     }
1407
1408     /* Copy periodic overlap in z */
1409 #pragma omp parallel for num_threads(pme->nthread) schedule(static)
1410     for (ix = 0; ix < pme->pmegrid_nx; ix++)
1411     {
1412         int iy, iz;
1413
1414         for (iy = 0; iy < pme->pmegrid_ny; iy++)
1415         {
1416             for (iz = 0; iz < overlap; iz++)
1417             {
1418                 pmegrid[(ix*pny+iy)*pnz+nz+iz] =
1419                     pmegrid[(ix*pny+iy)*pnz+iz];
1420             }
1421         }
1422     }
1423 }
1424
1425
1426 /* This has to be a macro to enable full compiler optimization with xlC (and probably others too) */
1427 #define DO_BSPLINE(order)                            \
1428     for (ithx = 0; (ithx < order); ithx++)                    \
1429     {                                                    \
1430         index_x = (i0+ithx)*pny*pnz;                     \
1431         valx    = qn*thx[ithx];                          \
1432                                                      \
1433         for (ithy = 0; (ithy < order); ithy++)                \
1434         {                                                \
1435             valxy    = valx*thy[ithy];                   \
1436             index_xy = index_x+(j0+ithy)*pnz;            \
1437                                                      \
1438             for (ithz = 0; (ithz < order); ithz++)            \
1439             {                                            \
1440                 index_xyz        = index_xy+(k0+ithz);   \
1441                 grid[index_xyz] += valxy*thz[ithz];      \
1442             }                                            \
1443         }                                                \
1444     }
1445
1446
1447 static void spread_q_bsplines_thread(pmegrid_t *pmegrid,
1448                                      pme_atomcomm_t *atc, splinedata_t *spline,
1449                                      pme_spline_work_t *work)
1450 {
1451
1452     /* spread charges from home atoms to local grid */
1453     real          *grid;
1454     pme_overlap_t *ol;
1455     int            b, i, nn, n, ithx, ithy, ithz, i0, j0, k0;
1456     int       *    idxptr;
1457     int            order, norder, index_x, index_xy, index_xyz;
1458     real           valx, valxy, qn;
1459     real          *thx, *thy, *thz;
1460     int            localsize, bndsize;
1461     int            pnx, pny, pnz, ndatatot;
1462     int            offx, offy, offz;
1463
1464     pnx = pmegrid->s[XX];
1465     pny = pmegrid->s[YY];
1466     pnz = pmegrid->s[ZZ];
1467
1468     offx = pmegrid->offset[XX];
1469     offy = pmegrid->offset[YY];
1470     offz = pmegrid->offset[ZZ];
1471
1472     ndatatot = pnx*pny*pnz;
1473     grid     = pmegrid->grid;
1474     for (i = 0; i < ndatatot; i++)
1475     {
1476         grid[i] = 0;
1477     }
1478
1479     order = pmegrid->order;
1480
1481     for (nn = 0; nn < spline->n; nn++)
1482     {
1483         n  = spline->ind[nn];
1484         qn = atc->q[n];
1485
1486         if (qn != 0)
1487         {
1488             idxptr = atc->idx[n];
1489             norder = nn*order;
1490
1491             i0   = idxptr[XX] - offx;
1492             j0   = idxptr[YY] - offy;
1493             k0   = idxptr[ZZ] - offz;
1494
1495             thx = spline->theta[XX] + norder;
1496             thy = spline->theta[YY] + norder;
1497             thz = spline->theta[ZZ] + norder;
1498
1499             switch (order)
1500             {
1501                 case 4:
1502 #ifdef PME_SSE_SPREAD_GATHER
1503 #ifdef PME_SSE_UNALIGNED
1504 #define PME_SPREAD_SSE_ORDER4
1505 #else
1506 #define PME_SPREAD_SSE_ALIGNED
1507 #define PME_ORDER 4
1508 #endif
1509 #include "pme_sse_single.h"
1510 #else
1511                     DO_BSPLINE(4);
1512 #endif
1513                     break;
1514                 case 5:
1515 #ifdef PME_SSE_SPREAD_GATHER
1516 #define PME_SPREAD_SSE_ALIGNED
1517 #define PME_ORDER 5
1518 #include "pme_sse_single.h"
1519 #else
1520                     DO_BSPLINE(5);
1521 #endif
1522                     break;
1523                 default:
1524                     DO_BSPLINE(order);
1525                     break;
1526             }
1527         }
1528     }
1529 }
1530
1531 static void set_grid_alignment(int *pmegrid_nz, int pme_order)
1532 {
1533 #ifdef PME_SSE_SPREAD_GATHER
1534     if (pme_order == 5
1535 #ifndef PME_SSE_UNALIGNED
1536         || pme_order == 4
1537 #endif
1538         )
1539     {
1540         /* Round nz up to a multiple of 4 to ensure alignment */
1541         *pmegrid_nz = ((*pmegrid_nz + 3) & ~3);
1542     }
1543 #endif
1544 }
1545
1546 static void set_gridsize_alignment(int gmx_unused *gridsize, int gmx_unused pme_order)
1547 {
1548 #ifdef PME_SSE_SPREAD_GATHER
1549 #ifndef PME_SSE_UNALIGNED
1550     if (pme_order == 4)
1551     {
1552         /* Add extra elements to ensured aligned operations do not go
1553          * beyond the allocated grid size.
1554          * Note that for pme_order=5, the pme grid z-size alignment
1555          * ensures that we will not go beyond the grid size.
1556          */
1557         *gridsize += 4;
1558     }
1559 #endif
1560 #endif
1561 }
1562
1563 static void pmegrid_init(pmegrid_t *grid,
1564                          int cx, int cy, int cz,
1565                          int x0, int y0, int z0,
1566                          int x1, int y1, int z1,
1567                          gmx_bool set_alignment,
1568                          int pme_order,
1569                          real *ptr)
1570 {
1571     int nz, gridsize;
1572
1573     grid->ci[XX]     = cx;
1574     grid->ci[YY]     = cy;
1575     grid->ci[ZZ]     = cz;
1576     grid->offset[XX] = x0;
1577     grid->offset[YY] = y0;
1578     grid->offset[ZZ] = z0;
1579     grid->n[XX]      = x1 - x0 + pme_order - 1;
1580     grid->n[YY]      = y1 - y0 + pme_order - 1;
1581     grid->n[ZZ]      = z1 - z0 + pme_order - 1;
1582     copy_ivec(grid->n, grid->s);
1583
1584     nz = grid->s[ZZ];
1585     set_grid_alignment(&nz, pme_order);
1586     if (set_alignment)
1587     {
1588         grid->s[ZZ] = nz;
1589     }
1590     else if (nz != grid->s[ZZ])
1591     {
1592         gmx_incons("pmegrid_init call with an unaligned z size");
1593     }
1594
1595     grid->order = pme_order;
1596     if (ptr == NULL)
1597     {
1598         gridsize = grid->s[XX]*grid->s[YY]*grid->s[ZZ];
1599         set_gridsize_alignment(&gridsize, pme_order);
1600         snew_aligned(grid->grid, gridsize, 16);
1601     }
1602     else
1603     {
1604         grid->grid = ptr;
1605     }
1606 }
1607
1608 static int div_round_up(int enumerator, int denominator)
1609 {
1610     return (enumerator + denominator - 1)/denominator;
1611 }
1612
1613 static void make_subgrid_division(const ivec n, int ovl, int nthread,
1614                                   ivec nsub)
1615 {
1616     int gsize_opt, gsize;
1617     int nsx, nsy, nsz;
1618     char *env;
1619
1620     gsize_opt = -1;
1621     for (nsx = 1; nsx <= nthread; nsx++)
1622     {
1623         if (nthread % nsx == 0)
1624         {
1625             for (nsy = 1; nsy <= nthread; nsy++)
1626             {
1627                 if (nsx*nsy <= nthread && nthread % (nsx*nsy) == 0)
1628                 {
1629                     nsz = nthread/(nsx*nsy);
1630
1631                     /* Determine the number of grid points per thread */
1632                     gsize =
1633                         (div_round_up(n[XX], nsx) + ovl)*
1634                         (div_round_up(n[YY], nsy) + ovl)*
1635                         (div_round_up(n[ZZ], nsz) + ovl);
1636
1637                     /* Minimize the number of grids points per thread
1638                      * and, secondarily, the number of cuts in minor dimensions.
1639                      */
1640                     if (gsize_opt == -1 ||
1641                         gsize < gsize_opt ||
1642                         (gsize == gsize_opt &&
1643                          (nsz < nsub[ZZ] || (nsz == nsub[ZZ] && nsy < nsub[YY]))))
1644                     {
1645                         nsub[XX]  = nsx;
1646                         nsub[YY]  = nsy;
1647                         nsub[ZZ]  = nsz;
1648                         gsize_opt = gsize;
1649                     }
1650                 }
1651             }
1652         }
1653     }
1654
1655     env = getenv("GMX_PME_THREAD_DIVISION");
1656     if (env != NULL)
1657     {
1658         sscanf(env, "%d %d %d", &nsub[XX], &nsub[YY], &nsub[ZZ]);
1659     }
1660
1661     if (nsub[XX]*nsub[YY]*nsub[ZZ] != nthread)
1662     {
1663         gmx_fatal(FARGS, "PME grid thread division (%d x %d x %d) does not match the total number of threads (%d)", nsub[XX], nsub[YY], nsub[ZZ], nthread);
1664     }
1665 }
1666
1667 static void pmegrids_init(pmegrids_t *grids,
1668                           int nx, int ny, int nz, int nz_base,
1669                           int pme_order,
1670                           gmx_bool bUseThreads,
1671                           int nthread,
1672                           int overlap_x,
1673                           int overlap_y)
1674 {
1675     ivec n, n_base, g0, g1;
1676     int t, x, y, z, d, i, tfac;
1677     int max_comm_lines = -1;
1678
1679     n[XX] = nx - (pme_order - 1);
1680     n[YY] = ny - (pme_order - 1);
1681     n[ZZ] = nz - (pme_order - 1);
1682
1683     copy_ivec(n, n_base);
1684     n_base[ZZ] = nz_base;
1685
1686     pmegrid_init(&grids->grid, 0, 0, 0, 0, 0, 0, n[XX], n[YY], n[ZZ], FALSE, pme_order,
1687                  NULL);
1688
1689     grids->nthread = nthread;
1690
1691     make_subgrid_division(n_base, pme_order-1, grids->nthread, grids->nc);
1692
1693     if (bUseThreads)
1694     {
1695         ivec nst;
1696         int gridsize;
1697
1698         for (d = 0; d < DIM; d++)
1699         {
1700             nst[d] = div_round_up(n[d], grids->nc[d]) + pme_order - 1;
1701         }
1702         set_grid_alignment(&nst[ZZ], pme_order);
1703
1704         if (debug)
1705         {
1706             fprintf(debug, "pmegrid thread local division: %d x %d x %d\n",
1707                     grids->nc[XX], grids->nc[YY], grids->nc[ZZ]);
1708             fprintf(debug, "pmegrid %d %d %d max thread pmegrid %d %d %d\n",
1709                     nx, ny, nz,
1710                     nst[XX], nst[YY], nst[ZZ]);
1711         }
1712
1713         snew(grids->grid_th, grids->nthread);
1714         t        = 0;
1715         gridsize = nst[XX]*nst[YY]*nst[ZZ];
1716         set_gridsize_alignment(&gridsize, pme_order);
1717         snew_aligned(grids->grid_all,
1718                      grids->nthread*gridsize+(grids->nthread+1)*GMX_CACHE_SEP,
1719                      16);
1720
1721         for (x = 0; x < grids->nc[XX]; x++)
1722         {
1723             for (y = 0; y < grids->nc[YY]; y++)
1724             {
1725                 for (z = 0; z < grids->nc[ZZ]; z++)
1726                 {
1727                     pmegrid_init(&grids->grid_th[t],
1728                                  x, y, z,
1729                                  (n[XX]*(x  ))/grids->nc[XX],
1730                                  (n[YY]*(y  ))/grids->nc[YY],
1731                                  (n[ZZ]*(z  ))/grids->nc[ZZ],
1732                                  (n[XX]*(x+1))/grids->nc[XX],
1733                                  (n[YY]*(y+1))/grids->nc[YY],
1734                                  (n[ZZ]*(z+1))/grids->nc[ZZ],
1735                                  TRUE,
1736                                  pme_order,
1737                                  grids->grid_all+GMX_CACHE_SEP+t*(gridsize+GMX_CACHE_SEP));
1738                     t++;
1739                 }
1740             }
1741         }
1742     }
1743     else
1744     {
1745         grids->grid_th = NULL;
1746     }
1747
1748     snew(grids->g2t, DIM);
1749     tfac = 1;
1750     for (d = DIM-1; d >= 0; d--)
1751     {
1752         snew(grids->g2t[d], n[d]);
1753         t = 0;
1754         for (i = 0; i < n[d]; i++)
1755         {
1756             /* The second check should match the parameters
1757              * of the pmegrid_init call above.
1758              */
1759             while (t + 1 < grids->nc[d] && i >= (n[d]*(t+1))/grids->nc[d])
1760             {
1761                 t++;
1762             }
1763             grids->g2t[d][i] = t*tfac;
1764         }
1765
1766         tfac *= grids->nc[d];
1767
1768         switch (d)
1769         {
1770             case XX: max_comm_lines = overlap_x;     break;
1771             case YY: max_comm_lines = overlap_y;     break;
1772             case ZZ: max_comm_lines = pme_order - 1; break;
1773         }
1774         grids->nthread_comm[d] = 0;
1775         while ((n[d]*grids->nthread_comm[d])/grids->nc[d] < max_comm_lines &&
1776                grids->nthread_comm[d] < grids->nc[d])
1777         {
1778             grids->nthread_comm[d]++;
1779         }
1780         if (debug != NULL)
1781         {
1782             fprintf(debug, "pmegrid thread grid communication range in %c: %d\n",
1783                     'x'+d, grids->nthread_comm[d]);
1784         }
1785         /* It should be possible to make grids->nthread_comm[d]==grids->nc[d]
1786          * work, but this is not a problematic restriction.
1787          */
1788         if (grids->nc[d] > 1 && grids->nthread_comm[d] > grids->nc[d])
1789         {
1790             gmx_fatal(FARGS, "Too many threads for PME (%d) compared to the number of grid lines, reduce the number of threads doing PME", grids->nthread);
1791         }
1792     }
1793 }
1794
1795
1796 static void pmegrids_destroy(pmegrids_t *grids)
1797 {
1798     int t;
1799
1800     if (grids->grid.grid != NULL)
1801     {
1802         sfree(grids->grid.grid);
1803
1804         if (grids->nthread > 0)
1805         {
1806             for (t = 0; t < grids->nthread; t++)
1807             {
1808                 sfree(grids->grid_th[t].grid);
1809             }
1810             sfree(grids->grid_th);
1811         }
1812     }
1813 }
1814
1815
1816 static void realloc_work(pme_work_t *work, int nkx)
1817 {
1818     if (nkx > work->nalloc)
1819     {
1820         work->nalloc = nkx;
1821         srenew(work->mhx, work->nalloc);
1822         srenew(work->mhy, work->nalloc);
1823         srenew(work->mhz, work->nalloc);
1824         srenew(work->m2, work->nalloc);
1825         /* Allocate an aligned pointer for SIMD operations, including extra
1826          * elements at the end for padding.
1827          */
1828 #ifdef PME_SIMD
1829 #define ALIGN_HERE  GMX_SIMD_WIDTH_HERE
1830 #else
1831 /* We can use any alignment, apart from 0, so we use 4 */
1832 #define ALIGN_HERE  4
1833 #endif
1834         sfree_aligned(work->denom);
1835         sfree_aligned(work->tmp1);
1836         sfree_aligned(work->eterm);
1837         snew_aligned(work->denom, work->nalloc+ALIGN_HERE, ALIGN_HERE*sizeof(real));
1838         snew_aligned(work->tmp1,  work->nalloc+ALIGN_HERE, ALIGN_HERE*sizeof(real));
1839         snew_aligned(work->eterm, work->nalloc+ALIGN_HERE, ALIGN_HERE*sizeof(real));
1840         srenew(work->m2inv, work->nalloc);
1841     }
1842 }
1843
1844
1845 static void free_work(pme_work_t *work)
1846 {
1847     sfree(work->mhx);
1848     sfree(work->mhy);
1849     sfree(work->mhz);
1850     sfree(work->m2);
1851     sfree_aligned(work->denom);
1852     sfree_aligned(work->tmp1);
1853     sfree_aligned(work->eterm);
1854     sfree(work->m2inv);
1855 }
1856
1857
1858 #ifdef PME_SIMD
1859 /* Calculate exponentials through SIMD */
1860 inline static void calc_exponentials(int start, int end, real f, real *d_aligned, real *r_aligned, real *e_aligned)
1861 {
1862     {
1863         const gmx_mm_pr two = gmx_set1_pr(2.0);
1864         gmx_mm_pr f_simd;
1865         gmx_mm_pr lu;
1866         gmx_mm_pr tmp_d1, d_inv, tmp_r, tmp_e;
1867         int kx;
1868         f_simd = gmx_load1_pr(&f);
1869         for (kx = 0; kx < end; kx += GMX_SIMD_WIDTH_HERE)
1870         {
1871             tmp_d1   = gmx_load_pr(d_aligned+kx);
1872             d_inv    = gmx_inv_pr(tmp_d1);
1873             tmp_r    = gmx_load_pr(r_aligned+kx);
1874             tmp_r    = gmx_exp_pr(tmp_r);
1875             tmp_e    = gmx_mul_pr(f_simd, d_inv);
1876             tmp_e    = gmx_mul_pr(tmp_e, tmp_r);
1877             gmx_store_pr(e_aligned+kx, tmp_e);
1878         }
1879     }
1880 }
1881 #else
1882 inline static void calc_exponentials(int start, int end, real f, real *d, real *r, real *e)
1883 {
1884     int kx;
1885     for (kx = start; kx < end; kx++)
1886     {
1887         d[kx] = 1.0/d[kx];
1888     }
1889     for (kx = start; kx < end; kx++)
1890     {
1891         r[kx] = exp(r[kx]);
1892     }
1893     for (kx = start; kx < end; kx++)
1894     {
1895         e[kx] = f*r[kx]*d[kx];
1896     }
1897 }
1898 #endif
1899
1900
1901 static int solve_pme_yzx(gmx_pme_t pme, t_complex *grid,
1902                          real ewaldcoeff, real vol,
1903                          gmx_bool bEnerVir,
1904                          int nthread, int thread)
1905 {
1906     /* do recip sum over local cells in grid */
1907     /* y major, z middle, x minor or continuous */
1908     t_complex *p0;
1909     int     kx, ky, kz, maxkx, maxky, maxkz;
1910     int     nx, ny, nz, iyz0, iyz1, iyz, iy, iz, kxstart, kxend;
1911     real    mx, my, mz;
1912     real    factor = M_PI*M_PI/(ewaldcoeff*ewaldcoeff);
1913     real    ets2, struct2, vfactor, ets2vf;
1914     real    d1, d2, energy = 0;
1915     real    by, bz;
1916     real    virxx = 0, virxy = 0, virxz = 0, viryy = 0, viryz = 0, virzz = 0;
1917     real    rxx, ryx, ryy, rzx, rzy, rzz;
1918     pme_work_t *work;
1919     real    *mhx, *mhy, *mhz, *m2, *denom, *tmp1, *eterm, *m2inv;
1920     real    mhxk, mhyk, mhzk, m2k;
1921     real    corner_fac;
1922     ivec    complex_order;
1923     ivec    local_ndata, local_offset, local_size;
1924     real    elfac;
1925
1926     elfac = ONE_4PI_EPS0/pme->epsilon_r;
1927
1928     nx = pme->nkx;
1929     ny = pme->nky;
1930     nz = pme->nkz;
1931
1932     /* Dimensions should be identical for A/B grid, so we just use A here */
1933     gmx_parallel_3dfft_complex_limits(pme->pfft_setupA,
1934                                       complex_order,
1935                                       local_ndata,
1936                                       local_offset,
1937                                       local_size);
1938
1939     rxx = pme->recipbox[XX][XX];
1940     ryx = pme->recipbox[YY][XX];
1941     ryy = pme->recipbox[YY][YY];
1942     rzx = pme->recipbox[ZZ][XX];
1943     rzy = pme->recipbox[ZZ][YY];
1944     rzz = pme->recipbox[ZZ][ZZ];
1945
1946     maxkx = (nx+1)/2;
1947     maxky = (ny+1)/2;
1948     maxkz = nz/2+1;
1949
1950     work  = &pme->work[thread];
1951     mhx   = work->mhx;
1952     mhy   = work->mhy;
1953     mhz   = work->mhz;
1954     m2    = work->m2;
1955     denom = work->denom;
1956     tmp1  = work->tmp1;
1957     eterm = work->eterm;
1958     m2inv = work->m2inv;
1959
1960     iyz0 = local_ndata[YY]*local_ndata[ZZ]* thread   /nthread;
1961     iyz1 = local_ndata[YY]*local_ndata[ZZ]*(thread+1)/nthread;
1962
1963     for (iyz = iyz0; iyz < iyz1; iyz++)
1964     {
1965         iy = iyz/local_ndata[ZZ];
1966         iz = iyz - iy*local_ndata[ZZ];
1967
1968         ky = iy + local_offset[YY];
1969
1970         if (ky < maxky)
1971         {
1972             my = ky;
1973         }
1974         else
1975         {
1976             my = (ky - ny);
1977         }
1978
1979         by = M_PI*vol*pme->bsp_mod[YY][ky];
1980
1981         kz = iz + local_offset[ZZ];
1982
1983         mz = kz;
1984
1985         bz = pme->bsp_mod[ZZ][kz];
1986
1987         /* 0.5 correction for corner points */
1988         corner_fac = 1;
1989         if (kz == 0 || kz == (nz+1)/2)
1990         {
1991             corner_fac = 0.5;
1992         }
1993
1994         p0 = grid + iy*local_size[ZZ]*local_size[XX] + iz*local_size[XX];
1995
1996         /* We should skip the k-space point (0,0,0) */
1997         if (local_offset[XX] > 0 || ky > 0 || kz > 0)
1998         {
1999             kxstart = local_offset[XX];
2000         }
2001         else
2002         {
2003             kxstart = local_offset[XX] + 1;
2004             p0++;
2005         }
2006         kxend = local_offset[XX] + local_ndata[XX];
2007
2008         if (bEnerVir)
2009         {
2010             /* More expensive inner loop, especially because of the storage
2011              * of the mh elements in array's.
2012              * Because x is the minor grid index, all mh elements
2013              * depend on kx for triclinic unit cells.
2014              */
2015
2016             /* Two explicit loops to avoid a conditional inside the loop */
2017             for (kx = kxstart; kx < maxkx; kx++)
2018             {
2019                 mx = kx;
2020
2021                 mhxk      = mx * rxx;
2022                 mhyk      = mx * ryx + my * ryy;
2023                 mhzk      = mx * rzx + my * rzy + mz * rzz;
2024                 m2k       = mhxk*mhxk + mhyk*mhyk + mhzk*mhzk;
2025                 mhx[kx]   = mhxk;
2026                 mhy[kx]   = mhyk;
2027                 mhz[kx]   = mhzk;
2028                 m2[kx]    = m2k;
2029                 denom[kx] = m2k*bz*by*pme->bsp_mod[XX][kx];
2030                 tmp1[kx]  = -factor*m2k;
2031             }
2032
2033             for (kx = maxkx; kx < kxend; kx++)
2034             {
2035                 mx = (kx - nx);
2036
2037                 mhxk      = mx * rxx;
2038                 mhyk      = mx * ryx + my * ryy;
2039                 mhzk      = mx * rzx + my * rzy + mz * rzz;
2040                 m2k       = mhxk*mhxk + mhyk*mhyk + mhzk*mhzk;
2041                 mhx[kx]   = mhxk;
2042                 mhy[kx]   = mhyk;
2043                 mhz[kx]   = mhzk;
2044                 m2[kx]    = m2k;
2045                 denom[kx] = m2k*bz*by*pme->bsp_mod[XX][kx];
2046                 tmp1[kx]  = -factor*m2k;
2047             }
2048
2049             for (kx = kxstart; kx < kxend; kx++)
2050             {
2051                 m2inv[kx] = 1.0/m2[kx];
2052             }
2053
2054             calc_exponentials(kxstart, kxend, elfac, denom, tmp1, eterm);
2055
2056             for (kx = kxstart; kx < kxend; kx++, p0++)
2057             {
2058                 d1      = p0->re;
2059                 d2      = p0->im;
2060
2061                 p0->re  = d1*eterm[kx];
2062                 p0->im  = d2*eterm[kx];
2063
2064                 struct2 = 2.0*(d1*d1+d2*d2);
2065
2066                 tmp1[kx] = eterm[kx]*struct2;
2067             }
2068
2069             for (kx = kxstart; kx < kxend; kx++)
2070             {
2071                 ets2     = corner_fac*tmp1[kx];
2072                 vfactor  = (factor*m2[kx] + 1.0)*2.0*m2inv[kx];
2073                 energy  += ets2;
2074
2075                 ets2vf   = ets2*vfactor;
2076                 virxx   += ets2vf*mhx[kx]*mhx[kx] - ets2;
2077                 virxy   += ets2vf*mhx[kx]*mhy[kx];
2078                 virxz   += ets2vf*mhx[kx]*mhz[kx];
2079                 viryy   += ets2vf*mhy[kx]*mhy[kx] - ets2;
2080                 viryz   += ets2vf*mhy[kx]*mhz[kx];
2081                 virzz   += ets2vf*mhz[kx]*mhz[kx] - ets2;
2082             }
2083         }
2084         else
2085         {
2086             /* We don't need to calculate the energy and the virial.
2087              * In this case the triclinic overhead is small.
2088              */
2089
2090             /* Two explicit loops to avoid a conditional inside the loop */
2091
2092             for (kx = kxstart; kx < maxkx; kx++)
2093             {
2094                 mx = kx;
2095
2096                 mhxk      = mx * rxx;
2097                 mhyk      = mx * ryx + my * ryy;
2098                 mhzk      = mx * rzx + my * rzy + mz * rzz;
2099                 m2k       = mhxk*mhxk + mhyk*mhyk + mhzk*mhzk;
2100                 denom[kx] = m2k*bz*by*pme->bsp_mod[XX][kx];
2101                 tmp1[kx]  = -factor*m2k;
2102             }
2103
2104             for (kx = maxkx; kx < kxend; kx++)
2105             {
2106                 mx = (kx - nx);
2107
2108                 mhxk      = mx * rxx;
2109                 mhyk      = mx * ryx + my * ryy;
2110                 mhzk      = mx * rzx + my * rzy + mz * rzz;
2111                 m2k       = mhxk*mhxk + mhyk*mhyk + mhzk*mhzk;
2112                 denom[kx] = m2k*bz*by*pme->bsp_mod[XX][kx];
2113                 tmp1[kx]  = -factor*m2k;
2114             }
2115
2116             calc_exponentials(kxstart, kxend, elfac, denom, tmp1, eterm);
2117
2118             for (kx = kxstart; kx < kxend; kx++, p0++)
2119             {
2120                 d1      = p0->re;
2121                 d2      = p0->im;
2122
2123                 p0->re  = d1*eterm[kx];
2124                 p0->im  = d2*eterm[kx];
2125             }
2126         }
2127     }
2128
2129     if (bEnerVir)
2130     {
2131         /* Update virial with local values.
2132          * The virial is symmetric by definition.
2133          * this virial seems ok for isotropic scaling, but I'm
2134          * experiencing problems on semiisotropic membranes.
2135          * IS THAT COMMENT STILL VALID??? (DvdS, 2001/02/07).
2136          */
2137         work->vir[XX][XX] = 0.25*virxx;
2138         work->vir[YY][YY] = 0.25*viryy;
2139         work->vir[ZZ][ZZ] = 0.25*virzz;
2140         work->vir[XX][YY] = work->vir[YY][XX] = 0.25*virxy;
2141         work->vir[XX][ZZ] = work->vir[ZZ][XX] = 0.25*virxz;
2142         work->vir[YY][ZZ] = work->vir[ZZ][YY] = 0.25*viryz;
2143
2144         /* This energy should be corrected for a charged system */
2145         work->energy = 0.5*energy;
2146     }
2147
2148     /* Return the loop count */
2149     return local_ndata[YY]*local_ndata[XX];
2150 }
2151
2152 static void get_pme_ener_vir(const gmx_pme_t pme, int nthread,
2153                              real *mesh_energy, matrix vir)
2154 {
2155     /* This function sums output over threads
2156      * and should therefore only be called after thread synchronization.
2157      */
2158     int thread;
2159
2160     *mesh_energy = pme->work[0].energy;
2161     copy_mat(pme->work[0].vir, vir);
2162
2163     for (thread = 1; thread < nthread; thread++)
2164     {
2165         *mesh_energy += pme->work[thread].energy;
2166         m_add(vir, pme->work[thread].vir, vir);
2167     }
2168 }
2169
2170 #define DO_FSPLINE(order)                      \
2171     for (ithx = 0; (ithx < order); ithx++)              \
2172     {                                              \
2173         index_x = (i0+ithx)*pny*pnz;               \
2174         tx      = thx[ithx];                       \
2175         dx      = dthx[ithx];                      \
2176                                                \
2177         for (ithy = 0; (ithy < order); ithy++)          \
2178         {                                          \
2179             index_xy = index_x+(j0+ithy)*pnz;      \
2180             ty       = thy[ithy];                  \
2181             dy       = dthy[ithy];                 \
2182             fxy1     = fz1 = 0;                    \
2183                                                \
2184             for (ithz = 0; (ithz < order); ithz++)      \
2185             {                                      \
2186                 gval  = grid[index_xy+(k0+ithz)];  \
2187                 fxy1 += thz[ithz]*gval;            \
2188                 fz1  += dthz[ithz]*gval;           \
2189             }                                      \
2190             fx += dx*ty*fxy1;                      \
2191             fy += tx*dy*fxy1;                      \
2192             fz += tx*ty*fz1;                       \
2193         }                                          \
2194     }
2195
2196
2197 static void gather_f_bsplines(gmx_pme_t pme, real *grid,
2198                               gmx_bool bClearF, pme_atomcomm_t *atc,
2199                               splinedata_t *spline,
2200                               real scale)
2201 {
2202     /* sum forces for local particles */
2203     int     nn, n, ithx, ithy, ithz, i0, j0, k0;
2204     int     index_x, index_xy;
2205     int     nx, ny, nz, pnx, pny, pnz;
2206     int *   idxptr;
2207     real    tx, ty, dx, dy, qn;
2208     real    fx, fy, fz, gval;
2209     real    fxy1, fz1;
2210     real    *thx, *thy, *thz, *dthx, *dthy, *dthz;
2211     int     norder;
2212     real    rxx, ryx, ryy, rzx, rzy, rzz;
2213     int     order;
2214
2215     pme_spline_work_t *work;
2216
2217     work = pme->spline_work;
2218
2219     order = pme->pme_order;
2220     thx   = spline->theta[XX];
2221     thy   = spline->theta[YY];
2222     thz   = spline->theta[ZZ];
2223     dthx  = spline->dtheta[XX];
2224     dthy  = spline->dtheta[YY];
2225     dthz  = spline->dtheta[ZZ];
2226     nx    = pme->nkx;
2227     ny    = pme->nky;
2228     nz    = pme->nkz;
2229     pnx   = pme->pmegrid_nx;
2230     pny   = pme->pmegrid_ny;
2231     pnz   = pme->pmegrid_nz;
2232
2233     rxx   = pme->recipbox[XX][XX];
2234     ryx   = pme->recipbox[YY][XX];
2235     ryy   = pme->recipbox[YY][YY];
2236     rzx   = pme->recipbox[ZZ][XX];
2237     rzy   = pme->recipbox[ZZ][YY];
2238     rzz   = pme->recipbox[ZZ][ZZ];
2239
2240     for (nn = 0; nn < spline->n; nn++)
2241     {
2242         n  = spline->ind[nn];
2243         qn = scale*atc->q[n];
2244
2245         if (bClearF)
2246         {
2247             atc->f[n][XX] = 0;
2248             atc->f[n][YY] = 0;
2249             atc->f[n][ZZ] = 0;
2250         }
2251         if (qn != 0)
2252         {
2253             fx     = 0;
2254             fy     = 0;
2255             fz     = 0;
2256             idxptr = atc->idx[n];
2257             norder = nn*order;
2258
2259             i0   = idxptr[XX];
2260             j0   = idxptr[YY];
2261             k0   = idxptr[ZZ];
2262
2263             /* Pointer arithmetic alert, next six statements */
2264             thx  = spline->theta[XX] + norder;
2265             thy  = spline->theta[YY] + norder;
2266             thz  = spline->theta[ZZ] + norder;
2267             dthx = spline->dtheta[XX] + norder;
2268             dthy = spline->dtheta[YY] + norder;
2269             dthz = spline->dtheta[ZZ] + norder;
2270
2271             switch (order)
2272             {
2273                 case 4:
2274 #ifdef PME_SSE_SPREAD_GATHER
2275 #ifdef PME_SSE_UNALIGNED
2276 #define PME_GATHER_F_SSE_ORDER4
2277 #else
2278 #define PME_GATHER_F_SSE_ALIGNED
2279 #define PME_ORDER 4
2280 #endif
2281 #include "pme_sse_single.h"
2282 #else
2283                     DO_FSPLINE(4);
2284 #endif
2285                     break;
2286                 case 5:
2287 #ifdef PME_SSE_SPREAD_GATHER
2288 #define PME_GATHER_F_SSE_ALIGNED
2289 #define PME_ORDER 5
2290 #include "pme_sse_single.h"
2291 #else
2292                     DO_FSPLINE(5);
2293 #endif
2294                     break;
2295                 default:
2296                     DO_FSPLINE(order);
2297                     break;
2298             }
2299
2300             atc->f[n][XX] += -qn*( fx*nx*rxx );
2301             atc->f[n][YY] += -qn*( fx*nx*ryx + fy*ny*ryy );
2302             atc->f[n][ZZ] += -qn*( fx*nx*rzx + fy*ny*rzy + fz*nz*rzz );
2303         }
2304     }
2305     /* Since the energy and not forces are interpolated
2306      * the net force might not be exactly zero.
2307      * This can be solved by also interpolating F, but
2308      * that comes at a cost.
2309      * A better hack is to remove the net force every
2310      * step, but that must be done at a higher level
2311      * since this routine doesn't see all atoms if running
2312      * in parallel. Don't know how important it is?  EL 990726
2313      */
2314 }
2315
2316
2317 static real gather_energy_bsplines(gmx_pme_t pme, real *grid,
2318                                    pme_atomcomm_t *atc)
2319 {
2320     splinedata_t *spline;
2321     int     n, ithx, ithy, ithz, i0, j0, k0;
2322     int     index_x, index_xy;
2323     int *   idxptr;
2324     real    energy, pot, tx, ty, qn, gval;
2325     real    *thx, *thy, *thz;
2326     int     norder;
2327     int     order;
2328
2329     spline = &atc->spline[0];
2330
2331     order = pme->pme_order;
2332
2333     energy = 0;
2334     for (n = 0; (n < atc->n); n++)
2335     {
2336         qn      = atc->q[n];
2337
2338         if (qn != 0)
2339         {
2340             idxptr = atc->idx[n];
2341             norder = n*order;
2342
2343             i0   = idxptr[XX];
2344             j0   = idxptr[YY];
2345             k0   = idxptr[ZZ];
2346
2347             /* Pointer arithmetic alert, next three statements */
2348             thx  = spline->theta[XX] + norder;
2349             thy  = spline->theta[YY] + norder;
2350             thz  = spline->theta[ZZ] + norder;
2351
2352             pot = 0;
2353             for (ithx = 0; (ithx < order); ithx++)
2354             {
2355                 index_x = (i0+ithx)*pme->pmegrid_ny*pme->pmegrid_nz;
2356                 tx      = thx[ithx];
2357
2358                 for (ithy = 0; (ithy < order); ithy++)
2359                 {
2360                     index_xy = index_x+(j0+ithy)*pme->pmegrid_nz;
2361                     ty       = thy[ithy];
2362
2363                     for (ithz = 0; (ithz < order); ithz++)
2364                     {
2365                         gval  = grid[index_xy+(k0+ithz)];
2366                         pot  += tx*ty*thz[ithz]*gval;
2367                     }
2368
2369                 }
2370             }
2371
2372             energy += pot*qn;
2373         }
2374     }
2375
2376     return energy;
2377 }
2378
2379 /* Macro to force loop unrolling by fixing order.
2380  * This gives a significant performance gain.
2381  */
2382 #define CALC_SPLINE(order)                     \
2383     {                                              \
2384         int j, k, l;                                 \
2385         real dr, div;                               \
2386         real data[PME_ORDER_MAX];                  \
2387         real ddata[PME_ORDER_MAX];                 \
2388                                                \
2389         for (j = 0; (j < DIM); j++)                     \
2390         {                                          \
2391             dr  = xptr[j];                         \
2392                                                \
2393             /* dr is relative offset from lower cell limit */ \
2394             data[order-1] = 0;                     \
2395             data[1]       = dr;                          \
2396             data[0]       = 1 - dr;                      \
2397                                                \
2398             for (k = 3; (k < order); k++)               \
2399             {                                      \
2400                 div       = 1.0/(k - 1.0);               \
2401                 data[k-1] = div*dr*data[k-2];      \
2402                 for (l = 1; (l < (k-1)); l++)           \
2403                 {                                  \
2404                     data[k-l-1] = div*((dr+l)*data[k-l-2]+(k-l-dr)* \
2405                                        data[k-l-1]);                \
2406                 }                                  \
2407                 data[0] = div*(1-dr)*data[0];      \
2408             }                                      \
2409             /* differentiate */                    \
2410             ddata[0] = -data[0];                   \
2411             for (k = 1; (k < order); k++)               \
2412             {                                      \
2413                 ddata[k] = data[k-1] - data[k];    \
2414             }                                      \
2415                                                \
2416             div           = 1.0/(order - 1);                 \
2417             data[order-1] = div*dr*data[order-2];  \
2418             for (l = 1; (l < (order-1)); l++)           \
2419             {                                      \
2420                 data[order-l-1] = div*((dr+l)*data[order-l-2]+    \
2421                                        (order-l-dr)*data[order-l-1]); \
2422             }                                      \
2423             data[0] = div*(1 - dr)*data[0];        \
2424                                                \
2425             for (k = 0; k < order; k++)                 \
2426             {                                      \
2427                 theta[j][i*order+k]  = data[k];    \
2428                 dtheta[j][i*order+k] = ddata[k];   \
2429             }                                      \
2430         }                                          \
2431     }
2432
2433 void make_bsplines(splinevec theta, splinevec dtheta, int order,
2434                    rvec fractx[], int nr, int ind[], real charge[],
2435                    gmx_bool bFreeEnergy)
2436 {
2437     /* construct splines for local atoms */
2438     int  i, ii;
2439     real *xptr;
2440
2441     for (i = 0; i < nr; i++)
2442     {
2443         /* With free energy we do not use the charge check.
2444          * In most cases this will be more efficient than calling make_bsplines
2445          * twice, since usually more than half the particles have charges.
2446          */
2447         ii = ind[i];
2448         if (bFreeEnergy || charge[ii] != 0.0)
2449         {
2450             xptr = fractx[ii];
2451             switch (order)
2452             {
2453                 case 4:  CALC_SPLINE(4);     break;
2454                 case 5:  CALC_SPLINE(5);     break;
2455                 default: CALC_SPLINE(order); break;
2456             }
2457         }
2458     }
2459 }
2460
2461
2462 void make_dft_mod(real *mod, real *data, int ndata)
2463 {
2464     int i, j;
2465     real sc, ss, arg;
2466
2467     for (i = 0; i < ndata; i++)
2468     {
2469         sc = ss = 0;
2470         for (j = 0; j < ndata; j++)
2471         {
2472             arg = (2.0*M_PI*i*j)/ndata;
2473             sc += data[j]*cos(arg);
2474             ss += data[j]*sin(arg);
2475         }
2476         mod[i] = sc*sc+ss*ss;
2477     }
2478     for (i = 0; i < ndata; i++)
2479     {
2480         if (mod[i] < 1e-7)
2481         {
2482             mod[i] = (mod[i-1]+mod[i+1])*0.5;
2483         }
2484     }
2485 }
2486
2487
2488 static void make_bspline_moduli(splinevec bsp_mod,
2489                                 int nx, int ny, int nz, int order)
2490 {
2491     int nmax = max(nx, max(ny, nz));
2492     real *data, *ddata, *bsp_data;
2493     int i, k, l;
2494     real div;
2495
2496     snew(data, order);
2497     snew(ddata, order);
2498     snew(bsp_data, nmax);
2499
2500     data[order-1] = 0;
2501     data[1]       = 0;
2502     data[0]       = 1;
2503
2504     for (k = 3; k < order; k++)
2505     {
2506         div       = 1.0/(k-1.0);
2507         data[k-1] = 0;
2508         for (l = 1; l < (k-1); l++)
2509         {
2510             data[k-l-1] = div*(l*data[k-l-2]+(k-l)*data[k-l-1]);
2511         }
2512         data[0] = div*data[0];
2513     }
2514     /* differentiate */
2515     ddata[0] = -data[0];
2516     for (k = 1; k < order; k++)
2517     {
2518         ddata[k] = data[k-1]-data[k];
2519     }
2520     div           = 1.0/(order-1);
2521     data[order-1] = 0;
2522     for (l = 1; l < (order-1); l++)
2523     {
2524         data[order-l-1] = div*(l*data[order-l-2]+(order-l)*data[order-l-1]);
2525     }
2526     data[0] = div*data[0];
2527
2528     for (i = 0; i < nmax; i++)
2529     {
2530         bsp_data[i] = 0;
2531     }
2532     for (i = 1; i <= order; i++)
2533     {
2534         bsp_data[i] = data[i-1];
2535     }
2536
2537     make_dft_mod(bsp_mod[XX], bsp_data, nx);
2538     make_dft_mod(bsp_mod[YY], bsp_data, ny);
2539     make_dft_mod(bsp_mod[ZZ], bsp_data, nz);
2540
2541     sfree(data);
2542     sfree(ddata);
2543     sfree(bsp_data);
2544 }
2545
2546
2547 /* Return the P3M optimal influence function */
2548 static double do_p3m_influence(double z, int order)
2549 {
2550     double z2, z4;
2551
2552     z2 = z*z;
2553     z4 = z2*z2;
2554
2555     /* The formula and most constants can be found in:
2556      * Ballenegger et al., JCTC 8, 936 (2012)
2557      */
2558     switch (order)
2559     {
2560         case 2:
2561             return 1.0 - 2.0*z2/3.0;
2562             break;
2563         case 3:
2564             return 1.0 - z2 + 2.0*z4/15.0;
2565             break;
2566         case 4:
2567             return 1.0 - 4.0*z2/3.0 + 2.0*z4/5.0 + 4.0*z2*z4/315.0;
2568             break;
2569         case 5:
2570             return 1.0 - 5.0*z2/3.0 + 7.0*z4/9.0 - 17.0*z2*z4/189.0 + 2.0*z4*z4/2835.0;
2571             break;
2572         case 6:
2573             return 1.0 - 2.0*z2 + 19.0*z4/15.0 - 256.0*z2*z4/945.0 + 62.0*z4*z4/4725.0 + 4.0*z2*z4*z4/155925.0;
2574             break;
2575         case 7:
2576             return 1.0 - 7.0*z2/3.0 + 28.0*z4/15.0 - 16.0*z2*z4/27.0 + 26.0*z4*z4/405.0 - 2.0*z2*z4*z4/1485.0 + 4.0*z4*z4*z4/6081075.0;
2577         case 8:
2578             return 1.0 - 8.0*z2/3.0 + 116.0*z4/45.0 - 344.0*z2*z4/315.0 + 914.0*z4*z4/4725.0 - 248.0*z4*z4*z2/22275.0 + 21844.0*z4*z4*z4/212837625.0 - 8.0*z4*z4*z4*z2/638512875.0;
2579             break;
2580     }
2581
2582     return 0.0;
2583 }
2584
2585 /* Calculate the P3M B-spline moduli for one dimension */
2586 static void make_p3m_bspline_moduli_dim(real *bsp_mod, int n, int order)
2587 {
2588     double zarg, zai, sinzai, infl;
2589     int    maxk, i;
2590
2591     if (order > 8)
2592     {
2593         gmx_fatal(FARGS, "The current P3M code only supports orders up to 8");
2594     }
2595
2596     zarg = M_PI/n;
2597
2598     maxk = (n + 1)/2;
2599
2600     for (i = -maxk; i < 0; i++)
2601     {
2602         zai          = zarg*i;
2603         sinzai       = sin(zai);
2604         infl         = do_p3m_influence(sinzai, order);
2605         bsp_mod[n+i] = infl*infl*pow(sinzai/zai, -2.0*order);
2606     }
2607     bsp_mod[0] = 1.0;
2608     for (i = 1; i < maxk; i++)
2609     {
2610         zai        = zarg*i;
2611         sinzai     = sin(zai);
2612         infl       = do_p3m_influence(sinzai, order);
2613         bsp_mod[i] = infl*infl*pow(sinzai/zai, -2.0*order);
2614     }
2615 }
2616
2617 /* Calculate the P3M B-spline moduli */
2618 static void make_p3m_bspline_moduli(splinevec bsp_mod,
2619                                     int nx, int ny, int nz, int order)
2620 {
2621     make_p3m_bspline_moduli_dim(bsp_mod[XX], nx, order);
2622     make_p3m_bspline_moduli_dim(bsp_mod[YY], ny, order);
2623     make_p3m_bspline_moduli_dim(bsp_mod[ZZ], nz, order);
2624 }
2625
2626
2627 static void setup_coordinate_communication(pme_atomcomm_t *atc)
2628 {
2629     int nslab, n, i;
2630     int fw, bw;
2631
2632     nslab = atc->nslab;
2633
2634     n = 0;
2635     for (i = 1; i <= nslab/2; i++)
2636     {
2637         fw = (atc->nodeid + i) % nslab;
2638         bw = (atc->nodeid - i + nslab) % nslab;
2639         if (n < nslab - 1)
2640         {
2641             atc->node_dest[n] = fw;
2642             atc->node_src[n]  = bw;
2643             n++;
2644         }
2645         if (n < nslab - 1)
2646         {
2647             atc->node_dest[n] = bw;
2648             atc->node_src[n]  = fw;
2649             n++;
2650         }
2651     }
2652 }
2653
2654 int gmx_pme_destroy(FILE *log, gmx_pme_t *pmedata)
2655 {
2656     int thread;
2657
2658     if (NULL != log)
2659     {
2660         fprintf(log, "Destroying PME data structures.\n");
2661     }
2662
2663     sfree((*pmedata)->nnx);
2664     sfree((*pmedata)->nny);
2665     sfree((*pmedata)->nnz);
2666
2667     pmegrids_destroy(&(*pmedata)->pmegridA);
2668
2669     sfree((*pmedata)->fftgridA);
2670     sfree((*pmedata)->cfftgridA);
2671     gmx_parallel_3dfft_destroy((*pmedata)->pfft_setupA);
2672
2673     if ((*pmedata)->pmegridB.grid.grid != NULL)
2674     {
2675         pmegrids_destroy(&(*pmedata)->pmegridB);
2676         sfree((*pmedata)->fftgridB);
2677         sfree((*pmedata)->cfftgridB);
2678         gmx_parallel_3dfft_destroy((*pmedata)->pfft_setupB);
2679     }
2680     for (thread = 0; thread < (*pmedata)->nthread; thread++)
2681     {
2682         free_work(&(*pmedata)->work[thread]);
2683     }
2684     sfree((*pmedata)->work);
2685
2686     sfree(*pmedata);
2687     *pmedata = NULL;
2688
2689     return 0;
2690 }
2691
2692 static int mult_up(int n, int f)
2693 {
2694     return ((n + f - 1)/f)*f;
2695 }
2696
2697
2698 static double pme_load_imbalance(gmx_pme_t pme)
2699 {
2700     int    nma, nmi;
2701     double n1, n2, n3;
2702
2703     nma = pme->nnodes_major;
2704     nmi = pme->nnodes_minor;
2705
2706     n1 = mult_up(pme->nkx, nma)*mult_up(pme->nky, nmi)*pme->nkz;
2707     n2 = mult_up(pme->nkx, nma)*mult_up(pme->nkz, nmi)*pme->nky;
2708     n3 = mult_up(pme->nky, nma)*mult_up(pme->nkz, nmi)*pme->nkx;
2709
2710     /* pme_solve is roughly double the cost of an fft */
2711
2712     return (n1 + n2 + 3*n3)/(double)(6*pme->nkx*pme->nky*pme->nkz);
2713 }
2714
2715 static void init_atomcomm(gmx_pme_t pme, pme_atomcomm_t *atc,
2716                           int dimind, gmx_bool bSpread)
2717 {
2718     int nk, k, s, thread;
2719
2720     atc->dimind    = dimind;
2721     atc->nslab     = 1;
2722     atc->nodeid    = 0;
2723     atc->pd_nalloc = 0;
2724 #ifdef GMX_MPI
2725     if (pme->nnodes > 1)
2726     {
2727         atc->mpi_comm = pme->mpi_comm_d[dimind];
2728         MPI_Comm_size(atc->mpi_comm, &atc->nslab);
2729         MPI_Comm_rank(atc->mpi_comm, &atc->nodeid);
2730     }
2731     if (debug)
2732     {
2733         fprintf(debug, "For PME atom communication in dimind %d: nslab %d rank %d\n", atc->dimind, atc->nslab, atc->nodeid);
2734     }
2735 #endif
2736
2737     atc->bSpread   = bSpread;
2738     atc->pme_order = pme->pme_order;
2739
2740     if (atc->nslab > 1)
2741     {
2742         /* These three allocations are not required for particle decomp. */
2743         snew(atc->node_dest, atc->nslab);
2744         snew(atc->node_src, atc->nslab);
2745         setup_coordinate_communication(atc);
2746
2747         snew(atc->count_thread, pme->nthread);
2748         for (thread = 0; thread < pme->nthread; thread++)
2749         {
2750             snew(atc->count_thread[thread], atc->nslab);
2751         }
2752         atc->count = atc->count_thread[0];
2753         snew(atc->rcount, atc->nslab);
2754         snew(atc->buf_index, atc->nslab);
2755     }
2756
2757     atc->nthread = pme->nthread;
2758     if (atc->nthread > 1)
2759     {
2760         snew(atc->thread_plist, atc->nthread);
2761     }
2762     snew(atc->spline, atc->nthread);
2763     for (thread = 0; thread < atc->nthread; thread++)
2764     {
2765         if (atc->nthread > 1)
2766         {
2767             snew(atc->thread_plist[thread].n, atc->nthread+2*GMX_CACHE_SEP);
2768             atc->thread_plist[thread].n += GMX_CACHE_SEP;
2769         }
2770         snew(atc->spline[thread].thread_one, pme->nthread);
2771         atc->spline[thread].thread_one[thread] = 1;
2772     }
2773 }
2774
2775 static void
2776 init_overlap_comm(pme_overlap_t *  ol,
2777                   int              norder,
2778 #ifdef GMX_MPI
2779                   MPI_Comm         comm,
2780 #endif
2781                   int              nnodes,
2782                   int              nodeid,
2783                   int              ndata,
2784                   int              commplainsize)
2785 {
2786     int lbnd, rbnd, maxlr, b, i;
2787     int exten;
2788     int nn, nk;
2789     pme_grid_comm_t *pgc;
2790     gmx_bool bCont;
2791     int fft_start, fft_end, send_index1, recv_index1;
2792 #ifdef GMX_MPI
2793     MPI_Status stat;
2794
2795     ol->mpi_comm = comm;
2796 #endif
2797
2798     ol->nnodes = nnodes;
2799     ol->nodeid = nodeid;
2800
2801     /* Linear translation of the PME grid won't affect reciprocal space
2802      * calculations, so to optimize we only interpolate "upwards",
2803      * which also means we only have to consider overlap in one direction.
2804      * I.e., particles on this node might also be spread to grid indices
2805      * that belong to higher nodes (modulo nnodes)
2806      */
2807
2808     snew(ol->s2g0, ol->nnodes+1);
2809     snew(ol->s2g1, ol->nnodes);
2810     if (debug)
2811     {
2812         fprintf(debug, "PME slab boundaries:");
2813     }
2814     for (i = 0; i < nnodes; i++)
2815     {
2816         /* s2g0 the local interpolation grid start.
2817          * s2g1 the local interpolation grid end.
2818          * Because grid overlap communication only goes forward,
2819          * the grid the slabs for fft's should be rounded down.
2820          */
2821         ol->s2g0[i] = ( i   *ndata + 0       )/nnodes;
2822         ol->s2g1[i] = ((i+1)*ndata + nnodes-1)/nnodes + norder - 1;
2823
2824         if (debug)
2825         {
2826             fprintf(debug, "  %3d %3d", ol->s2g0[i], ol->s2g1[i]);
2827         }
2828     }
2829     ol->s2g0[nnodes] = ndata;
2830     if (debug)
2831     {
2832         fprintf(debug, "\n");
2833     }
2834
2835     /* Determine with how many nodes we need to communicate the grid overlap */
2836     b = 0;
2837     do
2838     {
2839         b++;
2840         bCont = FALSE;
2841         for (i = 0; i < nnodes; i++)
2842         {
2843             if ((i+b <  nnodes && ol->s2g1[i] > ol->s2g0[i+b]) ||
2844                 (i+b >= nnodes && ol->s2g1[i] > ol->s2g0[i+b-nnodes] + ndata))
2845             {
2846                 bCont = TRUE;
2847             }
2848         }
2849     }
2850     while (bCont && b < nnodes);
2851     ol->noverlap_nodes = b - 1;
2852
2853     snew(ol->send_id, ol->noverlap_nodes);
2854     snew(ol->recv_id, ol->noverlap_nodes);
2855     for (b = 0; b < ol->noverlap_nodes; b++)
2856     {
2857         ol->send_id[b] = (ol->nodeid + (b + 1)) % ol->nnodes;
2858         ol->recv_id[b] = (ol->nodeid - (b + 1) + ol->nnodes) % ol->nnodes;
2859     }
2860     snew(ol->comm_data, ol->noverlap_nodes);
2861
2862     ol->send_size = 0;
2863     for (b = 0; b < ol->noverlap_nodes; b++)
2864     {
2865         pgc = &ol->comm_data[b];
2866         /* Send */
2867         fft_start        = ol->s2g0[ol->send_id[b]];
2868         fft_end          = ol->s2g0[ol->send_id[b]+1];
2869         if (ol->send_id[b] < nodeid)
2870         {
2871             fft_start += ndata;
2872             fft_end   += ndata;
2873         }
2874         send_index1       = ol->s2g1[nodeid];
2875         send_index1       = min(send_index1, fft_end);
2876         pgc->send_index0  = fft_start;
2877         pgc->send_nindex  = max(0, send_index1 - pgc->send_index0);
2878         ol->send_size    += pgc->send_nindex;
2879
2880         /* We always start receiving to the first index of our slab */
2881         fft_start        = ol->s2g0[ol->nodeid];
2882         fft_end          = ol->s2g0[ol->nodeid+1];
2883         recv_index1      = ol->s2g1[ol->recv_id[b]];
2884         if (ol->recv_id[b] > nodeid)
2885         {
2886             recv_index1 -= ndata;
2887         }
2888         recv_index1      = min(recv_index1, fft_end);
2889         pgc->recv_index0 = fft_start;
2890         pgc->recv_nindex = max(0, recv_index1 - pgc->recv_index0);
2891     }
2892
2893 #ifdef GMX_MPI
2894     /* Communicate the buffer sizes to receive */
2895     for (b = 0; b < ol->noverlap_nodes; b++)
2896     {
2897         MPI_Sendrecv(&ol->send_size, 1, MPI_INT, ol->send_id[b], b,
2898                      &ol->comm_data[b].recv_size, 1, MPI_INT, ol->recv_id[b], b,
2899                      ol->mpi_comm, &stat);
2900     }
2901 #endif
2902
2903     /* For non-divisible grid we need pme_order iso pme_order-1 */
2904     snew(ol->sendbuf, norder*commplainsize);
2905     snew(ol->recvbuf, norder*commplainsize);
2906 }
2907
2908 static void
2909 make_gridindex5_to_localindex(int n, int local_start, int local_range,
2910                               int **global_to_local,
2911                               real **fraction_shift)
2912 {
2913     int i;
2914     int * gtl;
2915     real * fsh;
2916
2917     snew(gtl, 5*n);
2918     snew(fsh, 5*n);
2919     for (i = 0; (i < 5*n); i++)
2920     {
2921         /* Determine the global to local grid index */
2922         gtl[i] = (i - local_start + n) % n;
2923         /* For coordinates that fall within the local grid the fraction
2924          * is correct, we don't need to shift it.
2925          */
2926         fsh[i] = 0;
2927         if (local_range < n)
2928         {
2929             /* Due to rounding issues i could be 1 beyond the lower or
2930              * upper boundary of the local grid. Correct the index for this.
2931              * If we shift the index, we need to shift the fraction by
2932              * the same amount in the other direction to not affect
2933              * the weights.
2934              * Note that due to this shifting the weights at the end of
2935              * the spline might change, but that will only involve values
2936              * between zero and values close to the precision of a real,
2937              * which is anyhow the accuracy of the whole mesh calculation.
2938              */
2939             /* With local_range=0 we should not change i=local_start */
2940             if (i % n != local_start)
2941             {
2942                 if (gtl[i] == n-1)
2943                 {
2944                     gtl[i] = 0;
2945                     fsh[i] = -1;
2946                 }
2947                 else if (gtl[i] == local_range)
2948                 {
2949                     gtl[i] = local_range - 1;
2950                     fsh[i] = 1;
2951                 }
2952             }
2953         }
2954     }
2955
2956     *global_to_local = gtl;
2957     *fraction_shift  = fsh;
2958 }
2959
2960 static pme_spline_work_t *make_pme_spline_work(int order)
2961 {
2962     pme_spline_work_t *work;
2963
2964 #ifdef PME_SSE_SPREAD_GATHER
2965     float  tmp[8];
2966     __m128 zero_SSE;
2967     int    of, i;
2968
2969     snew_aligned(work, 1, 16);
2970
2971     zero_SSE = _mm_setzero_ps();
2972
2973     /* Generate bit masks to mask out the unused grid entries,
2974      * as we only operate on order of the 8 grid entries that are
2975      * load into 2 SSE float registers.
2976      */
2977     for (of = 0; of < 8-(order-1); of++)
2978     {
2979         for (i = 0; i < 8; i++)
2980         {
2981             tmp[i] = (i >= of && i < of+order ? 1 : 0);
2982         }
2983         work->mask_SSE0[of] = _mm_loadu_ps(tmp);
2984         work->mask_SSE1[of] = _mm_loadu_ps(tmp+4);
2985         work->mask_SSE0[of] = _mm_cmpgt_ps(work->mask_SSE0[of], zero_SSE);
2986         work->mask_SSE1[of] = _mm_cmpgt_ps(work->mask_SSE1[of], zero_SSE);
2987     }
2988 #else
2989     work = NULL;
2990 #endif
2991
2992     return work;
2993 }
2994
2995 void gmx_pme_check_restrictions(int pme_order,
2996                                 int nkx, int nky, int nkz,
2997                                 int nnodes_major,
2998                                 int nnodes_minor,
2999                                 gmx_bool bUseThreads,
3000                                 gmx_bool bFatal,
3001                                 gmx_bool *bValidSettings)
3002 {
3003     if (pme_order > PME_ORDER_MAX)
3004     {
3005         if (!bFatal)
3006         {
3007             *bValidSettings = FALSE;
3008             return;
3009         }
3010         gmx_fatal(FARGS, "pme_order (%d) is larger than the maximum allowed value (%d). Modify and recompile the code if you really need such a high order.",
3011                   pme_order, PME_ORDER_MAX);
3012     }
3013
3014     if (nkx <= pme_order*(nnodes_major > 1 ? 2 : 1) ||
3015         nky <= pme_order*(nnodes_minor > 1 ? 2 : 1) ||
3016         nkz <= pme_order)
3017     {
3018         if (!bFatal)
3019         {
3020             *bValidSettings = FALSE;
3021             return;
3022         }
3023         gmx_fatal(FARGS, "The PME grid sizes need to be larger than pme_order (%d) and for dimensions with domain decomposition larger than 2*pme_order",
3024                   pme_order);
3025     }
3026
3027     /* Check for a limitation of the (current) sum_fftgrid_dd code.
3028      * We only allow multiple communication pulses in dim 1, not in dim 0.
3029      */
3030     if (bUseThreads && (nkx < nnodes_major*pme_order &&
3031                         nkx != nnodes_major*(pme_order - 1)))
3032     {
3033         if (!bFatal)
3034         {
3035             *bValidSettings = FALSE;
3036             return;
3037         }
3038         gmx_fatal(FARGS, "The number of PME grid lines per node along x is %g. But when using OpenMP threads, the number of grid lines per node along x should be >= pme_order (%d) or = pmeorder-1. To resolve this issue, use less nodes along x (and possibly more along y and/or z) by specifying -dd manually.",
3039                   nkx/(double)nnodes_major, pme_order);
3040     }
3041
3042     if (bValidSettings != NULL)
3043     {
3044         *bValidSettings = TRUE;
3045     }
3046
3047     return;
3048 }
3049
3050 int gmx_pme_init(gmx_pme_t *         pmedata,
3051                  t_commrec *         cr,
3052                  int                 nnodes_major,
3053                  int                 nnodes_minor,
3054                  t_inputrec *        ir,
3055                  int                 homenr,
3056                  gmx_bool            bFreeEnergy,
3057                  gmx_bool            bReproducible,
3058                  int                 nthread)
3059 {
3060     gmx_pme_t pme = NULL;
3061
3062     int  use_threads, sum_use_threads;
3063     ivec ndata;
3064
3065     if (debug)
3066     {
3067         fprintf(debug, "Creating PME data structures.\n");
3068     }
3069     snew(pme, 1);
3070
3071     pme->redist_init         = FALSE;
3072     pme->sum_qgrid_tmp       = NULL;
3073     pme->sum_qgrid_dd_tmp    = NULL;
3074     pme->buf_nalloc          = 0;
3075     pme->redist_buf_nalloc   = 0;
3076
3077     pme->nnodes              = 1;
3078     pme->bPPnode             = TRUE;
3079
3080     pme->nnodes_major        = nnodes_major;
3081     pme->nnodes_minor        = nnodes_minor;
3082
3083 #ifdef GMX_MPI
3084     if (nnodes_major*nnodes_minor > 1)
3085     {
3086         pme->mpi_comm = cr->mpi_comm_mygroup;
3087
3088         MPI_Comm_rank(pme->mpi_comm, &pme->nodeid);
3089         MPI_Comm_size(pme->mpi_comm, &pme->nnodes);
3090         if (pme->nnodes != nnodes_major*nnodes_minor)
3091         {
3092             gmx_incons("PME node count mismatch");
3093         }
3094     }
3095     else
3096     {
3097         pme->mpi_comm = MPI_COMM_NULL;
3098     }
3099 #endif
3100
3101     if (pme->nnodes == 1)
3102     {
3103 #ifdef GMX_MPI
3104         pme->mpi_comm_d[0] = MPI_COMM_NULL;
3105         pme->mpi_comm_d[1] = MPI_COMM_NULL;
3106 #endif
3107         pme->ndecompdim   = 0;
3108         pme->nodeid_major = 0;
3109         pme->nodeid_minor = 0;
3110 #ifdef GMX_MPI
3111         pme->mpi_comm_d[0] = pme->mpi_comm_d[1] = MPI_COMM_NULL;
3112 #endif
3113     }
3114     else
3115     {
3116         if (nnodes_minor == 1)
3117         {
3118 #ifdef GMX_MPI
3119             pme->mpi_comm_d[0] = pme->mpi_comm;
3120             pme->mpi_comm_d[1] = MPI_COMM_NULL;
3121 #endif
3122             pme->ndecompdim   = 1;
3123             pme->nodeid_major = pme->nodeid;
3124             pme->nodeid_minor = 0;
3125
3126         }
3127         else if (nnodes_major == 1)
3128         {
3129 #ifdef GMX_MPI
3130             pme->mpi_comm_d[0] = MPI_COMM_NULL;
3131             pme->mpi_comm_d[1] = pme->mpi_comm;
3132 #endif
3133             pme->ndecompdim   = 1;
3134             pme->nodeid_major = 0;
3135             pme->nodeid_minor = pme->nodeid;
3136         }
3137         else
3138         {
3139             if (pme->nnodes % nnodes_major != 0)
3140             {
3141                 gmx_incons("For 2D PME decomposition, #PME nodes must be divisible by the number of nodes in the major dimension");
3142             }
3143             pme->ndecompdim = 2;
3144
3145 #ifdef GMX_MPI
3146             MPI_Comm_split(pme->mpi_comm, pme->nodeid % nnodes_minor,
3147                            pme->nodeid, &pme->mpi_comm_d[0]);  /* My communicator along major dimension */
3148             MPI_Comm_split(pme->mpi_comm, pme->nodeid/nnodes_minor,
3149                            pme->nodeid, &pme->mpi_comm_d[1]);  /* My communicator along minor dimension */
3150
3151             MPI_Comm_rank(pme->mpi_comm_d[0], &pme->nodeid_major);
3152             MPI_Comm_size(pme->mpi_comm_d[0], &pme->nnodes_major);
3153             MPI_Comm_rank(pme->mpi_comm_d[1], &pme->nodeid_minor);
3154             MPI_Comm_size(pme->mpi_comm_d[1], &pme->nnodes_minor);
3155 #endif
3156         }
3157         pme->bPPnode = (cr->duty & DUTY_PP);
3158     }
3159
3160     pme->nthread = nthread;
3161
3162     /* Check if any of the PME MPI ranks uses threads */
3163     use_threads = (pme->nthread > 1 ? 1 : 0);
3164 #ifdef GMX_MPI
3165     if (pme->nnodes > 1)
3166     {
3167         MPI_Allreduce(&use_threads, &sum_use_threads, 1, MPI_INT,
3168                       MPI_SUM, pme->mpi_comm);
3169     }
3170     else
3171 #endif
3172     {
3173         sum_use_threads = use_threads;
3174     }
3175     pme->bUseThreads = (sum_use_threads > 0);
3176
3177     if (ir->ePBC == epbcSCREW)
3178     {
3179         gmx_fatal(FARGS, "pme does not (yet) work with pbc = screw");
3180     }
3181
3182     pme->bFEP        = ((ir->efep != efepNO) && bFreeEnergy);
3183     pme->nkx         = ir->nkx;
3184     pme->nky         = ir->nky;
3185     pme->nkz         = ir->nkz;
3186     pme->bP3M        = (ir->coulombtype == eelP3M_AD || getenv("GMX_PME_P3M") != NULL);
3187     pme->pme_order   = ir->pme_order;
3188     pme->epsilon_r   = ir->epsilon_r;
3189
3190     /* If we violate restrictions, generate a fatal error here */
3191     gmx_pme_check_restrictions(pme->pme_order,
3192                                pme->nkx, pme->nky, pme->nkz,
3193                                pme->nnodes_major,
3194                                pme->nnodes_minor,
3195                                pme->bUseThreads,
3196                                TRUE,
3197                                NULL);
3198
3199     if (pme->nnodes > 1)
3200     {
3201         double imbal;
3202
3203 #ifdef GMX_MPI
3204         MPI_Type_contiguous(DIM, mpi_type, &(pme->rvec_mpi));
3205         MPI_Type_commit(&(pme->rvec_mpi));
3206 #endif
3207
3208         /* Note that the charge spreading and force gathering, which usually
3209          * takes about the same amount of time as FFT+solve_pme,
3210          * is always fully load balanced
3211          * (unless the charge distribution is inhomogeneous).
3212          */
3213
3214         imbal = pme_load_imbalance(pme);
3215         if (imbal >= 1.2 && pme->nodeid_major == 0 && pme->nodeid_minor == 0)
3216         {
3217             fprintf(stderr,
3218                     "\n"
3219                     "NOTE: The load imbalance in PME FFT and solve is %d%%.\n"
3220                     "      For optimal PME load balancing\n"
3221                     "      PME grid_x (%d) and grid_y (%d) should be divisible by #PME_nodes_x (%d)\n"
3222                     "      and PME grid_y (%d) and grid_z (%d) should be divisible by #PME_nodes_y (%d)\n"
3223                     "\n",
3224                     (int)((imbal-1)*100 + 0.5),
3225                     pme->nkx, pme->nky, pme->nnodes_major,
3226                     pme->nky, pme->nkz, pme->nnodes_minor);
3227         }
3228     }
3229
3230     /* For non-divisible grid we need pme_order iso pme_order-1 */
3231     /* In sum_qgrid_dd x overlap is copied in place: take padding into account.
3232      * y is always copied through a buffer: we don't need padding in z,
3233      * but we do need the overlap in x because of the communication order.
3234      */
3235     init_overlap_comm(&pme->overlap[0], pme->pme_order,
3236 #ifdef GMX_MPI
3237                       pme->mpi_comm_d[0],
3238 #endif
3239                       pme->nnodes_major, pme->nodeid_major,
3240                       pme->nkx,
3241                       (div_round_up(pme->nky, pme->nnodes_minor)+pme->pme_order)*(pme->nkz+pme->pme_order-1));
3242
3243     /* Along overlap dim 1 we can send in multiple pulses in sum_fftgrid_dd.
3244      * We do this with an offset buffer of equal size, so we need to allocate
3245      * extra for the offset. That's what the (+1)*pme->nkz is for.
3246      */
3247     init_overlap_comm(&pme->overlap[1], pme->pme_order,
3248 #ifdef GMX_MPI
3249                       pme->mpi_comm_d[1],
3250 #endif
3251                       pme->nnodes_minor, pme->nodeid_minor,
3252                       pme->nky,
3253                       (div_round_up(pme->nkx, pme->nnodes_major)+pme->pme_order+1)*pme->nkz);
3254
3255     /* Double-check for a limitation of the (current) sum_fftgrid_dd code.
3256      * Note that gmx_pme_check_restrictions checked for this already.
3257      */
3258     if (pme->bUseThreads && pme->overlap[0].noverlap_nodes > 1)
3259     {
3260         gmx_incons("More than one communication pulse required for grid overlap communication along the major dimension while using threads");
3261     }
3262
3263     snew(pme->bsp_mod[XX], pme->nkx);
3264     snew(pme->bsp_mod[YY], pme->nky);
3265     snew(pme->bsp_mod[ZZ], pme->nkz);
3266
3267     /* The required size of the interpolation grid, including overlap.
3268      * The allocated size (pmegrid_n?) might be slightly larger.
3269      */
3270     pme->pmegrid_nx = pme->overlap[0].s2g1[pme->nodeid_major] -
3271         pme->overlap[0].s2g0[pme->nodeid_major];
3272     pme->pmegrid_ny = pme->overlap[1].s2g1[pme->nodeid_minor] -
3273         pme->overlap[1].s2g0[pme->nodeid_minor];
3274     pme->pmegrid_nz_base = pme->nkz;
3275     pme->pmegrid_nz      = pme->pmegrid_nz_base + pme->pme_order - 1;
3276     set_grid_alignment(&pme->pmegrid_nz, pme->pme_order);
3277
3278     pme->pmegrid_start_ix = pme->overlap[0].s2g0[pme->nodeid_major];
3279     pme->pmegrid_start_iy = pme->overlap[1].s2g0[pme->nodeid_minor];
3280     pme->pmegrid_start_iz = 0;
3281
3282     make_gridindex5_to_localindex(pme->nkx,
3283                                   pme->pmegrid_start_ix,
3284                                   pme->pmegrid_nx - (pme->pme_order-1),
3285                                   &pme->nnx, &pme->fshx);
3286     make_gridindex5_to_localindex(pme->nky,
3287                                   pme->pmegrid_start_iy,
3288                                   pme->pmegrid_ny - (pme->pme_order-1),
3289                                   &pme->nny, &pme->fshy);
3290     make_gridindex5_to_localindex(pme->nkz,
3291                                   pme->pmegrid_start_iz,
3292                                   pme->pmegrid_nz_base,
3293                                   &pme->nnz, &pme->fshz);
3294
3295     pmegrids_init(&pme->pmegridA,
3296                   pme->pmegrid_nx, pme->pmegrid_ny, pme->pmegrid_nz,
3297                   pme->pmegrid_nz_base,
3298                   pme->pme_order,
3299                   pme->bUseThreads,
3300                   pme->nthread,
3301                   pme->overlap[0].s2g1[pme->nodeid_major]-pme->overlap[0].s2g0[pme->nodeid_major+1],
3302                   pme->overlap[1].s2g1[pme->nodeid_minor]-pme->overlap[1].s2g0[pme->nodeid_minor+1]);
3303
3304     pme->spline_work = make_pme_spline_work(pme->pme_order);
3305
3306     ndata[0] = pme->nkx;
3307     ndata[1] = pme->nky;
3308     ndata[2] = pme->nkz;
3309
3310     /* This routine will allocate the grid data to fit the FFTs */
3311     gmx_parallel_3dfft_init(&pme->pfft_setupA, ndata,
3312                             &pme->fftgridA, &pme->cfftgridA,
3313                             pme->mpi_comm_d,
3314                             bReproducible, pme->nthread);
3315
3316     if (bFreeEnergy)
3317     {
3318         pmegrids_init(&pme->pmegridB,
3319                       pme->pmegrid_nx, pme->pmegrid_ny, pme->pmegrid_nz,
3320                       pme->pmegrid_nz_base,
3321                       pme->pme_order,
3322                       pme->bUseThreads,
3323                       pme->nthread,
3324                       pme->nkx % pme->nnodes_major != 0,
3325                       pme->nky % pme->nnodes_minor != 0);
3326
3327         gmx_parallel_3dfft_init(&pme->pfft_setupB, ndata,
3328                                 &pme->fftgridB, &pme->cfftgridB,
3329                                 pme->mpi_comm_d,
3330                                 bReproducible, pme->nthread);
3331     }
3332     else
3333     {
3334         pme->pmegridB.grid.grid = NULL;
3335         pme->fftgridB           = NULL;
3336         pme->cfftgridB          = NULL;
3337     }
3338
3339     if (!pme->bP3M)
3340     {
3341         /* Use plain SPME B-spline interpolation */
3342         make_bspline_moduli(pme->bsp_mod, pme->nkx, pme->nky, pme->nkz, pme->pme_order);
3343     }
3344     else
3345     {
3346         /* Use the P3M grid-optimized influence function */
3347         make_p3m_bspline_moduli(pme->bsp_mod, pme->nkx, pme->nky, pme->nkz, pme->pme_order);
3348     }
3349
3350     /* Use atc[0] for spreading */
3351     init_atomcomm(pme, &pme->atc[0], nnodes_major > 1 ? 0 : 1, TRUE);
3352     if (pme->ndecompdim >= 2)
3353     {
3354         init_atomcomm(pme, &pme->atc[1], 1, FALSE);
3355     }
3356
3357     if (pme->nnodes == 1)
3358     {
3359         pme->atc[0].n = homenr;
3360         pme_realloc_atomcomm_things(&pme->atc[0]);
3361     }
3362
3363     {
3364         int thread;
3365
3366         /* Use fft5d, order after FFT is y major, z, x minor */
3367
3368         snew(pme->work, pme->nthread);
3369         for (thread = 0; thread < pme->nthread; thread++)
3370         {
3371             realloc_work(&pme->work[thread], pme->nkx);
3372         }
3373     }
3374
3375     *pmedata = pme;
3376
3377     return 0;
3378 }
3379
3380 static void reuse_pmegrids(const pmegrids_t *old, pmegrids_t *new)
3381 {
3382     int d, t;
3383
3384     for (d = 0; d < DIM; d++)
3385     {
3386         if (new->grid.n[d] > old->grid.n[d])
3387         {
3388             return;
3389         }
3390     }
3391
3392     sfree_aligned(new->grid.grid);
3393     new->grid.grid = old->grid.grid;
3394
3395     if (new->grid_th != NULL && new->nthread == old->nthread)
3396     {
3397         sfree_aligned(new->grid_all);
3398         for (t = 0; t < new->nthread; t++)
3399         {
3400             new->grid_th[t].grid = old->grid_th[t].grid;
3401         }
3402     }
3403 }
3404
3405 int gmx_pme_reinit(gmx_pme_t *         pmedata,
3406                    t_commrec *         cr,
3407                    gmx_pme_t           pme_src,
3408                    const t_inputrec *  ir,
3409                    ivec                grid_size)
3410 {
3411     t_inputrec irc;
3412     int homenr;
3413     int ret;
3414
3415     irc     = *ir;
3416     irc.nkx = grid_size[XX];
3417     irc.nky = grid_size[YY];
3418     irc.nkz = grid_size[ZZ];
3419
3420     if (pme_src->nnodes == 1)
3421     {
3422         homenr = pme_src->atc[0].n;
3423     }
3424     else
3425     {
3426         homenr = -1;
3427     }
3428
3429     ret = gmx_pme_init(pmedata, cr, pme_src->nnodes_major, pme_src->nnodes_minor,
3430                        &irc, homenr, pme_src->bFEP, FALSE, pme_src->nthread);
3431
3432     if (ret == 0)
3433     {
3434         /* We can easily reuse the allocated pme grids in pme_src */
3435         reuse_pmegrids(&pme_src->pmegridA, &(*pmedata)->pmegridA);
3436         /* We would like to reuse the fft grids, but that's harder */
3437     }
3438
3439     return ret;
3440 }
3441
3442
3443 static void copy_local_grid(gmx_pme_t pme,
3444                             pmegrids_t *pmegrids, int thread, real *fftgrid)
3445 {
3446     ivec local_fft_ndata, local_fft_offset, local_fft_size;
3447     int  fft_my, fft_mz;
3448     int  nsx, nsy, nsz;
3449     ivec nf;
3450     int  offx, offy, offz, x, y, z, i0, i0t;
3451     int  d;
3452     pmegrid_t *pmegrid;
3453     real *grid_th;
3454
3455     gmx_parallel_3dfft_real_limits(pme->pfft_setupA,
3456                                    local_fft_ndata,
3457                                    local_fft_offset,
3458                                    local_fft_size);
3459     fft_my = local_fft_size[YY];
3460     fft_mz = local_fft_size[ZZ];
3461
3462     pmegrid = &pmegrids->grid_th[thread];
3463
3464     nsx = pmegrid->s[XX];
3465     nsy = pmegrid->s[YY];
3466     nsz = pmegrid->s[ZZ];
3467
3468     for (d = 0; d < DIM; d++)
3469     {
3470         nf[d] = min(pmegrid->n[d] - (pmegrid->order - 1),
3471                     local_fft_ndata[d] - pmegrid->offset[d]);
3472     }
3473
3474     offx = pmegrid->offset[XX];
3475     offy = pmegrid->offset[YY];
3476     offz = pmegrid->offset[ZZ];
3477
3478     /* Directly copy the non-overlapping parts of the local grids.
3479      * This also initializes the full grid.
3480      */
3481     grid_th = pmegrid->grid;
3482     for (x = 0; x < nf[XX]; x++)
3483     {
3484         for (y = 0; y < nf[YY]; y++)
3485         {
3486             i0  = ((offx + x)*fft_my + (offy + y))*fft_mz + offz;
3487             i0t = (x*nsy + y)*nsz;
3488             for (z = 0; z < nf[ZZ]; z++)
3489             {
3490                 fftgrid[i0+z] = grid_th[i0t+z];
3491             }
3492         }
3493     }
3494 }
3495
3496 static void
3497 reduce_threadgrid_overlap(gmx_pme_t pme,
3498                           const pmegrids_t *pmegrids, int thread,
3499                           real *fftgrid, real *commbuf_x, real *commbuf_y)
3500 {
3501     ivec local_fft_ndata, local_fft_offset, local_fft_size;
3502     int  fft_nx, fft_ny, fft_nz;
3503     int  fft_my, fft_mz;
3504     int  buf_my = -1;
3505     int  nsx, nsy, nsz;
3506     ivec ne;
3507     int  offx, offy, offz, x, y, z, i0, i0t;
3508     int  sx, sy, sz, fx, fy, fz, tx1, ty1, tz1, ox, oy, oz;
3509     gmx_bool bClearBufX, bClearBufY, bClearBufXY, bClearBuf;
3510     gmx_bool bCommX, bCommY;
3511     int  d;
3512     int  thread_f;
3513     const pmegrid_t *pmegrid, *pmegrid_g, *pmegrid_f;
3514     const real *grid_th;
3515     real *commbuf = NULL;
3516
3517     gmx_parallel_3dfft_real_limits(pme->pfft_setupA,
3518                                    local_fft_ndata,
3519                                    local_fft_offset,
3520                                    local_fft_size);
3521     fft_nx = local_fft_ndata[XX];
3522     fft_ny = local_fft_ndata[YY];
3523     fft_nz = local_fft_ndata[ZZ];
3524
3525     fft_my = local_fft_size[YY];
3526     fft_mz = local_fft_size[ZZ];
3527
3528     /* This routine is called when all thread have finished spreading.
3529      * Here each thread sums grid contributions calculated by other threads
3530      * to the thread local grid volume.
3531      * To minimize the number of grid copying operations,
3532      * this routines sums immediately from the pmegrid to the fftgrid.
3533      */
3534
3535     /* Determine which part of the full node grid we should operate on,
3536      * this is our thread local part of the full grid.
3537      */
3538     pmegrid = &pmegrids->grid_th[thread];
3539
3540     for (d = 0; d < DIM; d++)
3541     {
3542         ne[d] = min(pmegrid->offset[d]+pmegrid->n[d]-(pmegrid->order-1),
3543                     local_fft_ndata[d]);
3544     }
3545
3546     offx = pmegrid->offset[XX];
3547     offy = pmegrid->offset[YY];
3548     offz = pmegrid->offset[ZZ];
3549
3550
3551     bClearBufX  = TRUE;
3552     bClearBufY  = TRUE;
3553     bClearBufXY = TRUE;
3554
3555     /* Now loop over all the thread data blocks that contribute
3556      * to the grid region we (our thread) are operating on.
3557      */
3558     /* Note that ffy_nx/y is equal to the number of grid points
3559      * between the first point of our node grid and the one of the next node.
3560      */
3561     for (sx = 0; sx >= -pmegrids->nthread_comm[XX]; sx--)
3562     {
3563         fx     = pmegrid->ci[XX] + sx;
3564         ox     = 0;
3565         bCommX = FALSE;
3566         if (fx < 0)
3567         {
3568             fx    += pmegrids->nc[XX];
3569             ox    -= fft_nx;
3570             bCommX = (pme->nnodes_major > 1);
3571         }
3572         pmegrid_g = &pmegrids->grid_th[fx*pmegrids->nc[YY]*pmegrids->nc[ZZ]];
3573         ox       += pmegrid_g->offset[XX];
3574         if (!bCommX)
3575         {
3576             tx1 = min(ox + pmegrid_g->n[XX], ne[XX]);
3577         }
3578         else
3579         {
3580             tx1 = min(ox + pmegrid_g->n[XX], pme->pme_order);
3581         }
3582
3583         for (sy = 0; sy >= -pmegrids->nthread_comm[YY]; sy--)
3584         {
3585             fy     = pmegrid->ci[YY] + sy;
3586             oy     = 0;
3587             bCommY = FALSE;
3588             if (fy < 0)
3589             {
3590                 fy    += pmegrids->nc[YY];
3591                 oy    -= fft_ny;
3592                 bCommY = (pme->nnodes_minor > 1);
3593             }
3594             pmegrid_g = &pmegrids->grid_th[fy*pmegrids->nc[ZZ]];
3595             oy       += pmegrid_g->offset[YY];
3596             if (!bCommY)
3597             {
3598                 ty1 = min(oy + pmegrid_g->n[YY], ne[YY]);
3599             }
3600             else
3601             {
3602                 ty1 = min(oy + pmegrid_g->n[YY], pme->pme_order);
3603             }
3604
3605             for (sz = 0; sz >= -pmegrids->nthread_comm[ZZ]; sz--)
3606             {
3607                 fz = pmegrid->ci[ZZ] + sz;
3608                 oz = 0;
3609                 if (fz < 0)
3610                 {
3611                     fz += pmegrids->nc[ZZ];
3612                     oz -= fft_nz;
3613                 }
3614                 pmegrid_g = &pmegrids->grid_th[fz];
3615                 oz       += pmegrid_g->offset[ZZ];
3616                 tz1       = min(oz + pmegrid_g->n[ZZ], ne[ZZ]);
3617
3618                 if (sx == 0 && sy == 0 && sz == 0)
3619                 {
3620                     /* We have already added our local contribution
3621                      * before calling this routine, so skip it here.
3622                      */
3623                     continue;
3624                 }
3625
3626                 thread_f = (fx*pmegrids->nc[YY] + fy)*pmegrids->nc[ZZ] + fz;
3627
3628                 pmegrid_f = &pmegrids->grid_th[thread_f];
3629
3630                 grid_th = pmegrid_f->grid;
3631
3632                 nsx = pmegrid_f->s[XX];
3633                 nsy = pmegrid_f->s[YY];
3634                 nsz = pmegrid_f->s[ZZ];
3635
3636 #ifdef DEBUG_PME_REDUCE
3637                 printf("n%d t%d add %d  %2d %2d %2d  %2d %2d %2d  %2d-%2d %2d-%2d, %2d-%2d %2d-%2d, %2d-%2d %2d-%2d\n",
3638                        pme->nodeid, thread, thread_f,
3639                        pme->pmegrid_start_ix,
3640                        pme->pmegrid_start_iy,
3641                        pme->pmegrid_start_iz,
3642                        sx, sy, sz,
3643                        offx-ox, tx1-ox, offx, tx1,
3644                        offy-oy, ty1-oy, offy, ty1,
3645                        offz-oz, tz1-oz, offz, tz1);
3646 #endif
3647
3648                 if (!(bCommX || bCommY))
3649                 {
3650                     /* Copy from the thread local grid to the node grid */
3651                     for (x = offx; x < tx1; x++)
3652                     {
3653                         for (y = offy; y < ty1; y++)
3654                         {
3655                             i0  = (x*fft_my + y)*fft_mz;
3656                             i0t = ((x - ox)*nsy + (y - oy))*nsz - oz;
3657                             for (z = offz; z < tz1; z++)
3658                             {
3659                                 fftgrid[i0+z] += grid_th[i0t+z];
3660                             }
3661                         }
3662                     }
3663                 }
3664                 else
3665                 {
3666                     /* The order of this conditional decides
3667                      * where the corner volume gets stored with x+y decomp.
3668                      */
3669                     if (bCommY)
3670                     {
3671                         commbuf = commbuf_y;
3672                         buf_my  = ty1 - offy;
3673                         if (bCommX)
3674                         {
3675                             /* We index commbuf modulo the local grid size */
3676                             commbuf += buf_my*fft_nx*fft_nz;
3677
3678                             bClearBuf   = bClearBufXY;
3679                             bClearBufXY = FALSE;
3680                         }
3681                         else
3682                         {
3683                             bClearBuf  = bClearBufY;
3684                             bClearBufY = FALSE;
3685                         }
3686                     }
3687                     else
3688                     {
3689                         commbuf    = commbuf_x;
3690                         buf_my     = fft_ny;
3691                         bClearBuf  = bClearBufX;
3692                         bClearBufX = FALSE;
3693                     }
3694
3695                     /* Copy to the communication buffer */
3696                     for (x = offx; x < tx1; x++)
3697                     {
3698                         for (y = offy; y < ty1; y++)
3699                         {
3700                             i0  = (x*buf_my + y)*fft_nz;
3701                             i0t = ((x - ox)*nsy + (y - oy))*nsz - oz;
3702
3703                             if (bClearBuf)
3704                             {
3705                                 /* First access of commbuf, initialize it */
3706                                 for (z = offz; z < tz1; z++)
3707                                 {
3708                                     commbuf[i0+z]  = grid_th[i0t+z];
3709                                 }
3710                             }
3711                             else
3712                             {
3713                                 for (z = offz; z < tz1; z++)
3714                                 {
3715                                     commbuf[i0+z] += grid_th[i0t+z];
3716                                 }
3717                             }
3718                         }
3719                     }
3720                 }
3721             }
3722         }
3723     }
3724 }
3725
3726
3727 static void sum_fftgrid_dd(gmx_pme_t pme, real *fftgrid)
3728 {
3729     ivec local_fft_ndata, local_fft_offset, local_fft_size;
3730     pme_overlap_t *overlap;
3731     int  send_index0, send_nindex;
3732     int  recv_nindex;
3733 #ifdef GMX_MPI
3734     MPI_Status stat;
3735 #endif
3736     int  send_size_y, recv_size_y;
3737     int  ipulse, send_id, recv_id, datasize, gridsize, size_yx;
3738     real *sendptr, *recvptr;
3739     int  x, y, z, indg, indb;
3740
3741     /* Note that this routine is only used for forward communication.
3742      * Since the force gathering, unlike the charge spreading,
3743      * can be trivially parallelized over the particles,
3744      * the backwards process is much simpler and can use the "old"
3745      * communication setup.
3746      */
3747
3748     gmx_parallel_3dfft_real_limits(pme->pfft_setupA,
3749                                    local_fft_ndata,
3750                                    local_fft_offset,
3751                                    local_fft_size);
3752
3753     if (pme->nnodes_minor > 1)
3754     {
3755         /* Major dimension */
3756         overlap = &pme->overlap[1];
3757
3758         if (pme->nnodes_major > 1)
3759         {
3760             size_yx = pme->overlap[0].comm_data[0].send_nindex;
3761         }
3762         else
3763         {
3764             size_yx = 0;
3765         }
3766         datasize = (local_fft_ndata[XX] + size_yx)*local_fft_ndata[ZZ];
3767
3768         send_size_y = overlap->send_size;
3769
3770         for (ipulse = 0; ipulse < overlap->noverlap_nodes; ipulse++)
3771         {
3772             send_id       = overlap->send_id[ipulse];
3773             recv_id       = overlap->recv_id[ipulse];
3774             send_index0   =
3775                 overlap->comm_data[ipulse].send_index0 -
3776                 overlap->comm_data[0].send_index0;
3777             send_nindex   = overlap->comm_data[ipulse].send_nindex;
3778             /* We don't use recv_index0, as we always receive starting at 0 */
3779             recv_nindex   = overlap->comm_data[ipulse].recv_nindex;
3780             recv_size_y   = overlap->comm_data[ipulse].recv_size;
3781
3782             sendptr = overlap->sendbuf + send_index0*local_fft_ndata[ZZ];
3783             recvptr = overlap->recvbuf;
3784
3785 #ifdef GMX_MPI
3786             MPI_Sendrecv(sendptr, send_size_y*datasize, GMX_MPI_REAL,
3787                          send_id, ipulse,
3788                          recvptr, recv_size_y*datasize, GMX_MPI_REAL,
3789                          recv_id, ipulse,
3790                          overlap->mpi_comm, &stat);
3791 #endif
3792
3793             for (x = 0; x < local_fft_ndata[XX]; x++)
3794             {
3795                 for (y = 0; y < recv_nindex; y++)
3796                 {
3797                     indg = (x*local_fft_size[YY] + y)*local_fft_size[ZZ];
3798                     indb = (x*recv_size_y        + y)*local_fft_ndata[ZZ];
3799                     for (z = 0; z < local_fft_ndata[ZZ]; z++)
3800                     {
3801                         fftgrid[indg+z] += recvptr[indb+z];
3802                     }
3803                 }
3804             }
3805
3806             if (pme->nnodes_major > 1)
3807             {
3808                 /* Copy from the received buffer to the send buffer for dim 0 */
3809                 sendptr = pme->overlap[0].sendbuf;
3810                 for (x = 0; x < size_yx; x++)
3811                 {
3812                     for (y = 0; y < recv_nindex; y++)
3813                     {
3814                         indg = (x*local_fft_ndata[YY] + y)*local_fft_ndata[ZZ];
3815                         indb = ((local_fft_ndata[XX] + x)*recv_size_y + y)*local_fft_ndata[ZZ];
3816                         for (z = 0; z < local_fft_ndata[ZZ]; z++)
3817                         {
3818                             sendptr[indg+z] += recvptr[indb+z];
3819                         }
3820                     }
3821                 }
3822             }
3823         }
3824     }
3825
3826     /* We only support a single pulse here.
3827      * This is not a severe limitation, as this code is only used
3828      * with OpenMP and with OpenMP the (PME) domains can be larger.
3829      */
3830     if (pme->nnodes_major > 1)
3831     {
3832         /* Major dimension */
3833         overlap = &pme->overlap[0];
3834
3835         datasize = local_fft_ndata[YY]*local_fft_ndata[ZZ];
3836         gridsize = local_fft_size[YY] *local_fft_size[ZZ];
3837
3838         ipulse = 0;
3839
3840         send_id       = overlap->send_id[ipulse];
3841         recv_id       = overlap->recv_id[ipulse];
3842         send_nindex   = overlap->comm_data[ipulse].send_nindex;
3843         /* We don't use recv_index0, as we always receive starting at 0 */
3844         recv_nindex   = overlap->comm_data[ipulse].recv_nindex;
3845
3846         sendptr = overlap->sendbuf;
3847         recvptr = overlap->recvbuf;
3848
3849         if (debug != NULL)
3850         {
3851             fprintf(debug, "PME fftgrid comm %2d x %2d x %2d\n",
3852                     send_nindex, local_fft_ndata[YY], local_fft_ndata[ZZ]);
3853         }
3854
3855 #ifdef GMX_MPI
3856         MPI_Sendrecv(sendptr, send_nindex*datasize, GMX_MPI_REAL,
3857                      send_id, ipulse,
3858                      recvptr, recv_nindex*datasize, GMX_MPI_REAL,
3859                      recv_id, ipulse,
3860                      overlap->mpi_comm, &stat);
3861 #endif
3862
3863         for (x = 0; x < recv_nindex; x++)
3864         {
3865             for (y = 0; y < local_fft_ndata[YY]; y++)
3866             {
3867                 indg = (x*local_fft_size[YY]  + y)*local_fft_size[ZZ];
3868                 indb = (x*local_fft_ndata[YY] + y)*local_fft_ndata[ZZ];
3869                 for (z = 0; z < local_fft_ndata[ZZ]; z++)
3870                 {
3871                     fftgrid[indg+z] += recvptr[indb+z];
3872                 }
3873             }
3874         }
3875     }
3876 }
3877
3878
3879 static void spread_on_grid(gmx_pme_t pme,
3880                            pme_atomcomm_t *atc, pmegrids_t *grids,
3881                            gmx_bool bCalcSplines, gmx_bool bSpread,
3882                            real *fftgrid)
3883 {
3884     int nthread, thread;
3885 #ifdef PME_TIME_THREADS
3886     gmx_cycles_t c1, c2, c3, ct1a, ct1b, ct1c;
3887     static double cs1     = 0, cs2 = 0, cs3 = 0;
3888     static double cs1a[6] = {0, 0, 0, 0, 0, 0};
3889     static int cnt        = 0;
3890 #endif
3891
3892     nthread = pme->nthread;
3893     assert(nthread > 0);
3894
3895 #ifdef PME_TIME_THREADS
3896     c1 = omp_cyc_start();
3897 #endif
3898     if (bCalcSplines)
3899     {
3900 #pragma omp parallel for num_threads(nthread) schedule(static)
3901         for (thread = 0; thread < nthread; thread++)
3902         {
3903             int start, end;
3904
3905             start = atc->n* thread   /nthread;
3906             end   = atc->n*(thread+1)/nthread;
3907
3908             /* Compute fftgrid index for all atoms,
3909              * with help of some extra variables.
3910              */
3911             calc_interpolation_idx(pme, atc, start, end, thread);
3912         }
3913     }
3914 #ifdef PME_TIME_THREADS
3915     c1   = omp_cyc_end(c1);
3916     cs1 += (double)c1;
3917 #endif
3918
3919 #ifdef PME_TIME_THREADS
3920     c2 = omp_cyc_start();
3921 #endif
3922 #pragma omp parallel for num_threads(nthread) schedule(static)
3923     for (thread = 0; thread < nthread; thread++)
3924     {
3925         splinedata_t *spline;
3926         pmegrid_t *grid = NULL;
3927
3928         /* make local bsplines  */
3929         if (grids == NULL || !pme->bUseThreads)
3930         {
3931             spline = &atc->spline[0];
3932
3933             spline->n = atc->n;
3934
3935             if (bSpread)
3936             {
3937                 grid = &grids->grid;
3938             }
3939         }
3940         else
3941         {
3942             spline = &atc->spline[thread];
3943
3944             if (grids->nthread == 1)
3945             {
3946                 /* One thread, we operate on all charges */
3947                 spline->n = atc->n;
3948             }
3949             else
3950             {
3951                 /* Get the indices our thread should operate on */
3952                 make_thread_local_ind(atc, thread, spline);
3953             }
3954
3955             grid = &grids->grid_th[thread];
3956         }
3957
3958         if (bCalcSplines)
3959         {
3960             make_bsplines(spline->theta, spline->dtheta, pme->pme_order,
3961                           atc->fractx, spline->n, spline->ind, atc->q, pme->bFEP);
3962         }
3963
3964         if (bSpread)
3965         {
3966             /* put local atoms on grid. */
3967 #ifdef PME_TIME_SPREAD
3968             ct1a = omp_cyc_start();
3969 #endif
3970             spread_q_bsplines_thread(grid, atc, spline, pme->spline_work);
3971
3972             if (pme->bUseThreads)
3973             {
3974                 copy_local_grid(pme, grids, thread, fftgrid);
3975             }
3976 #ifdef PME_TIME_SPREAD
3977             ct1a          = omp_cyc_end(ct1a);
3978             cs1a[thread] += (double)ct1a;
3979 #endif
3980         }
3981     }
3982 #ifdef PME_TIME_THREADS
3983     c2   = omp_cyc_end(c2);
3984     cs2 += (double)c2;
3985 #endif
3986
3987     if (bSpread && pme->bUseThreads)
3988     {
3989 #ifdef PME_TIME_THREADS
3990         c3 = omp_cyc_start();
3991 #endif
3992 #pragma omp parallel for num_threads(grids->nthread) schedule(static)
3993         for (thread = 0; thread < grids->nthread; thread++)
3994         {
3995             reduce_threadgrid_overlap(pme, grids, thread,
3996                                       fftgrid,
3997                                       pme->overlap[0].sendbuf,
3998                                       pme->overlap[1].sendbuf);
3999         }
4000 #ifdef PME_TIME_THREADS
4001         c3   = omp_cyc_end(c3);
4002         cs3 += (double)c3;
4003 #endif
4004
4005         if (pme->nnodes > 1)
4006         {
4007             /* Communicate the overlapping part of the fftgrid.
4008              * For this communication call we need to check pme->bUseThreads
4009              * to have all ranks communicate here, regardless of pme->nthread.
4010              */
4011             sum_fftgrid_dd(pme, fftgrid);
4012         }
4013     }
4014
4015 #ifdef PME_TIME_THREADS
4016     cnt++;
4017     if (cnt % 20 == 0)
4018     {
4019         printf("idx %.2f spread %.2f red %.2f",
4020                cs1*1e-9, cs2*1e-9, cs3*1e-9);
4021 #ifdef PME_TIME_SPREAD
4022         for (thread = 0; thread < nthread; thread++)
4023         {
4024             printf(" %.2f", cs1a[thread]*1e-9);
4025         }
4026 #endif
4027         printf("\n");
4028     }
4029 #endif
4030 }
4031
4032
4033 static void dump_grid(FILE *fp,
4034                       int sx, int sy, int sz, int nx, int ny, int nz,
4035                       int my, int mz, const real *g)
4036 {
4037     int x, y, z;
4038
4039     for (x = 0; x < nx; x++)
4040     {
4041         for (y = 0; y < ny; y++)
4042         {
4043             for (z = 0; z < nz; z++)
4044             {
4045                 fprintf(fp, "%2d %2d %2d %6.3f\n",
4046                         sx+x, sy+y, sz+z, g[(x*my + y)*mz + z]);
4047             }
4048         }
4049     }
4050 }
4051
4052 static void dump_local_fftgrid(gmx_pme_t pme, const real *fftgrid)
4053 {
4054     ivec local_fft_ndata, local_fft_offset, local_fft_size;
4055
4056     gmx_parallel_3dfft_real_limits(pme->pfft_setupA,
4057                                    local_fft_ndata,
4058                                    local_fft_offset,
4059                                    local_fft_size);
4060
4061     dump_grid(stderr,
4062               pme->pmegrid_start_ix,
4063               pme->pmegrid_start_iy,
4064               pme->pmegrid_start_iz,
4065               pme->pmegrid_nx-pme->pme_order+1,
4066               pme->pmegrid_ny-pme->pme_order+1,
4067               pme->pmegrid_nz-pme->pme_order+1,
4068               local_fft_size[YY],
4069               local_fft_size[ZZ],
4070               fftgrid);
4071 }
4072
4073
4074 void gmx_pme_calc_energy(gmx_pme_t pme, int n, rvec *x, real *q, real *V)
4075 {
4076     pme_atomcomm_t *atc;
4077     pmegrids_t *grid;
4078
4079     if (pme->nnodes > 1)
4080     {
4081         gmx_incons("gmx_pme_calc_energy called in parallel");
4082     }
4083     if (pme->bFEP > 1)
4084     {
4085         gmx_incons("gmx_pme_calc_energy with free energy");
4086     }
4087
4088     atc            = &pme->atc_energy;
4089     atc->nthread   = 1;
4090     if (atc->spline == NULL)
4091     {
4092         snew(atc->spline, atc->nthread);
4093     }
4094     atc->nslab     = 1;
4095     atc->bSpread   = TRUE;
4096     atc->pme_order = pme->pme_order;
4097     atc->n         = n;
4098     pme_realloc_atomcomm_things(atc);
4099     atc->x         = x;
4100     atc->q         = q;
4101
4102     /* We only use the A-charges grid */
4103     grid = &pme->pmegridA;
4104
4105     /* Only calculate the spline coefficients, don't actually spread */
4106     spread_on_grid(pme, atc, NULL, TRUE, FALSE, pme->fftgridA);
4107
4108     *V = gather_energy_bsplines(pme, grid->grid.grid, atc);
4109 }
4110
4111
4112 static void reset_pmeonly_counters(gmx_wallcycle_t wcycle,
4113                                    t_nrnb *nrnb, t_inputrec *ir,
4114                                    gmx_large_int_t step)
4115 {
4116     /* Reset all the counters related to performance over the run */
4117     wallcycle_stop(wcycle, ewcRUN);
4118     wallcycle_reset_all(wcycle);
4119     init_nrnb(nrnb);
4120     if (ir->nsteps >= 0)
4121     {
4122         /* ir->nsteps is not used here, but we update it for consistency */
4123         ir->nsteps -= step - ir->init_step;
4124     }
4125     ir->init_step = step;
4126     wallcycle_start(wcycle, ewcRUN);
4127 }
4128
4129
4130 static void gmx_pmeonly_switch(int *npmedata, gmx_pme_t **pmedata,
4131                                ivec grid_size,
4132                                t_commrec *cr, t_inputrec *ir,
4133                                gmx_pme_t *pme_ret)
4134 {
4135     int ind;
4136     gmx_pme_t pme = NULL;
4137
4138     ind = 0;
4139     while (ind < *npmedata)
4140     {
4141         pme = (*pmedata)[ind];
4142         if (pme->nkx == grid_size[XX] &&
4143             pme->nky == grid_size[YY] &&
4144             pme->nkz == grid_size[ZZ])
4145         {
4146             *pme_ret = pme;
4147
4148             return;
4149         }
4150
4151         ind++;
4152     }
4153
4154     (*npmedata)++;
4155     srenew(*pmedata, *npmedata);
4156
4157     /* Generate a new PME data structure, copying part of the old pointers */
4158     gmx_pme_reinit(&((*pmedata)[ind]), cr, pme, ir, grid_size);
4159
4160     *pme_ret = (*pmedata)[ind];
4161 }
4162
4163
4164 int gmx_pmeonly(gmx_pme_t pme,
4165                 t_commrec *cr,    t_nrnb *nrnb,
4166                 gmx_wallcycle_t wcycle,
4167                 real ewaldcoeff,
4168                 t_inputrec *ir)
4169 {
4170     int npmedata;
4171     gmx_pme_t *pmedata;
4172     gmx_pme_pp_t pme_pp;
4173     int  ret;
4174     int  natoms;
4175     matrix box;
4176     rvec *x_pp      = NULL, *f_pp = NULL;
4177     real *chargeA   = NULL, *chargeB = NULL;
4178     real lambda     = 0;
4179     int  maxshift_x = 0, maxshift_y = 0;
4180     real energy, dvdlambda;
4181     matrix vir;
4182     float cycles;
4183     int  count;
4184     gmx_bool bEnerVir;
4185     gmx_large_int_t step, step_rel;
4186     ivec grid_switch;
4187
4188     /* This data will only use with PME tuning, i.e. switching PME grids */
4189     npmedata = 1;
4190     snew(pmedata, npmedata);
4191     pmedata[0] = pme;
4192
4193     pme_pp = gmx_pme_pp_init(cr);
4194
4195     init_nrnb(nrnb);
4196
4197     count = 0;
4198     do /****** this is a quasi-loop over time steps! */
4199     {
4200         /* The reason for having a loop here is PME grid tuning/switching */
4201         do
4202         {
4203             /* Domain decomposition */
4204             ret = gmx_pme_recv_q_x(pme_pp,
4205                                    &natoms,
4206                                    &chargeA, &chargeB, box, &x_pp, &f_pp,
4207                                    &maxshift_x, &maxshift_y,
4208                                    &pme->bFEP, &lambda,
4209                                    &bEnerVir,
4210                                    &step,
4211                                    grid_switch, &ewaldcoeff);
4212
4213             if (ret == pmerecvqxSWITCHGRID)
4214             {
4215                 /* Switch the PME grid to grid_switch */
4216                 gmx_pmeonly_switch(&npmedata, &pmedata, grid_switch, cr, ir, &pme);
4217             }
4218
4219             if (ret == pmerecvqxRESETCOUNTERS)
4220             {
4221                 /* Reset the cycle and flop counters */
4222                 reset_pmeonly_counters(wcycle, nrnb, ir, step);
4223             }
4224         }
4225         while (ret == pmerecvqxSWITCHGRID || ret == pmerecvqxRESETCOUNTERS);
4226
4227         if (ret == pmerecvqxFINISH)
4228         {
4229             /* We should stop: break out of the loop */
4230             break;
4231         }
4232
4233         step_rel = step - ir->init_step;
4234
4235         if (count == 0)
4236         {
4237             wallcycle_start(wcycle, ewcRUN);
4238         }
4239
4240         wallcycle_start(wcycle, ewcPMEMESH);
4241
4242         dvdlambda = 0;
4243         clear_mat(vir);
4244         gmx_pme_do(pme, 0, natoms, x_pp, f_pp, chargeA, chargeB, box,
4245                    cr, maxshift_x, maxshift_y, nrnb, wcycle, vir, ewaldcoeff,
4246                    &energy, lambda, &dvdlambda,
4247                    GMX_PME_DO_ALL_F | (bEnerVir ? GMX_PME_CALC_ENER_VIR : 0));
4248
4249         cycles = wallcycle_stop(wcycle, ewcPMEMESH);
4250
4251         gmx_pme_send_force_vir_ener(pme_pp,
4252                                     f_pp, vir, energy, dvdlambda,
4253                                     cycles);
4254
4255         count++;
4256     } /***** end of quasi-loop, we stop with the break above */
4257     while (TRUE);
4258
4259     return 0;
4260 }
4261
4262 int gmx_pme_do(gmx_pme_t pme,
4263                int start,       int homenr,
4264                rvec x[],        rvec f[],
4265                real *chargeA,   real *chargeB,
4266                matrix box, t_commrec *cr,
4267                int  maxshift_x, int maxshift_y,
4268                t_nrnb *nrnb,    gmx_wallcycle_t wcycle,
4269                matrix vir,      real ewaldcoeff,
4270                real *energy,    real lambda,
4271                real *dvdlambda, int flags)
4272 {
4273     int     q, d, i, j, ntot, npme;
4274     int     nx, ny, nz;
4275     int     n_d, local_ny;
4276     pme_atomcomm_t *atc = NULL;
4277     pmegrids_t *pmegrid = NULL;
4278     real    *grid       = NULL;
4279     real    *ptr;
4280     rvec    *x_d, *f_d;
4281     real    *charge = NULL, *q_d;
4282     real    energy_AB[2];
4283     matrix  vir_AB[2];
4284     gmx_bool bClearF;
4285     gmx_parallel_3dfft_t pfft_setup;
4286     real *  fftgrid;
4287     t_complex * cfftgrid;
4288     int     thread;
4289     const gmx_bool bCalcEnerVir = flags & GMX_PME_CALC_ENER_VIR;
4290     const gmx_bool bCalcF       = flags & GMX_PME_CALC_F;
4291
4292     assert(pme->nnodes > 0);
4293     assert(pme->nnodes == 1 || pme->ndecompdim > 0);
4294
4295     if (pme->nnodes > 1)
4296     {
4297         atc      = &pme->atc[0];
4298         atc->npd = homenr;
4299         if (atc->npd > atc->pd_nalloc)
4300         {
4301             atc->pd_nalloc = over_alloc_dd(atc->npd);
4302             srenew(atc->pd, atc->pd_nalloc);
4303         }
4304         atc->maxshift = (atc->dimind == 0 ? maxshift_x : maxshift_y);
4305     }
4306     else
4307     {
4308         /* This could be necessary for TPI */
4309         pme->atc[0].n = homenr;
4310     }
4311
4312     for (q = 0; q < (pme->bFEP ? 2 : 1); q++)
4313     {
4314         if (q == 0)
4315         {
4316             pmegrid    = &pme->pmegridA;
4317             fftgrid    = pme->fftgridA;
4318             cfftgrid   = pme->cfftgridA;
4319             pfft_setup = pme->pfft_setupA;
4320             charge     = chargeA+start;
4321         }
4322         else
4323         {
4324             pmegrid    = &pme->pmegridB;
4325             fftgrid    = pme->fftgridB;
4326             cfftgrid   = pme->cfftgridB;
4327             pfft_setup = pme->pfft_setupB;
4328             charge     = chargeB+start;
4329         }
4330         grid = pmegrid->grid.grid;
4331         /* Unpack structure */
4332         if (debug)
4333         {
4334             fprintf(debug, "PME: nnodes = %d, nodeid = %d\n",
4335                     cr->nnodes, cr->nodeid);
4336             fprintf(debug, "Grid = %p\n", (void*)grid);
4337             if (grid == NULL)
4338             {
4339                 gmx_fatal(FARGS, "No grid!");
4340             }
4341         }
4342         where();
4343
4344         m_inv_ur0(box, pme->recipbox);
4345
4346         if (pme->nnodes == 1)
4347         {
4348             atc = &pme->atc[0];
4349             if (DOMAINDECOMP(cr))
4350             {
4351                 atc->n = homenr;
4352                 pme_realloc_atomcomm_things(atc);
4353             }
4354             atc->x = x;
4355             atc->q = charge;
4356             atc->f = f;
4357         }
4358         else
4359         {
4360             wallcycle_start(wcycle, ewcPME_REDISTXF);
4361             for (d = pme->ndecompdim-1; d >= 0; d--)
4362             {
4363                 if (d == pme->ndecompdim-1)
4364                 {
4365                     n_d = homenr;
4366                     x_d = x + start;
4367                     q_d = charge;
4368                 }
4369                 else
4370                 {
4371                     n_d = pme->atc[d+1].n;
4372                     x_d = atc->x;
4373                     q_d = atc->q;
4374                 }
4375                 atc      = &pme->atc[d];
4376                 atc->npd = n_d;
4377                 if (atc->npd > atc->pd_nalloc)
4378                 {
4379                     atc->pd_nalloc = over_alloc_dd(atc->npd);
4380                     srenew(atc->pd, atc->pd_nalloc);
4381                 }
4382                 atc->maxshift = (atc->dimind == 0 ? maxshift_x : maxshift_y);
4383                 pme_calc_pidx_wrapper(n_d, pme->recipbox, x_d, atc);
4384                 where();
4385
4386                 /* Redistribute x (only once) and qA or qB */
4387                 if (DOMAINDECOMP(cr))
4388                 {
4389                     dd_pmeredist_x_q(pme, n_d, q == 0, x_d, q_d, atc);
4390                 }
4391                 else
4392                 {
4393                     pmeredist_pd(pme, TRUE, n_d, q == 0, x_d, q_d, atc);
4394                 }
4395             }
4396             where();
4397
4398             wallcycle_stop(wcycle, ewcPME_REDISTXF);
4399         }
4400
4401         if (debug)
4402         {
4403             fprintf(debug, "Node= %6d, pme local particles=%6d\n",
4404                     cr->nodeid, atc->n);
4405         }
4406
4407         if (flags & GMX_PME_SPREAD_Q)
4408         {
4409             wallcycle_start(wcycle, ewcPME_SPREADGATHER);
4410
4411             /* Spread the charges on a grid */
4412             spread_on_grid(pme, &pme->atc[0], pmegrid, q == 0, TRUE, fftgrid);
4413
4414             if (q == 0)
4415             {
4416                 inc_nrnb(nrnb, eNR_WEIGHTS, DIM*atc->n);
4417             }
4418             inc_nrnb(nrnb, eNR_SPREADQBSP,
4419                      pme->pme_order*pme->pme_order*pme->pme_order*atc->n);
4420
4421             if (!pme->bUseThreads)
4422             {
4423                 wrap_periodic_pmegrid(pme, grid);
4424
4425                 /* sum contributions to local grid from other nodes */
4426 #ifdef GMX_MPI
4427                 if (pme->nnodes > 1)
4428                 {
4429                     gmx_sum_qgrid_dd(pme, grid, GMX_SUM_QGRID_FORWARD);
4430                     where();
4431                 }
4432 #endif
4433
4434                 copy_pmegrid_to_fftgrid(pme, grid, fftgrid);
4435             }
4436
4437             wallcycle_stop(wcycle, ewcPME_SPREADGATHER);
4438
4439             /*
4440                dump_local_fftgrid(pme,fftgrid);
4441                exit(0);
4442              */
4443         }
4444
4445         /* Here we start a large thread parallel region */
4446 #pragma omp parallel num_threads(pme->nthread) private(thread)
4447         {
4448             thread = gmx_omp_get_thread_num();
4449             if (flags & GMX_PME_SOLVE)
4450             {
4451                 int loop_count;
4452
4453                 /* do 3d-fft */
4454                 if (thread == 0)
4455                 {
4456                     wallcycle_start(wcycle, ewcPME_FFT);
4457                 }
4458                 gmx_parallel_3dfft_execute(pfft_setup, GMX_FFT_REAL_TO_COMPLEX,
4459                                            thread, wcycle);
4460                 if (thread == 0)
4461                 {
4462                     wallcycle_stop(wcycle, ewcPME_FFT);
4463                 }
4464                 where();
4465
4466                 /* solve in k-space for our local cells */
4467                 if (thread == 0)
4468                 {
4469                     wallcycle_start(wcycle, ewcPME_SOLVE);
4470                 }
4471                 loop_count =
4472                     solve_pme_yzx(pme, cfftgrid, ewaldcoeff,
4473                                   box[XX][XX]*box[YY][YY]*box[ZZ][ZZ],
4474                                   bCalcEnerVir,
4475                                   pme->nthread, thread);
4476                 if (thread == 0)
4477                 {
4478                     wallcycle_stop(wcycle, ewcPME_SOLVE);
4479                     where();
4480                     inc_nrnb(nrnb, eNR_SOLVEPME, loop_count);
4481                 }
4482             }
4483
4484             if (bCalcF)
4485             {
4486                 /* do 3d-invfft */
4487                 if (thread == 0)
4488                 {
4489                     where();
4490                     wallcycle_start(wcycle, ewcPME_FFT);
4491                 }
4492                 gmx_parallel_3dfft_execute(pfft_setup, GMX_FFT_COMPLEX_TO_REAL,
4493                                            thread, wcycle);
4494                 if (thread == 0)
4495                 {
4496                     wallcycle_stop(wcycle, ewcPME_FFT);
4497
4498                     where();
4499
4500                     if (pme->nodeid == 0)
4501                     {
4502                         ntot  = pme->nkx*pme->nky*pme->nkz;
4503                         npme  = ntot*log((real)ntot)/log(2.0);
4504                         inc_nrnb(nrnb, eNR_FFT, 2*npme);
4505                     }
4506
4507                     wallcycle_start(wcycle, ewcPME_SPREADGATHER);
4508                 }
4509
4510                 copy_fftgrid_to_pmegrid(pme, fftgrid, grid, pme->nthread, thread);
4511             }
4512         }
4513         /* End of thread parallel section.
4514          * With MPI we have to synchronize here before gmx_sum_qgrid_dd.
4515          */
4516
4517         if (bCalcF)
4518         {
4519             /* distribute local grid to all nodes */
4520 #ifdef GMX_MPI
4521             if (pme->nnodes > 1)
4522             {
4523                 gmx_sum_qgrid_dd(pme, grid, GMX_SUM_QGRID_BACKWARD);
4524             }
4525 #endif
4526             where();
4527
4528             unwrap_periodic_pmegrid(pme, grid);
4529
4530             /* interpolate forces for our local atoms */
4531
4532             where();
4533
4534             /* If we are running without parallelization,
4535              * atc->f is the actual force array, not a buffer,
4536              * therefore we should not clear it.
4537              */
4538             bClearF = (q == 0 && PAR(cr));
4539 #pragma omp parallel for num_threads(pme->nthread) schedule(static)
4540             for (thread = 0; thread < pme->nthread; thread++)
4541             {
4542                 gather_f_bsplines(pme, grid, bClearF, atc,
4543                                   &atc->spline[thread],
4544                                   pme->bFEP ? (q == 0 ? 1.0-lambda : lambda) : 1.0);
4545             }
4546
4547             where();
4548
4549             inc_nrnb(nrnb, eNR_GATHERFBSP,
4550                      pme->pme_order*pme->pme_order*pme->pme_order*pme->atc[0].n);
4551             wallcycle_stop(wcycle, ewcPME_SPREADGATHER);
4552         }
4553
4554         if (bCalcEnerVir)
4555         {
4556             /* This should only be called on the master thread
4557              * and after the threads have synchronized.
4558              */
4559             get_pme_ener_vir(pme, pme->nthread, &energy_AB[q], vir_AB[q]);
4560         }
4561     } /* of q-loop */
4562
4563     if (bCalcF && pme->nnodes > 1)
4564     {
4565         wallcycle_start(wcycle, ewcPME_REDISTXF);
4566         for (d = 0; d < pme->ndecompdim; d++)
4567         {
4568             atc = &pme->atc[d];
4569             if (d == pme->ndecompdim - 1)
4570             {
4571                 n_d = homenr;
4572                 f_d = f + start;
4573             }
4574             else
4575             {
4576                 n_d = pme->atc[d+1].n;
4577                 f_d = pme->atc[d+1].f;
4578             }
4579             if (DOMAINDECOMP(cr))
4580             {
4581                 dd_pmeredist_f(pme, atc, n_d, f_d,
4582                                d == pme->ndecompdim-1 && pme->bPPnode);
4583             }
4584             else
4585             {
4586                 pmeredist_pd(pme, FALSE, n_d, TRUE, f_d, NULL, atc);
4587             }
4588         }
4589
4590         wallcycle_stop(wcycle, ewcPME_REDISTXF);
4591     }
4592     where();
4593
4594     if (bCalcEnerVir)
4595     {
4596         if (!pme->bFEP)
4597         {
4598             *energy = energy_AB[0];
4599             m_add(vir, vir_AB[0], vir);
4600         }
4601         else
4602         {
4603             *energy     = (1.0-lambda)*energy_AB[0] + lambda*energy_AB[1];
4604             *dvdlambda += energy_AB[1] - energy_AB[0];
4605             for (i = 0; i < DIM; i++)
4606             {
4607                 for (j = 0; j < DIM; j++)
4608                 {
4609                     vir[i][j] += (1.0-lambda)*vir_AB[0][i][j] +
4610                         lambda*vir_AB[1][i][j];
4611                 }
4612             }
4613         }
4614     }
4615     else
4616     {
4617         *energy = 0;
4618     }
4619
4620     if (debug)
4621     {
4622         fprintf(debug, "PME mesh energy: %g\n", *energy);
4623     }
4624
4625     return 0;
4626 }