src/cpu/stft.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // ─── CPU STFT / iSTFT (brosoundml CHUNK 2) ───────────────────────────────── | ||
| 2 | // | ||
| 3 | // Short-time Fourier transform and its inverse, plus their adjoints. CPU | ||
| 4 | // backend, FP32-only. No external libraries — the per-frame DFT reuses the | ||
| 5 | // hand-rolled mixed-radix + Bluestein engine from detail/cpu/fft_core.h | ||
| 6 | // (so n_fft = 400 and prime sizes both work, exactly as in fft.cpp). | ||
| 7 | // | ||
| 8 | // Ops implemented here: | ||
| 9 | // stft / stft_backward real signal <-> complex spectrogram | ||
| 10 | // istft / istft_backward complex spectrogram <-> real signal (COLA OLA) | ||
| 11 | // | ||
| 12 | // ── Layout (see the doc comments in ops.h for the full contract) ──────────── | ||
| 13 | // signal: REAL (N, signal_len) — N batched signals, one / row. | ||
| 14 | // window: REAL (1, win_length) — caller-supplied. | ||
| 15 | // spec: interleaved-complex (N*frames, 2*bins), bins = n_fft/2+1. Each | ||
| 16 | // frame is a row; the N signals' frame blocks are stacked in order. | ||
| 17 | // | ||
| 18 | // ── Frame model ───────────────────────────────────────────────────────────── | ||
| 19 | // Frame f of signal b takes n_fft samples starting at padded position | ||
| 20 | // f*hop_length, multiplies the central win_length of them by `window`, and | ||
| 21 | // rfft's the n_fft buffer. The window sits centred in the n_fft buffer | ||
| 22 | // (pad = (n_fft-win_length)/2 zeros each side). When center == true the | ||
| 23 | // signal is reflect-padded by n_fft/2 each side first; otherwise the raw | ||
| 24 | // signal is used. `padded_index` below maps a padded position back to a raw | ||
| 25 | // signal index (reflecting at the borders when center == true) so the forward | ||
| 26 | // op and its adjoint share one indexing rule and stay exact transposes. | ||
| 27 | // | ||
| 28 | // ── Normalisation ─────────────────────────────────────────────────────────── | ||
| 29 | // rfft uses the "backward" convention (forward unscaled). normalized == true | ||
| 30 | // multiplies the forward spectrum by 1/sqrt(n_fft) (istft divides by it). | ||
| 31 | // | ||
| 32 | // ── Gradient design ───────────────────────────────────────────────────────── | ||
| 33 | // stft and istft are linear but NOT mutual adjoints (window + COLA). Each | ||
| 34 | // backward op is the exact transpose of its own forward linear map — see the | ||
| 35 | // ops.h header note. They are the minimal correct set for the | ||
| 36 | // multi-resolution STFT loss. | ||
| 37 | |||
| 38 | #include <brotensor/detail/cpu/fft_core.h> | ||
| 39 | #include <brotensor/tensor.h> | ||
| 40 | |||
| 41 | #include <algorithm> | ||
| 42 | #include <cmath> | ||
| 43 | #include <cstddef> | ||
| 44 | #include <stdexcept> | ||
| 45 | #include <string> | ||
| 46 | #include <vector> | ||
| 47 | |||
| 48 | namespace brotensor::detail::cpu { | ||
| 49 | |||
| 50 | using fftcore::Cd; | ||
| 51 | using fftcore::dft_1d; | ||
| 52 | |||
| 53 | namespace { | ||
| 54 | |||
| 55 | ✗ | [[noreturn]] void fail(const char* op, const std::string& reason) { | |
| 56 | ✗ | throw std::runtime_error(std::string("brotensor: ") + op + ": " + reason); | |
| 57 | ✗ | } | |
| 58 | |||
| 59 | 12838 | void require_fp32_host(const char* op, const ::brotensor::Tensor& t, | |
| 60 | const char* name) { | ||
| 61 |
1/2✓ Branch 0 taken 12838 times.
✗ Branch 1 not taken.
|
12838 | if (t.device != ::brotensor::Device::CPU) { |
| 62 | ✗ | fail(op, std::string(name) + " must be a CPU tensor"); | |
| 63 | } | ||
| 64 |
1/2✓ Branch 0 taken 12838 times.
✗ Branch 1 not taken.
|
12838 | if (t.dtype != ::brotensor::Dtype::FP32) { |
| 65 | ✗ | fail(op, std::string(name) + " must be FP32 (CPU is FP32-only)"); | |
| 66 | } | ||
| 67 | 12838 | } | |
| 68 | |||
| 69 | // Common parameter validation + derived sizes for all four ops. | ||
| 70 | 6419 | struct StftGeom { | |
| 71 | 6419 | int bins = 0; // n_fft/2 + 1 | |
| 72 | 6419 | int frames = 0; // frames per signal | |
| 73 | 6419 | int padded_len = 0; // signal length the frame loop indexes into | |
| 74 | 6419 | int pad_lo = 0; // (n_fft - win_length) / 2 — window offset in buffer | |
| 75 | }; | ||
| 76 | |||
| 77 | 6419 | StftGeom check_geom(const char* op, int N, int signal_len, int n_fft, | |
| 78 | int hop_length, int win_length, bool center) { | ||
| 79 |
1/4✓ Branch 0 taken 6419 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
6419 | if (N < 0) fail(op, "N must be >= 0"); |
| 80 |
1/4✓ Branch 0 taken 6419 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
6419 | if (n_fft < 1) fail(op, "n_fft must be >= 1"); |
| 81 |
1/4✓ Branch 0 taken 6419 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
6419 | if (hop_length < 1) fail(op, "hop_length must be >= 1"); |
| 82 |
1/2✓ Branch 0 taken 6419 times.
✗ Branch 1 not taken.
|
6419 | if (win_length < 1 || win_length > n_fft) { |
| 83 | ✗ | fail(op, "win_length must satisfy 1 <= win_length <= n_fft"); | |
| 84 | } | ||
| 85 |
1/4✓ Branch 0 taken 6419 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
6419 | if (signal_len < 1) fail(op, "signal_len must be >= 1"); |
| 86 | |||
| 87 | 6419 | StftGeom g; | |
| 88 | 6419 | g.bins = n_fft / 2 + 1; | |
| 89 | 6419 | g.pad_lo = (n_fft - win_length) / 2; | |
| 90 | |||
| 91 |
2/2✓ Branch 0 taken 3661 times.
✓ Branch 1 taken 2758 times.
|
6419 | if (center) { |
| 92 | // Reflect padding by n_fft/2 each side. numpy/torch 'reflect' mode | ||
| 93 | // needs at least 2 samples (the reflected index must stay in range); | ||
| 94 | // require enough signal to fill the n_fft/2 pad. | ||
| 95 |
1/2✓ Branch 0 taken 3661 times.
✗ Branch 1 not taken.
|
3661 | if (signal_len < n_fft / 2 + 1) { |
| 96 | ✗ | fail(op, "center=true needs signal_len >= n_fft/2 + 1"); | |
| 97 | } | ||
| 98 | 3661 | g.padded_len = signal_len + n_fft; | |
| 99 | 3661 | g.frames = 1 + signal_len / hop_length; | |
| 100 | 3661 | } else { | |
| 101 |
1/2✓ Branch 0 taken 2758 times.
✗ Branch 1 not taken.
|
2758 | if (signal_len < n_fft) { |
| 102 | ✗ | fail(op, "center=false needs signal_len >= n_fft"); | |
| 103 | } | ||
| 104 | 2758 | g.padded_len = signal_len; | |
| 105 | 2758 | g.frames = 1 + (signal_len - n_fft) / hop_length; | |
| 106 | } | ||
| 107 | 6419 | return g; | |
| 108 | ✗ | } | |
| 109 | |||
| 110 | // Map a padded position p in [0, padded_len) to a raw signal index in | ||
| 111 | // [0, signal_len). center == false is the identity; center == true reflects | ||
| 112 | // at the borders (numpy 'reflect': edge sample not repeated). | ||
| 113 | // | ||
| 114 | // Reflection over [0, L-1] with period 2*(L-1): fold q into that range. | ||
| 115 | 189520 | inline int reflect_index(int q, int L) { | |
| 116 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 189520 times.
|
189520 | if (L == 1) return 0; |
| 117 | 189520 | const int period = 2 * (L - 1); | |
| 118 | 189520 | int m = q % period; | |
| 119 |
2/2✓ Branch 0 taken 178644 times.
✓ Branch 1 taken 10876 times.
|
189520 | if (m < 0) m += period; |
| 120 |
2/2✓ Branch 0 taken 171702 times.
✓ Branch 1 taken 17818 times.
|
189520 | return (m < L) ? m : period - m; |
| 121 | 189520 | } | |
| 122 | |||
| 123 | 326720 | inline int padded_index(int p, int signal_len, int n_fft, bool center) { | |
| 124 |
2/2✓ Branch 0 taken 137200 times.
✓ Branch 1 taken 189520 times.
|
326720 | if (!center) return p; |
| 125 | 189520 | return reflect_index(p - n_fft / 2, signal_len); | |
| 126 | 326720 | } | |
| 127 | |||
| 128 | } // namespace | ||
| 129 | |||
| 130 | // ════════════════════════════════════════════════════════════════════════════ | ||
| 131 | // stft — real signal -> complex spectrogram | ||
| 132 | // ════════════════════════════════════════════════════════════════════════════ | ||
| 133 | 1286 | void stft(const ::brotensor::Tensor& signal, const ::brotensor::Tensor& window, | |
| 134 | int N, int n_fft, int hop_length, int win_length, | ||
| 135 | bool center, bool normalized, ::brotensor::Tensor& spec) { | ||
| 136 | 1286 | require_fp32_host("stft", signal, "signal"); | |
| 137 | 1286 | require_fp32_host("stft", window, "window"); | |
| 138 |
1/2✓ Branch 0 taken 1286 times.
✗ Branch 1 not taken.
|
1286 | if (signal.rows != N) { |
| 139 | ✗ | fail("stft", "signal.rows must equal N"); | |
| 140 | } | ||
| 141 | 1286 | const int signal_len = signal.cols; | |
| 142 |
1/2✓ Branch 0 taken 1286 times.
✗ Branch 1 not taken.
|
1286 | if (window.rows != 1 || window.cols != win_length) { |
| 143 | ✗ | fail("stft", "window must be a (1, win_length) tensor"); | |
| 144 | } | ||
| 145 | 2572 | const StftGeom g = check_geom("stft", N, signal_len, n_fft, hop_length, | |
| 146 | 1286 | win_length, center); | |
| 147 | 1286 | const int out_rows = N * g.frames; | |
| 148 | 1286 | const int out_cols = 2 * g.bins; | |
| 149 |
3/4✓ Branch 0 taken 1256 times.
✓ Branch 1 taken 30 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 1256 times.
|
1286 | if (spec.rows != out_rows || spec.cols != out_cols) { |
| 150 | 30 | spec.resize(out_rows, out_cols); | |
| 151 | 30 | } | |
| 152 |
1/2✓ Branch 0 taken 1286 times.
✗ Branch 1 not taken.
|
1286 | if (out_rows == 0) return; |
| 153 | |||
| 154 | 1286 | const float* sig = signal.host_f32(); | |
| 155 | 1286 | const float* win = window.host_f32(); | |
| 156 | 1286 | float* sp = spec.host_f32_mut(); | |
| 157 |
2/2✓ Branch 0 taken 637 times.
✓ Branch 1 taken 649 times.
|
1286 | const double norm = normalized |
| 158 | 637 | ? 1.0 / std::sqrt(static_cast<double>(n_fft)) | |
| 159 | : 1.0; | ||
| 160 | |||
| 161 | 1286 | std::vector<Cd> buf(static_cast<std::size_t>(n_fft)), out; | |
| 162 |
2/2✓ Branch 0 taken 2574 times.
✓ Branch 1 taken 1286 times.
|
3860 | for (int b = 0; b < N; ++b) { |
| 163 | 2574 | const float* srow = sig + static_cast<std::size_t>(b) * signal_len; | |
| 164 |
2/2✓ Branch 0 taken 2574 times.
✓ Branch 1 taken 16368 times.
|
18942 | for (int f = 0; f < g.frames; ++f) { |
| 165 | // Build the windowed n_fft frame buffer (real, im = 0). | ||
| 166 |
2/2✓ Branch 0 taken 324816 times.
✓ Branch 1 taken 16368 times.
|
341184 | for (int i = 0; i < n_fft; ++i) buf[static_cast<std::size_t>(i)] = Cd{}; |
| 167 | 16368 | const int base = f * hop_length; // padded-position start | |
| 168 |
2/2✓ Branch 0 taken 16368 times.
✓ Branch 1 taken 323804 times.
|
340172 | for (int j = 0; j < win_length; ++j) { |
| 169 | 323804 | const int i = g.pad_lo + j; | |
| 170 | 323804 | const int p = base + i; | |
| 171 |
1/2✓ Branch 0 taken 323804 times.
✗ Branch 1 not taken.
|
323804 | const int s = padded_index(p, signal_len, n_fft, center); |
| 172 | 323804 | buf[static_cast<std::size_t>(i)] = | |
| 173 | 647608 | {static_cast<double>(srow[s]) * | |
| 174 | 323804 | static_cast<double>(win[j]), | |
| 175 | 0.0}; | ||
| 176 | 323804 | } | |
| 177 |
1/2✓ Branch 0 taken 16368 times.
✗ Branch 1 not taken.
|
16368 | dft_1d(buf, out, -1); // unscaled forward DFT |
| 178 | 32736 | float* dst = sp + static_cast<std::size_t>( | |
| 179 | 32736 | static_cast<std::size_t>(b) * g.frames + f) * | |
| 180 | 16368 | out_cols; | |
| 181 |
2/2✓ Branch 0 taken 178776 times.
✓ Branch 1 taken 16368 times.
|
195144 | for (int k = 0; k < g.bins; ++k) { |
| 182 | 178776 | dst[2 * k] = static_cast<float>( | |
| 183 | 178776 | out[static_cast<std::size_t>(k)].re * norm); | |
| 184 | 178776 | dst[2 * k + 1] = static_cast<float>( | |
| 185 | 178776 | out[static_cast<std::size_t>(k)].im * norm); | |
| 186 | 178776 | } | |
| 187 | 16368 | } | |
| 188 | 2574 | } | |
| 189 | 1286 | } | |
| 190 | |||
| 191 | // ════════════════════════════════════════════════════════════════════════════ | ||
| 192 | // stft_backward — adjoint of stft | ||
| 193 | // ════════════════════════════════════════════════════════════════════════════ | ||
| 194 | // | ||
| 195 | // stft is the linear map spec = R * W * P * signal where P scatters the | ||
| 196 | // signal into frame buffers (with reflect padding folded in), W multiplies by | ||
| 197 | // the window, and R is the truncated forward DFT. Its adjoint applied to | ||
| 198 | // dSpec is P^T * W^T * R^T * dSpec : | ||
| 199 | // * R^T per frame is exactly rfft_backward's adjoint (the +1-sign unscaled | ||
| 200 | // DFT of the zero-padded n_fft spectrum, real part); | ||
| 201 | // * W^T is the same window multiply (diagonal — self-transpose); | ||
| 202 | // * P^T accumulates each frame sample back into the signal (the same index | ||
| 203 | // map, summed), so overlapping frames add — NO COLA division here (that | ||
| 204 | // belongs to istft, a different map). | ||
| 205 | // dSignal is *overwritten* (zeroed then accumulated). | ||
| 206 | 11 | void stft_backward(const ::brotensor::Tensor& dSpec, | |
| 207 | const ::brotensor::Tensor& window, | ||
| 208 | int N, int signal_len, int n_fft, int hop_length, | ||
| 209 | int win_length, bool center, bool normalized, | ||
| 210 | ::brotensor::Tensor& dSignal) { | ||
| 211 | 11 | require_fp32_host("stft_backward", dSpec, "dSpec"); | |
| 212 | 11 | require_fp32_host("stft_backward", window, "window"); | |
| 213 |
1/2✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
|
11 | if (window.rows != 1 || window.cols != win_length) { |
| 214 | ✗ | fail("stft_backward", "window must be a (1, win_length) tensor"); | |
| 215 | } | ||
| 216 | 22 | const StftGeom g = check_geom("stft_backward", N, signal_len, n_fft, | |
| 217 | 11 | hop_length, win_length, center); | |
| 218 | 11 | const int exp_rows = N * g.frames; | |
| 219 | 11 | const int exp_cols = 2 * g.bins; | |
| 220 |
1/2✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
|
11 | if (dSpec.rows != exp_rows || dSpec.cols != exp_cols) { |
| 221 | ✗ | fail("stft_backward", "dSpec shape must match the stft output shape"); | |
| 222 | } | ||
| 223 |
3/4✓ Branch 0 taken 8 times.
✓ Branch 1 taken 3 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
|
11 | if (dSignal.rows != N || dSignal.cols != signal_len) { |
| 224 | 3 | dSignal.resize(N, signal_len); | |
| 225 | 3 | } | |
| 226 |
1/2✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
|
11 | if (dSignal.size() != 0) { |
| 227 | 11 | float* z = dSignal.host_f32_mut(); | |
| 228 |
2/2✓ Branch 0 taken 1120 times.
✓ Branch 1 taken 11 times.
|
1131 | for (int i = 0; i < dSignal.size(); ++i) z[i] = 0.0f; |
| 229 | 11 | } | |
| 230 |
1/2✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
|
11 | if (exp_rows == 0) return; |
| 231 | |||
| 232 | 11 | const float* gp = dSpec.host_f32(); | |
| 233 | 11 | const float* win = window.host_f32(); | |
| 234 | 11 | float* dsig = dSignal.host_f32_mut(); | |
| 235 |
2/2✓ Branch 0 taken 5 times.
✓ Branch 1 taken 6 times.
|
11 | const double norm = normalized |
| 236 | 5 | ? 1.0 / std::sqrt(static_cast<double>(n_fft)) | |
| 237 | : 1.0; | ||
| 238 | |||
| 239 | // Per frame: spec[k] = norm * (truncated DFT)[k]. The adjoint of the | ||
| 240 | // truncated forward DFT is: zero-pad dSpec to length n_fft, run an | ||
| 241 | // unscaled +1-sign DFT, take the real part (== rfft_backward's core). | ||
| 242 | 11 | std::vector<Cd> spec(static_cast<std::size_t>(n_fft)), tbuf; | |
| 243 |
2/2✓ Branch 0 taken 23 times.
✓ Branch 1 taken 11 times.
|
34 | for (int b = 0; b < N; ++b) { |
| 244 | 23 | float* drow = dsig + static_cast<std::size_t>(b) * signal_len; | |
| 245 |
2/2✓ Branch 0 taken 23 times.
✓ Branch 1 taken 251 times.
|
274 | for (int f = 0; f < g.frames; ++f) { |
| 246 | 502 | const float* grow = gp + static_cast<std::size_t>( | |
| 247 | 502 | static_cast<std::size_t>(b) * | |
| 248 | 502 | g.frames + | |
| 249 | 502 | f) * exp_cols; | |
| 250 |
2/2✓ Branch 0 taken 3168 times.
✓ Branch 1 taken 251 times.
|
3419 | for (int k = 0; k < n_fft; ++k) spec[static_cast<std::size_t>(k)] = Cd{}; |
| 251 |
2/2✓ Branch 0 taken 1835 times.
✓ Branch 1 taken 251 times.
|
2086 | for (int k = 0; k < g.bins; ++k) { |
| 252 | 1835 | spec[static_cast<std::size_t>(k)] = | |
| 253 | 3670 | {static_cast<double>(grow[2 * k]) * norm, | |
| 254 | 1835 | static_cast<double>(grow[2 * k + 1]) * norm}; | |
| 255 | 1835 | } | |
| 256 |
1/2✓ Branch 0 taken 251 times.
✗ Branch 1 not taken.
|
251 | dft_1d(spec, tbuf, +1); // adjoint of truncated forward DFT |
| 257 | // W^T (window) then P^T (scatter-add into the signal). | ||
| 258 | 251 | const int base = f * hop_length; | |
| 259 |
2/2✓ Branch 0 taken 251 times.
✓ Branch 1 taken 2916 times.
|
3167 | for (int j = 0; j < win_length; ++j) { |
| 260 | 2916 | const int i = g.pad_lo + j; | |
| 261 | 2916 | const int p = base + i; | |
| 262 |
1/2✓ Branch 0 taken 2916 times.
✗ Branch 1 not taken.
|
2916 | const int s = padded_index(p, signal_len, n_fft, center); |
| 263 | 2916 | drow[s] += static_cast<float>( | |
| 264 | 5832 | tbuf[static_cast<std::size_t>(i)].re * | |
| 265 | 2916 | static_cast<double>(win[j])); | |
| 266 | 2916 | } | |
| 267 | 251 | } | |
| 268 | 23 | } | |
| 269 | 11 | } | |
| 270 | |||
| 271 | // ════════════════════════════════════════════════════════════════════════════ | ||
| 272 | // istft — complex spectrogram -> real signal (windowed overlap-add + COLA) | ||
| 273 | // ════════════════════════════════════════════════════════════════════════════ | ||
| 274 | // | ||
| 275 | // Per frame: irfft the n_fft spectrum, multiply by the window, scatter-add | ||
| 276 | // into the output. Then divide each output sample by the overlap-added | ||
| 277 | // squared window (the COLA envelope) so a COLA-satisfying window+hop makes | ||
| 278 | // istft(stft(x)) == x. Samples with a ~0 envelope (edges with no frame | ||
| 279 | // coverage) stay 0. | ||
| 280 | 5111 | void istft(const ::brotensor::Tensor& spec, const ::brotensor::Tensor& window, | |
| 281 | int N, int signal_len, int n_fft, int hop_length, int win_length, | ||
| 282 | bool center, bool normalized, ::brotensor::Tensor& signal) { | ||
| 283 | 5111 | require_fp32_host("istft", spec, "spec"); | |
| 284 | 5111 | require_fp32_host("istft", window, "window"); | |
| 285 |
1/2✓ Branch 0 taken 5111 times.
✗ Branch 1 not taken.
|
5111 | if (window.rows != 1 || window.cols != win_length) { |
| 286 | ✗ | fail("istft", "window must be a (1, win_length) tensor"); | |
| 287 | } | ||
| 288 | 10222 | const StftGeom g = check_geom("istft", N, signal_len, n_fft, hop_length, | |
| 289 | 5111 | win_length, center); | |
| 290 | 5111 | const int exp_rows = N * g.frames; | |
| 291 | 5111 | const int exp_cols = 2 * g.bins; | |
| 292 |
1/2✓ Branch 0 taken 5111 times.
✗ Branch 1 not taken.
|
5111 | if (spec.rows != exp_rows || spec.cols != exp_cols) { |
| 293 | ✗ | fail("istft", "spec shape must match the stft output shape"); | |
| 294 | } | ||
| 295 |
3/4✓ Branch 0 taken 5096 times.
✓ Branch 1 taken 15 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 5096 times.
|
5111 | if (signal.rows != N || signal.cols != signal_len) { |
| 296 | 15 | signal.resize(N, signal_len); | |
| 297 | 15 | } | |
| 298 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 5111 times.
|
5111 | if (signal.size() != 0) { |
| 299 | 5111 | float* z = signal.host_f32_mut(); | |
| 300 |
2/2✓ Branch 0 taken 451096 times.
✓ Branch 1 taken 5111 times.
|
456207 | for (int i = 0; i < signal.size(); ++i) z[i] = 0.0f; |
| 301 | 5111 | } | |
| 302 |
1/2✓ Branch 0 taken 5111 times.
✗ Branch 1 not taken.
|
5111 | if (exp_rows == 0) return; |
| 303 | |||
| 304 | 5111 | const float* sp = spec.host_f32(); | |
| 305 | 5111 | const float* win = window.host_f32(); | |
| 306 | 5111 | float* sig = signal.host_f32_mut(); | |
| 307 | // istft inverts stft's optional 1/sqrt(n_fft): multiply the spectrum by | ||
| 308 | // sqrt(n_fft) so the irfft 1/n_fft scaling lands at the right amplitude. | ||
| 309 |
2/2✓ Branch 0 taken 2555 times.
✓ Branch 1 taken 2556 times.
|
5111 | const double norm = normalized ? std::sqrt(static_cast<double>(n_fft)) |
| 310 | : 1.0; | ||
| 311 | 5111 | const double invN = 1.0 / static_cast<double>(n_fft); | |
| 312 | |||
| 313 | // COLA envelope: overlap-added squared window, in padded coordinates. | ||
| 314 | 5111 | std::vector<double> env(static_cast<std::size_t>(g.padded_len), 0.0); | |
| 315 |
2/2✓ Branch 0 taken 60119 times.
✓ Branch 1 taken 5111 times.
|
65230 | for (int f = 0; f < g.frames; ++f) { |
| 316 | 60119 | const int base = f * hop_length; | |
| 317 |
2/2✓ Branch 0 taken 811508 times.
✓ Branch 1 taken 60119 times.
|
871627 | for (int j = 0; j < win_length; ++j) { |
| 318 | 811508 | const int p = base + g.pad_lo + j; | |
| 319 | 811508 | const double w = static_cast<double>(win[j]); | |
| 320 | 811508 | env[static_cast<std::size_t>(p)] += w * w; | |
| 321 | 811508 | } | |
| 322 | 60119 | } | |
| 323 | |||
| 324 | // Per signal: overlap-add the windowed irfft frames, then COLA-divide. | ||
| 325 |
1/2✓ Branch 0 taken 5111 times.
✗ Branch 1 not taken.
|
5111 | std::vector<Cd> full(static_cast<std::size_t>(n_fft)), out; |
| 326 |
1/2✓ Branch 0 taken 5111 times.
✗ Branch 1 not taken.
|
5111 | std::vector<double> acc(static_cast<std::size_t>(g.padded_len)); |
| 327 |
2/2✓ Branch 0 taken 10223 times.
✓ Branch 1 taken 5111 times.
|
15334 | for (int b = 0; b < N; ++b) { |
| 328 |
1/2✓ Branch 0 taken 10223 times.
✗ Branch 1 not taken.
|
10223 | std::fill(acc.begin(), acc.end(), 0.0); |
| 329 |
2/2✓ Branch 0 taken 120259 times.
✓ Branch 1 taken 10223 times.
|
130482 | for (int f = 0; f < g.frames; ++f) { |
| 330 | 240518 | const float* srow = sp + static_cast<std::size_t>( | |
| 331 | 240518 | static_cast<std::size_t>(b) * | |
| 332 | 240518 | g.frames + | |
| 333 | 240518 | f) * exp_cols; | |
| 334 | // Rebuild the Hermitian-symmetric n_fft spectrum, irfft it. | ||
| 335 |
2/2✓ Branch 0 taken 932019 times.
✓ Branch 1 taken 120259 times.
|
1052278 | for (int k = 0; k < g.bins; ++k) { |
| 336 | 932019 | full[static_cast<std::size_t>(k)] = | |
| 337 | 1864038 | {static_cast<double>(srow[2 * k]) * norm, | |
| 338 | 932019 | static_cast<double>(srow[2 * k + 1]) * norm}; | |
| 339 | 932019 | } | |
| 340 |
2/2✓ Branch 0 taken 691501 times.
✓ Branch 1 taken 120259 times.
|
811760 | for (int k = 1; k < n_fft - g.bins + 1; ++k) { |
| 341 | 691501 | const Cd c = full[static_cast<std::size_t>(k)]; | |
| 342 | 691501 | full[static_cast<std::size_t>(n_fft - k)] = {c.re, -c.im}; | |
| 343 | 691501 | } | |
| 344 |
1/2✓ Branch 0 taken 120259 times.
✗ Branch 1 not taken.
|
120259 | dft_1d(full, out, +1); // inverse DFT, still needs *1/n_fft |
| 345 | 120259 | const int base = f * hop_length; | |
| 346 |
2/2✓ Branch 0 taken 1623268 times.
✓ Branch 1 taken 120259 times.
|
1743527 | for (int j = 0; j < win_length; ++j) { |
| 347 | 1623268 | const int i = g.pad_lo + j; | |
| 348 | 1623268 | const int p = base + i; | |
| 349 | 1623268 | const double t = out[static_cast<std::size_t>(i)].re * invN; | |
| 350 | 1623268 | acc[static_cast<std::size_t>(p)] += | |
| 351 | 1623268 | t * static_cast<double>(win[j]); | |
| 352 | 1623268 | } | |
| 353 | 120259 | } | |
| 354 | // COLA-divide and strip centre padding back to the raw signal. | ||
| 355 | 10223 | float* drow = sig + static_cast<std::size_t>(b) * signal_len; | |
| 356 |
2/2✓ Branch 0 taken 6009 times.
✓ Branch 1 taken 4214 times.
|
10223 | const int shift = center ? n_fft / 2 : 0; |
| 357 |
2/2✓ Branch 0 taken 451096 times.
✓ Branch 1 taken 10223 times.
|
461319 | for (int n = 0; n < signal_len; ++n) { |
| 358 | 451096 | const int p = n + shift; | |
| 359 | 451096 | const double e = env[static_cast<std::size_t>(p)]; | |
| 360 |
2/2✓ Branch 0 taken 437444 times.
✓ Branch 1 taken 13652 times.
|
451096 | drow[n] = (e > 1e-10) |
| 361 | 874888 | ? static_cast<float>(acc[static_cast<std::size_t>(p)] | |
| 362 | 437444 | / e) | |
| 363 | : 0.0f; | ||
| 364 | 451096 | } | |
| 365 | 10223 | } | |
| 366 | 5111 | } | |
| 367 | |||
| 368 | // ════════════════════════════════════════════════════════════════════════════ | ||
| 369 | // istft_backward — adjoint of istft | ||
| 370 | // ════════════════════════════════════════════════════════════════════════════ | ||
| 371 | // | ||
| 372 | // istft is the linear map signal = D * E^{-1} * P * W * I * spec where I is | ||
| 373 | // the per-frame inverse DFT (1/n_fft scaled), W the window multiply, P the | ||
| 374 | // overlap-add scatter, E^{-1} the per-sample COLA division, and D the | ||
| 375 | // centre-padding strip. The COLA envelope E depends only on the window/hop, | ||
| 376 | // not on the spectrum, so E^{-1} is a (data-independent) diagonal — and its | ||
| 377 | // transpose is itself. The adjoint applied to dSignal is therefore | ||
| 378 | // I^T * W^T * P^T * E^{-1} * D^T * dSignal : | ||
| 379 | // * D^T scatters dSignal back into padded coordinates; | ||
| 380 | // * E^{-1} divides by the same COLA envelope (diagonal, self-transpose); | ||
| 381 | // * P^T gathers each frame's window_length samples; | ||
| 382 | // * W^T is the window multiply again; | ||
| 383 | // * I^T is the adjoint of the inverse DFT — which is irfft_backward's core | ||
| 384 | // (forward-sign DFT of the gathered frame, 1/n_fft scaling, and the | ||
| 385 | // interior-bin doubling from the Hermitian fold). | ||
| 386 | // dSpec is *overwritten*. | ||
| 387 | 11 | void istft_backward(const ::brotensor::Tensor& dSignal, | |
| 388 | const ::brotensor::Tensor& window, | ||
| 389 | int N, int signal_len, int n_fft, int hop_length, | ||
| 390 | int win_length, bool center, bool normalized, | ||
| 391 | ::brotensor::Tensor& dSpec) { | ||
| 392 | 11 | require_fp32_host("istft_backward", dSignal, "dSignal"); | |
| 393 | 11 | require_fp32_host("istft_backward", window, "window"); | |
| 394 |
1/2✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
|
11 | if (dSignal.rows != N || dSignal.cols != signal_len) { |
| 395 | ✗ | fail("istft_backward", "dSignal must be a (N, signal_len) tensor"); | |
| 396 | } | ||
| 397 |
1/2✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
|
11 | if (window.rows != 1 || window.cols != win_length) { |
| 398 | ✗ | fail("istft_backward", "window must be a (1, win_length) tensor"); | |
| 399 | } | ||
| 400 | 22 | const StftGeom g = check_geom("istft_backward", N, signal_len, n_fft, | |
| 401 | 11 | hop_length, win_length, center); | |
| 402 | 11 | const int out_rows = N * g.frames; | |
| 403 | 11 | const int out_cols = 2 * g.bins; | |
| 404 |
3/4✓ Branch 0 taken 8 times.
✓ Branch 1 taken 3 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
|
11 | if (dSpec.rows != out_rows || dSpec.cols != out_cols) { |
| 405 | 3 | dSpec.resize(out_rows, out_cols); | |
| 406 | 3 | } | |
| 407 |
1/2✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
|
11 | if (out_rows == 0) return; |
| 408 | |||
| 409 | 11 | const float* dsig = dSignal.host_f32(); | |
| 410 | 11 | const float* win = window.host_f32(); | |
| 411 | 11 | float* gp = dSpec.host_f32_mut(); | |
| 412 |
2/2✓ Branch 0 taken 5 times.
✓ Branch 1 taken 6 times.
|
11 | const double norm = normalized ? std::sqrt(static_cast<double>(n_fft)) |
| 413 | : 1.0; | ||
| 414 | 11 | const double invN = 1.0 / static_cast<double>(n_fft); | |
| 415 | |||
| 416 | // Same COLA envelope as istft (window/hop only). | ||
| 417 | 11 | std::vector<double> env(static_cast<std::size_t>(g.padded_len), 0.0); | |
| 418 |
2/2✓ Branch 0 taken 159 times.
✓ Branch 1 taken 11 times.
|
170 | for (int f = 0; f < g.frames; ++f) { |
| 419 | 159 | const int base = f * hop_length; | |
| 420 |
2/2✓ Branch 0 taken 1844 times.
✓ Branch 1 taken 159 times.
|
2003 | for (int j = 0; j < win_length; ++j) { |
| 421 | 1844 | const int p = base + g.pad_lo + j; | |
| 422 | 1844 | const double w = static_cast<double>(win[j]); | |
| 423 | 1844 | env[static_cast<std::size_t>(p)] += w * w; | |
| 424 | 1844 | } | |
| 425 | 159 | } | |
| 426 | |||
| 427 | 11 | const bool even = (n_fft % 2 == 0); | |
| 428 |
1/2✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
|
11 | std::vector<double> gacc(static_cast<std::size_t>(g.padded_len)); |
| 429 |
1/2✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
|
11 | std::vector<Cd> frame(static_cast<std::size_t>(n_fft)), spec; |
| 430 |
2/2✓ Branch 0 taken 23 times.
✓ Branch 1 taken 11 times.
|
34 | for (int b = 0; b < N; ++b) { |
| 431 | // D^T then E^{-1}: scatter dSignal into padded coords, COLA-divide. | ||
| 432 |
1/2✓ Branch 0 taken 23 times.
✗ Branch 1 not taken.
|
23 | std::fill(gacc.begin(), gacc.end(), 0.0); |
| 433 | 23 | const float* drow = dsig + static_cast<std::size_t>(b) * signal_len; | |
| 434 |
2/2✓ Branch 0 taken 13 times.
✓ Branch 1 taken 10 times.
|
23 | const int shift = center ? n_fft / 2 : 0; |
| 435 |
2/2✓ Branch 0 taken 1120 times.
✓ Branch 1 taken 23 times.
|
1143 | for (int n = 0; n < signal_len; ++n) { |
| 436 | 1120 | const int p = n + shift; | |
| 437 | 1120 | const double e = env[static_cast<std::size_t>(p)]; | |
| 438 | 1120 | gacc[static_cast<std::size_t>(p)] = | |
| 439 |
2/2✓ Branch 0 taken 1092 times.
✓ Branch 1 taken 28 times.
|
1120 | (e > 1e-10) ? static_cast<double>(drow[n]) / e : 0.0; |
| 440 | 1120 | } | |
| 441 |
2/2✓ Branch 0 taken 23 times.
✓ Branch 1 taken 339 times.
|
362 | for (int f = 0; f < g.frames; ++f) { |
| 442 | // P^T (gather) then W^T (window): the frame's n_fft time buffer. | ||
| 443 |
2/2✓ Branch 0 taken 4192 times.
✓ Branch 1 taken 339 times.
|
4531 | for (int k = 0; k < n_fft; ++k) frame[static_cast<std::size_t>(k)] = Cd{}; |
| 444 | 339 | const int base = f * hop_length; | |
| 445 |
2/2✓ Branch 0 taken 3940 times.
✓ Branch 1 taken 339 times.
|
4279 | for (int j = 0; j < win_length; ++j) { |
| 446 | 3940 | const int i = g.pad_lo + j; | |
| 447 | 3940 | const int p = base + i; | |
| 448 | 3940 | frame[static_cast<std::size_t>(i)] = | |
| 449 | 7880 | {gacc[static_cast<std::size_t>(p)] * | |
| 450 | 3940 | static_cast<double>(win[j]), | |
| 451 | 0.0}; | ||
| 452 | 3940 | } | |
| 453 | // I^T: adjoint of the 1/n_fft inverse DFT — forward-sign DFT, | ||
| 454 | // 1/n_fft scaling, interior-bin doubling for the Hermitian fold. | ||
| 455 |
1/2✓ Branch 0 taken 339 times.
✗ Branch 1 not taken.
|
339 | dft_1d(frame, spec, -1); |
| 456 | 678 | float* grow = gp + static_cast<std::size_t>( | |
| 457 | 678 | static_cast<std::size_t>(b) * g.frames + | |
| 458 | 678 | f) * out_cols; | |
| 459 |
2/2✓ Branch 0 taken 2435 times.
✓ Branch 1 taken 339 times.
|
2774 | for (int k = 0; k < g.bins; ++k) { |
| 460 | 2435 | double s = 2.0; | |
| 461 |
2/2✓ Branch 0 taken 2096 times.
✓ Branch 1 taken 339 times.
|
2435 | if (k == 0) s = 1.0; |
| 462 |
3/4✓ Branch 0 taken 2435 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2096 times.
✓ Branch 3 taken 339 times.
|
2435 | if (even && k == n_fft / 2) s = 1.0; |
| 463 | 2435 | const double scale = s * invN * norm; | |
| 464 | 2435 | grow[2 * k] = static_cast<float>( | |
| 465 | 2435 | scale * spec[static_cast<std::size_t>(k)].re); | |
| 466 | 2435 | grow[2 * k + 1] = static_cast<float>( | |
| 467 | 2435 | scale * spec[static_cast<std::size_t>(k)].im); | |
| 468 | 2435 | } | |
| 469 | 339 | } | |
| 470 | 23 | } | |
| 471 | 11 | } | |
| 472 | |||
| 473 | } // namespace brotensor::detail::cpu | ||
| 474 |