GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 88.2% 149 / 0 / 169
Functions: 93.8% 15 / 0 / 16
Branches: 72.6% 77 / 0 / 106
include/brotensor/detail/cpu/fft_core.h
Line Branch Exec Source
1 #pragma once
2
3 // ─── CPU FFT core — shared internals ───────────────────────────────────────
4 //
5 // Hand-rolled mixed-radix + Bluestein DFT, factored out of src/cpu/fft.cpp so
6 // the STFT / iSTFT ops (and any later spectral op) can reuse the exact same
7 // transform instead of duplicating it.
8 //
9 // This is a CPU-backend-private header (brotensor::detail::cpu). It is
10 // header-only and double-precision internally for accuracy, matching the
11 // fft.cpp original. No external libraries.
12 //
13 // * Cd — tiny double-precision complex helper.
14 // * dft_1d — unscaled length-N DFT, `sign` = -1 forward / +1
15 // inverse. The "backward" 1/N inverse scaling is
16 // applied by callers that want an inverse.
17 // * load/store_complex_row — interleaved-complex (R, 2*C) tensor row I/O.
18 //
19 // The mixed-radix engine handles sizes whose prime factorisation uses only
20 // the radices 2/3/5/7 (covers Whisper's n_fft = 400 = 2^4 * 5^2); anything
21 // with a large or prime factor falls back to a Bluestein chirp-z transform,
22 // so dft_1d is correct for every length >= 1.
23 //
24 // The twiddle table (mixed-radix path) and chirp table (Bluestein path) are
25 // each a pure function of (N, sign), so they're cached per-thread across
26 // calls (see cached_twiddles / cached_chirp below) — a multi-frame STFT that
27 // calls dft_1d once per frame at a fixed n_fft reuses the same tables
28 // instead of rebuilding them every frame.
29
30 #include <cmath>
31 #include <cstddef>
32 #include <cstdint>
33 #include <unordered_map>
34 #include <vector>
35
36 namespace brotensor::detail::cpu::fftcore {
37
38 constexpr double kPi = 3.14159265358979323846;
39
40 // ─── tiny complex helper (double precision internally for accuracy) ────────
41 5590050 struct Cd {
42 5590050 double re = 0.0;
43 5590050 double im = 0.0;
44 };
45
46 22794620 inline Cd operator+(Cd a, Cd b) { return {a.re + b.re, a.im + b.im}; }
47 inline Cd operator-(Cd a, Cd b) { return {a.re - b.re, a.im - b.im}; }
48 23055344 inline Cd operator*(Cd a, Cd b) {
49 23055344 return {a.re * b.re - a.im * b.im, a.re * b.im + a.im * b.re};
50 }
51
52 // ─── twiddle-factor table: w[k] = exp(sign * 2*pi*i * k / N) ───────────────
53 108 inline std::vector<Cd> make_twiddles(int N, int sign) {
54 108 std::vector<Cd> w(static_cast<std::size_t>(N));
55 216 const double s = static_cast<double>(sign) * 2.0 * kPi
56 108 / static_cast<double>(N);
57
2/2
✓ Branch 0 taken 108 times.
✓ Branch 1 taken 10116 times.
 10224 for (int k = 0; k < N; ++k) {
58 10116 w[static_cast<std::size_t>(k)] = {std::cos(s * k), std::sin(s * k)};
59 10116 }
60 108 return w;
61 108 }
62
63 // ─── (N, sign) -> precomputed-table cache ──────────────────────────────────
64 //
65 // The mixed-radix recursion and its Bluestein fallback each build a
66 // trigonometric table (the twiddle table / the chirp table) that is a pure
67 // function of (N, sign) — once computed its values never change. A single
68 // STFT/iSTFT call runs `dft_1d` once per frame/row with the *same* n_fft, so
69 // the recursion revisits the exact same set of (N, sign) sub-problems (N's
70 // factorisation tree is deterministic) on every frame. Without caching that
71 // means thousands of redundant heap allocations and sin/cos evaluations per
72 // call. Cache the tables keyed on (N, sign) instead of rebuilding them.
73 //
74 // The cache is `thread_local` rather than behind a mutex: this codebase's
75 // convention (see CLAUDE.md — "no mutexes") is single-owner/thread-local
76 // state, and the CPU backend already uses exactly this
77 // `thread_local static` reusable-scratch pattern elsewhere (mirrored from
78 // src/cuda/resblock.cu and src/metal/resblock.mm, which cache per-thread
79 // scratch tensors across calls). A thread_local table means each thread that
80 // calls into the CPU FFT backend gets its own private cache — no shared
81 // mutable state, so nothing to race on, without needing any locking
82 // machinery that single-threaded CPU-backend use doesn't need.
83 //
84 // No eviction policy: the cache grows to the number of distinct (N, sign)
85 // pairs a thread has ever computed, which in practice is bounded by the
86 // handful of n_fft / signal-length configurations a process actually uses
87 // (plus the O(log N) sub-sizes each one's recursion visits) — not by the
88 // number of frames or calls, so it saturates quickly and stays small.
89 2121533 inline std::uint64_t table_cache_key(int N, int sign) {
90 4243066 return (static_cast<std::uint64_t>(static_cast<std::uint32_t>(N)) << 1)
91 2121533 | (sign > 0 ? 1u : 0u);
92 }
93
94 // Cached twiddle table for fft_recursive's per-level combine step.
95 2120443 inline const std::vector<Cd>& cached_twiddles(int N, int sign) {
96
2/2
✓ Branch 0 taken 2120439 times.
✓ Branch 1 taken 4 times.
 2120443 thread_local std::unordered_map<std::uint64_t, std::vector<Cd>> cache;
97 2120443 const std::uint64_t key = table_cache_key(N, sign);
98 2120443 auto it = cache.find(key);
99
2/2
✓ Branch 0 taken 2120335 times.
✓ Branch 1 taken 108 times.
 2120443 if (it != cache.end()) return it->second;
100
2/4
✓ Branch 0 taken 108 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 108 times.
✗ Branch 3 not taken.
 108 return cache.emplace(key, make_twiddles(N, sign)).first->second;
101 2120443 }
102
103 // Cached chirp table for Bluestein's chirp-z transform (exp(sign * i*pi*n^2/N)).
104 1090 inline const std::vector<Cd>& cached_chirp(int N, int sign) {
105
2/2
✓ Branch 0 taken 1088 times.
✓ Branch 1 taken 2 times.
 1090 thread_local std::unordered_map<std::uint64_t, std::vector<Cd>> cache;
106 1090 const std::uint64_t key = table_cache_key(N, sign);
107 1090 auto it = cache.find(key);
108
2/2
✓ Branch 0 taken 1084 times.
✓ Branch 1 taken 6 times.
 1090 if (it != cache.end()) return it->second;
109 6 std::vector<Cd> chirp(static_cast<std::size_t>(N));
110
2/2
✓ Branch 0 taken 934 times.
✓ Branch 1 taken 6 times.
 940 for (int n = 0; n < N; ++n) {
111 934 const long long n2 = (static_cast<long long>(n) * n) % (2LL * N);
112 1868 const double ang = static_cast<double>(sign) * kPi
113 934 * static_cast<double>(n2) / static_cast<double>(N);
114 934 chirp[static_cast<std::size_t>(n)] = {std::cos(ang), std::sin(ang)};
115 934 }
116
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
 6 return cache.emplace(key, std::move(chirp)).first->second;
117 1090 }
118
119 // ─── naive O(N^2) DFT — base case + Bluestein small kernels ────────────────
120 ✗ inline void dft_naive(const std::vector<Cd>& in, std::vector<Cd>& out,
121 int sign) {
122 ✗ const int N = static_cast<int>(in.size());
123 ✗ out.assign(static_cast<std::size_t>(N), Cd{});
124 ✗ const double s = static_cast<double>(sign) * 2.0 * kPi
125 ✗ / static_cast<double>(N);
126 ✗ for (int k = 0; k < N; ++k) {
127 ✗ Cd acc{};
128 ✗ for (int n = 0; n < N; ++n) {
129 ✗ const double a = s * (static_cast<long long>(k) * n % N);
130 ✗ const Cd tw{std::cos(a), std::sin(a)};
131 ✗ acc = acc + in[static_cast<std::size_t>(n)] * tw;
132 ✗ }
133 ✗ out[static_cast<std::size_t>(k)] = acc;
134 ✗ }
135 ✗ }
136
137 // Smallest supported radix (2,3,5,7) dividing N, or 0 if none does.
138 2120443 inline int small_radix(int N) {
139
1/2
✓ Branch 0 taken 2207157 times.
✗ Branch 1 not taken.
 2207157 for (int r : {2, 3, 5, 7}) {
140
2/2
✓ Branch 0 taken 2120443 times.
✓ Branch 1 taken 86714 times.
 2207157 if (N % r == 0) return r;
141 }
142 ✗ return 0;
143 2120443 }
144
145 // True iff N factors entirely into the supported small radices.
146 1972443 inline bool is_smooth(int N) {
147 1972443 int m = N;
148
2/2
✓ Branch 0 taken 7889772 times.
✓ Branch 1 taken 1972443 times.
 9862215 for (int r : {2, 3, 5, 7}) {
149
2/2
✓ Branch 0 taken 1108480 times.
✓ Branch 1 taken 7889772 times.
 8998252 while (m % r == 0) m /= r;
150 }
151 1972443 return m == 1;
152 }
153
154
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 138979 times.
 138979 inline bool is_power_of_two(int N) { return N > 0 && (N & (N - 1)) == 0; }
155
156 // Forward declaration: Bluestein needs the recursive (power-of-two) FFT.
157 inline void fft_recursive(const std::vector<Cd>& in, std::vector<Cd>& out,
158 int sign);
159
160 // ─── Bluestein chirp-z transform — arbitrary-length fallback ───────────────
161 1090 inline void bluestein(const std::vector<Cd>& in, std::vector<Cd>& out,
162 int sign) {
163 1090 const int N = static_cast<int>(in.size());
164
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1090 times.
 1090 if (N <= 1) { out = in; return; }
165
166 1090 int M = 1;
167
2/2
✓ Branch 0 taken 7634 times.
✓ Branch 1 taken 1090 times.
 8724 while (M < 2 * N - 1) M <<= 1;
168
169 1090 const std::vector<Cd>& chirp = cached_chirp(N, sign);
170
171 1090 std::vector<Cd> a(static_cast<std::size_t>(M), Cd{});
172
2/2
✓ Branch 0 taken 59002 times.
✓ Branch 1 taken 1090 times.
 60092 for (int n = 0; n < N; ++n) {
173 59002 a[static_cast<std::size_t>(n)] =
174 59002 in[static_cast<std::size_t>(n)] * chirp[static_cast<std::size_t>(n)];
175 59002 }
176
177
1/2
✓ Branch 0 taken 1090 times.
✗ Branch 1 not taken.
 1090 std::vector<Cd> b(static_cast<std::size_t>(M), Cd{});
178 1090 b[0] = {chirp[0].re, -chirp[0].im};
179
2/2
✓ Branch 0 taken 57912 times.
✓ Branch 1 taken 1090 times.
 59002 for (int k = 1; k < N; ++k) {
180 115824 const Cd c{chirp[static_cast<std::size_t>(k)].re,
181 57912 -chirp[static_cast<std::size_t>(k)].im};
182 57912 b[static_cast<std::size_t>(k)] = c;
183 57912 b[static_cast<std::size_t>(M - k)] = c;
184 57912 }
185
186 1090 std::vector<Cd> fa, fb;
187
1/2
✓ Branch 0 taken 1090 times.
✗ Branch 1 not taken.
 1090 fft_recursive(a, fa, -1);
188
1/2
✓ Branch 0 taken 1090 times.
✗ Branch 1 not taken.
 1090 fft_recursive(b, fb, -1);
189
1/2
✓ Branch 0 taken 1090 times.
✗ Branch 1 not taken.
 1090 std::vector<Cd> prod(static_cast<std::size_t>(M));
190
2/2
✓ Branch 0 taken 142720 times.
✓ Branch 1 taken 1090 times.
 143810 for (int i = 0; i < M; ++i) {
191 142720 prod[static_cast<std::size_t>(i)] =
192 142720 fa[static_cast<std::size_t>(i)] * fb[static_cast<std::size_t>(i)];
193 142720 }
194 1090 std::vector<Cd> conv;
195
1/2
✓ Branch 0 taken 1090 times.
✗ Branch 1 not taken.
 1090 fft_recursive(prod, conv, +1);
196 1090 const double invM = 1.0 / static_cast<double>(M);
197
198
1/2
✓ Branch 0 taken 1090 times.
✗ Branch 1 not taken.
 1090 out.assign(static_cast<std::size_t>(N), Cd{});
199
2/2
✓ Branch 0 taken 59002 times.
✓ Branch 1 taken 1090 times.
 60092 for (int k = 0; k < N; ++k) {
200 59002 Cd c = conv[static_cast<std::size_t>(k)];
201 59002 c.re *= invM;
202 59002 c.im *= invM;
203 59002 out[static_cast<std::size_t>(k)] =
204 59002 chirp[static_cast<std::size_t>(k)] * c;
205 59002 }
206 1090 }
207
208 // ─── mixed-radix recursive FFT ─────────────────────────────────────────────
209 4512009 inline void fft_recursive(const std::vector<Cd>& in, std::vector<Cd>& out,
210 int sign) {
211 4512009 const int N = static_cast<int>(in.size());
212
2/2
✓ Branch 0 taken 2391566 times.
✓ Branch 1 taken 2120443 times.
 4512009 if (N <= 1) { out = in; return; }
213
214 2120443 const int r = small_radix(N);
215
1/2
✓ Branch 0 taken 2120443 times.
✗ Branch 1 not taken.
 2120443 if (r == 0) {
216 ✗ bluestein(in, out, sign);
217 ✗ return;
218 }
219
220 2120443 const int m = N / r;
221
222
3/4
✓ Branch 0 taken 1970570 times.
✓ Branch 1 taken 149873 times.
✓ Branch 2 taken 1970570 times.
✗ Branch 3 not taken.
 2120443 if (N <= 8 && !is_smooth(m)) {
223 ✗ dft_naive(in, out, sign);
224 ✗ return;
225 }
226
227
1/2
✓ Branch 0 taken 2120443 times.
✗ Branch 1 not taken.
 2120443 std::vector<std::vector<Cd>> subs(
228 2120443 static_cast<std::size_t>(r),
229 2120443 std::vector<Cd>(static_cast<std::size_t>(m)));
230
2/2
✓ Branch 0 taken 5320754 times.
✓ Branch 1 taken 2120443 times.
 7441197 for (int t = 0; t < m; ++t) {
231
2/2
✓ Branch 0 taken 10857340 times.
✓ Branch 1 taken 5320754 times.
 16178094 for (int j = 0; j < r; ++j) {
232 10857340 subs[static_cast<std::size_t>(j)][static_cast<std::size_t>(t)] =
233 10857340 in[static_cast<std::size_t>(t * r + j)];
234 10857340 }
235 5320754 }
236
237
1/2
✓ Branch 0 taken 2120443 times.
✗ Branch 1 not taken.
 2120443 std::vector<std::vector<Cd>> subF(static_cast<std::size_t>(r));
238
2/2
✓ Branch 0 taken 2120443 times.
✓ Branch 1 taken 4370850 times.
 6491293 for (int j = 0; j < r; ++j) {
239
2/4
✓ Branch 0 taken 4370850 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4370850 times.
✗ Branch 3 not taken.
 8741700 fft_recursive(subs[static_cast<std::size_t>(j)],
240 4370850 subF[static_cast<std::size_t>(j)], sign);
241 4370850 }
242
243
1/2
✓ Branch 0 taken 2120443 times.
✗ Branch 1 not taken.
 2120443 const std::vector<Cd>& w = cached_twiddles(N, sign);
244
1/2
✓ Branch 0 taken 2120443 times.
✗ Branch 1 not taken.
 2120443 out.assign(static_cast<std::size_t>(N), Cd{});
245
2/2
✓ Branch 0 taken 2120443 times.
✓ Branch 1 taken 10857340 times.
 12977783 for (int k = 0; k < N; ++k) {
246 10857340 const int km = k % m;
247 10857340 Cd acc{};
248
2/2
✓ Branch 0 taken 10857340 times.
✓ Branch 1 taken 22794620 times.
 33651960 for (int j = 0; j < r; ++j) {
249 22794620 const int idx = (j * k) % N;
250
3/6
✓ Branch 0 taken 22794620 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 22794620 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 22794620 times.
✗ Branch 5 not taken.
 45589240 acc = acc + w[static_cast<std::size_t>(idx)] *
251 45589240 subF[static_cast<std::size_t>(j)]
252 22794620 [static_cast<std::size_t>(km)];
253 22794620 }
254 10857340 out[static_cast<std::size_t>(k)] = acc;
255 10857340 }
256 4512009 }
257
258 // ─── public single-signal transform (unscaled) ─────────────────────────────
259 //
260 // dft_1d runs the unscaled transform with the requested sign. The 1/N inverse
261 // scaling for the "backward" convention is applied by inverse callers.
262 138981 inline void dft_1d(const std::vector<Cd>& in, std::vector<Cd>& out, int sign) {
263 138981 const int N = static_cast<int>(in.size());
264
1/2
✓ Branch 0 taken 138981 times.
✗ Branch 1 not taken.
 138981 if (N == 0) { out.clear(); return; }
265
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 138979 times.
 138981 if (N == 1) { out = in; return; }
266
4/4
✓ Branch 0 taken 1873 times.
✓ Branch 1 taken 137106 times.
✓ Branch 2 taken 783 times.
✓ Branch 3 taken 1090 times.
 138979 if (is_power_of_two(N) || is_smooth(N)) {
267 137889 fft_recursive(in, out, sign);
268 137889 } else {
269 1090 bluestein(in, out, sign);
270 }
271 138981 }
272
273 // ─── interleaved-complex tensor row helpers ────────────────────────────────
274 364 inline void load_complex_row(const float* base, int row, int cols2,
275 std::vector<Cd>& dst) {
276 364 const int C = cols2 / 2;
277 364 dst.assign(static_cast<std::size_t>(C), Cd{});
278 364 const float* p = base + static_cast<std::size_t>(row) * cols2;
279
2/2
✓ Branch 0 taken 14104 times.
✓ Branch 1 taken 364 times.
 14468 for (int c = 0; c < C; ++c) {
280 28208 dst[static_cast<std::size_t>(c)] = {static_cast<double>(p[2 * c]),
281 14104 static_cast<double>(p[2 * c + 1])};
282 14104 }
283 364 }
284
285 364 inline void store_complex_row(float* base, int row, int cols2,
286 const std::vector<Cd>& src) {
287 364 const int C = cols2 / 2;
288 364 float* p = base + static_cast<std::size_t>(row) * cols2;
289
2/2
✓ Branch 0 taken 14104 times.
✓ Branch 1 taken 364 times.
 14468 for (int c = 0; c < C; ++c) {
290 14104 p[2 * c] = static_cast<float>(src[static_cast<std::size_t>(c)].re);
291 14104 p[2 * c + 1] = static_cast<float>(src[static_cast<std::size_t>(c)].im);
292 14104 }
293 364 }
294
295 } // namespace brotensor::detail::cpu::fftcore
296