src/gromacs/fft/gpu_3dfft_cufft.cu

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2016,2017,2018,2019,2020,2021, 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
  36 /*! \internal \file
  37  *  \brief Implements GPU 3D FFT routines for CUDA.
  38  *
  39  *  \author Aleksei Iupinov <a.yupinov@gmail.com>
  40  *  \author Mark Abraham <mark.j.abraham@gmail.com>
  41  *  \ingroup module_fft
  42  */
  43
  44 #include "gmxpre.h"
  45
  46 #include "gpu_3dfft_cufft.h"
  47
  48 #include "gromacs/gpu_utils/device_stream.h"
  49 #include "gromacs/utility/arrayref.h"
  50 #include "gromacs/utility/fatalerror.h"
  51 #include "gromacs/utility/gmxassert.h"
  52
  53 namespace gmx
  54 {
  55 static void handleCufftError(cufftResult_t status, const char* msg)
  56 {
  57     if (status != CUFFT_SUCCESS)
  58     {
  59         gmx_fatal(FARGS, "%s (error code %d)\n", msg, status);
  60     }
  61 }
  62
  63 Gpu3dFft::ImplCuFft::ImplCuFft(bool allocateGrids,
  64                                MPI_Comm /*comm*/,
  65                                ArrayRef<const int> gridSizesInXForEachRank,
  66                                ArrayRef<const int> gridSizesInYForEachRank,
  67                                const int /*nz*/,
  68                                bool /*performOutOfPlaceFFT*/,
  69                                const DeviceContext& /*context*/,
  70                                const DeviceStream&  pmeStream,
  71                                ivec                 realGridSize,
  72                                ivec                 realGridSizePadded,
  73                                ivec                 complexGridSizePadded,
  74                                DeviceBuffer<float>* realGrid,
  75                                DeviceBuffer<float>* complexGrid) :
  76     realGrid_(reinterpret_cast<cufftReal*>(*realGrid)),
  77     complexGrid_(reinterpret_cast<cufftComplex*>(*complexGrid))
  78 {
  79     GMX_RELEASE_ASSERT(allocateGrids == false, "Grids needs to be pre-allocated");
  80     GMX_RELEASE_ASSERT(gridSizesInXForEachRank.size() == 1 && gridSizesInYForEachRank.size() == 1,
  81                        "FFT decomposition not implemented with cuFFT backend");
  82
  83     const int complexGridSizePaddedTotal =
  84             complexGridSizePadded[XX] * complexGridSizePadded[YY] * complexGridSizePadded[ZZ];
  85     const int realGridSizePaddedTotal =
  86             realGridSizePadded[XX] * realGridSizePadded[YY] * realGridSizePadded[ZZ];
  87
  88     GMX_RELEASE_ASSERT(realGrid_, "Bad (null) input real-space grid");
  89     GMX_RELEASE_ASSERT(complexGrid_, "Bad (null) input complex grid");
  90
  91     cufftResult_t result;
  92     /* Commented code for a simple 3D grid with no padding */
  93     /*
  94        result = cufftPlan3d(&planR2C_, realGridSize[XX], realGridSize[YY], realGridSize[ZZ],
  95        CUFFT_R2C); handleCufftError(result, "cufftPlan3d R2C plan failure");
  96
  97        result = cufftPlan3d(&planC2R_, realGridSize[XX], realGridSize[YY], realGridSize[ZZ],
  98        CUFFT_C2R); handleCufftError(result, "cufftPlan3d C2R plan failure");
  99      */
 100
 101     const int rank = 3, batch = 1;
 102     result = cufftPlanMany(&planR2C_,
 103                            rank,
 104                            realGridSize,
 105                            realGridSizePadded,
 106                            1,
 107                            realGridSizePaddedTotal,
 108                            complexGridSizePadded,
 109                            1,
 110                            complexGridSizePaddedTotal,
 111                            CUFFT_R2C,
 112                            batch);
 113     handleCufftError(result, "cufftPlanMany R2C plan failure");
 114
 115     result = cufftPlanMany(&planC2R_,
 116                            rank,
 117                            realGridSize,
 118                            complexGridSizePadded,
 119                            1,
 120                            complexGridSizePaddedTotal,
 121                            realGridSizePadded,
 122                            1,
 123                            realGridSizePaddedTotal,
 124                            CUFFT_C2R,
 125                            batch);
 126     handleCufftError(result, "cufftPlanMany C2R plan failure");
 127
 128     cudaStream_t stream = pmeStream.stream();
 129     GMX_RELEASE_ASSERT(stream, "Can not use the default CUDA stream for PME cuFFT");
 130
 131     result = cufftSetStream(planR2C_, stream);
 132     handleCufftError(result, "cufftSetStream R2C failure");
 133
 134     result = cufftSetStream(planC2R_, stream);
 135     handleCufftError(result, "cufftSetStream C2R failure");
 136 }
 137
 138 Gpu3dFft::ImplCuFft::~ImplCuFft()
 139 {
 140     cufftResult_t result;
 141     result = cufftDestroy(planR2C_);
 142     handleCufftError(result, "cufftDestroy R2C failure");
 143     result = cufftDestroy(planC2R_);
 144     handleCufftError(result, "cufftDestroy C2R failure");
 145 }
 146
 147 void Gpu3dFft::ImplCuFft::perform3dFft(gmx_fft_direction dir, CommandEvent* /*timingEvent*/)
 148 {
 149     cufftResult_t result;
 150     if (dir == GMX_FFT_REAL_TO_COMPLEX)
 151     {
 152         result = cufftExecR2C(planR2C_, realGrid_, complexGrid_);
 153         handleCufftError(result, "cuFFT R2C execution failure");
 154     }
 155     else
 156     {
 157         result = cufftExecC2R(planC2R_, complexGrid_, realGrid_);
 158         handleCufftError(result, "cuFFT C2R execution failure");
 159     }
 160 }
 161
 162 } // namespace gmx