471f57c25b43e66eff7b77cd6f0c7e94a8d4bfc0
[alexxy/gromacs.git] / src / gromacs / fft / fft5d.cpp
1 /*
2  * This file is part of the GROMACS molecular simulation package.
3  *
4  * Copyright (c) 2009,2010,2012,2013,2014,2015,2016, by the GROMACS development team, led by
5  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
6  * and including many others, as listed in the AUTHORS file in the
7  * top-level source directory and at http://www.gromacs.org.
8  *
9  * GROMACS is free software; you can redistribute it and/or
10  * modify it under the terms of the GNU Lesser General Public License
11  * as published by the Free Software Foundation; either version 2.1
12  * of the License, or (at your option) any later version.
13  *
14  * GROMACS is distributed in the hope that it will be useful,
15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
17  * Lesser General Public License for more details.
18  *
19  * You should have received a copy of the GNU Lesser General Public
20  * License along with GROMACS; if not, see
21  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
22  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
23  *
24  * If you want to redistribute modifications to GROMACS, please
25  * consider that scientific software is very special. Version
26  * control is crucial - bugs must be traceable. We will be happy to
27  * consider code for inclusion in the official distribution, but
28  * derived work must not be called official GROMACS. Details are found
29  * in the README & COPYING files - if they are missing, get the
30  * official version at http://www.gromacs.org.
31  *
32  * To help us fund GROMACS development, we humbly ask that you cite
33  * the research papers on the package. Check out http://www.gromacs.org.
34  */
35 #include "gmxpre.h"
36
37 #include "fft5d.h"
38
39 #include "config.h"
40
41 #include <assert.h>
42 #include <float.h>
43 #include <math.h>
44 #include <stdio.h>
45 #include <stdlib.h>
46 #include <string.h>
47
48 #include <algorithm>
49
50 #include "gromacs/utility/exceptions.h"
51 #include "gromacs/utility/fatalerror.h"
52 #include "gromacs/utility/gmxmpi.h"
53 #include "gromacs/utility/smalloc.h"
54
55 #ifdef NOGMX
56 #define GMX_PARALLEL_ENV_INITIALIZED 1
57 #else
58 #if GMX_MPI
59 #define GMX_PARALLEL_ENV_INITIALIZED 1
60 #else
61 #define GMX_PARALLEL_ENV_INITIALIZED 0
62 #endif
63 #endif
64
65 #if GMX_OPENMP
66 /* TODO: Do we still need this? Are we still planning ot use fftw + OpenMP? */
67 #define FFT5D_THREADS
68 /* requires fftw compiled with openmp */
69 /* #define FFT5D_FFTW_THREADS (now set by cmake) */
70 #endif
71
72 #ifndef __FLT_EPSILON__
73 #define __FLT_EPSILON__ FLT_EPSILON
74 #define __DBL_EPSILON__ DBL_EPSILON
75 #endif
76
77 #ifdef NOGMX
78 FILE* debug = 0;
79 #endif
80
81 #if GMX_FFT_FFTW3
82
83 #include "gromacs/utility/exceptions.h"
84 #include "gromacs/utility/mutex.h"
85 /* none of the fftw3 calls, except execute(), are thread-safe, so
86    we need to serialize them with this mutex. */
87 static gmx::Mutex big_fftw_mutex;
88 #define FFTW_LOCK try { big_fftw_mutex.lock(); } GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR
89 #define FFTW_UNLOCK try { big_fftw_mutex.unlock(); } GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR
90 #endif /* GMX_FFT_FFTW3 */
91
92 #if GMX_MPI
93 /* largest factor smaller than sqrt */
94 static int lfactor(int z)
95 {
96     int i = static_cast<int>(sqrt(static_cast<double>(z)));
97     while (z%i != 0)
98     {
99         i--;
100     }
101     return i;
102 }
103 #endif
104
105 /* largest factor */
106 static int l2factor(int z)
107 {
108     int i;
109     if (z == 1)
110     {
111         i = 1;
112     }
113     else
114     {
115         i = z/2;
116         while (z%i != 0)
117         {
118             i--;
119         }
120     }
121     return i;
122 }
123
124 /* largest prime factor: WARNING: slow recursion, only use for small numbers */
125 static int lpfactor(int z)
126 {
127     int f = l2factor(z);
128     if (f == 1)
129     {
130         return z;
131     }
132     return std::max(lpfactor(f), lpfactor(z/f));
133 }
134
135 #if !GMX_MPI
136 #if HAVE_GETTIMEOFDAY
137 #include <sys/time.h>
138 double MPI_Wtime()
139 {
140     struct timeval tv;
141     gettimeofday(&tv, 0);
142     return tv.tv_sec+tv.tv_usec*1e-6;
143 }
144 #else
145 double MPI_Wtime()
146 {
147     return 0.0;
148 }
149 #endif
150 #endif
151
152 static int vmax(int* a, int s)
153 {
154     int i, max = 0;
155     for (i = 0; i < s; i++)
156     {
157         if (a[i] > max)
158         {
159             max = a[i];
160         }
161     }
162     return max;
163 }
164
165
166 /* NxMxK the size of the data
167  * comm communicator to use for fft5d
168  * P0 number of processor in 1st axes (can be null for automatic)
169  * lin is allocated by fft5d because size of array is only known after planning phase
170  * rlout2 is only used as intermediate buffer - only returned after allocation to reuse for back transform - should not be used by caller
171  */
172 fft5d_plan fft5d_plan_3d(int NG, int MG, int KG, MPI_Comm comm[2], int flags, t_complex** rlin, t_complex** rlout, t_complex** rlout2, t_complex** rlout3, int nthreads)
173 {
174
175     int        P[2], bMaster, prank[2], i, t;
176     int        rNG, rMG, rKG;
177     int       *N0 = 0, *N1 = 0, *M0 = 0, *M1 = 0, *K0 = 0, *K1 = 0, *oN0 = 0, *oN1 = 0, *oM0 = 0, *oM1 = 0, *oK0 = 0, *oK1 = 0;
178     int        N[3], M[3], K[3], pN[3], pM[3], pK[3], oM[3], oK[3], *iNin[3] = {0}, *oNin[3] = {0}, *iNout[3] = {0}, *oNout[3] = {0};
179     int        C[3], rC[3], nP[2];
180     int        lsize;
181     t_complex *lin = 0, *lout = 0, *lout2 = 0, *lout3 = 0;
182     fft5d_plan plan;
183     int        s;
184
185     /* comm, prank and P are in the order of the decomposition (plan->cart is in the order of transposes) */
186 #if GMX_MPI
187     if (GMX_PARALLEL_ENV_INITIALIZED && comm[0] != MPI_COMM_NULL)
188     {
189         MPI_Comm_size(comm[0], &P[0]);
190         MPI_Comm_rank(comm[0], &prank[0]);
191     }
192     else
193 #endif
194     {
195         P[0]     = 1;
196         prank[0] = 0;
197     }
198 #if GMX_MPI
199     if (GMX_PARALLEL_ENV_INITIALIZED && comm[1] != MPI_COMM_NULL)
200     {
201         MPI_Comm_size(comm[1], &P[1]);
202         MPI_Comm_rank(comm[1], &prank[1]);
203     }
204     else
205 #endif
206     {
207         P[1]     = 1;
208         prank[1] = 0;
209     }
210
211     bMaster = (prank[0] == 0 && prank[1] == 0);
212
213
214     if (debug)
215     {
216         fprintf(debug, "FFT5D: Using %dx%d rank grid, rank %d,%d\n",
217                 P[0], P[1], prank[0], prank[1]);
218     }
219
220     if (bMaster)
221     {
222         if (debug)
223         {
224             fprintf(debug, "FFT5D: N: %d, M: %d, K: %d, P: %dx%d, real2complex: %d, backward: %d, order yz: %d, debug %d\n",
225                     NG, MG, KG, P[0], P[1], (flags&FFT5D_REALCOMPLEX) > 0, (flags&FFT5D_BACKWARD) > 0, (flags&FFT5D_ORDER_YZ) > 0, (flags&FFT5D_DEBUG) > 0);
226         }
227         /* The check below is not correct, one prime factor 11 or 13 is ok.
228            if (fft5d_fmax(fft5d_fmax(lpfactor(NG),lpfactor(MG)),lpfactor(KG))>7) {
229             printf("WARNING: FFT very slow with prime factors larger 7\n");
230             printf("Change FFT size or in case you cannot change it look at\n");
231             printf("http://www.fftw.org/fftw3_doc/Generating-your-own-code.html\n");
232            }
233          */
234     }
235
236     if (NG == 0 || MG == 0 || KG == 0)
237     {
238         if (bMaster)
239         {
240             printf("FFT5D: FATAL: Datasize cannot be zero in any dimension\n");
241         }
242         return 0;
243     }
244
245     rNG = NG; rMG = MG; rKG = KG;
246
247     if (flags&FFT5D_REALCOMPLEX)
248     {
249         if (!(flags&FFT5D_BACKWARD))
250         {
251             NG = NG/2+1;
252         }
253         else
254         {
255             if (!(flags&FFT5D_ORDER_YZ))
256             {
257                 MG = MG/2+1;
258             }
259             else
260             {
261                 KG = KG/2+1;
262             }
263         }
264     }
265
266
267     /*for transpose we need to know the size for each processor not only our own size*/
268
269     N0  = (int*)malloc(P[0]*sizeof(int)); N1 = (int*)malloc(P[1]*sizeof(int));
270     M0  = (int*)malloc(P[0]*sizeof(int)); M1 = (int*)malloc(P[1]*sizeof(int));
271     K0  = (int*)malloc(P[0]*sizeof(int)); K1 = (int*)malloc(P[1]*sizeof(int));
272     oN0 = (int*)malloc(P[0]*sizeof(int)); oN1 = (int*)malloc(P[1]*sizeof(int));
273     oM0 = (int*)malloc(P[0]*sizeof(int)); oM1 = (int*)malloc(P[1]*sizeof(int));
274     oK0 = (int*)malloc(P[0]*sizeof(int)); oK1 = (int*)malloc(P[1]*sizeof(int));
275
276     for (i = 0; i < P[0]; i++)
277     {
278         #define EVENDIST
279         #ifndef EVENDIST
280         oN0[i] = i*ceil((double)NG/P[0]);
281         oM0[i] = i*ceil((double)MG/P[0]);
282         oK0[i] = i*ceil((double)KG/P[0]);
283         #else
284         oN0[i] = (NG*i)/P[0];
285         oM0[i] = (MG*i)/P[0];
286         oK0[i] = (KG*i)/P[0];
287         #endif
288     }
289     for (i = 0; i < P[1]; i++)
290     {
291         #ifndef EVENDIST
292         oN1[i] = i*ceil((double)NG/P[1]);
293         oM1[i] = i*ceil((double)MG/P[1]);
294         oK1[i] = i*ceil((double)KG/P[1]);
295         #else
296         oN1[i] = (NG*i)/P[1];
297         oM1[i] = (MG*i)/P[1];
298         oK1[i] = (KG*i)/P[1];
299         #endif
300     }
301     for (i = 0; i < P[0]-1; i++)
302     {
303         N0[i] = oN0[i+1]-oN0[i];
304         M0[i] = oM0[i+1]-oM0[i];
305         K0[i] = oK0[i+1]-oK0[i];
306     }
307     N0[P[0]-1] = NG-oN0[P[0]-1];
308     M0[P[0]-1] = MG-oM0[P[0]-1];
309     K0[P[0]-1] = KG-oK0[P[0]-1];
310     for (i = 0; i < P[1]-1; i++)
311     {
312         N1[i] = oN1[i+1]-oN1[i];
313         M1[i] = oM1[i+1]-oM1[i];
314         K1[i] = oK1[i+1]-oK1[i];
315     }
316     N1[P[1]-1] = NG-oN1[P[1]-1];
317     M1[P[1]-1] = MG-oM1[P[1]-1];
318     K1[P[1]-1] = KG-oK1[P[1]-1];
319
320     /* for step 1-3 the local N,M,K sizes of the transposed system
321        C: contiguous dimension, and nP: number of processor in subcommunicator
322        for that step */
323
324
325     pM[0] = M0[prank[0]];
326     oM[0] = oM0[prank[0]];
327     pK[0] = K1[prank[1]];
328     oK[0] = oK1[prank[1]];
329     C[0]  = NG;
330     rC[0] = rNG;
331     if (!(flags&FFT5D_ORDER_YZ))
332     {
333         N[0]     = vmax(N1, P[1]);
334         M[0]     = M0[prank[0]];
335         K[0]     = vmax(K1, P[1]);
336         pN[0]    = N1[prank[1]];
337         iNout[0] = N1;
338         oNout[0] = oN1;
339         nP[0]    = P[1];
340         C[1]     = KG;
341         rC[1]    = rKG;
342         N[1]     = vmax(K0, P[0]);
343         pN[1]    = K0[prank[0]];
344         iNin[1]  = K1;
345         oNin[1]  = oK1;
346         iNout[1] = K0;
347         oNout[1] = oK0;
348         M[1]     = vmax(M0, P[0]);
349         pM[1]    = M0[prank[0]];
350         oM[1]    = oM0[prank[0]];
351         K[1]     = N1[prank[1]];
352         pK[1]    = N1[prank[1]];
353         oK[1]    = oN1[prank[1]];
354         nP[1]    = P[0];
355         C[2]     = MG;
356         rC[2]    = rMG;
357         iNin[2]  = M0;
358         oNin[2]  = oM0;
359         M[2]     = vmax(K0, P[0]);
360         pM[2]    = K0[prank[0]];
361         oM[2]    = oK0[prank[0]];
362         K[2]     = vmax(N1, P[1]);
363         pK[2]    = N1[prank[1]];
364         oK[2]    = oN1[prank[1]];
365         free(N0); free(oN0); /*these are not used for this order*/
366         free(M1); free(oM1); /*the rest is freed in destroy*/
367     }
368     else
369     {
370         N[0]     = vmax(N0, P[0]);
371         M[0]     = vmax(M0, P[0]);
372         K[0]     = K1[prank[1]];
373         pN[0]    = N0[prank[0]];
374         iNout[0] = N0;
375         oNout[0] = oN0;
376         nP[0]    = P[0];
377         C[1]     = MG;
378         rC[1]    = rMG;
379         N[1]     = vmax(M1, P[1]);
380         pN[1]    = M1[prank[1]];
381         iNin[1]  = M0;
382         oNin[1]  = oM0;
383         iNout[1] = M1;
384         oNout[1] = oM1;
385         M[1]     = N0[prank[0]];
386         pM[1]    = N0[prank[0]];
387         oM[1]    = oN0[prank[0]];
388         K[1]     = vmax(K1, P[1]);
389         pK[1]    = K1[prank[1]];
390         oK[1]    = oK1[prank[1]];
391         nP[1]    = P[1];
392         C[2]     = KG;
393         rC[2]    = rKG;
394         iNin[2]  = K1;
395         oNin[2]  = oK1;
396         M[2]     = vmax(N0, P[0]);
397         pM[2]    = N0[prank[0]];
398         oM[2]    = oN0[prank[0]];
399         K[2]     = vmax(M1, P[1]);
400         pK[2]    = M1[prank[1]];
401         oK[2]    = oM1[prank[1]];
402         free(N1); free(oN1); /*these are not used for this order*/
403         free(K0); free(oK0); /*the rest is freed in destroy*/
404     }
405     N[2] = pN[2] = -1;       /*not used*/
406
407     /*
408        Difference between x-y-z regarding 2d decomposition is whether they are
409        distributed along axis 1, 2 or both
410      */
411
412     /* int lsize = fmax(N[0]*M[0]*K[0]*nP[0],N[1]*M[1]*K[1]*nP[1]); */
413     lsize = std::max(N[0]*M[0]*K[0]*nP[0], std::max(N[1]*M[1]*K[1]*nP[1], C[2]*M[2]*K[2]));
414     /* int lsize = fmax(C[0]*M[0]*K[0],fmax(C[1]*M[1]*K[1],C[2]*M[2]*K[2])); */
415     if (!(flags&FFT5D_NOMALLOC))
416     {
417         snew_aligned(lin, lsize, 32);
418         snew_aligned(lout, lsize, 32);
419         if (nthreads > 1)
420         {
421             /* We need extra transpose buffers to avoid OpenMP barriers */
422             snew_aligned(lout2, lsize, 32);
423             snew_aligned(lout3, lsize, 32);
424         }
425         else
426         {
427             /* We can reuse the buffers to avoid cache misses */
428             lout2 = lin;
429             lout3 = lout;
430         }
431     }
432     else
433     {
434         lin  = *rlin;
435         lout = *rlout;
436         if (nthreads > 1)
437         {
438             lout2 = *rlout2;
439             lout3 = *rlout3;
440         }
441         else
442         {
443             lout2 = lin;
444             lout3 = lout;
445         }
446     }
447
448     plan = (fft5d_plan)calloc(1, sizeof(struct fft5d_plan_t));
449
450
451     if (debug)
452     {
453         fprintf(debug, "Running on %d threads\n", nthreads);
454     }
455
456 #if GMX_FFT_FFTW3
457     /* Don't add more stuff here! We have already had at least one bug because we are reimplementing
458      * the low-level FFT interface instead of using the Gromacs FFT module. If we need more
459      * generic functionality it is far better to extend the interface so we can use it for
460      * all FFT libraries instead of writing FFTW-specific code here.
461      */
462
463     /*if not FFTW - then we don't do a 3d plan but instead use only 1D plans */
464     /* It is possible to use the 3d plan with OMP threads - but in that case it is not allowed to be called from
465      * within a parallel region. For now deactivated. If it should be supported it has to made sure that
466      * that the execute of the 3d plan is in a master/serial block (since it contains it own parallel region)
467      * and that the 3d plan is faster than the 1d plan.
468      */
469     if ((!(flags&FFT5D_INPLACE)) && (!(P[0] > 1 || P[1] > 1)) && nthreads == 1) /*don't do 3d plan in parallel or if in_place requested */
470     {
471         int fftwflags = FFTW_DESTROY_INPUT;
472         FFTW(iodim) dims[3];
473         int inNG = NG, outMG = MG, outKG = KG;
474
475         FFTW_LOCK;
476
477         fftwflags |= (flags & FFT5D_NOMEASURE) ? FFTW_ESTIMATE : FFTW_MEASURE;
478
479         if (flags&FFT5D_REALCOMPLEX)
480         {
481             if (!(flags&FFT5D_BACKWARD))        /*input pointer is not complex*/
482             {
483                 inNG *= 2;
484             }
485             else                                /*output pointer is not complex*/
486             {
487                 if (!(flags&FFT5D_ORDER_YZ))
488                 {
489                     outMG *= 2;
490                 }
491                 else
492                 {
493                     outKG *= 2;
494                 }
495             }
496         }
497
498         if (!(flags&FFT5D_BACKWARD))
499         {
500             dims[0].n  = KG;
501             dims[1].n  = MG;
502             dims[2].n  = rNG;
503
504             dims[0].is = inNG*MG;         /*N M K*/
505             dims[1].is = inNG;
506             dims[2].is = 1;
507             if (!(flags&FFT5D_ORDER_YZ))
508             {
509                 dims[0].os = MG;           /*M K N*/
510                 dims[1].os = 1;
511                 dims[2].os = MG*KG;
512             }
513             else
514             {
515                 dims[0].os = 1;           /*K N M*/
516                 dims[1].os = KG*NG;
517                 dims[2].os = KG;
518             }
519         }
520         else
521         {
522             if (!(flags&FFT5D_ORDER_YZ))
523             {
524                 dims[0].n  = NG;
525                 dims[1].n  = KG;
526                 dims[2].n  = rMG;
527
528                 dims[0].is = 1;
529                 dims[1].is = NG*MG;
530                 dims[2].is = NG;
531
532                 dims[0].os = outMG*KG;
533                 dims[1].os = outMG;
534                 dims[2].os = 1;
535             }
536             else
537             {
538                 dims[0].n  = MG;
539                 dims[1].n  = NG;
540                 dims[2].n  = rKG;
541
542                 dims[0].is = NG;
543                 dims[1].is = 1;
544                 dims[2].is = NG*MG;
545
546                 dims[0].os = outKG*NG;
547                 dims[1].os = outKG;
548                 dims[2].os = 1;
549             }
550         }
551 #ifdef FFT5D_THREADS
552 #ifdef FFT5D_FFTW_THREADS
553         FFTW(plan_with_nthreads)(nthreads);
554 #endif
555 #endif
556         if ((flags&FFT5D_REALCOMPLEX) && !(flags&FFT5D_BACKWARD))
557         {
558             plan->p3d = FFTW(plan_guru_dft_r2c)(/*rank*/ 3, dims,
559                                                          /*howmany*/ 0, /*howmany_dims*/ 0,
560                                                          (real*)lin, (FFTW(complex) *) lout,
561                                                          /*flags*/ fftwflags);
562         }
563         else if ((flags&FFT5D_REALCOMPLEX) && (flags&FFT5D_BACKWARD))
564         {
565             plan->p3d = FFTW(plan_guru_dft_c2r)(/*rank*/ 3, dims,
566                                                          /*howmany*/ 0, /*howmany_dims*/ 0,
567                                                          (FFTW(complex) *) lin, (real*)lout,
568                                                          /*flags*/ fftwflags);
569         }
570         else
571         {
572             plan->p3d = FFTW(plan_guru_dft)(/*rank*/ 3, dims,
573                                                      /*howmany*/ 0, /*howmany_dims*/ 0,
574                                                      (FFTW(complex) *) lin, (FFTW(complex) *) lout,
575                                                      /*sign*/ (flags&FFT5D_BACKWARD) ? 1 : -1, /*flags*/ fftwflags);
576         }
577 #ifdef FFT5D_THREADS
578 #ifdef FFT5D_FFTW_THREADS
579         FFTW(plan_with_nthreads)(1);
580 #endif
581 #endif
582         FFTW_UNLOCK;
583     }
584     if (!plan->p3d) /* for decomposition and if 3d plan did not work */
585     {
586 #endif              /* GMX_FFT_FFTW3 */
587     for (s = 0; s < 3; s++)
588     {
589         if (debug)
590         {
591             fprintf(debug, "FFT5D: Plan s %d rC %d M %d pK %d C %d lsize %d\n",
592                     s, rC[s], M[s], pK[s], C[s], lsize);
593         }
594         plan->p1d[s] = (gmx_fft_t*)malloc(sizeof(gmx_fft_t)*nthreads);
595
596         /* Make sure that the init routines are only called by one thread at a time and in order
597            (later is only important to not confuse valgrind)
598          */
599 #pragma omp parallel for num_threads(nthreads) schedule(static) ordered
600         for (t = 0; t < nthreads; t++)
601         {
602 #pragma omp ordered
603             {
604                 try
605                 {
606                     int tsize = ((t+1)*pM[s]*pK[s]/nthreads)-(t*pM[s]*pK[s]/nthreads);
607
608                     if ((flags&FFT5D_REALCOMPLEX) && ((!(flags&FFT5D_BACKWARD) && s == 0) || ((flags&FFT5D_BACKWARD) && s == 2)))
609                     {
610                         gmx_fft_init_many_1d_real( &plan->p1d[s][t], rC[s], tsize, (flags&FFT5D_NOMEASURE) ? GMX_FFT_FLAG_CONSERVATIVE : 0 );
611                     }
612                     else
613                     {
614                         gmx_fft_init_many_1d     ( &plan->p1d[s][t],  C[s], tsize, (flags&FFT5D_NOMEASURE) ? GMX_FFT_FLAG_CONSERVATIVE : 0 );
615                     }
616                 }
617                 GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR;
618             }
619         }
620     }
621
622 #if GMX_FFT_FFTW3
623 }
624 #endif
625     if ((flags&FFT5D_ORDER_YZ))   /*plan->cart is in the order of transposes */
626     {
627         plan->cart[0] = comm[0]; plan->cart[1] = comm[1];
628     }
629     else
630     {
631         plan->cart[1] = comm[0]; plan->cart[0] = comm[1];
632     }
633 #ifdef FFT5D_MPI_TRANSPOSE
634     FFTW_LOCK;
635     for (s = 0; s < 2; s++)
636     {
637         if ((s == 0 && !(flags&FFT5D_ORDER_YZ)) || (s == 1 && (flags&FFT5D_ORDER_YZ)))
638         {
639             plan->mpip[s] = FFTW(mpi_plan_many_transpose)(nP[s], nP[s], N[s]*K[s]*pM[s]*2, 1, 1, (real*)lout2, (real*)lout3, plan->cart[s], FFTW_PATIENT);
640         }
641         else
642         {
643             plan->mpip[s] = FFTW(mpi_plan_many_transpose)(nP[s], nP[s], N[s]*pK[s]*M[s]*2, 1, 1, (real*)lout2, (real*)lout3, plan->cart[s], FFTW_PATIENT);
644         }
645     }
646     FFTW_UNLOCK;
647 #endif
648
649
650     plan->lin   = lin;
651     plan->lout  = lout;
652     plan->lout2 = lout2;
653     plan->lout3 = lout3;
654
655     plan->NG = NG; plan->MG = MG; plan->KG = KG;
656
657     for (s = 0; s < 3; s++)
658     {
659         plan->N[s]    = N[s]; plan->M[s] = M[s]; plan->K[s] = K[s]; plan->pN[s] = pN[s]; plan->pM[s] = pM[s]; plan->pK[s] = pK[s];
660         plan->oM[s]   = oM[s]; plan->oK[s] = oK[s];
661         plan->C[s]    = C[s]; plan->rC[s] = rC[s];
662         plan->iNin[s] = iNin[s]; plan->oNin[s] = oNin[s]; plan->iNout[s] = iNout[s]; plan->oNout[s] = oNout[s];
663     }
664     for (s = 0; s < 2; s++)
665     {
666         plan->P[s] = nP[s]; plan->coor[s] = prank[s];
667     }
668
669 /*    plan->fftorder=fftorder;
670     plan->direction=direction;
671     plan->realcomplex=realcomplex;
672  */
673     plan->flags    = flags;
674     plan->nthreads = nthreads;
675     *rlin          = lin;
676     *rlout         = lout;
677     *rlout2        = lout2;
678     *rlout3        = lout3;
679     return plan;
680 }
681
682
683 enum order {
684     XYZ,
685     XZY,
686     YXZ,
687     YZX,
688     ZXY,
689     ZYX
690 };
691
692
693
694 /*here x,y,z and N,M,K is in rotated coordinate system!!
695    x (and N) is mayor (consecutive) dimension, y (M) middle and z (K) major
696    maxN,maxM,maxK is max size of local data
697    pN, pM, pK is local size specific to current processor (only different to max if not divisible)
698    NG, MG, KG is size of global data*/
699 static void splitaxes(t_complex* lout, const t_complex* lin,
700                       int maxN, int maxM, int maxK, int pM,
701                       int P, int NG, int *N, int* oN, int starty, int startz, int endy, int endz)
702 {
703     int x, y, z, i;
704     int in_i, out_i, in_z, out_z, in_y, out_y;
705     int s_y, e_y;
706
707     for (z = startz; z < endz+1; z++) /*3. z l*/
708     {
709         if (z == startz)
710         {
711             s_y = starty;
712         }
713         else
714         {
715             s_y = 0;
716         }
717         if (z == endz)
718         {
719             e_y = endy;
720         }
721         else
722         {
723             e_y = pM;
724         }
725         out_z  = z*maxN*maxM;
726         in_z   = z*NG*pM;
727
728         for (i = 0; i < P; i++) /*index cube along long axis*/
729         {
730             out_i  = out_z  + i*maxN*maxM*maxK;
731             in_i   = in_z + oN[i];
732             for (y = s_y; y < e_y; y++)   /*2. y k*/
733             {
734                 out_y  = out_i  + y*maxN;
735                 in_y   = in_i + y*NG;
736                 for (x = 0; x < N[i]; x++)       /*1. x j*/
737                 {
738                     lout[out_y+x] = lin[in_y+x]; /*in=z*NG*pM+oN[i]+y*NG+x*/
739                     /*after split important that each processor chunk i has size maxN*maxM*maxK and thus being the same size*/
740                     /*before split data contiguos - thus if different processor get different amount oN is different*/
741                 }
742             }
743         }
744     }
745 }
746
747 /*make axis contiguous again (after AllToAll) and also do local transpose*/
748 /*transpose mayor and major dimension
749    variables see above
750    the major, middle, minor order is only correct for x,y,z (N,M,K) for the input
751    N,M,K local dimensions
752    KG global size*/
753 static void joinAxesTrans13(t_complex* lout, const t_complex* lin,
754                             int maxN, int maxM, int maxK, int pM,
755                             int P, int KG, int* K, int* oK, int starty, int startx, int endy, int endx)
756 {
757     int i, x, y, z;
758     int out_i, in_i, out_x, in_x, out_z, in_z;
759     int s_y, e_y;
760
761     for (x = startx; x < endx+1; x++) /*1.j*/
762     {
763         if (x == startx)
764         {
765             s_y = starty;
766         }
767         else
768         {
769             s_y = 0;
770         }
771         if (x == endx)
772         {
773             e_y = endy;
774         }
775         else
776         {
777             e_y = pM;
778         }
779
780         out_x  = x*KG*pM;
781         in_x   = x;
782
783         for (i = 0; i < P; i++) /*index cube along long axis*/
784         {
785             out_i  = out_x  + oK[i];
786             in_i   = in_x + i*maxM*maxN*maxK;
787             for (z = 0; z < K[i]; z++) /*3.l*/
788             {
789                 out_z  = out_i  + z;
790                 in_z   = in_i + z*maxM*maxN;
791                 for (y = s_y; y < e_y; y++)              /*2.k*/
792                 {
793                     lout[out_z+y*KG] = lin[in_z+y*maxN]; /*out=x*KG*pM+oK[i]+z+y*KG*/
794                 }
795             }
796         }
797     }
798 }
799
800 /*make axis contiguous again (after AllToAll) and also do local transpose
801    tranpose mayor and middle dimension
802    variables see above
803    the minor, middle, major order is only correct for x,y,z (N,M,K) for the input
804    N,M,K local size
805    MG, global size*/
806 static void joinAxesTrans12(t_complex* lout, const t_complex* lin, int maxN, int maxM, int maxK, int pN,
807                             int P, int MG, int* M, int* oM, int startx, int startz, int endx, int endz)
808 {
809     int i, z, y, x;
810     int out_i, in_i, out_z, in_z, out_x, in_x;
811     int s_x, e_x;
812
813     for (z = startz; z < endz+1; z++)
814     {
815         if (z == startz)
816         {
817             s_x = startx;
818         }
819         else
820         {
821             s_x = 0;
822         }
823         if (z == endz)
824         {
825             e_x = endx;
826         }
827         else
828         {
829             e_x = pN;
830         }
831         out_z  = z*MG*pN;
832         in_z   = z*maxM*maxN;
833
834         for (i = 0; i < P; i++) /*index cube along long axis*/
835         {
836             out_i  = out_z  + oM[i];
837             in_i   = in_z + i*maxM*maxN*maxK;
838             for (x = s_x; x < e_x; x++)
839             {
840                 out_x  = out_i  + x*MG;
841                 in_x   = in_i + x;
842                 for (y = 0; y < M[i]; y++)
843                 {
844                     lout[out_x+y] = lin[in_x+y*maxN]; /*out=z*MG*pN+oM[i]+x*MG+y*/
845                 }
846             }
847         }
848     }
849 }
850
851
852 static void rotate_offsets(int x[])
853 {
854     int t = x[0];
855 /*    x[0]=x[2];
856     x[2]=x[1];
857     x[1]=t;*/
858     x[0] = x[1];
859     x[1] = x[2];
860     x[2] = t;
861 }
862
863 /*compute the offset to compare or print transposed local data in original input coordinates
864    xs matrix dimension size, xl dimension length, xc decomposition offset
865    s: step in computation = number of transposes*/
866 static void compute_offsets(fft5d_plan plan, int xs[], int xl[], int xc[], int NG[], int s)
867 {
868 /*    int direction = plan->direction;
869     int fftorder = plan->fftorder;*/
870
871     int  o = 0;
872     int  pos[3], i;
873     int *pM = plan->pM, *pK = plan->pK, *oM = plan->oM, *oK = plan->oK,
874     *C      = plan->C, *rC = plan->rC;
875
876     NG[0] = plan->NG; NG[1] = plan->MG; NG[2] = plan->KG;
877
878     if (!(plan->flags&FFT5D_ORDER_YZ))
879     {
880         switch (s)
881         {
882             case 0: o = XYZ; break;
883             case 1: o = ZYX; break;
884             case 2: o = YZX; break;
885             default: assert(0);
886         }
887     }
888     else
889     {
890         switch (s)
891         {
892             case 0: o = XYZ; break;
893             case 1: o = YXZ; break;
894             case 2: o = ZXY; break;
895             default: assert(0);
896         }
897     }
898
899     switch (o)
900     {
901         case XYZ: pos[0] = 1; pos[1] = 2; pos[2] = 3; break;
902         case XZY: pos[0] = 1; pos[1] = 3; pos[2] = 2; break;
903         case YXZ: pos[0] = 2; pos[1] = 1; pos[2] = 3; break;
904         case YZX: pos[0] = 3; pos[1] = 1; pos[2] = 2; break;
905         case ZXY: pos[0] = 2; pos[1] = 3; pos[2] = 1; break;
906         case ZYX: pos[0] = 3; pos[1] = 2; pos[2] = 1; break;
907     }
908     /*if (debug) printf("pos: %d %d %d\n",pos[0],pos[1],pos[2]);*/
909
910     /*xs, xl give dimension size and data length in local transposed coordinate system
911        for 0(/1/2): x(/y/z) in original coordinate system*/
912     for (i = 0; i < 3; i++)
913     {
914         switch (pos[i])
915         {
916             case 1: xs[i] = 1;         xc[i] = 0;     xl[i] = C[s]; break;
917             case 2: xs[i] = C[s];      xc[i] = oM[s]; xl[i] = pM[s]; break;
918             case 3: xs[i] = C[s]*pM[s]; xc[i] = oK[s]; xl[i] = pK[s]; break;
919         }
920     }
921     /*input order is different for test program to match FFTW order
922        (important for complex to real)*/
923     if (plan->flags&FFT5D_BACKWARD)
924     {
925         rotate_offsets(xs);
926         rotate_offsets(xl);
927         rotate_offsets(xc);
928         rotate_offsets(NG);
929         if (plan->flags&FFT5D_ORDER_YZ)
930         {
931             rotate_offsets(xs);
932             rotate_offsets(xl);
933             rotate_offsets(xc);
934             rotate_offsets(NG);
935         }
936     }
937     if ((plan->flags&FFT5D_REALCOMPLEX) && ((!(plan->flags&FFT5D_BACKWARD) && s == 0) || ((plan->flags&FFT5D_BACKWARD) && s == 2)))
938     {
939         xl[0] = rC[s];
940     }
941 }
942
943 static void print_localdata(const t_complex* lin, const char* txt, int s, fft5d_plan plan)
944 {
945     int  x, y, z, l;
946     int *coor = plan->coor;
947     int  xs[3], xl[3], xc[3], NG[3];
948     int  ll = (plan->flags&FFT5D_REALCOMPLEX) ? 1 : 2;
949     compute_offsets(plan, xs, xl, xc, NG, s);
950     fprintf(debug, txt, coor[0], coor[1]);
951     /*printf("xs: %d %d %d, xl: %d %d %d\n",xs[0],xs[1],xs[2],xl[0],xl[1],xl[2]);*/
952     for (z = 0; z < xl[2]; z++)
953     {
954         for (y = 0; y < xl[1]; y++)
955         {
956             fprintf(debug, "%d %d: ", coor[0], coor[1]);
957             for (x = 0; x < xl[0]; x++)
958             {
959                 for (l = 0; l < ll; l++)
960                 {
961                     fprintf(debug, "%f ", ((real*)lin)[(z*xs[2]+y*xs[1])*2+(x*xs[0])*ll+l]);
962                 }
963                 fprintf(debug, ",");
964             }
965             fprintf(debug, "\n");
966         }
967     }
968 }
969
970 void fft5d_execute(fft5d_plan plan, int thread, fft5d_time times)
971 {
972     t_complex  *lin   = plan->lin;
973     t_complex  *lout  = plan->lout;
974     t_complex  *lout2 = plan->lout2;
975     t_complex  *lout3 = plan->lout3;
976     t_complex  *fftout, *joinin;
977
978     gmx_fft_t **p1d = plan->p1d;
979 #ifdef FFT5D_MPI_TRANSPOSE
980     FFTW(plan) *mpip = plan->mpip;
981 #endif
982 #if GMX_MPI
983     MPI_Comm *cart = plan->cart;
984 #endif
985 #ifdef NOGMX
986     double time_fft = 0, time_local = 0, time_mpi[2] = {0}, time = 0;
987 #endif
988     int   *N = plan->N, *M = plan->M, *K = plan->K, *pN = plan->pN, *pM = plan->pM, *pK = plan->pK,
989     *C       = plan->C, *P = plan->P, **iNin = plan->iNin, **oNin = plan->oNin, **iNout = plan->iNout, **oNout = plan->oNout;
990     int    s = 0, tstart, tend, bParallelDim;
991
992
993 #if GMX_FFT_FFTW3
994     if (plan->p3d)
995     {
996         if (thread == 0)
997         {
998 #ifdef NOGMX
999             if (times != 0)
1000             {
1001                 time = MPI_Wtime();
1002             }
1003 #endif
1004             FFTW(execute)(plan->p3d);
1005 #ifdef NOGMX
1006             if (times != 0)
1007             {
1008                 times->fft += MPI_Wtime()-time;
1009             }
1010 #endif
1011         }
1012         return;
1013     }
1014 #endif
1015
1016     s = 0;
1017
1018     /*lin: x,y,z*/
1019     if ((plan->flags&FFT5D_DEBUG) && thread == 0)
1020     {
1021         print_localdata(lin, "%d %d: copy in lin\n", s, plan);
1022     }
1023
1024     for (s = 0; s < 2; s++)  /*loop over first two FFT steps (corner rotations)*/
1025
1026     {
1027 #if GMX_MPI
1028         if (GMX_PARALLEL_ENV_INITIALIZED && cart[s] != MPI_COMM_NULL && P[s] > 1)
1029         {
1030             bParallelDim = 1;
1031         }
1032         else
1033 #endif
1034         {
1035             bParallelDim = 0;
1036         }
1037
1038         /* ---------- START FFT ------------ */
1039 #ifdef NOGMX
1040         if (times != 0 && thread == 0)
1041         {
1042             time = MPI_Wtime();
1043         }
1044 #endif
1045
1046         if (bParallelDim || plan->nthreads == 1)
1047         {
1048             fftout = lout;
1049         }
1050         else
1051         {
1052             if (s == 0)
1053             {
1054                 fftout = lout3;
1055             }
1056             else
1057             {
1058                 fftout = lout2;
1059             }
1060         }
1061
1062         tstart = (thread*pM[s]*pK[s]/plan->nthreads)*C[s];
1063         if ((plan->flags&FFT5D_REALCOMPLEX) && !(plan->flags&FFT5D_BACKWARD) && s == 0)
1064         {
1065             gmx_fft_many_1d_real(p1d[s][thread], (plan->flags&FFT5D_BACKWARD) ? GMX_FFT_COMPLEX_TO_REAL : GMX_FFT_REAL_TO_COMPLEX, lin+tstart, fftout+tstart);
1066         }
1067         else
1068         {
1069             gmx_fft_many_1d(     p1d[s][thread], (plan->flags&FFT5D_BACKWARD) ? GMX_FFT_BACKWARD : GMX_FFT_FORWARD,               lin+tstart, fftout+tstart);
1070
1071         }
1072
1073 #ifdef NOGMX
1074         if (times != NULL && thread == 0)
1075         {
1076             time_fft += MPI_Wtime()-time;
1077         }
1078 #endif
1079         if ((plan->flags&FFT5D_DEBUG) && thread == 0)
1080         {
1081             print_localdata(lout, "%d %d: FFT %d\n", s, plan);
1082         }
1083         /* ---------- END FFT ------------ */
1084
1085         /* ---------- START SPLIT + TRANSPOSE------------ (if parallel in in this dimension)*/
1086         if (bParallelDim)
1087         {
1088 #ifdef NOGMX
1089             if (times != NULL && thread == 0)
1090             {
1091                 time = MPI_Wtime();
1092             }
1093 #endif
1094             /*prepare for A
1095                llToAll
1096                1. (most outer) axes (x) is split into P[s] parts of size N[s]
1097                for sending*/
1098             if (pM[s] > 0)
1099             {
1100                 tend    = ((thread+1)*pM[s]*pK[s]/plan->nthreads);
1101                 tstart /= C[s];
1102                 splitaxes(lout2, lout, N[s], M[s], K[s], pM[s], P[s], C[s], iNout[s], oNout[s], tstart%pM[s], tstart/pM[s], tend%pM[s], tend/pM[s]);
1103             }
1104 #pragma omp barrier /*barrier required before AllToAll (all input has to be their) - before timing to make timing more acurate*/
1105 #ifdef NOGMX
1106             if (times != NULL && thread == 0)
1107             {
1108                 time_local += MPI_Wtime()-time;
1109             }
1110 #endif
1111
1112             /* ---------- END SPLIT , START TRANSPOSE------------ */
1113
1114             if (thread == 0)
1115             {
1116 #ifdef NOGMX
1117                 if (times != 0)
1118                 {
1119                     time = MPI_Wtime();
1120                 }
1121 #else
1122                 wallcycle_start(times, ewcPME_FFTCOMM);
1123 #endif
1124 #ifdef FFT5D_MPI_TRANSPOSE
1125                 FFTW(execute)(mpip[s]);
1126 #else
1127 #if GMX_MPI
1128                 if ((s == 0 && !(plan->flags&FFT5D_ORDER_YZ)) || (s == 1 && (plan->flags&FFT5D_ORDER_YZ)))
1129                 {
1130                     MPI_Alltoall((real *)lout2, N[s]*pM[s]*K[s]*sizeof(t_complex)/sizeof(real), GMX_MPI_REAL, (real *)lout3, N[s]*pM[s]*K[s]*sizeof(t_complex)/sizeof(real), GMX_MPI_REAL, cart[s]);
1131                 }
1132                 else
1133                 {
1134                     MPI_Alltoall((real *)lout2, N[s]*M[s]*pK[s]*sizeof(t_complex)/sizeof(real), GMX_MPI_REAL, (real *)lout3, N[s]*M[s]*pK[s]*sizeof(t_complex)/sizeof(real), GMX_MPI_REAL, cart[s]);
1135                 }
1136 #else
1137                 gmx_incons("fft5d MPI call without MPI configuration");
1138 #endif /*GMX_MPI*/
1139 #endif /*FFT5D_MPI_TRANSPOSE*/
1140 #ifdef NOGMX
1141                 if (times != 0)
1142                 {
1143                     time_mpi[s] = MPI_Wtime()-time;
1144                 }
1145 #else
1146                 wallcycle_stop(times, ewcPME_FFTCOMM);
1147 #endif
1148             } /*master*/
1149         }     /* bPrallelDim */
1150 #pragma omp barrier  /*both needed for parallel and non-parallel dimension (either have to wait on data from AlltoAll or from last FFT*/
1151
1152         /* ---------- END SPLIT + TRANSPOSE------------ */
1153
1154         /* ---------- START JOIN ------------ */
1155 #ifdef NOGMX
1156         if (times != NULL && thread == 0)
1157         {
1158             time = MPI_Wtime();
1159         }
1160 #endif
1161
1162         if (bParallelDim)
1163         {
1164             joinin = lout3;
1165         }
1166         else
1167         {
1168             joinin = fftout;
1169         }
1170         /*bring back in matrix form
1171            thus make  new 1. axes contiguos
1172            also local transpose 1 and 2/3
1173            runs on thread used for following FFT (thus needing a barrier before but not afterwards)
1174          */
1175         if ((s == 0 && !(plan->flags&FFT5D_ORDER_YZ)) || (s == 1 && (plan->flags&FFT5D_ORDER_YZ)))
1176         {
1177             if (pM[s] > 0)
1178             {
1179                 tstart = ( thread   *pM[s]*pN[s]/plan->nthreads);
1180                 tend   = ((thread+1)*pM[s]*pN[s]/plan->nthreads);
1181                 joinAxesTrans13(lin, joinin, N[s], pM[s], K[s], pM[s], P[s], C[s+1], iNin[s+1], oNin[s+1], tstart%pM[s], tstart/pM[s], tend%pM[s], tend/pM[s]);
1182             }
1183         }
1184         else
1185         {
1186             if (pN[s] > 0)
1187             {
1188                 tstart = ( thread   *pK[s]*pN[s]/plan->nthreads);
1189                 tend   = ((thread+1)*pK[s]*pN[s]/plan->nthreads);
1190                 joinAxesTrans12(lin, joinin, N[s], M[s], pK[s], pN[s], P[s], C[s+1], iNin[s+1], oNin[s+1], tstart%pN[s], tstart/pN[s], tend%pN[s], tend/pN[s]);
1191             }
1192         }
1193
1194 #ifdef NOGMX
1195         if (times != NULL && thread == 0)
1196         {
1197             time_local += MPI_Wtime()-time;
1198         }
1199 #endif
1200         if ((plan->flags&FFT5D_DEBUG) && thread == 0)
1201         {
1202             print_localdata(lin, "%d %d: tranposed %d\n", s+1, plan);
1203         }
1204         /* ---------- END JOIN ------------ */
1205
1206         /*if (debug) print_localdata(lin, "%d %d: transposed x-z\n", N1, M0, K, ZYX, coor);*/
1207     }  /* for(s=0;s<2;s++) */
1208 #ifdef NOGMX
1209     if (times != NULL && thread == 0)
1210     {
1211         time = MPI_Wtime();
1212     }
1213 #endif
1214
1215     if (plan->flags&FFT5D_INPLACE)
1216     {
1217         lout = lin;                          /*in place currently not supported*/
1218
1219     }
1220     /*  ----------- FFT ----------- */
1221     tstart = (thread*pM[s]*pK[s]/plan->nthreads)*C[s];
1222     if ((plan->flags&FFT5D_REALCOMPLEX) && (plan->flags&FFT5D_BACKWARD))
1223     {
1224         gmx_fft_many_1d_real(p1d[s][thread], (plan->flags&FFT5D_BACKWARD) ? GMX_FFT_COMPLEX_TO_REAL : GMX_FFT_REAL_TO_COMPLEX, lin+tstart, lout+tstart);
1225     }
1226     else
1227     {
1228         gmx_fft_many_1d(     p1d[s][thread], (plan->flags&FFT5D_BACKWARD) ? GMX_FFT_BACKWARD : GMX_FFT_FORWARD,               lin+tstart, lout+tstart);
1229     }
1230     /* ------------ END FFT ---------*/
1231
1232 #ifdef NOGMX
1233     if (times != NULL && thread == 0)
1234     {
1235         time_fft += MPI_Wtime()-time;
1236
1237         times->fft   += time_fft;
1238         times->local += time_local;
1239         times->mpi2  += time_mpi[1];
1240         times->mpi1  += time_mpi[0];
1241     }
1242 #endif
1243
1244     if ((plan->flags&FFT5D_DEBUG) && thread == 0)
1245     {
1246         print_localdata(lout, "%d %d: FFT %d\n", s, plan);
1247     }
1248 }
1249
1250 void fft5d_destroy(fft5d_plan plan)
1251 {
1252     int s, t;
1253
1254     for (s = 0; s < 3; s++)
1255     {
1256         if (plan->p1d[s])
1257         {
1258             for (t = 0; t < plan->nthreads; t++)
1259             {
1260                 gmx_many_fft_destroy(plan->p1d[s][t]);
1261             }
1262             free(plan->p1d[s]);
1263         }
1264         if (plan->iNin[s])
1265         {
1266             free(plan->iNin[s]);
1267             plan->iNin[s] = 0;
1268         }
1269         if (plan->oNin[s])
1270         {
1271             free(plan->oNin[s]);
1272             plan->oNin[s] = 0;
1273         }
1274         if (plan->iNout[s])
1275         {
1276             free(plan->iNout[s]);
1277             plan->iNout[s] = 0;
1278         }
1279         if (plan->oNout[s])
1280         {
1281             free(plan->oNout[s]);
1282             plan->oNout[s] = 0;
1283         }
1284     }
1285 #if GMX_FFT_FFTW3
1286     FFTW_LOCK;
1287 #ifdef FFT5D_MPI_TRANSPOS
1288     for (s = 0; s < 2; s++)
1289     {
1290         FFTW(destroy_plan)(plan->mpip[s]);
1291     }
1292 #endif /* FFT5D_MPI_TRANSPOS */
1293     if (plan->p3d)
1294     {
1295         FFTW(destroy_plan)(plan->p3d);
1296     }
1297     FFTW_UNLOCK;
1298 #endif /* GMX_FFT_FFTW3 */
1299
1300     if (!(plan->flags&FFT5D_NOMALLOC))
1301     {
1302         sfree_aligned(plan->lin);
1303         sfree_aligned(plan->lout);
1304         if (plan->nthreads > 1)
1305         {
1306             sfree_aligned(plan->lout2);
1307             sfree_aligned(plan->lout3);
1308         }
1309     }
1310
1311 #ifdef FFT5D_THREADS
1312 #ifdef FFT5D_FFTW_THREADS
1313     /*FFTW(cleanup_threads)();*/
1314 #endif
1315 #endif
1316
1317     free(plan);
1318 }
1319
1320 /*Is this better than direct access of plan? enough data?
1321    here 0,1 reference divided by which processor grid dimension (not FFT step!)*/
1322 void fft5d_local_size(fft5d_plan plan, int* N1, int* M0, int* K0, int* K1, int** coor)
1323 {
1324     *N1 = plan->N[0];
1325     *M0 = plan->M[0];
1326     *K1 = plan->K[0];
1327     *K0 = plan->N[1];
1328
1329     *coor = plan->coor;
1330 }
1331
1332
1333 /*same as fft5d_plan_3d but with cartesian coordinator and automatic splitting
1334    of processor dimensions*/
1335 fft5d_plan fft5d_plan_3d_cart(int NG, int MG, int KG, MPI_Comm comm, int P0, int flags, t_complex** rlin, t_complex** rlout, t_complex** rlout2, t_complex** rlout3, int nthreads)
1336 {
1337     MPI_Comm cart[2] = {0};
1338 #if GMX_MPI
1339     int      size = 1, prank = 0;
1340     int      P[2];
1341     int      coor[2];
1342     int      wrap[] = {0, 0};
1343     MPI_Comm gcart;
1344     int      rdim1[] = {0, 1}, rdim2[] = {1, 0};
1345
1346     MPI_Comm_size(comm, &size);
1347     MPI_Comm_rank(comm, &prank);
1348
1349     if (P0 == 0)
1350     {
1351         P0 = lfactor(size);
1352     }
1353     if (size%P0 != 0)
1354     {
1355         if (prank == 0)
1356         {
1357             printf("FFT5D: WARNING: Number of ranks %d not evenly divisible by %d\n", size, P0);
1358         }
1359         P0 = lfactor(size);
1360     }
1361
1362     P[0] = P0; P[1] = size/P0; /*number of processors in the two dimensions*/
1363
1364     /*Difference between x-y-z regarding 2d decomposition is whether they are
1365        distributed along axis 1, 2 or both*/
1366
1367     MPI_Cart_create(comm, 2, P, wrap, 1, &gcart); /*parameter 4: value 1: reorder*/
1368     MPI_Cart_get(gcart, 2, P, wrap, coor);
1369     MPI_Cart_sub(gcart, rdim1, &cart[0]);
1370     MPI_Cart_sub(gcart, rdim2, &cart[1]);
1371 #else
1372     (void)P0;
1373     (void)comm;
1374 #endif
1375     return fft5d_plan_3d(NG, MG, KG, cart, flags, rlin, rlout, rlout2, rlout3, nthreads);
1376 }
1377
1378
1379
1380 /*prints in original coordinate system of data (as the input to FFT)*/
1381 void fft5d_compare_data(const t_complex* lin, const t_complex* in, fft5d_plan plan, int bothLocal, int normalize)
1382 {
1383     int  xs[3], xl[3], xc[3], NG[3];
1384     int  x, y, z, l;
1385     int *coor = plan->coor;
1386     int  ll   = 2; /*compare ll values per element (has to be 2 for complex)*/
1387     if ((plan->flags&FFT5D_REALCOMPLEX) && (plan->flags&FFT5D_BACKWARD))
1388     {
1389         ll = 1;
1390     }
1391
1392     compute_offsets(plan, xs, xl, xc, NG, 2);
1393     if (plan->flags&FFT5D_DEBUG)
1394     {
1395         printf("Compare2\n");
1396     }
1397     for (z = 0; z < xl[2]; z++)
1398     {
1399         for (y = 0; y < xl[1]; y++)
1400         {
1401             if (plan->flags&FFT5D_DEBUG)
1402             {
1403                 printf("%d %d: ", coor[0], coor[1]);
1404             }
1405             for (x = 0; x < xl[0]; x++)
1406             {
1407                 for (l = 0; l < ll; l++)   /*loop over real/complex parts*/
1408                 {
1409                     real a, b;
1410                     a = ((real*)lin)[(z*xs[2]+y*xs[1])*2+x*xs[0]*ll+l];
1411                     if (normalize)
1412                     {
1413                         a /= plan->rC[0]*plan->rC[1]*plan->rC[2];
1414                     }
1415                     if (!bothLocal)
1416                     {
1417                         b = ((real*)in)[((z+xc[2])*NG[0]*NG[1]+(y+xc[1])*NG[0])*2+(x+xc[0])*ll+l];
1418                     }
1419                     else
1420                     {
1421                         b = ((real*)in)[(z*xs[2]+y*xs[1])*2+x*xs[0]*ll+l];
1422                     }
1423                     if (plan->flags&FFT5D_DEBUG)
1424                     {
1425                         printf("%f %f, ", a, b);
1426                     }
1427                     else
1428                     {
1429                         if (fabs(a-b) > 2*NG[0]*NG[1]*NG[2]*GMX_REAL_EPS)
1430                         {
1431                             printf("result incorrect on %d,%d at %d,%d,%d: FFT5D:%f reference:%f\n", coor[0], coor[1], x, y, z, a, b);
1432                         }
1433 /*                        assert(fabs(a-b)<2*NG[0]*NG[1]*NG[2]*GMX_REAL_EPS);*/
1434                     }
1435                 }
1436                 if (plan->flags&FFT5D_DEBUG)
1437                 {
1438                     printf(",");
1439                 }
1440             }
1441             if (plan->flags&FFT5D_DEBUG)
1442             {
1443                 printf("\n");
1444             }
1445         }
1446     }
1447
1448 }