src/fft_cufft.cpp

   1 #include "fft_cufft.h"
   2 #include <cublas_v2.h>
   3
   4 cuFFT::cuFFT()
   5 {
   6     CublasErrorCheck(cublasCreate(&cublas));
   7 }
   8
   9 void cuFFT::init(unsigned width, unsigned height, unsigned num_of_feats, unsigned num_of_scales)
  10 {
  11     Fft::init(width, height, num_of_feats, num_of_scales);
  12
  13     std::cout << "FFT: cuFFT" << std::endl;
  14
  15     // FFT forward
  16     {
  17         int rank = 2;
  18         int n[] = {int(m_height), int(m_width)};
  19         int howmany = IF_BIG_BATCH(m_num_of_scales, 1);
  20         int idist = m_height * m_width, odist = m_height * (m_width / 2 + 1);
  21         int istride = 1, ostride = 1;
  22         int *inembed = n, onembed[] = {int(m_height), int(m_width) / 2 + 1};
  23
  24         CufftErrorCheck(cufftPlanMany(&plan_f, rank, n, inembed, istride, idist, onembed, ostride, odist, CUFFT_R2C, howmany));
  25         CufftErrorCheck(cufftSetStream(plan_f, cudaStreamPerThread));
  26     }
  27
  28     // FFT forward window
  29     {
  30         int rank = 2;
  31         int n[] = {int(m_height), int(m_width)};
  32         int howmany = m_num_of_feats * IF_BIG_BATCH(m_num_of_scales, 1);
  33         int idist = m_height * m_width, odist = m_height * (m_width / 2 + 1);
  34         int istride = 1, ostride = 1;
  35         int *inembed = n, onembed[] = {int(m_height), int(m_width) / 2 + 1};
  36
  37         CufftErrorCheck(cufftPlanMany(&plan_fw, rank, n, inembed, istride, idist, onembed, ostride, odist, CUFFT_R2C, howmany));
  38         CufftErrorCheck(cufftSetStream(plan_fw, cudaStreamPerThread));
  39     }
  40     // FFT inverse all channels
  41     {
  42         int rank = 2;
  43         int n[] = {int(m_height), int(m_width)};
  44         int howmany = m_num_of_feats * IF_BIG_BATCH(m_num_of_scales, 1);
  45         int idist = int(m_height * (m_width / 2 + 1)), odist = 1;
  46         int istride = 1, ostride = m_num_of_feats * IF_BIG_BATCH(m_num_of_scales, 1);
  47         int inembed[] = {int(m_height), int(m_width / 2 + 1)}, *onembed = n;
  48
  49         CufftErrorCheck(cufftPlanMany(&plan_i_features, rank, n, inembed, istride, idist, onembed, ostride, odist, CUFFT_C2R, howmany));
  50         CufftErrorCheck(cufftSetStream(plan_i_features, cudaStreamPerThread));
  51     }
  52     // FFT inverse one channel
  53     {
  54         int rank = 2;
  55         int n[] = {int(m_height), int(m_width)};
  56         int howmany = IF_BIG_BATCH(m_num_of_scales, 1);
  57         int idist = m_height * (m_width / 2 + 1), odist = 1;
  58         int istride = 1, ostride = IF_BIG_BATCH(m_num_of_scales, 1);
  59         int inembed[] = {int(m_height), int(m_width / 2 + 1)}, *onembed = n;
  60
  61         CufftErrorCheck(cufftPlanMany(&plan_i_1ch, rank, n, inembed, istride, idist, onembed, ostride, odist, CUFFT_C2R, howmany));
  62         CufftErrorCheck(cufftSetStream(plan_i_1ch, cudaStreamPerThread));
  63     }
  64 }
  65
  66 void cuFFT::set_window(const MatDynMem &window)
  67 {
  68     Fft::set_window(window);
  69     m_window = window;
  70 }
  71
  72 void cuFFT::forward(const MatDynMem &real_input, ComplexMat &complex_result)
  73 {
  74     Fft::forward(real_input, complex_result);
  75     auto in = static_cast<cufftReal *>(const_cast<MatDynMem&>(real_input).deviceMem());
  76
  77     CufftErrorCheck(cufftExecR2C(plan_f, in, complex_result.get_p_data()));
  78 }
  79
  80 void cuFFT::forward_window(MatDynMem &feat, ComplexMat &complex_result, MatDynMem &temp)
  81 {
  82     Fft::forward_window(feat, complex_result, temp);
  83
  84     uint n_channels = feat.size[0];
  85     cufftReal *temp_data = temp.deviceMem();
  86
  87     for (uint i = 0; i < n_channels; ++i) {
  88         cv::Mat feat_plane(feat.dims - 1, feat.size + 1, feat.cv::Mat::type(), feat.ptr<void>(i));
  89         cv::Mat temp_plane(temp.dims - 1, temp.size + 1, temp.cv::Mat::type(), temp.ptr(i));
  90         temp_plane = feat_plane.mul(m_window);
  91     }
  92     CufftErrorCheck(cufftExecR2C(plan_fw, temp_data, complex_result.get_p_data()));
  93 }
  94
  95 void cuFFT::inverse(ComplexMat &complex_input, MatDynMem &real_result)
  96 {
  97     Fft::inverse(complex_input, real_result);
  98
  99     uint n_channels = complex_input.n_channels;
 100     cufftComplex *in = reinterpret_cast<cufftComplex *>(complex_input.get_p_data());
 101     cufftReal *out = real_result.deviceMem();
 102     float alpha = 1.0 / (m_width * m_height);
 103
 104     if (n_channels == 1) {
 105         CufftErrorCheck(cufftExecC2R(plan_i_1ch, in, out));
 106     } else {
 107         CufftErrorCheck(cufftExecC2R(plan_i_features, in, out));
 108     }
 109     // TODO: Investigate whether this scalling is needed or not
 110     CublasErrorCheck(cublasSscal(cublas, real_result.total(), &alpha, out, 1));
 111 }
 112
 113 cuFFT::~cuFFT()
 114 {
 115     CublasErrorCheck(cublasDestroy(cublas));
 116
 117     CufftErrorCheck(cufftDestroy(plan_f));
 118     CufftErrorCheck(cufftDestroy(plan_fw));
 119     CufftErrorCheck(cufftDestroy(plan_i_1ch));
 120     CufftErrorCheck(cufftDestroy(plan_i_features));
 121 }