src/fft_cufft.h

   1 #ifndef FFT_CUDA_H
   2 #define FFT_CUDA_H
   3
   4 #include <cufft.h>
   5 #include <cuda_runtime.h>
   6 #include <cublas_v2.h>
   7
   8 #include "fft.h"
   9 #include "cuda_error_check.hpp"
  10 #include "pragmas.h"
  11
  12 struct ThreadCtx;
  13
  14 class cuFFT : public Fft
  15 {
  16 public:
  17     cuFFT();
  18     void init(unsigned width, unsigned height, unsigned num_of_feats, unsigned num_of_scales);
  19     void set_window(const MatDynMem &window);
  20     void forward(const MatScales &real_input, ComplexMat &complex_result);
  21     void forward_window(MatScaleFeats &patch_feats_in, ComplexMat &complex_result, MatScaleFeats &tmp);
  22     void inverse(ComplexMat &complex_input, MatScales &real_result);
  23     ~cuFFT();
  24
  25 protected:
  26     cufftHandle create_plan_fwd(uint howmany) const;
  27     cufftHandle create_plan_inv(uint howmany) const;
  28
  29 private:
  30     cv::Mat m_window;
  31     cufftHandle plan_f, plan_fw, plan_i_1ch;
  32 #ifdef BIG_BATCH
  33     cufftHandle plan_f_all_scales, plan_fw_all_scales, plan_i_all_scales;
  34 #endif
  35     cublasHandle_t cublas;
  36 };
  37
  38 #endif // FFT_CUDA_H