src/fft_fftw.cpp

   1 #include "fft_fftw.h"
   2
   3 #include "fft.h"
   4
   5 #ifdef OPENMP
   6 #include <omp.h>
   7 #endif
   8
   9 #if !defined(ASYNC) && !defined(OPENMP) && !defined(CUFFTW)
  10 #define FFTW_PLAN_WITH_THREADS() fftw_plan_with_nthreads(4);
  11 #else
  12 #define FFTW_PLAN_WITH_THREADS()
  13 #endif
  14
  15 Fftw::Fftw(){}
  16
  17 void Fftw::init(unsigned width, unsigned height, unsigned num_of_feats, unsigned num_of_scales)
  18 {
  19     Fft::init(width, height, num_of_feats, num_of_scales);
  20
  21 #if (!defined(ASYNC) && !defined(CUFFTW)) && defined(OPENMP)
  22     fftw_init_threads();
  23 #endif // OPENMP
  24
  25 #ifndef CUFFTW
  26     std::cout << "FFT: FFTW" << std::endl;
  27 #else
  28     std::cout << "FFT: cuFFTW" << std::endl;
  29 #endif
  30     fftwf_cleanup();
  31     // FFT forward one scale
  32     {
  33         cv::Mat in_f = cv::Mat::zeros(int(m_height), int(m_width), CV_32FC1);
  34         ComplexMat out_f(int(m_height), m_width / 2 + 1, 1);
  35         plan_f = fftwf_plan_dft_r2c_2d(int(m_height), int(m_width), reinterpret_cast<float *>(in_f.data),
  36                                        reinterpret_cast<fftwf_complex *>(out_f.get_p_data()), FFTW_PATIENT);
  37     }
  38 #ifdef BIG_BATCH
  39     // FFT forward all scales
  40     if (m_num_of_scales > 1 && BIG_BATCH_MODE) {
  41         cv::Mat in_f_all = cv::Mat::zeros(m_height * m_num_of_scales, m_width, CV_32F);
  42         ComplexMat out_f_all(m_height, m_width / 2 + 1, m_num_of_scales);
  43         float *in = reinterpret_cast<float *>(in_f_all.data);
  44         fftwf_complex *out = reinterpret_cast<fftwf_complex *>(out_f_all.get_p_data());
  45         int rank = 2;
  46         int n[] = {(int)m_height, (int)m_width};
  47         int howmany = m_num_of_scales;
  48         int idist = m_height * m_width, odist = m_height * (m_width / 2 + 1);
  49         int istride = 1, ostride = 1;
  50         int *inembed = NULL, *onembed = NULL;
  51
  52         FFTW_PLAN_WITH_THREADS();
  53         plan_f_all_scales = fftwf_plan_many_dft_r2c(rank, n, howmany, in, inembed, istride, idist, out, onembed,
  54                                                     ostride, odist, FFTW_PATIENT);
  55     }
  56 #endif
  57     // FFT forward window one scale
  58     {
  59         cv::Mat in_fw = cv::Mat::zeros(int(m_height * m_num_of_feats), int(m_width), CV_32F);
  60         ComplexMat out_fw(int(m_height), m_width / 2 + 1, int(m_num_of_feats));
  61         float *in = reinterpret_cast<float *>(in_fw.data);
  62         fftwf_complex *out = reinterpret_cast<fftwf_complex *>(out_fw.get_p_data());
  63         int rank = 2;
  64         int n[] = {int(m_height), int(m_width)};
  65         int howmany = int(m_num_of_feats);
  66         int idist = int(m_height * m_width), odist = int(m_height * (m_width / 2 + 1));
  67         int istride = 1, ostride = 1;
  68         int *inembed = nullptr, *onembed = nullptr;
  69
  70         FFTW_PLAN_WITH_THREADS();
  71         plan_fw = fftwf_plan_many_dft_r2c(rank, n, howmany, in, inembed, istride, idist, out, onembed, ostride, odist,
  72                                           FFTW_PATIENT);
  73     }
  74 #ifdef BIG_BATCH
  75     // FFT forward window all scales all feats
  76     if (m_num_of_scales > 1 && BIG_BATCH_MODE) {
  77         cv::Mat in_all = cv::Mat::zeros(m_height * (m_num_of_scales * m_num_of_feats), m_width, CV_32F);
  78         ComplexMat out_all(m_height, m_width / 2 + 1, m_num_of_scales * m_num_of_feats);
  79         float *in = reinterpret_cast<float *>(in_all.data);
  80         fftwf_complex *out = reinterpret_cast<fftwf_complex *>(out_all.get_p_data());
  81         int rank = 2;
  82         int n[] = {(int)m_height, (int)m_width};
  83         int howmany = m_num_of_scales * m_num_of_feats;
  84         int idist = m_height * m_width, odist = m_height * (m_width / 2 + 1);
  85         int istride = 1, ostride = 1;
  86         int *inembed = NULL, *onembed = NULL;
  87
  88         FFTW_PLAN_WITH_THREADS();
  89         plan_fw_all_scales = fftwf_plan_many_dft_r2c(rank, n, howmany, in, inembed, istride, idist, out, onembed,
  90                                                      ostride, odist, FFTW_PATIENT);
  91     }
  92 #endif
  93     // FFT inverse one scale
  94     {
  95         ComplexMat in_i(m_height, m_width, m_num_of_feats);
  96         cv::Mat out_i = cv::Mat::zeros(int(m_height), int(m_width), CV_32FC(int(m_num_of_feats)));
  97         fftwf_complex *in = reinterpret_cast<fftwf_complex *>(in_i.get_p_data());
  98         float *out = reinterpret_cast<float *>(out_i.data);
  99         int rank = 2;
 100         int n[] = {int(m_height), int(m_width)};
 101         int howmany = int(m_num_of_feats);
 102         int idist = int(m_height * (m_width / 2 + 1)), odist = 1;
 103         int istride = 1, ostride = int(m_num_of_feats);
 104         int inembed[] = {int(m_height), int(m_width / 2 + 1)}, *onembed = n;
 105
 106         FFTW_PLAN_WITH_THREADS();
 107         plan_i_features = fftwf_plan_many_dft_c2r(rank, n, howmany, in, inembed, istride, idist, out, onembed, ostride,
 108                                                   odist, FFTW_PATIENT);
 109     }
 110     // FFT inverse all scales
 111 #ifdef BIG_BATCH
 112     if (m_num_of_scales > 1 && BIG_BATCH_MODE) {
 113         ComplexMat in_i_all(m_height, m_width, m_num_of_feats * m_num_of_scales);
 114         cv::Mat out_i_all = cv::Mat::zeros(m_height, m_width, CV_32FC(m_num_of_feats * m_num_of_scales));
 115         fftwf_complex *in = reinterpret_cast<fftwf_complex *>(in_i_all.get_p_data());
 116         float *out = reinterpret_cast<float *>(out_i_all.data);
 117         int rank = 2;
 118         int n[] = {(int)m_height, (int)m_width};
 119         int howmany = m_num_of_feats * m_num_of_scales;
 120         int idist = m_height * (m_width / 2 + 1), odist = 1;
 121         int istride = 1, ostride = m_num_of_feats * m_num_of_scales;
 122         int inembed[] = {(int)m_height, (int)m_width / 2 + 1}, *onembed = n;
 123
 124         FFTW_PLAN_WITH_THREADS();
 125         plan_i_features_all_scales = fftwf_plan_many_dft_c2r(rank, n, howmany, in, inembed, istride, idist, out,
 126                                                              onembed, ostride, odist, FFTW_PATIENT);
 127     }
 128 #endif
 129     // FFT inver one channel one scale
 130     {
 131         ComplexMat in_i1(int(m_height), int(m_width), 1);
 132         cv::Mat out_i1 = cv::Mat::zeros(int(m_height), int(m_width), CV_32FC1);
 133         fftwf_complex *in = reinterpret_cast<fftwf_complex *>(in_i1.get_p_data());
 134         float *out = reinterpret_cast<float *>(out_i1.data);
 135         int rank = 2;
 136         int n[] = {int(m_height), int(m_width)};
 137         int howmany = 1;
 138         int idist = int(m_height * (m_width / 2 + 1)), odist = 1;
 139         int istride = 1, ostride = 1;
 140         int inembed[] = {int(m_height), int(m_width) / 2 + 1}, *onembed = n;
 141
 142         FFTW_PLAN_WITH_THREADS();
 143         plan_i_1ch = fftwf_plan_many_dft_c2r(rank, n, howmany, in, inembed, istride, idist, out, onembed, ostride,
 144                                              odist, FFTW_PATIENT);
 145     }
 146 #ifdef BIG_BATCH
 147     // FFT inver one channel all scales
 148     if (m_num_of_scales > 1 && BIG_BATCH_MODE) {
 149         ComplexMat in_i1_all(m_height, m_width, m_num_of_scales);
 150         cv::Mat out_i1_all = cv::Mat::zeros(m_height, m_width, CV_32FC(m_num_of_scales));
 151         fftwf_complex *in = reinterpret_cast<fftwf_complex *>(in_i1_all.get_p_data());
 152         float *out = reinterpret_cast<float *>(out_i1_all.data);
 153         int rank = 2;
 154         int n[] = {(int)m_height, (int)m_width};
 155         int howmany = m_num_of_scales;
 156         int idist = m_height * (m_width / 2 + 1), odist = 1;
 157         int istride = 1, ostride = m_num_of_scales;
 158         int inembed[] = {(int)m_height, (int)m_width / 2 + 1}, *onembed = n;
 159
 160         FFTW_PLAN_WITH_THREADS();
 161         plan_i_1ch_all_scales = fftwf_plan_many_dft_c2r(rank, n, howmany, in, inembed, istride, idist, out, onembed,
 162                                                         ostride, odist, FFTW_PATIENT);
 163     }
 164 #endif
 165 }
 166
 167 void Fftw::set_window(const MatDynMem &window)
 168 {
 169     Fft::set_window(window);
 170     m_window = window;
 171 }
 172
 173 void Fftw::forward(MatDynMem &&real_input, ComplexMat & complex_result)
 174 {
 175     Fft::forward(real_input, complex_result);
 176
 177     if (BIG_BATCH_MODE && real_input.rows == int(m_height * m_num_of_scales)) {
 178         fftwf_execute_dft_r2c(plan_f_all_scales, reinterpret_cast<float *>(real_input.data),
 179                               reinterpret_cast<fftwf_complex *>(complex_result.get_p_data()));
 180     } else {
 181         fftwf_execute_dft_r2c(plan_f, reinterpret_cast<float *>(real_input.data),
 182                               reinterpret_cast<fftwf_complex *>(complex_result.get_p_data()));
 183     }
 184     return;
 185 }
 186
 187 void Fftw::forward_window(MatDynMem &feat, ComplexMat & complex_result, MatDynMem &temp)
 188 {
 189     Fft::forward_window(feat, complex_result, temp);
 190
 191     int n_channels = feat.size[0];
 192     for (int i = 0; i < n_channels; ++i) {
 193         cv::Mat feat_plane(feat.dims - 1, feat.size + 1, feat.cv::Mat::type(), feat.ptr<void>(i));
 194         cv::Mat temp_plane(temp.dims - 1, temp.size + 1, temp.cv::Mat::type(), temp.ptr(i));
 195         temp_plane = feat_plane.mul(m_window);
 196     }
 197
 198     float *in = temp.ptr<float>();
 199     fftwf_complex *out = reinterpret_cast<fftwf_complex *>(complex_result.get_p_data());
 200
 201     if (n_channels <= int(m_num_of_feats))
 202         fftwf_execute_dft_r2c(plan_fw, in, out);
 203     else
 204         fftwf_execute_dft_r2c(plan_fw_all_scales, in, out);
 205     return;
 206 }
 207
 208 void Fftw::inverse(ComplexMat &  complex_input, MatDynMem & real_result)
 209 {
 210     Fft::inverse(complex_input, real_result);
 211
 212     int n_channels = complex_input.n_channels;
 213     fftwf_complex *in = reinterpret_cast<fftwf_complex *>(complex_input.get_p_data());
 214     float *out = real_result.ptr<float>();
 215
 216     if (n_channels == 1)
 217         fftwf_execute_dft_c2r(plan_i_1ch, in, out);
 218     else if (BIG_BATCH_MODE && n_channels == int(m_num_of_scales))
 219         fftwf_execute_dft_c2r(plan_i_1ch_all_scales, in, out);
 220     else if (BIG_BATCH_MODE && n_channels == int(m_num_of_feats) * int(m_num_of_scales))
 221         fftwf_execute_dft_c2r(plan_i_features_all_scales, in, out);
 222     else
 223         fftwf_execute_dft_c2r(plan_i_features, in, out);
 224
 225     real_result = real_result / (m_width * m_height);
 226 }
 227
 228 Fftw::~Fftw()
 229 {
 230     fftwf_destroy_plan(plan_f);
 231     fftwf_destroy_plan(plan_fw);
 232     fftwf_destroy_plan(plan_i_features);
 233     fftwf_destroy_plan(plan_i_1ch);
 234
 235     if (BIG_BATCH_MODE) {
 236         fftwf_destroy_plan(plan_f_all_scales);
 237         fftwf_destroy_plan(plan_i_features_all_scales);
 238         fftwf_destroy_plan(plan_fw_all_scales);
 239         fftwf_destroy_plan(plan_i_1ch_all_scales);
 240     }
 241 }