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