Merge release-4-6 into master
[alexxy/gromacs.git] / src / gromacs / mdlib / nbnxn_atomdata.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  * Written by David van der Spoel, Erik Lindahl, Berk Hess, and others.
11  * Copyright (c) 1991-2000, University of Groningen, The Netherlands.
12  * Copyright (c) 2001-2012, The GROMACS development team,
13  * check out http://www.gromacs.org for more information.
14
15  * This program is free software; you can redistribute it and/or
16  * modify it under the terms of the GNU General Public License
17  * as published by the Free Software Foundation; either version 2
18  * of the License, or (at your option) any later version.
19  *
20  * If you want to redistribute modifications, please consider that
21  * scientific software is very special. Version control is crucial -
22  * bugs must be traceable. We will be happy to consider code for
23  * inclusion in the official distribution, but derived work must not
24  * be called official GROMACS. Details are found in the README & COPYING
25  * files - if they are missing, get the official version at www.gromacs.org.
26  *
27  * To help us fund GROMACS development, we humbly ask that you cite
28  * the papers on the package - you can find them in the top README file.
29  *
30  * For more info, check our website at http://www.gromacs.org
31  */
32
33 #ifdef HAVE_CONFIG_H
34 #include <config.h>
35 #endif
36
37 #include <math.h>
38 #include <string.h>
39 #include "smalloc.h"
40 #include "macros.h"
41 #include "vec.h"
42 #include "nbnxn_consts.h"
43 #include "nbnxn_internal.h"
44 #include "nbnxn_search.h"
45 #include "nbnxn_atomdata.h"
46 #include "gmx_omp_nthreads.h"
47
48 /* Default nbnxn allocation routine, allocates 32 byte aligned,
49  * which works for plain C and aligned SSE and AVX loads/stores.
50  */
51 void nbnxn_alloc_aligned(void **ptr,size_t nbytes)
52 {
53     *ptr = save_malloc_aligned("ptr",__FILE__,__LINE__,nbytes,1,32);
54 }
55
56 /* Free function for memory allocated with nbnxn_alloc_aligned */
57 void nbnxn_free_aligned(void *ptr)
58 {
59     sfree_aligned(ptr);
60 }
61
62 /* Reallocation wrapper function for nbnxn data structures */
63 void nbnxn_realloc_void(void **ptr,
64                         int nbytes_copy,int nbytes_new,
65                         nbnxn_alloc_t *ma,
66                         nbnxn_free_t  *mf)
67 {
68     void *ptr_new;
69
70     ma(&ptr_new,nbytes_new);
71
72     if (nbytes_new > 0 && ptr_new == NULL)
73     {
74         gmx_fatal(FARGS, "Allocation of %d bytes failed", nbytes_new);
75     }
76
77     if (nbytes_copy > 0)
78     {
79         if (nbytes_new < nbytes_copy)
80         {
81             gmx_incons("In nbnxn_realloc_void: new size less than copy size");
82         }
83         memcpy(ptr_new,*ptr,nbytes_copy);
84     }
85     if (*ptr != NULL)
86     {
87         mf(*ptr);
88     }
89     *ptr = ptr_new;
90 }
91
92 /* Reallocate the nbnxn_atomdata_t for a size of n atoms */
93 void nbnxn_atomdata_realloc(nbnxn_atomdata_t *nbat,int n)
94 {
95     int t;
96
97     nbnxn_realloc_void((void **)&nbat->type,
98                        nbat->natoms*sizeof(*nbat->type),
99                        n*sizeof(*nbat->type),
100                        nbat->alloc,nbat->free);
101     nbnxn_realloc_void((void **)&nbat->lj_comb,
102                        nbat->natoms*2*sizeof(*nbat->lj_comb),
103                        n*2*sizeof(*nbat->lj_comb),
104                        nbat->alloc,nbat->free);
105     if (nbat->XFormat != nbatXYZQ)
106     {
107         nbnxn_realloc_void((void **)&nbat->q,
108                            nbat->natoms*sizeof(*nbat->q),
109                            n*sizeof(*nbat->q),
110                            nbat->alloc,nbat->free);
111     }
112     if (nbat->nenergrp > 1)
113     {
114         nbnxn_realloc_void((void **)&nbat->energrp,
115                            nbat->natoms/nbat->na_c*sizeof(*nbat->energrp),
116                            n/nbat->na_c*sizeof(*nbat->energrp),
117                            nbat->alloc,nbat->free);
118     }
119     nbnxn_realloc_void((void **)&nbat->x,
120                        nbat->natoms*nbat->xstride*sizeof(*nbat->x),
121                        n*nbat->xstride*sizeof(*nbat->x),
122                        nbat->alloc,nbat->free);
123     for(t=0; t<nbat->nout; t++)
124     {
125         /* Allocate one element extra for possible signaling with CUDA */
126         nbnxn_realloc_void((void **)&nbat->out[t].f,
127                            nbat->natoms*nbat->fstride*sizeof(*nbat->out[t].f),
128                            n*nbat->fstride*sizeof(*nbat->out[t].f),
129                            nbat->alloc,nbat->free);
130     }
131     nbat->nalloc = n;
132 }
133
134 /* Initializes an nbnxn_atomdata_output_t data structure */
135 static void nbnxn_atomdata_output_init(nbnxn_atomdata_output_t *out,
136                                        int nb_kernel_type,
137                                        int nenergrp,int stride,
138                                        nbnxn_alloc_t *ma)
139 {
140     int cj_size;
141
142     out->f = NULL;
143     ma((void **)&out->fshift,SHIFTS*DIM*sizeof(*out->fshift));
144     out->nV = nenergrp*nenergrp;
145     ma((void **)&out->Vvdw,out->nV*sizeof(*out->Vvdw));
146     ma((void **)&out->Vc  ,out->nV*sizeof(*out->Vc  ));
147
148     if (nb_kernel_type == nbk4xN_X86_SIMD128 ||
149         nb_kernel_type == nbk4xN_X86_SIMD256)
150     {
151         cj_size = nbnxn_kernel_to_cj_size(nb_kernel_type);
152         out->nVS = nenergrp*nenergrp*stride*(cj_size>>1)*cj_size;
153         ma((void **)&out->VSvdw,out->nVS*sizeof(*out->VSvdw));
154         ma((void **)&out->VSc  ,out->nVS*sizeof(*out->VSc  ));
155     }
156     else
157     {
158         out->nVS = 0;
159     }
160 }
161
162 static void copy_int_to_nbat_int(const int *a,int na,int na_round,
163                                  const int *in,int fill,int *innb)
164 {
165     int i,j;
166
167     j = 0;
168     for(i=0; i<na; i++)
169     {
170         innb[j++] = in[a[i]];
171     }
172     /* Complete the partially filled last cell with fill */
173     for(; i<na_round; i++)
174     {
175         innb[j++] = fill;
176     }
177 }
178
179 static void clear_nbat_real(int na,int nbatFormat,real *xnb,int a0)
180 {
181     int a,d,j,c;
182
183     switch (nbatFormat)
184     {
185     case nbatXYZ:
186         for(a=0; a<na; a++)
187         {
188             for(d=0; d<DIM; d++)
189             {
190                 xnb[(a0+a)*STRIDE_XYZ+d] = 0;
191             }
192         }
193         break;
194     case nbatXYZQ:
195         for(a=0; a<na; a++)
196         {
197             for(d=0; d<DIM; d++)
198             {
199                 xnb[(a0+a)*STRIDE_XYZQ+d] = 0;
200             }
201         }
202         break;
203     case nbatX4:
204         j = X4_IND_A(a0);
205         c = a0 & (PACK_X4-1);
206         for(a=0; a<na; a++)
207         {
208             xnb[j+XX*PACK_X4] = 0;
209             xnb[j+YY*PACK_X4] = 0;
210             xnb[j+ZZ*PACK_X4] = 0;
211             j++;
212             c++;
213             if (c == PACK_X4)
214             {
215                 j += (DIM-1)*PACK_X4;
216                 c  = 0;
217             }
218         }
219         break;
220     case nbatX8:
221         j = X8_IND_A(a0);
222         c = a0 & (PACK_X8-1);
223         for(a=0; a<na; a++)
224         {
225             xnb[j+XX*PACK_X8] = 0;
226             xnb[j+YY*PACK_X8] = 0;
227             xnb[j+ZZ*PACK_X8] = 0;
228             j++;
229             c++;
230             if (c == PACK_X8)
231             {
232                 j += (DIM-1)*PACK_X8;
233                 c  = 0;
234             }
235         }
236         break;
237     }
238 }
239
240 void copy_rvec_to_nbat_real(const int *a,int na,int na_round,
241                             rvec *x,int nbatFormat,real *xnb,int a0,
242                             int cx,int cy,int cz)
243 {
244     int i,j,c;
245
246 /* We might need to place filler particles to fill up the cell to na_round.
247  * The coefficients (LJ and q) for such particles are zero.
248  * But we might still get NaN as 0*NaN when distances are too small.
249  * We hope that -107 nm is far away enough from to zero
250  * to avoid accidental short distances to particles shifted down for pbc.
251  */
252 #define NBAT_FAR_AWAY 107
253
254     switch (nbatFormat)
255     {
256     case nbatXYZ:
257         j = a0*STRIDE_XYZ;
258         for(i=0; i<na; i++)
259         {
260             xnb[j++] = x[a[i]][XX];
261             xnb[j++] = x[a[i]][YY];
262             xnb[j++] = x[a[i]][ZZ];
263         }
264         /* Complete the partially filled last cell with copies of the last element.
265          * This simplifies the bounding box calculation and avoid
266          * numerical issues with atoms that are coincidentally close.
267          */
268         for(; i<na_round; i++)
269         {
270             xnb[j++] = -NBAT_FAR_AWAY*(1 + cx);
271             xnb[j++] = -NBAT_FAR_AWAY*(1 + cy);
272             xnb[j++] = -NBAT_FAR_AWAY*(1 + cz + i);
273         }
274         break;
275     case nbatXYZQ:
276         j = a0*STRIDE_XYZQ;
277         for(i=0; i<na; i++)
278         {
279             xnb[j++] = x[a[i]][XX];
280             xnb[j++] = x[a[i]][YY];
281             xnb[j++] = x[a[i]][ZZ];
282             j++;
283         }
284         /* Complete the partially filled last cell with particles far apart */
285         for(; i<na_round; i++)
286         {
287             xnb[j++] = -NBAT_FAR_AWAY*(1 + cx);
288             xnb[j++] = -NBAT_FAR_AWAY*(1 + cy);
289             xnb[j++] = -NBAT_FAR_AWAY*(1 + cz + i);
290             j++;
291         }
292         break;
293     case nbatX4:
294         j = X4_IND_A(a0);
295         c = a0 & (PACK_X4-1);
296         for(i=0; i<na; i++)
297         {
298             xnb[j+XX*PACK_X4] = x[a[i]][XX];
299             xnb[j+YY*PACK_X4] = x[a[i]][YY];
300             xnb[j+ZZ*PACK_X4] = x[a[i]][ZZ];
301             j++;
302             c++;
303             if (c == PACK_X4)
304             {
305                 j += (DIM-1)*PACK_X4;
306                 c  = 0;
307             }
308         }
309         /* Complete the partially filled last cell with particles far apart */
310         for(; i<na_round; i++)
311         {
312             xnb[j+XX*PACK_X4] = -NBAT_FAR_AWAY*(1 + cx);
313             xnb[j+YY*PACK_X4] = -NBAT_FAR_AWAY*(1 + cy);
314             xnb[j+ZZ*PACK_X4] = -NBAT_FAR_AWAY*(1 + cz + i);
315             j++;
316             c++;
317             if (c == PACK_X4)
318             {
319                 j += (DIM-1)*PACK_X4;
320                 c  = 0;
321             }
322         }
323         break;
324     case nbatX8:
325         j = X8_IND_A(a0);
326         c = a0 & (PACK_X8 - 1);
327         for(i=0; i<na; i++)
328         {
329             xnb[j+XX*PACK_X8] = x[a[i]][XX];
330             xnb[j+YY*PACK_X8] = x[a[i]][YY];
331             xnb[j+ZZ*PACK_X8] = x[a[i]][ZZ];
332             j++;
333             c++;
334             if (c == PACK_X8)
335             {
336                 j += (DIM-1)*PACK_X8;
337                 c  = 0;
338             }
339         }
340         /* Complete the partially filled last cell with particles far apart */
341         for(; i<na_round; i++)
342         {
343             xnb[j+XX*PACK_X8] = -NBAT_FAR_AWAY*(1 + cx);
344             xnb[j+YY*PACK_X8] = -NBAT_FAR_AWAY*(1 + cy);
345             xnb[j+ZZ*PACK_X8] = -NBAT_FAR_AWAY*(1 + cz + i);
346             j++;
347             c++;
348             if (c == PACK_X8)
349             {
350                 j += (DIM-1)*PACK_X8;
351                 c  = 0;
352             }
353         }
354         break;
355     default:
356         gmx_incons("Unsupported stride");
357     }
358 }
359
360 /* Determines the combination rule (or none) to be used, stores it,
361  * and sets the LJ parameters required with the rule.
362  */
363 static void set_combination_rule_data(nbnxn_atomdata_t *nbat)
364 {
365     int  nt,i,j;
366     real c6,c12;
367
368     nt = nbat->ntype;
369
370     switch (nbat->comb_rule)
371     {
372     case  ljcrGEOM:
373         nbat->comb_rule = ljcrGEOM;
374
375         for(i=0; i<nt; i++)
376         {
377             /* Copy the diagonal from the nbfp matrix */
378             nbat->nbfp_comb[i*2  ] = sqrt(nbat->nbfp[(i*nt+i)*2  ]);
379             nbat->nbfp_comb[i*2+1] = sqrt(nbat->nbfp[(i*nt+i)*2+1]);
380         }
381         break;
382     case ljcrLB:
383         for(i=0; i<nt; i++)
384         {
385             /* Get 6*C6 and 12*C12 from the diagonal of the nbfp matrix */
386             c6  = nbat->nbfp[(i*nt+i)*2  ];
387             c12 = nbat->nbfp[(i*nt+i)*2+1];
388             if (c6 > 0 && c12 > 0)
389             {
390                 /* We store 0.5*2^1/6*sigma and sqrt(4*3*eps),
391                  * so we get 6*C6 and 12*C12 after combining.
392                  */
393                 nbat->nbfp_comb[i*2  ] = 0.5*pow(c12/c6,1.0/6.0);
394                 nbat->nbfp_comb[i*2+1] = sqrt(c6*c6/c12);
395             }
396             else
397             {
398                 nbat->nbfp_comb[i*2  ] = 0;
399                 nbat->nbfp_comb[i*2+1] = 0;
400             }
401         }
402         break;
403     case ljcrNONE:
404         /* In nbfp_s4 we use a stride of 4 for storing two parameters */
405         nbat->alloc((void **)&nbat->nbfp_s4,nt*nt*4*sizeof(*nbat->nbfp_s4));
406         for(i=0; i<nt; i++)
407         {
408             for(j=0; j<nt; j++)
409             {
410                 nbat->nbfp_s4[(i*nt+j)*4+0] = nbat->nbfp[(i*nt+j)*2+0];
411                 nbat->nbfp_s4[(i*nt+j)*4+1] = nbat->nbfp[(i*nt+j)*2+1];
412                 nbat->nbfp_s4[(i*nt+j)*4+2] = 0;
413                 nbat->nbfp_s4[(i*nt+j)*4+3] = 0;
414             }
415         }
416         break;
417     default:
418         gmx_incons("Unknown combination rule");
419         break;
420     }
421 }
422
423 /* Initializes an nbnxn_atomdata_t data structure */
424 void nbnxn_atomdata_init(FILE *fp,
425                          nbnxn_atomdata_t *nbat,
426                          int nb_kernel_type,
427                          int ntype,const real *nbfp,
428                          int n_energygroups,
429                          int nout,
430                          nbnxn_alloc_t *alloc,
431                          nbnxn_free_t  *free)
432 {
433     int  i,j;
434     real c6,c12,tol;
435     char *ptr;
436     gmx_bool simple,bCombGeom,bCombLB;
437
438     if (alloc == NULL)
439     {
440         nbat->alloc = nbnxn_alloc_aligned;
441     }
442     else
443     {
444         nbat->alloc = alloc;
445     }
446     if (free == NULL)
447     {
448         nbat->free = nbnxn_free_aligned;
449     }
450     else
451     {
452         nbat->free = free;
453     }
454
455     if (debug)
456     {
457         fprintf(debug,"There are %d atom types in the system, adding one for nbnxn_atomdata_t\n",ntype);
458     }
459     nbat->ntype = ntype + 1;
460     nbat->alloc((void **)&nbat->nbfp,
461                 nbat->ntype*nbat->ntype*2*sizeof(*nbat->nbfp));
462     nbat->alloc((void **)&nbat->nbfp_comb,nbat->ntype*2*sizeof(*nbat->nbfp_comb));
463
464     /* A tolerance of 1e-5 seems reasonable for (possibly hand-typed)
465      * force-field floating point parameters.
466      */
467     tol = 1e-5;
468     ptr = getenv("GMX_LJCOMB_TOL");
469     if (ptr != NULL)
470     {
471         double dbl;
472
473         sscanf(ptr,"%lf",&dbl);
474         tol = dbl;
475     }
476     bCombGeom = TRUE;
477     bCombLB   = TRUE;
478
479     /* Temporarily fill nbat->nbfp_comb with sigma and epsilon
480      * to check for the LB rule.
481      */
482     for(i=0; i<ntype; i++)
483     {
484         c6  = nbfp[(i*ntype+i)*2  ];
485         c12 = nbfp[(i*ntype+i)*2+1];
486         if (c6 > 0 && c12 > 0)
487         {
488             nbat->nbfp_comb[i*2  ] = pow(c12/c6,1.0/6.0);
489             nbat->nbfp_comb[i*2+1] = 0.25*c6*c6/c12;
490         }
491         else if (c6 == 0 && c12 == 0)
492         {
493             nbat->nbfp_comb[i*2  ] = 0;
494             nbat->nbfp_comb[i*2+1] = 0;
495         }
496         else
497         {
498             /* Can not use LB rule with only dispersion or repulsion */
499             bCombLB = FALSE;
500         }
501     }
502
503     for(i=0; i<nbat->ntype; i++)
504     {
505         for(j=0; j<nbat->ntype; j++)
506         {
507             if (i < ntype && j < ntype)
508             {
509                 /* fr->nbfp has been updated, so that array too now stores c6/c12 including
510                  * the 6.0/12.0 prefactors to save 2 flops in the most common case (force-only).
511                  */
512                 c6  = nbfp[(i*ntype+j)*2  ];
513                 c12 = nbfp[(i*ntype+j)*2+1];
514                 nbat->nbfp[(i*nbat->ntype+j)*2  ] = c6;
515                 nbat->nbfp[(i*nbat->ntype+j)*2+1] = c12;
516                 c6  /= 6.0;
517                 c12 /= 12.0;
518                 
519                 bCombGeom = bCombGeom &&
520                     gmx_within_tol(c6*c6  ,nbfp[(i*ntype+i)*2  ]*nbfp[(j*ntype+j)*2  ],tol) &&
521                     gmx_within_tol(c12*c12,nbfp[(i*ntype+i)*2+1]*nbfp[(j*ntype+j)*2+1],tol);
522
523                 bCombLB = bCombLB &&
524                     ((c6 == 0 && c12 == 0 &&
525                       (nbat->nbfp_comb[i*2+1] == 0 || nbat->nbfp_comb[j*2+1] == 0)) ||
526                      (c6 > 0 && c12 > 0 &&
527                       gmx_within_tol(pow(c12/c6,1.0/6.0),0.5*(nbat->nbfp_comb[i*2]+nbat->nbfp_comb[j*2]),tol) &&
528                       gmx_within_tol(0.25*c6*c6/c12,sqrt(nbat->nbfp_comb[i*2+1]*nbat->nbfp_comb[j*2+1]),tol)));
529             }
530             else
531             {
532                 /* Add zero parameters for the additional dummy atom type */
533                 nbat->nbfp[(i*nbat->ntype+j)*2  ] = 0;
534                 nbat->nbfp[(i*nbat->ntype+j)*2+1] = 0;
535             }
536         }
537     }
538     if (debug)
539     {
540         fprintf(debug,"Combination rules: geometric %d Lorentz-Berthelot %d\n",
541                 bCombGeom,bCombLB);
542     }
543
544     simple = nbnxn_kernel_pairlist_simple(nb_kernel_type);
545
546     if (simple)
547     {
548         /* We prefer the geometic combination rule,
549          * as that gives a slightly faster kernel than the LB rule.
550          */
551         if (bCombGeom)
552         {
553             nbat->comb_rule = ljcrGEOM;
554         }
555         else if (bCombLB)
556         {
557             nbat->comb_rule = ljcrLB;
558         }
559         else
560         {
561             nbat->comb_rule = ljcrNONE;
562
563             nbat->free(nbat->nbfp_comb);
564         }
565
566         if (fp)
567         {
568             if (nbat->comb_rule == ljcrNONE)
569             {
570                 fprintf(fp,"Using full Lennard-Jones parameter combination matrix\n\n");
571             }
572             else
573             {
574                 fprintf(fp,"Using %s Lennard-Jones combination rule\n\n",
575                         nbat->comb_rule==ljcrGEOM ? "geometric" : "Lorentz-Berthelot");
576             }
577         }
578
579         set_combination_rule_data(nbat);
580     }
581     else
582     {
583         nbat->comb_rule = ljcrNONE;
584
585         nbat->free(nbat->nbfp_comb);
586     }
587
588     nbat->natoms  = 0;
589     nbat->type    = NULL;
590     nbat->lj_comb = NULL;
591     if (simple)
592     {
593         switch (nb_kernel_type)
594         {
595         case nbk4xN_X86_SIMD128:
596             nbat->XFormat = nbatX4;
597             break;
598         case nbk4xN_X86_SIMD256:
599 #ifndef GMX_DOUBLE
600             nbat->XFormat = nbatX8;
601 #else
602             nbat->XFormat = nbatX4;
603 #endif
604             break;
605         default:
606             nbat->XFormat = nbatXYZ;
607             break;
608         }
609
610         nbat->FFormat = nbat->XFormat;
611     }
612     else
613     {
614         nbat->XFormat = nbatXYZQ;
615         nbat->FFormat = nbatXYZ;
616     }
617     nbat->q       = NULL;
618     nbat->nenergrp = n_energygroups;
619     if (!simple)
620     {
621         /* Energy groups not supported yet for super-sub lists */
622         if (n_energygroups > 1 && fp != NULL)
623         {
624             fprintf(fp,"\nNOTE: With GPUs, reporting energy group contributions is not supported\n\n");
625         }
626         nbat->nenergrp = 1;
627     }
628     /* Temporary storage goes as #grp^3*simd_width^2/2, so limit to 64 */
629     if (nbat->nenergrp > 64)
630     {
631         gmx_fatal(FARGS,"With NxN kernels not more than 64 energy groups are supported\n");
632     }
633     nbat->neg_2log = 1;
634     while (nbat->nenergrp > (1<<nbat->neg_2log))
635     {
636         nbat->neg_2log++;
637     }
638     nbat->energrp = NULL;
639     nbat->alloc((void **)&nbat->shift_vec,SHIFTS*sizeof(*nbat->shift_vec));
640     nbat->xstride = (nbat->XFormat == nbatXYZQ ? STRIDE_XYZQ : DIM);
641     nbat->fstride = (nbat->FFormat == nbatXYZQ ? STRIDE_XYZQ : DIM);
642     nbat->x       = NULL;
643     nbat->nout    = nout;
644     snew(nbat->out,nbat->nout);
645     nbat->nalloc  = 0;
646     for(i=0; i<nbat->nout; i++)
647     {
648         nbnxn_atomdata_output_init(&nbat->out[i],
649                                    nb_kernel_type,
650                                    nbat->nenergrp,1<<nbat->neg_2log,
651                                    nbat->alloc);
652     }
653 }
654
655 static void copy_lj_to_nbat_lj_comb_x4(const real *ljparam_type,
656                                        const int *type,int na,
657                                        real *ljparam_at)
658 {
659     int is,k,i;
660
661     /* The LJ params follow the combination rule:
662      * copy the params for the type array to the atom array.
663      */
664     for(is=0; is<na; is+=PACK_X4)
665     {
666         for(k=0; k<PACK_X4; k++)
667         {
668             i = is + k;
669             ljparam_at[is*2        +k] = ljparam_type[type[i]*2  ];
670             ljparam_at[is*2+PACK_X4+k] = ljparam_type[type[i]*2+1];
671         }
672     }
673 }
674
675 static void copy_lj_to_nbat_lj_comb_x8(const real *ljparam_type,
676                                        const int *type,int na,
677                                        real *ljparam_at)
678 {
679     int is,k,i;
680
681     /* The LJ params follow the combination rule:
682      * copy the params for the type array to the atom array.
683      */
684     for(is=0; is<na; is+=PACK_X8)
685     {
686         for(k=0; k<PACK_X8; k++)
687         {
688             i = is + k;
689             ljparam_at[is*2        +k] = ljparam_type[type[i]*2  ];
690             ljparam_at[is*2+PACK_X8+k] = ljparam_type[type[i]*2+1];
691         }
692     }
693 }
694
695 /* Sets the atom type and LJ data in nbnxn_atomdata_t */
696 static void nbnxn_atomdata_set_atomtypes(nbnxn_atomdata_t *nbat,
697                                          int ngrid,
698                                          const nbnxn_search_t nbs,
699                                          const int *type)
700 {
701     int g,i,ncz,ash;
702     const nbnxn_grid_t *grid;
703
704     for(g=0; g<ngrid; g++)
705     {
706         grid = &nbs->grid[g];
707
708         /* Loop over all columns and copy and fill */
709         for(i=0; i<grid->ncx*grid->ncy; i++)
710         {
711             ncz = grid->cxy_ind[i+1] - grid->cxy_ind[i];
712             ash = (grid->cell0 + grid->cxy_ind[i])*grid->na_sc;
713
714             copy_int_to_nbat_int(nbs->a+ash,grid->cxy_na[i],ncz*grid->na_sc,
715                                  type,nbat->ntype-1,nbat->type+ash);
716
717             if (nbat->comb_rule != ljcrNONE)
718             {
719                 if (nbat->XFormat == nbatX4)
720                 {
721                     copy_lj_to_nbat_lj_comb_x4(nbat->nbfp_comb,
722                                                nbat->type+ash,ncz*grid->na_sc,
723                                                nbat->lj_comb+ash*2);
724                 }
725                 else if (nbat->XFormat == nbatX8)
726                 {
727                     copy_lj_to_nbat_lj_comb_x8(nbat->nbfp_comb,
728                                                nbat->type+ash,ncz*grid->na_sc,
729                                                nbat->lj_comb+ash*2);
730                 }
731             }
732         }
733     }
734 }
735
736 /* Sets the charges in nbnxn_atomdata_t *nbat */
737 static void nbnxn_atomdata_set_charges(nbnxn_atomdata_t *nbat,
738                                        int ngrid,
739                                        const nbnxn_search_t nbs,
740                                        const real *charge)
741 {
742     int  g,cxy,ncz,ash,na,na_round,i,j;
743     real *q;
744     const nbnxn_grid_t *grid;
745
746     for(g=0; g<ngrid; g++)
747     {
748         grid = &nbs->grid[g];
749
750         /* Loop over all columns and copy and fill */
751         for(cxy=0; cxy<grid->ncx*grid->ncy; cxy++)
752         {
753             ash = (grid->cell0 + grid->cxy_ind[cxy])*grid->na_sc;
754             na  = grid->cxy_na[cxy];
755             na_round = (grid->cxy_ind[cxy+1] - grid->cxy_ind[cxy])*grid->na_sc;
756
757             if (nbat->XFormat == nbatXYZQ)
758             {
759                 q = nbat->x + ash*STRIDE_XYZQ + ZZ + 1;
760                 for(i=0; i<na; i++)
761                 {
762                     *q = charge[nbs->a[ash+i]];
763                     q += STRIDE_XYZQ;
764                 }
765                 /* Complete the partially filled last cell with zeros */
766                 for(; i<na_round; i++)
767                 {
768                     *q = 0;
769                     q += STRIDE_XYZQ;
770                 }
771             }
772             else
773             {
774                 q = nbat->q + ash;
775                 for(i=0; i<na; i++)
776                 {
777                     *q = charge[nbs->a[ash+i]];
778                     q++;
779                 }
780                 /* Complete the partially filled last cell with zeros */
781                 for(; i<na_round; i++)
782                 {
783                     *q = 0;
784                     q++;
785                 }
786             }
787         }
788     }
789 }
790
791 /* Copies the energy group indices to a reordered and packed array */
792 static void copy_egp_to_nbat_egps(const int *a,int na,int na_round,
793                                   int na_c,int bit_shift,
794                                   const int *in,int *innb)
795 {
796     int i,j,sa,at;
797     int comb;
798
799     j = 0;
800     for(i=0; i<na; i+=na_c)
801     {
802         /* Store na_c energy group numbers into one int */
803         comb = 0;
804         for(sa=0; sa<na_c; sa++)
805         {
806             at = a[i+sa];
807             if (at >= 0)
808             {
809                 comb |= (GET_CGINFO_GID(in[at]) << (sa*bit_shift));
810             }
811         }
812         innb[j++] = comb;
813     }
814     /* Complete the partially filled last cell with fill */
815     for(; i<na_round; i+=na_c)
816     {
817         innb[j++] = 0;
818     }
819 }
820
821 /* Set the energy group indices for atoms in nbnxn_atomdata_t */
822 static void nbnxn_atomdata_set_energygroups(nbnxn_atomdata_t *nbat,
823                                             int ngrid,
824                                             const nbnxn_search_t nbs,
825                                             const int *atinfo)
826 {
827     int g,i,ncz,ash;
828     const nbnxn_grid_t *grid;
829
830     for(g=0; g<ngrid; g++)
831     {
832         grid = &nbs->grid[g];
833
834         /* Loop over all columns and copy and fill */
835         for(i=0; i<grid->ncx*grid->ncy; i++)
836         {
837             ncz = grid->cxy_ind[i+1] - grid->cxy_ind[i];
838             ash = (grid->cell0 + grid->cxy_ind[i])*grid->na_sc;
839
840             copy_egp_to_nbat_egps(nbs->a+ash,grid->cxy_na[i],ncz*grid->na_sc,
841                                   nbat->na_c,nbat->neg_2log,
842                                   atinfo,nbat->energrp+(ash>>grid->na_c_2log));
843         }
844     }
845 }
846
847 /* Sets all required atom parameter data in nbnxn_atomdata_t */
848 void nbnxn_atomdata_set(nbnxn_atomdata_t *nbat,
849                         int locality,
850                         const nbnxn_search_t nbs,
851                         const t_mdatoms *mdatoms,
852                         const int *atinfo)
853 {
854     int ngrid;
855
856     if (locality == eatLocal)
857     {
858         ngrid = 1;
859     }
860     else
861     {
862         ngrid = nbs->ngrid;
863     }
864
865     nbnxn_atomdata_set_atomtypes(nbat,ngrid,nbs,mdatoms->typeA);
866
867     nbnxn_atomdata_set_charges(nbat,ngrid,nbs,mdatoms->chargeA);
868
869     if (nbat->nenergrp > 1)
870     {
871         nbnxn_atomdata_set_energygroups(nbat,ngrid,nbs,atinfo);
872     }
873 }
874
875 /* Copies the shift vector array to nbnxn_atomdata_t */
876 void nbnxn_atomdata_copy_shiftvec(gmx_bool bDynamicBox,
877                                    rvec *shift_vec,
878                                    nbnxn_atomdata_t *nbat)
879 {
880     int i;
881
882     nbat->bDynamicBox = bDynamicBox;
883     for(i=0; i<SHIFTS; i++)
884     {
885         copy_rvec(shift_vec[i],nbat->shift_vec[i]);
886     }
887 }
888
889 /* Copies (and reorders) the coordinates to nbnxn_atomdata_t */
890 void nbnxn_atomdata_copy_x_to_nbat_x(const nbnxn_search_t nbs,
891                                       int locality,
892                                       gmx_bool FillLocal,
893                                       rvec *x,
894                                       nbnxn_atomdata_t *nbat)
895 {
896     int g0=0,g1=0;
897     int nth,th;
898
899     switch (locality)
900     {
901     case eatAll:
902         g0 = 0;
903         g1 = nbs->ngrid;
904         break;
905     case eatLocal:
906         g0 = 0;
907         g1 = 1;
908         break;
909     case eatNonlocal:
910         g0 = 1;
911         g1 = nbs->ngrid;
912         break;
913     }
914
915     if (FillLocal)
916     {
917         nbat->natoms_local = nbs->grid[0].nc*nbs->grid[0].na_sc;
918     }
919
920     nth = gmx_omp_nthreads_get(emntPairsearch);
921
922 #pragma omp parallel for num_threads(nth) schedule(static)
923     for(th=0; th<nth; th++)
924     {
925         int g;
926
927         for(g=g0; g<g1; g++)
928         {
929             const nbnxn_grid_t *grid;
930             int cxy0,cxy1,cxy;
931
932             grid = &nbs->grid[g];
933
934             cxy0 = (grid->ncx*grid->ncy* th   +nth-1)/nth;
935             cxy1 = (grid->ncx*grid->ncy*(th+1)+nth-1)/nth;
936
937             for(cxy=cxy0; cxy<cxy1; cxy++)
938             {
939                 int na,ash,na_fill;
940
941                 na  = grid->cxy_na[cxy];
942                 ash = (grid->cell0 + grid->cxy_ind[cxy])*grid->na_sc;
943
944                 if (g == 0 && FillLocal)
945                 {
946                     na_fill =
947                         (grid->cxy_ind[cxy+1] - grid->cxy_ind[cxy])*grid->na_sc;
948                 }
949                 else
950                 {
951                     /* We fill only the real particle locations.
952                      * We assume the filling entries at the end have been
953                      * properly set before during ns.
954                      */
955                     na_fill = na;
956                 }
957                 copy_rvec_to_nbat_real(nbs->a+ash,na,na_fill,x,
958                                        nbat->XFormat,nbat->x,ash,
959                                        0,0,0);
960             }
961         }
962     }
963 }
964
965 static void
966 nbnxn_atomdata_reduce_reals(real * gmx_restrict dest,
967                             nbnxn_atomdata_output_t * gmx_restrict src,
968                             int nsrc,
969                             int i0, int i1)
970 {
971     int i,s;
972
973     for(i=i0; i<i1; i++)
974     {
975         for(s=0; s<nsrc; s++)
976         {
977             dest[i] += src[s].f[i];
978         }
979     }
980 }
981
982 static void
983 nbnxn_atomdata_reduce_reals_x86_simd(real * gmx_restrict dest,
984                                      nbnxn_atomdata_output_t * gmx_restrict src,
985                                      int nsrc,
986                                      int i0, int i1)
987 {
988 #ifdef NBNXN_SEARCH_SSE
989 #ifdef GMX_X86_AVX_256
990 #define GMX_MM256_HERE
991 #else
992 #define GMX_MM128_HERE
993 #endif
994 #include "gmx_x86_simd_macros.h"
995
996     int       i,s;
997     gmx_mm_pr dest_SSE,src_SSE;
998
999     if ((i0 & (GMX_X86_SIMD_WIDTH_HERE-1)) ||
1000         (i1 & (GMX_X86_SIMD_WIDTH_HERE-1)))
1001     {
1002         gmx_incons("bounds not a multiple of GMX_X86_SIMD_WIDTH_HERE in nbnxn_atomdata_reduce_reals_x86_simd");
1003     }
1004
1005     for(i=i0; i<i1; i+=GMX_X86_SIMD_WIDTH_HERE)
1006     {
1007         dest_SSE = gmx_load_pr(dest+i);
1008         for(s=0; s<nsrc; s++)
1009         {
1010             src_SSE  = gmx_load_pr(src[s].f+i);
1011             dest_SSE = gmx_add_pr(dest_SSE,src_SSE);
1012         }
1013         gmx_store_pr(dest+i,dest_SSE);
1014     }
1015
1016 #undef GMX_MM128_HERE
1017 #undef GMX_MM256_HERE
1018 #endif
1019 }
1020
1021 /* Add part of the force array(s) from nbnxn_atomdata_t to f */
1022 static void
1023 nbnxn_atomdata_add_nbat_f_to_f_part(const nbnxn_search_t nbs,
1024                                     const nbnxn_atomdata_t *nbat,
1025                                     nbnxn_atomdata_output_t *out,
1026                                     int nfa,
1027                                     int a0,int a1,
1028                                     rvec *f)
1029 {
1030     int  a,i,fa;
1031     const int  *cell;
1032     const real *fnb;
1033
1034     cell = nbs->cell;
1035
1036     /* Loop over all columns and copy and fill */
1037     switch (nbat->FFormat)
1038     {
1039     case nbatXYZ:
1040     case nbatXYZQ:
1041         if (nfa == 1)
1042         {
1043             fnb = out[0].f;
1044
1045             for(a=a0; a<a1; a++)
1046             {
1047                 i = cell[a]*nbat->fstride;
1048
1049                 f[a][XX] += fnb[i];
1050                 f[a][YY] += fnb[i+1];
1051                 f[a][ZZ] += fnb[i+2];
1052             }
1053         }
1054         else
1055         {
1056             for(a=a0; a<a1; a++)
1057             {
1058                 i = cell[a]*nbat->fstride;
1059
1060                 for(fa=0; fa<nfa; fa++)
1061                 {
1062                     f[a][XX] += out[fa].f[i];
1063                     f[a][YY] += out[fa].f[i+1];
1064                     f[a][ZZ] += out[fa].f[i+2];
1065                 }
1066             }
1067         }
1068         break;
1069     case nbatX4:
1070         if (nfa == 1)
1071         {
1072             fnb = out[0].f;
1073
1074             for(a=a0; a<a1; a++)
1075             {
1076                 i = X4_IND_A(cell[a]);
1077
1078                 f[a][XX] += fnb[i+XX*PACK_X4];
1079                 f[a][YY] += fnb[i+YY*PACK_X4];
1080                 f[a][ZZ] += fnb[i+ZZ*PACK_X4];
1081             }
1082         }
1083         else
1084         {
1085             for(a=a0; a<a1; a++)
1086             {
1087                 i = X4_IND_A(cell[a]);
1088                 
1089                 for(fa=0; fa<nfa; fa++)
1090                 {
1091                     f[a][XX] += out[fa].f[i+XX*PACK_X4];
1092                     f[a][YY] += out[fa].f[i+YY*PACK_X4];
1093                     f[a][ZZ] += out[fa].f[i+ZZ*PACK_X4];
1094                 }
1095             }
1096         }
1097         break;
1098     case nbatX8:
1099         if (nfa == 1)
1100         {
1101             fnb = out[0].f;
1102
1103             for(a=a0; a<a1; a++)
1104             {
1105                 i = X8_IND_A(cell[a]);
1106
1107                 f[a][XX] += fnb[i+XX*PACK_X8];
1108                 f[a][YY] += fnb[i+YY*PACK_X8];
1109                 f[a][ZZ] += fnb[i+ZZ*PACK_X8];
1110             }
1111         }
1112         else
1113         {
1114             for(a=a0; a<a1; a++)
1115             {
1116                 i = X8_IND_A(cell[a]);
1117                 
1118                 for(fa=0; fa<nfa; fa++)
1119                 {
1120                     f[a][XX] += out[fa].f[i+XX*PACK_X8];
1121                     f[a][YY] += out[fa].f[i+YY*PACK_X8];
1122                     f[a][ZZ] += out[fa].f[i+ZZ*PACK_X8];
1123                 }
1124             }
1125         }
1126         break;
1127     }
1128 }
1129
1130 /* Add the force array(s) from nbnxn_atomdata_t to f */
1131 void nbnxn_atomdata_add_nbat_f_to_f(const nbnxn_search_t nbs,
1132                                     int locality,
1133                                     const nbnxn_atomdata_t *nbat,
1134                                     rvec *f)
1135 {
1136     int a0=0,na=0;
1137     int nth,th;
1138     gmx_bool bStreamingReduce;
1139
1140     nbs_cycle_start(&nbs->cc[enbsCCreducef]);
1141
1142     switch (locality)
1143     {
1144     case eatAll:
1145         a0 = 0;
1146         na = nbs->natoms_nonlocal;
1147         break;
1148     case eatLocal:
1149         a0 = 0;
1150         na = nbs->natoms_local;
1151         break;
1152     case eatNonlocal:
1153         a0 = nbs->natoms_local;
1154         na = nbs->natoms_nonlocal - nbs->natoms_local;
1155         break;
1156     }
1157
1158     nth = gmx_omp_nthreads_get(emntNonbonded);
1159
1160     /* Using the two-step streaming reduction is probably always faster */
1161     bStreamingReduce = (nbat->nout > 1);
1162
1163     if (bStreamingReduce)
1164     {
1165         /* Reduce the force thread output buffers into buffer 0, before adding
1166          * them to the, differently ordered, "real" force buffer.
1167          */
1168 #pragma omp parallel for num_threads(nth) schedule(static)
1169         for(th=0; th<nth; th++)
1170         {
1171             int g0,g1;
1172             int b0,b1,nb;
1173             int blocksize,i0,i1;
1174
1175             /* For which grids should we reduce the force output? */
1176             g0 = ((locality==eatLocal || locality==eatAll) ? 0 : 1);
1177             g1 = (locality==eatLocal ? 1 : nbs->ngrid);
1178
1179             /* Get the grid cell bounds */
1180             b0 = nbs->grid[g0].cell0;
1181             b1 = nbs->grid[g1-1].cell0 + nbs->grid[g1-1].nc;
1182             blocksize = nbs->grid[g0].na_sc*nbat->fstride;
1183             /* The simple grid size in atoms is a multiple of na_cj.
1184              * With float-AVX256 we use this and make blocksize a multiple of 8.
1185              */
1186             if (nbs->grid[0].bSimple && nbs->grid[0].na_cj > nbs->grid[0].na_c)
1187             {
1188                 blocksize *= 2;
1189                 b0 /= 2;
1190                 b1 /= 2;
1191             }
1192             nb = b1 - b0;
1193
1194             /* Calculate the index range for our thread */
1195             i0 = (b0 + (nb* th   )/nth)*blocksize;
1196             i1 = (b0 + (nb*(th+1))/nth)*blocksize;
1197
1198 #ifdef NBNXN_SEARCH_SSE
1199             nbnxn_atomdata_reduce_reals_x86_simd(
1200 #else
1201             nbnxn_atomdata_reduce_reals(    
1202 #endif
1203                                         nbat->out[0].f,
1204                                         nbat->out+1,nbat->nout - 1,
1205                                         i0,i1);
1206         }
1207     }
1208
1209 #pragma omp parallel for num_threads(nth) schedule(static)
1210     for(th=0; th<nth; th++)
1211     {
1212         nbnxn_atomdata_add_nbat_f_to_f_part(nbs,nbat,
1213                                             nbat->out,
1214                                             bStreamingReduce ? 1 : nbat->nout,
1215                                             a0+((th+0)*na)/nth,
1216                                             a0+((th+1)*na)/nth,
1217                                             f);
1218     }
1219
1220     nbs_cycle_stop(&nbs->cc[enbsCCreducef]);
1221 }
1222
1223 /* Adds the shift forces from nbnxn_atomdata_t to fshift */
1224 void nbnxn_atomdata_add_nbat_fshift_to_fshift(const nbnxn_atomdata_t *nbat,
1225                                               rvec *fshift)
1226 {
1227     const nbnxn_atomdata_output_t *out;
1228     int  th;
1229     int  s;
1230     rvec sum;
1231
1232     out = nbat->out;
1233     
1234     for(s=0; s<SHIFTS; s++)
1235     {
1236         clear_rvec(sum);
1237         for(th=0; th<nbat->nout; th++)
1238         {
1239             sum[XX] += out[th].fshift[s*DIM+XX];
1240             sum[YY] += out[th].fshift[s*DIM+YY];
1241             sum[ZZ] += out[th].fshift[s*DIM+ZZ];
1242         }
1243         rvec_inc(fshift[s],sum);
1244     }
1245 }