src/cpu/fft.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // ─── CPU spectral / FFT core (brosoundml CHUNK 1) ────────────────────────── | ||
| 2 | // | ||
| 3 | // Hand-rolled FFT for the CPU backend. No external libraries. FP32-only, | ||
| 4 | // matching the CPU backend's FP32-only contract. | ||
| 5 | // | ||
| 6 | // Ops implemented here: | ||
| 7 | // complex_mul / complex_abs / complex_angle / complex_from_polar | ||
| 8 | // complex_mul_backward / complex_abs_backward | ||
| 9 | // fft / ifft — complex -> complex, one signal per tensor row | ||
| 10 | // rfft / irfft — real (R,L) <-> complex (R, 2*(L/2+1)) | ||
| 11 | // rfft_backward / irfft_backward — adjoints of rfft / irfft | ||
| 12 | // | ||
| 13 | // ── Complex layout ───────────────────────────────────────────────────────── | ||
| 14 | // A complex tensor is a regular FP32 tensor with the bin axis stored | ||
| 15 | // interleaved [re, im, re, im, ...]. A complex spectrum of C bins over R rows | ||
| 16 | // is an (R, 2*C) FP32 tensor. There is no new Dtype. | ||
| 17 | // | ||
| 18 | // ── FFT algorithm ────────────────────────────────────────────────────────── | ||
| 19 | // Mixed-radix Cooley-Tukey (radix 2/3/5/7) handles sizes whose prime | ||
| 20 | // factorisation only uses small primes — this covers Whisper's n_fft = 400 | ||
| 21 | // (= 2^4 * 5^2). Any remaining factor (a large or genuinely prime factor) is | ||
| 22 | // transformed with a Bluestein chirp-z transform, which reduces an | ||
| 23 | // arbitrary-length DFT to a power-of-two convolution. The whole thing is | ||
| 24 | // therefore correct for *every* length >= 1, including primes. | ||
| 25 | // | ||
| 26 | // ── Normalisation ────────────────────────────────────────────────────────── | ||
| 27 | // "backward" convention (numpy default): the forward transform is unscaled, | ||
| 28 | // the inverse transform is scaled by 1/N. | ||
| 29 | // | ||
| 30 | // ── Gradient design (the linear-transform adjoints) ──────────────────────── | ||
| 31 | // fft / ifft / rfft / irfft are all linear maps, so the backward of each is | ||
| 32 | // the adjoint (conjugate transpose) of its forward matrix applied to the | ||
| 33 | // upstream gradient. We deliberately keep the vtable minimal: | ||
| 34 | // | ||
| 35 | // * fft / ifft are complex->complex and self-similar. The adjoint of the | ||
| 36 | // length-N forward DFT matrix F is F^H = conj(F) = N * F^{-1}. With this | ||
| 37 | // library's transforms that means: | ||
| 38 | // grad_x(fft) = N * ifft(grad_y) | ||
| 39 | // grad_x(ifft) = (1/N) * fft(grad_y) | ||
| 40 | // Both adjoints are an *existing* transform composed with a scalar, so we | ||
| 41 | // do NOT add fft_backward / ifft_backward rows. Training code spells the | ||
| 42 | // gradient as `ifft(g); scale_inplace(g, N)` (or the ifft dual). This is | ||
| 43 | // documented on the fft / ifft declarations in ops.h. | ||
| 44 | // | ||
| 45 | // * rfft / irfft are NOT mutual adjoints. rfft maps a real length-L signal | ||
| 46 | // to its non-redundant half-spectrum (L/2+1 bins); irfft maps a | ||
| 47 | // half-spectrum back to a real signal assuming Hermitian symmetry AND | ||
| 48 | // applies the 1/L inverse scaling. rfft_backward is the plain adjoint of | ||
| 49 | // the truncated DFT matrix (no bin weighting — rfft does no folding). | ||
| 50 | // irfft_backward carries the 1/L scaling and the interior-bin weighting | ||
| 51 | // that irfft's Hermitian folding implies. Getting that weighting wrong is | ||
| 52 | // a silent training bug, so both are explicit ops rather than something | ||
| 53 | // callers reconstruct. They are the minimal correct set for the gradient | ||
| 54 | // path of a multi-resolution STFT loss. | ||
| 55 | // | ||
| 56 | // This file ports cleanly to CUDA / Metal later (the math is backend-neutral); | ||
| 57 | // only the host_f32 accessors are CPU-specific. | ||
| 58 | // | ||
| 59 | // The mixed-radix + Bluestein transform engine itself lives in the shared | ||
| 60 | // header detail/cpu/fft_core.h so the STFT / iSTFT ops (stft.cpp) reuse the | ||
| 61 | // exact same DFT instead of copy-pasting it. | ||
| 62 | |||
| 63 | #include <brotensor/detail/cpu/fft_core.h> | ||
| 64 | #include <brotensor/tensor.h> | ||
| 65 | |||
| 66 | #include <cmath> | ||
| 67 | #include <cstddef> | ||
| 68 | #include <stdexcept> | ||
| 69 | #include <string> | ||
| 70 | #include <vector> | ||
| 71 | |||
| 72 | namespace brotensor::detail::cpu { | ||
| 73 | |||
| 74 | // Pull the shared FFT-core internals (Cd, dft_1d, complex-row I/O) into scope. | ||
| 75 | using fftcore::Cd; | ||
| 76 | using fftcore::dft_1d; | ||
| 77 | using fftcore::load_complex_row; | ||
| 78 | using fftcore::store_complex_row; | ||
| 79 | |||
| 80 | namespace { | ||
| 81 | |||
| 82 | ✗ | [[noreturn]] void fail(const char* op, const std::string& reason) { | |
| 83 | ✗ | throw std::runtime_error(std::string("brotensor: ") + op + ": " + reason); | |
| 84 | ✗ | } | |
| 85 | |||
| 86 | 1230 | void require_fp32_host(const char* op, const ::brotensor::Tensor& t, | |
| 87 | const char* name) { | ||
| 88 |
1/2✓ Branch 0 taken 1230 times.
✗ Branch 1 not taken.
|
1230 | if (t.device != ::brotensor::Device::CPU) { |
| 89 | ✗ | fail(op, std::string(name) + " must be a CPU tensor"); | |
| 90 | } | ||
| 91 |
1/2✓ Branch 0 taken 1230 times.
✗ Branch 1 not taken.
|
1230 | if (t.dtype != ::brotensor::Dtype::FP32) { |
| 92 | ✗ | fail(op, std::string(name) + " must be FP32 (CPU is FP32-only)"); | |
| 93 | } | ||
| 94 | 1230 | } | |
| 95 | |||
| 96 | } // namespace | ||
| 97 | |||
| 98 | // ════════════════════════════════════════════════════════════════════════════ | ||
| 99 | // Complex elementwise ops | ||
| 100 | // ════════════════════════════════════════════════════════════════════════════ | ||
| 101 | |||
| 102 | // y = a * b, complex elementwise. a, b, y are interleaved-complex (R, 2*C). | ||
| 103 | 52 | void complex_mul(const ::brotensor::Tensor& a, const ::brotensor::Tensor& b, | |
| 104 | ::brotensor::Tensor& y) { | ||
| 105 | 52 | require_fp32_host("complex_mul", a, "a"); | |
| 106 | 52 | require_fp32_host("complex_mul", b, "b"); | |
| 107 |
1/2✓ Branch 0 taken 52 times.
✗ Branch 1 not taken.
|
52 | if (a.rows != b.rows || a.cols != b.cols) { |
| 108 | ✗ | fail("complex_mul", "a and b must have identical shape"); | |
| 109 | } | ||
| 110 |
1/2✓ Branch 0 taken 52 times.
✗ Branch 1 not taken.
|
52 | if (a.cols % 2 != 0) { |
| 111 | ✗ | fail("complex_mul", "cols must be even (interleaved [re,im] layout)"); | |
| 112 | } | ||
| 113 |
3/4✓ Branch 0 taken 50 times.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 50 times.
|
52 | if (y.rows != a.rows || y.cols != a.cols) y.resize(a.rows, a.cols); |
| 114 | 52 | const int n = a.size(); | |
| 115 |
1/2✓ Branch 0 taken 52 times.
✗ Branch 1 not taken.
|
52 | if (n == 0) return; |
| 116 | 52 | const float* ap = a.host_f32(); | |
| 117 | 52 | const float* bp = b.host_f32(); | |
| 118 | 52 | float* yp = y.host_f32_mut(); | |
| 119 |
2/2✓ Branch 0 taken 52 times.
✓ Branch 1 taken 312 times.
|
364 | for (int i = 0; i < n; i += 2) { |
| 120 | 312 | const float ar = ap[i], ai = ap[i + 1]; | |
| 121 | 312 | const float br = bp[i], bi = bp[i + 1]; | |
| 122 | 312 | yp[i] = ar * br - ai * bi; | |
| 123 | 312 | yp[i + 1] = ar * bi + ai * br; | |
| 124 | 312 | } | |
| 125 | 52 | } | |
| 126 | |||
| 127 | // Backward of complex_mul. y = a * b ⇒ (Wirtinger / real-pair gradient) | ||
| 128 | // dA = dY * conj(b), dB = dY * conj(a). | ||
| 129 | // dA and dB are *accumulated into* — the caller zeros them (mirrors the | ||
| 130 | // accumulation contract used by linear_backward / matmul_backward). | ||
| 131 | 3 | void complex_mul_backward(const ::brotensor::Tensor& a, | |
| 132 | const ::brotensor::Tensor& b, | ||
| 133 | const ::brotensor::Tensor& dY, | ||
| 134 | ::brotensor::Tensor& dA, ::brotensor::Tensor& dB) { | ||
| 135 | 3 | require_fp32_host("complex_mul_backward", a, "a"); | |
| 136 | 3 | require_fp32_host("complex_mul_backward", b, "b"); | |
| 137 | 3 | require_fp32_host("complex_mul_backward", dY, "dY"); | |
| 138 | 3 | require_fp32_host("complex_mul_backward", dA, "dA"); | |
| 139 | 3 | require_fp32_host("complex_mul_backward", dB, "dB"); | |
| 140 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | if (a.rows != b.rows || a.cols != b.cols) { |
| 141 | ✗ | fail("complex_mul_backward", "a and b must have identical shape"); | |
| 142 | } | ||
| 143 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | if (dY.rows != a.rows || dY.cols != a.cols) { |
| 144 | ✗ | fail("complex_mul_backward", "dY must match a / b shape"); | |
| 145 | } | ||
| 146 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | if (dA.rows != a.rows || dA.cols != a.cols) { |
| 147 | ✗ | fail("complex_mul_backward", "dA must be pre-sized to a's shape"); | |
| 148 | } | ||
| 149 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | if (dB.rows != a.rows || dB.cols != a.cols) { |
| 150 | ✗ | fail("complex_mul_backward", "dB must be pre-sized to b's shape"); | |
| 151 | } | ||
| 152 | 3 | const int n = a.size(); | |
| 153 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | if (n == 0) return; |
| 154 | 3 | const float* ap = a.host_f32(); | |
| 155 | 3 | const float* bp = b.host_f32(); | |
| 156 | 3 | const float* gp = dY.host_f32(); | |
| 157 | 3 | float* dap = dA.host_f32_mut(); | |
| 158 | 3 | float* dbp = dB.host_f32_mut(); | |
| 159 |
2/2✓ Branch 0 taken 3 times.
✓ Branch 1 taken 36 times.
|
39 | for (int i = 0; i < n; i += 2) { |
| 160 | 36 | const float gr = gp[i], gi = gp[i + 1]; | |
| 161 | 36 | const float ar = ap[i], ai = ap[i + 1]; | |
| 162 | 36 | const float br = bp[i], bi = bp[i + 1]; | |
| 163 | // dA = dY * conj(b): (gr+igi)(br-ibi) | ||
| 164 | 36 | dap[i] += gr * br + gi * bi; | |
| 165 | 36 | dap[i + 1] += gi * br - gr * bi; | |
| 166 | // dB = dY * conj(a): (gr+igi)(ar-iai) | ||
| 167 | 36 | dbp[i] += gr * ar + gi * ai; | |
| 168 | 36 | dbp[i + 1] += gi * ar - gr * ai; | |
| 169 | 36 | } | |
| 170 | 3 | } | |
| 171 | |||
| 172 | // y = |z|, real magnitude per complex bin. Input z is interleaved-complex | ||
| 173 | // (R, 2*C); output y is REAL (R, C). | ||
| 174 | 35 | void complex_abs(const ::brotensor::Tensor& z, ::brotensor::Tensor& y) { | |
| 175 | 35 | require_fp32_host("complex_abs", z, "z"); | |
| 176 |
1/2✓ Branch 0 taken 35 times.
✗ Branch 1 not taken.
|
35 | if (z.cols % 2 != 0) { |
| 177 | ✗ | fail("complex_abs", "z.cols must be even (interleaved [re,im] layout)"); | |
| 178 | } | ||
| 179 | 35 | const int C = z.cols / 2; | |
| 180 |
3/4✓ Branch 0 taken 34 times.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✓ Branch 3 taken 34 times.
|
35 | if (y.rows != z.rows || y.cols != C) y.resize(z.rows, C); |
| 181 |
1/2✓ Branch 0 taken 35 times.
✗ Branch 1 not taken.
|
35 | if (z.size() == 0) return; |
| 182 | 35 | const float* zp = z.host_f32(); | |
| 183 | 35 | float* yp = y.host_f32_mut(); | |
| 184 |
2/2✓ Branch 0 taken 70 times.
✓ Branch 1 taken 35 times.
|
105 | for (int r = 0; r < z.rows; ++r) { |
| 185 | 70 | const float* zr = zp + static_cast<std::size_t>(r) * z.cols; | |
| 186 | 70 | float* yr = yp + static_cast<std::size_t>(r) * C; | |
| 187 |
2/2✓ Branch 0 taken 287 times.
✓ Branch 1 taken 70 times.
|
357 | for (int c = 0; c < C; ++c) { |
| 188 | 287 | const float re = zr[2 * c], im = zr[2 * c + 1]; | |
| 189 | 287 | yr[c] = std::sqrt(re * re + im * im); | |
| 190 | 287 | } | |
| 191 | 70 | } | |
| 192 | 35 | } | |
| 193 | |||
| 194 | // Backward of complex_abs. With r = |z| = sqrt(re^2 + im^2): | ||
| 195 | // d|z|/d(re) = re / r, d|z|/d(im) = im / r. | ||
| 196 | // dZ is interleaved-complex (R, 2*C), *overwritten* (matches the GPU | ||
| 197 | // activation-backward convention — backward writes dZ directly). At r == 0 | ||
| 198 | // the gradient is set to 0 (the magnitude is non-differentiable there; | ||
| 199 | // 0 is the conventional choice). | ||
| 200 | 2 | void complex_abs_backward(const ::brotensor::Tensor& z, | |
| 201 | const ::brotensor::Tensor& dY, | ||
| 202 | ::brotensor::Tensor& dZ) { | ||
| 203 | 2 | require_fp32_host("complex_abs_backward", z, "z"); | |
| 204 | 2 | require_fp32_host("complex_abs_backward", dY, "dY"); | |
| 205 |
1/2✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
|
2 | if (z.cols % 2 != 0) { |
| 206 | ✗ | fail("complex_abs_backward", | |
| 207 | ✗ | "z.cols must be even (interleaved [re,im] layout)"); | |
| 208 | } | ||
| 209 | 2 | const int C = z.cols / 2; | |
| 210 |
1/2✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
|
2 | if (dY.rows != z.rows || dY.cols != C) { |
| 211 | ✗ | fail("complex_abs_backward", "dY must be the real (R, C) magnitude grad"); | |
| 212 | } | ||
| 213 |
3/4✓ Branch 0 taken 1 time.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
|
2 | if (dZ.rows != z.rows || dZ.cols != z.cols) dZ.resize(z.rows, z.cols); |
| 214 |
1/2✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
|
2 | if (z.size() == 0) return; |
| 215 | 2 | const float* zp = z.host_f32(); | |
| 216 | 2 | const float* gp = dY.host_f32(); | |
| 217 | 2 | float* dzp = dZ.host_f32_mut(); | |
| 218 |
2/2✓ Branch 0 taken 5 times.
✓ Branch 1 taken 2 times.
|
7 | for (int r = 0; r < z.rows; ++r) { |
| 219 | 5 | const float* zr = zp + static_cast<std::size_t>(r) * z.cols; | |
| 220 | 5 | const float* gr = gp + static_cast<std::size_t>(r) * C; | |
| 221 | 5 | float* dzr = dzp + static_cast<std::size_t>(r) * z.cols; | |
| 222 |
2/2✓ Branch 0 taken 29 times.
✓ Branch 1 taken 5 times.
|
34 | for (int c = 0; c < C; ++c) { |
| 223 | 29 | const float re = zr[2 * c], im = zr[2 * c + 1]; | |
| 224 | 29 | const float mag = std::sqrt(re * re + im * im); | |
| 225 |
1/2✓ Branch 0 taken 29 times.
✗ Branch 1 not taken.
|
29 | if (mag > 0.0f) { |
| 226 | 29 | const float inv = gr[c] / mag; | |
| 227 | 29 | dzr[2 * c] = re * inv; | |
| 228 | 29 | dzr[2 * c + 1] = im * inv; | |
| 229 | 29 | } else { | |
| 230 | ✗ | dzr[2 * c] = 0.0f; | |
| 231 | ✗ | dzr[2 * c + 1] = 0.0f; | |
| 232 | } | ||
| 233 | 29 | } | |
| 234 | 5 | } | |
| 235 | 2 | } | |
| 236 | |||
| 237 | // y = atan2(im, re), the phase angle per complex bin, in radians (-pi, pi]. | ||
| 238 | // Input z is interleaved-complex (R, 2*C); output y is REAL (R, C). No | ||
| 239 | // backward — phase is rarely used in a differentiable loss and atan2 is | ||
| 240 | // non-differentiable at the origin; add a backward later if a consumer needs | ||
| 241 | // one. | ||
| 242 | 2 | void complex_angle(const ::brotensor::Tensor& z, ::brotensor::Tensor& y) { | |
| 243 | 2 | require_fp32_host("complex_angle", z, "z"); | |
| 244 |
1/2✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
|
2 | if (z.cols % 2 != 0) { |
| 245 | ✗ | fail("complex_angle", "z.cols must be even (interleaved [re,im] layout)"); | |
| 246 | } | ||
| 247 | 2 | const int C = z.cols / 2; | |
| 248 |
3/4✓ Branch 0 taken 1 time.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
|
2 | if (y.rows != z.rows || y.cols != C) y.resize(z.rows, C); |
| 249 |
1/2✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
|
2 | if (z.size() == 0) return; |
| 250 | 2 | const float* zp = z.host_f32(); | |
| 251 | 2 | float* yp = y.host_f32_mut(); | |
| 252 |
2/2✓ Branch 0 taken 6 times.
✓ Branch 1 taken 2 times.
|
8 | for (int r = 0; r < z.rows; ++r) { |
| 253 | 6 | const float* zr = zp + static_cast<std::size_t>(r) * z.cols; | |
| 254 | 6 | float* yr = yp + static_cast<std::size_t>(r) * C; | |
| 255 |
2/2✓ Branch 0 taken 22 times.
✓ Branch 1 taken 6 times.
|
28 | for (int c = 0; c < C; ++c) { |
| 256 | 22 | yr[c] = std::atan2(zr[2 * c + 1], zr[2 * c]); | |
| 257 | 22 | } | |
| 258 | 6 | } | |
| 259 | 2 | } | |
| 260 | |||
| 261 | // y = mag * exp(i*phase) — build a complex tensor from polar components. | ||
| 262 | // mag and phase are REAL (R, C); output y is interleaved-complex (R, 2*C). | ||
| 263 | // y.re = mag * cos(phase), y.im = mag * sin(phase). | ||
| 264 | // Inverse of (complex_abs, complex_angle) taken together. | ||
| 265 | 2 | void complex_from_polar(const ::brotensor::Tensor& mag, | |
| 266 | const ::brotensor::Tensor& phase, | ||
| 267 | ::brotensor::Tensor& y) { | ||
| 268 | 2 | require_fp32_host("complex_from_polar", mag, "mag"); | |
| 269 | 2 | require_fp32_host("complex_from_polar", phase, "phase"); | |
| 270 |
1/2✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
|
2 | if (mag.rows != phase.rows || mag.cols != phase.cols) { |
| 271 | ✗ | fail("complex_from_polar", "mag and phase must have identical shape"); | |
| 272 | } | ||
| 273 | 2 | const int C = mag.cols; | |
| 274 |
3/4✓ Branch 0 taken 1 time.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
|
2 | if (y.rows != mag.rows || y.cols != 2 * C) y.resize(mag.rows, 2 * C); |
| 275 |
1/2✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
|
2 | if (mag.size() == 0) return; |
| 276 | 2 | const float* mp = mag.host_f32(); | |
| 277 | 2 | const float* pp = phase.host_f32(); | |
| 278 | 2 | float* yp = y.host_f32_mut(); | |
| 279 |
2/2✓ Branch 0 taken 5 times.
✓ Branch 1 taken 2 times.
|
7 | for (int r = 0; r < mag.rows; ++r) { |
| 280 | 5 | const float* mr = mp + static_cast<std::size_t>(r) * C; | |
| 281 | 5 | const float* pr = pp + static_cast<std::size_t>(r) * C; | |
| 282 | 5 | float* yr = yp + static_cast<std::size_t>(r) * (2 * C); | |
| 283 |
2/2✓ Branch 0 taken 38 times.
✓ Branch 1 taken 5 times.
|
43 | for (int c = 0; c < C; ++c) { |
| 284 | 38 | yr[2 * c] = mr[c] * std::cos(pr[c]); | |
| 285 | 38 | yr[2 * c + 1] = mr[c] * std::sin(pr[c]); | |
| 286 | 38 | } | |
| 287 | 5 | } | |
| 288 | 2 | } | |
| 289 | |||
| 290 | // ════════════════════════════════════════════════════════════════════════════ | ||
| 291 | // Complex <-> complex FFT / IFFT | ||
| 292 | // ════════════════════════════════════════════════════════════════════════════ | ||
| 293 | |||
| 294 | // fft: forward DFT, one signal per row. x and y are interleaved-complex | ||
| 295 | // (R, 2*N). "backward" normalisation — the forward transform is unscaled. | ||
| 296 | 335 | void fft(const ::brotensor::Tensor& x, ::brotensor::Tensor& y) { | |
| 297 | 335 | require_fp32_host("fft", x, "x"); | |
| 298 |
1/2✓ Branch 0 taken 335 times.
✗ Branch 1 not taken.
|
335 | if (x.cols % 2 != 0) { |
| 299 | ✗ | fail("fft", "x.cols must be even (interleaved [re,im] layout)"); | |
| 300 | } | ||
| 301 |
3/4✓ Branch 0 taken 331 times.
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 331 times.
|
335 | if (y.rows != x.rows || y.cols != x.cols) y.resize(x.rows, x.cols); |
| 302 |
1/2✓ Branch 0 taken 335 times.
✗ Branch 1 not taken.
|
335 | if (x.size() == 0) return; |
| 303 | 335 | const float* xp = x.host_f32(); | |
| 304 | 335 | float* yp = y.host_f32_mut(); | |
| 305 | 335 | std::vector<Cd> in, out; | |
| 306 |
2/2✓ Branch 0 taken 340 times.
✓ Branch 1 taken 335 times.
|
675 | for (int r = 0; r < x.rows; ++r) { |
| 307 |
1/2✓ Branch 0 taken 340 times.
✗ Branch 1 not taken.
|
340 | load_complex_row(xp, r, x.cols, in); |
| 308 |
1/2✓ Branch 0 taken 340 times.
✗ Branch 1 not taken.
|
340 | dft_1d(in, out, -1); |
| 309 |
1/2✓ Branch 0 taken 340 times.
✗ Branch 1 not taken.
|
340 | store_complex_row(yp, r, x.cols, out); |
| 310 | 340 | } | |
| 311 | 335 | } | |
| 312 | |||
| 313 | // ifft: inverse DFT, one signal per row. x and y are interleaved-complex | ||
| 314 | // (R, 2*N). "backward" normalisation — the inverse transform is scaled by 1/N. | ||
| 315 | 21 | void ifft(const ::brotensor::Tensor& x, ::brotensor::Tensor& y) { | |
| 316 | 21 | require_fp32_host("ifft", x, "x"); | |
| 317 |
1/2✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
|
21 | if (x.cols % 2 != 0) { |
| 318 | ✗ | fail("ifft", "x.cols must be even (interleaved [re,im] layout)"); | |
| 319 | } | ||
| 320 |
3/4✓ Branch 0 taken 19 times.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 19 times.
|
21 | if (y.rows != x.rows || y.cols != x.cols) y.resize(x.rows, x.cols); |
| 321 |
1/2✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
|
21 | if (x.size() == 0) return; |
| 322 | 21 | const int N = x.cols / 2; | |
| 323 |
1/2✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
|
21 | const double inv = (N > 0) ? 1.0 / static_cast<double>(N) : 1.0; |
| 324 | 21 | const float* xp = x.host_f32(); | |
| 325 | 21 | float* yp = y.host_f32_mut(); | |
| 326 | 21 | std::vector<Cd> in, out; | |
| 327 |
2/2✓ Branch 0 taken 24 times.
✓ Branch 1 taken 21 times.
|
45 | for (int r = 0; r < x.rows; ++r) { |
| 328 |
1/2✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
|
24 | load_complex_row(xp, r, x.cols, in); |
| 329 |
1/2✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
|
24 | dft_1d(in, out, +1); |
| 330 |
2/2✓ Branch 0 taken 1091 times.
✓ Branch 1 taken 24 times.
|
1115 | for (auto& v : out) { v.re *= inv; v.im *= inv; } |
| 331 |
1/2✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
|
24 | store_complex_row(yp, r, x.cols, out); |
| 332 | 24 | } | |
| 333 | 21 | } | |
| 334 | |||
| 335 | // ════════════════════════════════════════════════════════════════════════════ | ||
| 336 | // Real <-> complex rfft / irfft | ||
| 337 | // ════════════════════════════════════════════════════════════════════════════ | ||
| 338 | |||
| 339 | // rfft: real-input FFT. x is REAL (R, L); y is the non-redundant | ||
| 340 | // half-spectrum, interleaved-complex (R, 2*(L/2+1)). "backward" normalisation | ||
| 341 | // (forward unscaled). Only bins 0 .. L/2 are stored — the remaining bins are | ||
| 342 | // the conjugates of these by Hermitian symmetry of a real signal. | ||
| 343 | 332 | void rfft(const ::brotensor::Tensor& x, ::brotensor::Tensor& y) { | |
| 344 | 332 | require_fp32_host("rfft", x, "x"); | |
| 345 | 332 | const int L = x.cols; | |
| 346 |
1/2✓ Branch 0 taken 332 times.
✗ Branch 1 not taken.
|
332 | if (L == 0) { |
| 347 | ✗ | fail("rfft", "signal length L (x.cols) must be >= 1"); | |
| 348 | } | ||
| 349 | 332 | const int C = L / 2 + 1; | |
| 350 |
3/4✓ Branch 0 taken 329 times.
✓ Branch 1 taken 3 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 329 times.
|
332 | if (y.rows != x.rows || y.cols != 2 * C) y.resize(x.rows, 2 * C); |
| 351 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 332 times.
|
332 | if (x.size() == 0) return; |
| 352 | 332 | const float* xp = x.host_f32(); | |
| 353 | 332 | float* yp = y.host_f32_mut(); | |
| 354 | 332 | std::vector<Cd> in(static_cast<std::size_t>(L)), out; | |
| 355 |
2/2✓ Branch 0 taken 332 times.
✓ Branch 1 taken 653 times.
|
985 | for (int r = 0; r < x.rows; ++r) { |
| 356 | 653 | const float* xr = xp + static_cast<std::size_t>(r) * L; | |
| 357 |
2/2✓ Branch 0 taken 25723 times.
✓ Branch 1 taken 653 times.
|
26376 | for (int n = 0; n < L; ++n) { |
| 358 | 25723 | in[static_cast<std::size_t>(n)] = {static_cast<double>(xr[n]), 0.0}; | |
| 359 | 25723 | } | |
| 360 |
1/2✓ Branch 0 taken 653 times.
✗ Branch 1 not taken.
|
653 | dft_1d(in, out, -1); |
| 361 | 653 | float* yr = yp + static_cast<std::size_t>(r) * (2 * C); | |
| 362 |
2/2✓ Branch 0 taken 13269 times.
✓ Branch 1 taken 653 times.
|
13922 | for (int c = 0; c < C; ++c) { |
| 363 | 13269 | yr[2 * c] = static_cast<float>(out[static_cast<std::size_t>(c)].re); | |
| 364 | 13269 | yr[2 * c + 1] = static_cast<float>(out[static_cast<std::size_t>(c)].im); | |
| 365 | 13269 | } | |
| 366 | 653 | } | |
| 367 | 332 | } | |
| 368 | |||
| 369 | // irfft: inverse real FFT. x is a half-spectrum, interleaved-complex | ||
| 370 | // (R, 2*(L/2+1)); y is the reconstructed REAL signal (R, L). "backward" | ||
| 371 | // normalisation (scaled by 1/L). The full Hermitian-symmetric spectrum is | ||
| 372 | // rebuilt from the stored half (bin L-k = conj(bin k)) before the inverse | ||
| 373 | // transform; the DC and (for even L) Nyquist bins have no conjugate partner. | ||
| 374 | // | ||
| 375 | // `L` must be passed explicitly: a half-spectrum with C bins is ambiguous | ||
| 376 | // between L = 2*(C-1) (even) and L = 2*C-1 (odd). | ||
| 377 | 364 | void irfft(const ::brotensor::Tensor& x, int L, ::brotensor::Tensor& y) { | |
| 378 | 364 | require_fp32_host("irfft", x, "x"); | |
| 379 |
1/2✓ Branch 0 taken 364 times.
✗ Branch 1 not taken.
|
364 | if (x.cols % 2 != 0) { |
| 380 | ✗ | fail("irfft", "x.cols must be even (interleaved [re,im] layout)"); | |
| 381 | } | ||
| 382 | 364 | const int C = x.cols / 2; | |
| 383 |
1/2✓ Branch 0 taken 364 times.
✗ Branch 1 not taken.
|
364 | if (L <= 0) { |
| 384 | ✗ | fail("irfft", "output length L must be >= 1"); | |
| 385 | } | ||
| 386 |
1/2✓ Branch 0 taken 364 times.
✗ Branch 1 not taken.
|
364 | if (C != L / 2 + 1) { |
| 387 | ✗ | fail("irfft", "half-spectrum bin count must equal L/2+1"); | |
| 388 | } | ||
| 389 |
3/4✓ Branch 0 taken 361 times.
✓ Branch 1 taken 3 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 361 times.
|
364 | if (y.rows != x.rows || y.cols != L) y.resize(x.rows, L); |
| 390 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 364 times.
|
364 | if (x.size() == 0) return; |
| 391 | 364 | const double inv = 1.0 / static_cast<double>(L); | |
| 392 | 364 | const float* xp = x.host_f32(); | |
| 393 | 364 | float* yp = y.host_f32_mut(); | |
| 394 | 364 | std::vector<Cd> full(static_cast<std::size_t>(L)), out; | |
| 395 |
2/2✓ Branch 0 taken 364 times.
✓ Branch 1 taken 717 times.
|
1081 | for (int r = 0; r < x.rows; ++r) { |
| 396 | 717 | const float* xr = xp + static_cast<std::size_t>(r) * x.cols; | |
| 397 | // Stored half-spectrum into bins 0 .. C-1. | ||
| 398 |
2/2✓ Branch 0 taken 13709 times.
✓ Branch 1 taken 717 times.
|
14426 | for (int c = 0; c < C; ++c) { |
| 399 | 27418 | full[static_cast<std::size_t>(c)] = {static_cast<double>(xr[2 * c]), | |
| 400 | 13709 | static_cast<double>(xr[2 * c + 1])}; | |
| 401 | 13709 | } | |
| 402 | // Hermitian mirror: bin L-k = conj(bin k) for k = 1 .. L-C. | ||
| 403 |
2/2✓ Branch 0 taken 12782 times.
✓ Branch 1 taken 717 times.
|
13499 | for (int k = 1; k < L - C + 1; ++k) { |
| 404 | 12782 | const Cd c = full[static_cast<std::size_t>(k)]; | |
| 405 | 12782 | full[static_cast<std::size_t>(L - k)] = {c.re, -c.im}; | |
| 406 | 12782 | } | |
| 407 |
1/2✓ Branch 0 taken 717 times.
✗ Branch 1 not taken.
|
717 | dft_1d(full, out, +1); |
| 408 | 717 | float* yr = yp + static_cast<std::size_t>(r) * L; | |
| 409 |
2/2✓ Branch 0 taken 26491 times.
✓ Branch 1 taken 717 times.
|
27208 | for (int n = 0; n < L; ++n) { |
| 410 | 26491 | yr[n] = static_cast<float>(out[static_cast<std::size_t>(n)].re * inv); | |
| 411 | 26491 | } | |
| 412 | 717 | } | |
| 413 | 364 | } | |
| 414 | |||
| 415 | // ── rfft_backward — adjoint of rfft ──────────────────────────────────────── | ||
| 416 | // | ||
| 417 | // rfft maps a real length-L signal x to its non-redundant half-spectrum Y | ||
| 418 | // (C = L/2+1 complex bins): Y[k] = sum_n x[n] * exp(-i 2*pi k n / L). This is | ||
| 419 | // just the first C rows of the length-L DFT matrix applied to a real vector. | ||
| 420 | // | ||
| 421 | // For a real-valued loss formed directly on the spectrum, | ||
| 422 | // loss = sum_{k=0}^{C-1} ( dY[k].re * Y[k].re + dY[k].im * Y[k].im ), | ||
| 423 | // the gradient w.r.t. the real signal is the plain conjugate transpose of | ||
| 424 | // that truncated DFT matrix — NO conjugate-pair weighting (the doubling lives | ||
| 425 | // in irfft / irfft_backward, which fold the Hermitian half back; rfft does no | ||
| 426 | // folding so its adjoint does none either): | ||
| 427 | // | ||
| 428 | // dX[n] = sum_{k=0}^{C-1} ( dY[k].re * cos(2*pi k n / L) | ||
| 429 | // - dY[k].im * sin(2*pi k n / L) ) | ||
| 430 | // = Re( sum_{k=0}^{C-1} dY[k] * exp(+i 2*pi k n / L) ). | ||
| 431 | // | ||
| 432 | // Computed by zero-padding dY to length L and running an inverse-sign | ||
| 433 | // (sign = +1) unscaled DFT, then taking the real part. dX is *overwritten*. | ||
| 434 | 7 | void rfft_backward(const ::brotensor::Tensor& dY, int L, | |
| 435 | ::brotensor::Tensor& dX) { | ||
| 436 | 7 | require_fp32_host("rfft_backward", dY, "dY"); | |
| 437 |
1/2✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
|
7 | if (dY.cols % 2 != 0) { |
| 438 | ✗ | fail("rfft_backward", "dY.cols must be even (interleaved [re,im] layout)"); | |
| 439 | } | ||
| 440 | 7 | const int C = dY.cols / 2; | |
| 441 |
1/2✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
|
7 | if (L <= 0) { |
| 442 | ✗ | fail("rfft_backward", "signal length L must be >= 1"); | |
| 443 | } | ||
| 444 |
1/2✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
|
7 | if (C != L / 2 + 1) { |
| 445 | ✗ | fail("rfft_backward", "dY bin count must equal L/2+1"); | |
| 446 | } | ||
| 447 |
3/4✓ Branch 0 taken 5 times.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 5 times.
|
7 | if (dX.rows != dY.rows || dX.cols != L) dX.resize(dY.rows, L); |
| 448 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 7 times.
|
7 | if (dY.size() == 0) return; |
| 449 | 7 | const float* gp = dY.host_f32(); | |
| 450 | 7 | float* dxp = dX.host_f32_mut(); | |
| 451 | 7 | std::vector<Cd> spec(static_cast<std::size_t>(L)), out; | |
| 452 |
2/2✓ Branch 0 taken 7 times.
✓ Branch 1 taken 15 times.
|
22 | for (int r = 0; r < dY.rows; ++r) { |
| 453 | 15 | const float* gr = gp + static_cast<std::size_t>(r) * dY.cols; | |
| 454 | // Zero-padded length-L spectrum: bins 0..C-1 carry dY[k], rest zero. | ||
| 455 |
2/2✓ Branch 0 taken 198 times.
✓ Branch 1 taken 15 times.
|
213 | for (int k = 0; k < L; ++k) spec[static_cast<std::size_t>(k)] = Cd{}; |
| 456 |
2/2✓ Branch 0 taken 111 times.
✓ Branch 1 taken 15 times.
|
126 | for (int k = 0; k < C; ++k) { |
| 457 | 111 | spec[static_cast<std::size_t>(k)] = | |
| 458 | 222 | {static_cast<double>(gr[2 * k]), | |
| 459 | 111 | static_cast<double>(gr[2 * k + 1])}; | |
| 460 | 111 | } | |
| 461 | // dX[n] = Re( sum_k spec[k] * exp(+i 2pi k n / L) ). | ||
| 462 |
1/2✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
|
15 | dft_1d(spec, out, +1); |
| 463 | 15 | float* dxr = dxp + static_cast<std::size_t>(r) * L; | |
| 464 |
2/2✓ Branch 0 taken 198 times.
✓ Branch 1 taken 15 times.
|
213 | for (int n = 0; n < L; ++n) { |
| 465 | 198 | dxr[n] = static_cast<float>(out[static_cast<std::size_t>(n)].re); | |
| 466 | 198 | } | |
| 467 | 15 | } | |
| 468 | 7 | } | |
| 469 | |||
| 470 | // ── irfft_backward — adjoint of irfft ────────────────────────────────────── | ||
| 471 | // | ||
| 472 | // irfft maps a half-spectrum X (C = L/2+1 complex bins) to a real signal y | ||
| 473 | // of length L, with the 1/L inverse scaling. The adjoint maps the upstream | ||
| 474 | // real gradient dY (real (R, L)) back to the half-spectrum gradient dX | ||
| 475 | // (interleaved-complex (R, 2*C)). | ||
| 476 | // | ||
| 477 | // Forward (per output sample n): | ||
| 478 | // y[n] = (1/L) * [ X[0].re | ||
| 479 | // + sum_{k=1}^{C-1} 2 * Re( X[k] * exp(i 2pi k n / L) ) | ||
| 480 | // - (L even ? X[L/2].re : 0) ] (Nyquist counted once) | ||
| 481 | // so the adjoint, per stored bin k: | ||
| 482 | // dX[k].re = (s_k / L) * sum_n dY[n] * cos(2pi k n / L) | ||
| 483 | // dX[k].im = -(s_k / L) * sum_n dY[n] * sin(2pi k n / L) | ||
| 484 | // s_k = 1 for k = 0 and k = L/2 (L even); 2 otherwise. | ||
| 485 | // | ||
| 486 | // dX is *overwritten*. This op is the transpose of rfft_backward. | ||
| 487 | 7 | void irfft_backward(const ::brotensor::Tensor& dY, | |
| 488 | ::brotensor::Tensor& dX) { | ||
| 489 | 7 | require_fp32_host("irfft_backward", dY, "dY"); | |
| 490 | 7 | const int L = dY.cols; | |
| 491 |
1/2✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
|
7 | if (L == 0) { |
| 492 | ✗ | fail("irfft_backward", "dY length L (dY.cols) must be >= 1"); | |
| 493 | } | ||
| 494 | 7 | const int C = L / 2 + 1; | |
| 495 |
3/4✓ Branch 0 taken 5 times.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 5 times.
|
7 | if (dX.rows != dY.rows || dX.cols != 2 * C) dX.resize(dY.rows, 2 * C); |
| 496 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 7 times.
|
7 | if (dY.size() == 0) return; |
| 497 | 7 | const bool even = (L % 2 == 0); | |
| 498 | 7 | const double invL = 1.0 / static_cast<double>(L); | |
| 499 | 7 | const float* gp = dY.host_f32(); | |
| 500 | 7 | float* dxp = dX.host_f32_mut(); | |
| 501 | // dX[k] = (s_k / L) * conj( forward_DFT(dY)[k] ), for k = 0 .. C-1. | ||
| 502 | 7 | std::vector<Cd> in(static_cast<std::size_t>(L)), spec; | |
| 503 |
2/2✓ Branch 0 taken 7 times.
✓ Branch 1 taken 15 times.
|
22 | for (int r = 0; r < dY.rows; ++r) { |
| 504 | 15 | const float* gr = gp + static_cast<std::size_t>(r) * L; | |
| 505 |
2/2✓ Branch 0 taken 198 times.
✓ Branch 1 taken 15 times.
|
213 | for (int n = 0; n < L; ++n) { |
| 506 | 198 | in[static_cast<std::size_t>(n)] = {static_cast<double>(gr[n]), 0.0}; | |
| 507 | 198 | } | |
| 508 |
1/2✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
|
15 | dft_1d(in, spec, -1); // spec[k] = sum_n dY[n] exp(-i 2pi k n / L) |
| 509 | 15 | float* dxr = dxp + static_cast<std::size_t>(r) * (2 * C); | |
| 510 |
2/2✓ Branch 0 taken 111 times.
✓ Branch 1 taken 15 times.
|
126 | for (int k = 0; k < C; ++k) { |
| 511 | 111 | double s = 2.0; | |
| 512 |
2/2✓ Branch 0 taken 96 times.
✓ Branch 1 taken 15 times.
|
111 | if (k == 0) s = 1.0; |
| 513 |
4/4✓ Branch 0 taken 39 times.
✓ Branch 1 taken 72 times.
✓ Branch 2 taken 30 times.
✓ Branch 3 taken 9 times.
|
111 | if (even && k == L / 2) s = 1.0; |
| 514 | 111 | const double scale = s * invL; | |
| 515 | // cos sum = spec[k].re; -sin sum = spec[k].im (already conj-form). | ||
| 516 | 111 | dxr[2 * k] = static_cast<float>(scale * spec[static_cast<std::size_t>(k)].re); | |
| 517 | 111 | dxr[2 * k + 1] = static_cast<float>(scale * spec[static_cast<std::size_t>(k)].im); | |
| 518 | 111 | } | |
| 519 | 15 | } | |
| 520 | 7 | } | |
| 521 | |||
| 522 | } // namespace brotensor::detail::cpu | ||
| 523 |