pspg/field/FFT.h

#ifndef PSPG_FFT_H
#define PSPG_FFT_H
 
/*
* PSCF Package 
*
* Copyright 2016 - 2022, The Regents of the University of Minnesota
* Distributed under the terms of the GNU General Public License.
*/
 
#include <pspg/field/RDField.h>
#include <pspg/field/RDFieldDft.h>
#include <pspg/math/GpuResources.h>
#include <pscf/math/IntVec.h>
#include <util/global.h>
 
#include <cufft.h>
#include <cuda.h>
#include <cuda_runtime.h>
 
namespace Pscf {
namespace Pspg {
 
   using namespace Util;
   using namespace Pscf;
 
   template <int D>
   class FFT 
   {
 
   public:
 
      FFT();
 
      virtual ~FFT();
 
      void setup(IntVec<D> const & meshDimensions);
 
      void setup(RDField<D>& rDField, RDFieldDft<D>& kDField);
 
      void forwardTransform(RDField<D> & rField, RDFieldDft<D>& kField) 
      const;
 
      void forwardTransformSafe(RDField<D> const & rField, 
                                RDFieldDft<D>& kField) const;
 
      void inverseTransform(RDFieldDft<D> & kField, RDField<D>& rField) 
      const;
 
      void inverseTransformSafe(RDFieldDft<D> const & kField, 
                                RDField<D>& rField) const;
 
      const IntVec<D>& meshDimensions() const;
 
      bool isSetup() const;
 
      cufftHandle& fPlan();
 
      cufftHandle& iPlan();
 
   private:
 
      // Vector containing number of grid points in each direction.
      IntVec<D> meshDimensions_;
 
      // Private r-space array for performing safe transforms.
      mutable RDField<D> rFieldCopy_;
 
      // Private k-space array for performing safe transforms.
      mutable RDFieldDft<D> kFieldCopy_;
 
      // Number of points in r-space grid
      int rSize_;
 
      // Number of points in k-space grid
      int kSize_;
 
      // Pointer to a plan for a forward transform.
      cufftHandle fPlan_;
 
      // Pointer to a plan for an inverse transform.
      cufftHandle iPlan_;
 
      // Have array dimension and plan been initialized?
      bool isSetup_;
 
      void makePlans(RDField<D>& rDField, RDFieldDft<D>& kDField);
 
   };
 
   // Declarations of explicit specializations
 
   template <>
   void FFT<1>::makePlans(RDField<1>& rField, RDFieldDft<1>& kField);
 
   template <>
   void FFT<2>::makePlans(RDField<2>& rField, RDFieldDft<2>& kField);
 
   template <>
   void FFT<3>::makePlans(RDField<3>& rField, RDFieldDft<3>& kField);
 
   /*
   * Return the dimensions of the grid for which this was allocated.
   */
   template <int D>
   inline const IntVec<D>& FFT<D>::meshDimensions() const
   {  return meshDimensions_; }
 
   template <int D>
   inline bool FFT<D>::isSetup() const
   { return isSetup_; }
 
   template <int D>
   inline cufftHandle& FFT<D>::fPlan() 
   { return fPlan_; }
 
   template <int D>
   inline cufftHandle& FFT<D>::iPlan()
   { return iPlan_; }
 
 
   #ifndef PSPG_FFT_TPP
   // Suppress implicit instantiation
   extern template class FFT<1>;
   extern template class FFT<2>;
   extern template class FFT<3>;
   #endif
 
static __global__ 
void scaleRealData(cudaReal* data, cudaReal scale, int size) {
   //write code that will scale
   int nThreads = blockDim.x * gridDim.x;
   int startId = blockIdx.x * blockDim.x + threadIdx.x;
   for(int i = startId; i < size; i += nThreads ) {
      data[i] *= scale;
   }
}
 
}
}
 
//#include "FFT.tpp"
#endif