diff --git a/include/armadillo_bits/op_fft_meat.hpp b/include/armadillo_bits/op_fft_meat.hpp index 74fef1d0..4f5d93af 100644 --- a/include/armadillo_bits/op_fft_meat.hpp +++ b/include/armadillo_bits/op_fft_meat.hpp @@ -43,12 +43,14 @@ class fft_engine_wrapper } inline - fft_engine_wrapper(const uword N) + fft_engine_wrapper(const uword N_samples, const uword N_exec) { arma_extra_debug_sigprint(); - worker_kissfft = (N < threshold) ? new fft_engine_kissfft(N) : nullptr; - worker_fftw3 = (N >= threshold) ? new fft_engine_fftw3 (N) : nullptr; + const bool use_fftw3 = N_samples >= (threshold / N_exec); + + worker_kissfft = (use_fftw3 == false) ? new fft_engine_kissfft(N_samples) : nullptr; + worker_fftw3 = (use_fftw3 == true ) ? new fft_engine_fftw3 (N_samples) : nullptr; } inline @@ -94,7 +96,8 @@ op_fft_real::apply( Mat< std::complex >& out, const mtOp< const uword N_user = (in.aux_uword_b == 0) ? in.aux_uword_a : N_orig; #if defined(ARMA_USE_FFTW3) - fft_engine_wrapper worker(N_user); + const uword N_exec = (is_vec) ? uword(1) : n_cols; + fft_engine_wrapper worker(N_user, N_exec); #else fft_engine_kissfft worker(N_user); #endif @@ -198,7 +201,8 @@ op_fft_cx::apply_noalias(Mat& out, const Mat& X, const uword a, const uw const uword N_user = (b == 0) ? a : N_orig; #if defined(ARMA_USE_FFTW3) - fft_engine_wrapper worker(N_user); + const uword N_exec = (is_vec) ? uword(1) : n_cols; + fft_engine_wrapper worker(N_user, N_exec); #else fft_engine_kissfft worker(N_user); #endif