mel: precompute FFT twiddle factors and reuse scratch

src/fft.cpp

@@ -7,75 +7,120 @@
 #ifndef M_PI
 #define M_PI 3.14159265358979323846
 #endif
-#include <cstdint>
+#include <memory>
 #include <vector>
 
 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<double>& re, std::vector<double>& im) {
-    const int n = static_cast<int>(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<float>& re, std::vector<float>& 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<int> rev_;
+    std::vector<double> tw_r_, tw_i_;
+    std::vector<int> stage_off_;
+    mutable std::vector<double> d_re_, d_im_;  // reusable scratch (no per-call alloc)
+};
 
 void rfft(const std::vector<float>& in, std::vector<float>& re, std::vector<float>& im) {
     const int n = static_cast<int>(in.size());
     assert(n > 0 && (n & (n - 1)) == 0);
 
-    // Build complex buffer from real input (imag = 0)
-    std::vector<double> buf_re(n), buf_im(n);
-    for (int i = 0; i < n; ++i) {
-        buf_re[i] = static_cast<double>(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<float>(buf_re[k]);
-        im[k] = static_cast<float>(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<FftPlan> cached;
+    if (cached_n != n) {
+        cached = std::make_unique<FftPlan>(n);
+        cached_n = n;
     }
+    cached->apply(in.data(), re, im);
 }
 
 } // namespace pk

tests/test_fft.cpp

@@ -2,17 +2,36 @@
 #include <vector>
 #include <cmath>
 #include <cstdio>
-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<float> x(N);
-    // pure cosine at bin k=8: magnitude should peak at bin 8
-    for (int i=0;i<N;++i) x[i]=std::cos(2.0*M_PI*8*i/N);
-    std::vector<float> 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<float> 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;
 }