FFTBatched.h

#ifndef PRDC_CUDA_FFT_BATCHED_H
#define PRDC_CUDA_FFT_BATCHED_H
 
/*
* PSCF - Polymer Self-Consistent Field 
*
* Copyright 2015 - 2025, The Regents of the University of Minnesota
* Distributed under the terms of the GNU General Public License.
*/
 
//#include "RField.h"
//#include "RFieldDft.h"
#include <pscf/cuda/cudaTypes.h>
#include <pscf/cuda/DeviceArray.h>
#include <pscf/math/IntVec.h>
#include <util/global.h>
 
#include <cufft.h>
 
namespace Pscf {
namespace Prdc {
namespace Cuda {
 
   using namespace Util;

   template <int D>
   class FFTBatched
   {
 
   public:

      FFTBatched();

      virtual ~FFTBatched();

      void setup(IntVec<D> const & meshDimensions, int batchSize);
      
      void forwardTransform(DeviceArray<cudaReal> const & rFields, 
                            DeviceArray<cudaComplex>& kFields) const;

      void inverseTransformUnsafe(DeviceArray<cudaComplex>& kFields, 
                                  DeviceArray<cudaReal>& rFields) const;

      const IntVec<D>& meshDimensions() const;

      void resetBatchSize(int batchSize);

      bool isSetup() const;
 
   private:

      IntVec<D> meshDimensions_;

      IntVec<D> kMeshDimensions_;

      int batchSize_;

      int rSize_;

      int kSize_;

      cufftHandle fPlan_;

      cufftHandle iPlan_;

      bool isSetup_;

      void makePlans(int batchSize);
 
   };
 
   // Return the dimensions of the grid for which this was allocated.
   template <int D>
   inline const IntVec<D>& FFTBatched<D>::meshDimensions() const
   {  return meshDimensions_; }
 
   // Has the setup method been called?
   template <int D>
   inline bool FFTBatched<D>::isSetup() const
   { return isSetup_; }
 
   #ifndef PRDC_CUDA_FFT_BATCHED_TPP
   // Explicit instantiation declarations
   extern template class FFTBatched<1>;
   extern template class FFTBatched<2>;
   extern template class FFTBatched<3>;
   #endif
 
}
}
}
 
#endif