diff --git a/src/fft.cpp b/src/fft.cpp index 4571a60..843708b 100644 --- a/src/fft.cpp +++ b/src/fft.cpp @@ -7,75 +7,120 @@ #ifndef M_PI #define M_PI 3.14159265358979323846 #endif -#include +#include #include namespace pk { -// Iterative radix-2 Cooley-Tukey FFT (in-place, complex double buffer). -// n must be a power of 2. -static void fft_inplace(std::vector& re, std::vector& im) { - const int n = static_cast(re.size()); - assert(n > 0 && (n & (n - 1)) == 0); // must be power of 2 - - // Bit-reversal permutation - for (int i = 1, j = 0; i < n; ++i) { - int bit = n >> 1; - for (; j & bit; bit >>= 1) j ^= bit; - j ^= bit; - if (i < j) { - std::swap(re[i], re[j]); - std::swap(im[i], im[j]); +// Precomputed radix-2 Cooley-Tukey plan for a FIXED length n (a power of two). +// The twiddle sequences are generated with the exact iterative complex-multiply +// used by the original per-frame FFT, so the transform is bit-identical to it, +// but the cos/sin-per-stage and per-butterfly twiddle advance are paid once at +// plan construction instead of on every frame. +class FftPlan { +public: + explicit FftPlan(int n) : n_(n) { + assert(n > 0 && (n & (n - 1)) == 0 && "n must be a power of 2"); + // Bit-reversal permutation. + rev_.resize((size_t)n); + for (int i = 1, j = 0; i < n; ++i) { + int bit = n >> 1; + for (; j & bit; bit >>= 1) j ^= bit; + j ^= bit; + rev_[i] = j; } - } - - // Butterfly stages - for (int len = 2; len <= n; len <<= 1) { - double ang = -2.0 * M_PI / len; - double wr = std::cos(ang); - double wi = std::sin(ang); - for (int i = 0; i < n; i += len) { + // Per-stage twiddle sequences, generated bit-identically to the original + // loop (start at (1,0) and advance by complex-multiplying cos/sin(-2π/len)). + stage_off_.push_back(0); + for (int len = 2; len <= n; len <<= 1) { + double ang = -2.0 * M_PI / len; + double wr = std::cos(ang); + double wi = std::sin(ang); double cur_wr = 1.0, cur_wi = 0.0; for (int k = 0; k < len / 2; ++k) { - int u = i + k; - int v = i + k + len / 2; - double tr = cur_wr * re[v] - cur_wi * im[v]; - double ti = cur_wr * im[v] + cur_wi * re[v]; - re[v] = re[u] - tr; - im[v] = im[u] - ti; - re[u] = re[u] + tr; - im[u] = im[u] + ti; - // advance twiddle factor + tw_r_.push_back(cur_wr); + tw_i_.push_back(cur_wi); double new_wr = cur_wr * wr - cur_wi * wi; double new_wi = cur_wr * wi + cur_wi * wr; cur_wr = new_wr; cur_wi = new_wi; } + stage_off_.push_back((int)tw_r_.size()); } } -} + + // Transform `in` into the full n complex bins, keeping the entire transform + // in double (bit-identical to the original per-frame FFT); copies bins + // 0..n_bins-1 into the float outputs. + void apply(const float* in, std::vector& re, std::vector& im) const { + const int n = n_; + d_re_.resize((size_t)n); + d_im_.resize((size_t)n); + for (int i = 0; i < n; ++i) { + d_re_[(size_t)rev_[i]] = (double)in[i]; + d_im_[(size_t)rev_[i]] = 0.0; + } + const double* tw_r = tw_r_.data(); + const double* tw_i = tw_i_.data(); + double* re0 = d_re_.data(); + double* im0 = d_im_.data(); + int st = 0; + for (int len = 2; len <= n; len <<= 1, ++st) { + const int half = len >> 1; + const double* wbase_r = tw_r + stage_off_[st]; + const double* wbase_i = tw_i + stage_off_[st]; + for (int i = 0; i < n; i += len) { + const double* wr = wbase_r; + const double* wi = wbase_i; + for (int k = 0; k < half; ++k) { + const int u = i + k; + const int v = u + half; + const double tr = wr[k] * re0[v] - wi[k] * im0[v]; + const double ti = wr[k] * im0[v] + wi[k] * re0[v]; + re0[v] = re0[u] - tr; + im0[v] = im0[u] - ti; + re0[u] = re0[u] + tr; + im0[u] = im0[u] + ti; + } + } + } + const int n_bins = n / 2 + 1; + re.resize((size_t)n_bins); + im.resize((size_t)n_bins); + for (int b = 0; b < n_bins; ++b) { + re[(size_t)b] = (float)d_re_[b]; + im[(size_t)b] = (float)d_im_[b]; + } + } + + int n() const { return n_; } + +private: + int n_; + std::vector rev_; + std::vector tw_r_, tw_i_; + std::vector stage_off_; + mutable std::vector d_re_, d_im_; // reusable scratch (no per-call alloc) +}; void rfft(const std::vector& in, std::vector& re, std::vector& im) { const int n = static_cast(in.size()); assert(n > 0 && (n & (n - 1)) == 0); - // Build complex buffer from real input (imag = 0) - std::vector buf_re(n), buf_im(n); - for (int i = 0; i < n; ++i) { - buf_re[i] = static_cast(in[i]); - buf_im[i] = 0.0; - } - - fft_inplace(buf_re, buf_im); - - // Copy bins 0..n/2 into output (n/2 + 1 bins) - const int n_bins = n / 2 + 1; - re.resize(n_bins); - im.resize(n_bins); - for (int k = 0; k < n_bins; ++k) { - re[k] = static_cast(buf_re[k]); - im[k] = static_cast(buf_im[k]); + // Cache the plan for the LAST-seen n, rebuilding when n changes. rfft is a + // public function and may be called with different power-of-two lengths on + // the same thread (e.g. different models with different n_fft), so the plan + // MUST be keyed by n — caching only the first length would run a + // differently-sized input against the wrong transform. thread_local so + // concurrent rfft callers (e.g. a multi-request server) never race on the + // reusable double scratch. + static thread_local int cached_n = 0; + static thread_local std::unique_ptr cached; + if (cached_n != n) { + cached = std::make_unique(n); + cached_n = n; } + cached->apply(in.data(), re, im); } } // namespace pk diff --git a/tests/test_fft.cpp b/tests/test_fft.cpp index 3c52d2c..25ddf7d 100644 --- a/tests/test_fft.cpp +++ b/tests/test_fft.cpp @@ -2,17 +2,36 @@ #include #include #include -int main() { - const int N=512; + +// Run rfft on a pure cosine at bin k of length N and check the magnitude peak +// lands at bin k. Returns 0 on success, else the (wrong) peak bin. +static int check_fft(int N, int k) { std::vector x(N); - // pure cosine at bin k=8: magnitude should peak at bin 8 - for (int i=0;i re(N/2+1), im(N/2+1); + for (int i = 0; i < N; ++i) x[i] = std::cos(2.0 * M_PI * k * i / N); + std::vector re(N / 2 + 1), im(N / 2 + 1); pk::rfft(x, re, im); - // find peak bin - int peak=0; double best=-1; - for (int k=0;k<=N/2;++k){ double m=re[k]*re[k]+im[k]*im[k]; if(m>best){best=m;peak=k;} } - if (peak!=8){ std::fprintf(stderr,"peak at %d, expected 8\n",peak); return 1; } - std::printf("fft ok: peak bin=%d\n", peak); + int peak = 0; + double best = -1; + for (int b = 0; b <= N / 2; ++b) { + double m = (double)re[b] * re[b] + (double)im[b] * im[b]; + if (m > best) { best = m; peak = b; } + } + return peak == k ? 0 : peak; +} + +int main() { + // rfft caches a per-length FftPlan thread-locally. Call it with several + // different power-of-two lengths on the SAME thread, in both orders, so a + // plan keyed only to the first-seen length would produce a wrong peak (or + // read past a shorter input) and fail here. Regression for the rfft plan + // cache being keyed by n. + if (int p = check_fft(512, 8)) { std::fprintf(stderr, "512@8 peak=%d\n", p); return 1; } + if (int p = check_fft(256, 4)) { std::fprintf(stderr, "256@4 peak=%d\n", p); return 1; } + if (int p = check_fft(512, 16)) { std::fprintf(stderr, "512@16 peak=%d\n", p); return 1; } + if (int p = check_fft(128, 3)) { std::fprintf(stderr, "128@3 peak=%d\n", p); return 1; } + if (int p = check_fft(1024, 20)){ std::fprintf(stderr, "1024@20 peak=%d\n", p); return 1; } + if (int p = check_fft(256, 7)) { std::fprintf(stderr, "256@7 peak=%d\n", p); return 1; } + + std::printf("fft ok: multi-length, both orders\n"); return 0; }