src/gromacs/fft/parallel_3dfft.cpp

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 1991-2005 David van der Spoel, Erik Lindahl, University of Groningen.
   5  * Copyright (c) 2013,2014,2017,2018,2019,2020,2021, by the GROMACS development team, led by
   6  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   7  * and including many others, as listed in the AUTHORS file in the
   8  * top-level source directory and at http://www.gromacs.org.
   9  *
  10  * GROMACS is free software; you can redistribute it and/or
  11  * modify it under the terms of the GNU Lesser General Public License
  12  * as published by the Free Software Foundation; either version 2.1
  13  * of the License, or (at your option) any later version.
  14  *
  15  * GROMACS is distributed in the hope that it will be useful,
  16  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  17  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  18  * Lesser General Public License for more details.
  19  *
  20  * You should have received a copy of the GNU Lesser General Public
  21  * License along with GROMACS; if not, see
  22  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
  23  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  24  *
  25  * If you want to redistribute modifications to GROMACS, please
  26  * consider that scientific software is very special. Version
  27  * control is crucial - bugs must be traceable. We will be happy to
  28  * consider code for inclusion in the official distribution, but
  29  * derived work must not be called official GROMACS. Details are found
  30  * in the README & COPYING files - if they are missing, get the
  31  * official version at http://www.gromacs.org.
  32  *
  33  * To help us fund GROMACS development, we humbly ask that you cite
  34  * the research papers on the package. Check out http://www.gromacs.org.
  35  */
  36 #include "gmxpre.h"
  37
  38 #include "parallel_3dfft.h"
  39
  40 #include <cerrno>
  41 #include <cstdlib>
  42 #include <cstring>
  43
  44 #include "gromacs/fft/fft.h"
  45 #include "gromacs/fft/fft5d.h"
  46 #include "gromacs/math/gmxcomplex.h"
  47 #include "gromacs/utility/fatalerror.h"
  48 #include "gromacs/utility/gmxmpi.h"
  49 #include "gromacs/utility/smalloc.h"
  50
  51 struct gmx_parallel_3dfft
  52 {
  53     fft5d_plan p1, p2;
  54 };
  55
  56 int gmx_parallel_3dfft_init(gmx_parallel_3dfft_t* pfft_setup,
  57                             const ivec            ndata,
  58                             real**                real_data,
  59                             t_complex**           complex_data,
  60                             MPI_Comm              comm[2],
  61                             gmx_bool              bReproducible,
  62                             int                   nthreads,
  63                             gmx::PinningPolicy    realGridAllocation)
  64 {
  65     int        rN = ndata[2], M = ndata[1], K = ndata[0];
  66     int        flags   = FFT5D_REALCOMPLEX | FFT5D_ORDER_YZ; /* FFT5D_DEBUG */
  67     MPI_Comm   rcomm[] = { comm[1], comm[0] };
  68     int        Nb, Mb, Kb;  /* dimension for backtransform (in starting order) */
  69     t_complex *buf1, *buf2; /*intermediate buffers - used internally.*/
  70
  71     snew(*pfft_setup, 1);
  72     if (bReproducible)
  73     {
  74         flags |= FFT5D_NOMEASURE;
  75     }
  76
  77     if (!(flags & FFT5D_ORDER_YZ))
  78     {
  79         Nb = M;
  80         Mb = K;
  81         Kb = rN;
  82     }
  83     else
  84     {
  85         Nb = K;
  86         Mb = rN;
  87         Kb = M; /* currently always true because ORDER_YZ always set */
  88     }
  89
  90     (*pfft_setup)->p1 = fft5d_plan_3d(
  91             rN, M, K, rcomm, flags, reinterpret_cast<t_complex**>(real_data), complex_data, &buf1, &buf2, nthreads, realGridAllocation);
  92
  93     (*pfft_setup)->p2 = fft5d_plan_3d(Nb,
  94                                       Mb,
  95                                       Kb,
  96                                       rcomm,
  97                                       (flags | FFT5D_BACKWARD | FFT5D_NOMALLOC) ^ FFT5D_ORDER_YZ,
  98                                       complex_data,
  99                                       reinterpret_cast<t_complex**>(real_data),
 100                                       &buf1,
 101                                       &buf2,
 102                                       nthreads);
 103
 104     return static_cast<int>((*pfft_setup)->p1 != nullptr && (*pfft_setup)->p2 != nullptr);
 105 }
 106
 107
 108 static int fft5d_limits(fft5d_plan p, ivec local_ndata, ivec local_offset, ivec local_size)
 109 {
 110     local_offset[2] = 0;
 111     local_offset[1] = p->oM[0]; /*=p->coor[0]*p->MG/p->P[0]; */
 112     local_offset[0] = p->oK[0]; /*=p->coor[1]*p->KG/p->P[1]; */
 113
 114     local_ndata[2] = p->rC[0];
 115     local_ndata[1] = p->pM[0];
 116     local_ndata[0] = p->pK[0];
 117
 118     if ((!(p->flags & FFT5D_BACKWARD)) && (p->flags & FFT5D_REALCOMPLEX))
 119     {
 120         // C is length in multiples of complex local_size in multiples of real
 121         local_size[2] = p->C[0] * 2;
 122     }
 123     else
 124     {
 125         local_size[2] = p->C[0];
 126     }
 127     local_size[1] = p->pM[0];
 128     local_size[0] = p->pK[0];
 129     return 0;
 130 }
 131
 132 int gmx_parallel_3dfft_real_limits(gmx_parallel_3dfft_t pfft_setup, ivec local_ndata, ivec local_offset, ivec local_size)
 133 {
 134     return fft5d_limits(pfft_setup->p1, local_ndata, local_offset, local_size);
 135 }
 136
 137 static void reorder_ivec_yzx(ivec v)
 138 {
 139     int tmp;
 140
 141     tmp   = v[0];
 142     v[XX] = v[2];
 143     v[ZZ] = v[1];
 144     v[YY] = tmp;
 145 }
 146
 147 int gmx_parallel_3dfft_complex_limits(gmx_parallel_3dfft_t pfft_setup,
 148                                       ivec                 complex_order,
 149                                       ivec                 local_ndata,
 150                                       ivec                 local_offset,
 151                                       ivec                 local_size)
 152 {
 153     int ret;
 154
 155     /* For now everything is in-order, but prepare to save communication by avoiding transposes */
 156     complex_order[0] = 0;
 157     complex_order[1] = 1;
 158     complex_order[2] = 2;
 159
 160     ret = fft5d_limits(pfft_setup->p2, local_ndata, local_offset, local_size);
 161
 162     reorder_ivec_yzx(local_ndata);
 163     reorder_ivec_yzx(local_offset);
 164     reorder_ivec_yzx(local_size);
 165
 166     return ret;
 167 }
 168
 169
 170 int gmx_parallel_3dfft_execute(gmx_parallel_3dfft_t   pfft_setup,
 171                                enum gmx_fft_direction dir,
 172                                int                    thread,
 173                                gmx_wallcycle*         wcycle)
 174 {
 175     if (((pfft_setup->p1->flags & FFT5D_REALCOMPLEX) == 0)
 176         ^ (dir == GMX_FFT_FORWARD || dir == GMX_FFT_BACKWARD))
 177     {
 178         gmx_fatal(FARGS, "Invalid transform. Plan and execution don't match regarding reel/complex");
 179     }
 180     if (dir == GMX_FFT_FORWARD || dir == GMX_FFT_REAL_TO_COMPLEX)
 181     {
 182         fft5d_execute(pfft_setup->p1, thread, wcycle);
 183     }
 184     else
 185     {
 186         fft5d_execute(pfft_setup->p2, thread, wcycle);
 187     }
 188     return 0;
 189 }
 190
 191 int gmx_parallel_3dfft_destroy(gmx_parallel_3dfft_t pfft_setup)
 192 {
 193     if (pfft_setup)
 194     {
 195         fft5d_destroy(pfft_setup->p2);
 196         fft5d_destroy(pfft_setup->p1);
 197         sfree(pfft_setup);
 198     }
 199     return 0;
 200 }