From cbcd0eb2ebcb7aebeb4fd8544674553c6cf61fe7 Mon Sep 17 00:00:00 2001 From: ruv Date: Thu, 18 Jun 2026 00:34:20 -0400 Subject: [PATCH] perf(photonlayer-core): fold Fraunhofer fftshift into checkerboard premult + precompute FFT twiddle tables MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit OPT-A (bit-identical): replace `fft_2d + fftshift_2d` in both Fraunhofer paths (free `fraunhofer()` and `Propagator::propagate_into`) with a ±1 checkerboard premultiply `(-1)^(x+y)` before the transform. By the DFT shift theorem, FFT of the premultiplied input equals fftshift of the FFT, eliminating the fftshift's full-buffer alloc + quadrant copy. True negate (`Complex::ZERO - c`) is exact ±1.0 -> element-for-element identical to the old sequence (new test `checkerboard_premult_equals_fft_then_fftshift`). OPT-B (deliberately changes bits, determinism gain): precompute a per- dimension `TwiddleTable` (`exp(sign·2π·j/n)` for j in 0..n/2) and INDEX it by stride per butterfly instead of accumulating `w *= wlen`. Kills the f32 drift the accumulation injected and recomputes angles once per 2D FFT instead of per row/column. Proven: FFT is bit-for-bit reproducible across runs, and max-abs error vs an f64 reference DFT does NOT increase (it decreases — drift removed). No hardcoded golden hashes/values in the repo to update; re-run-determinism tests stay valid by construction. Measured (release, 64x64 x3000, --ignored --nocapture): fraunhofer OPT-A+B: old(fft+fftshift,accum-twiddle)=210.5ms -> new(checkerboard+table)=116.1ms = 1.81x, max_diff_vs_old=5.7e-6 (f32 noise). M1 cached-propagator benchmark still 2.00x and bit-identical. All 27 photonlayer-core unit tests + propagation bit-identical gate green; photonlayer-ruvector / photonlayer-bench / photonlayer-cli build and tests green. Determinism invariant preserved (scalar cos/sin FFT, no FMA/SIMD/RFFT). Co-Authored-By: claude-flow --- crates/photonlayer-core/src/fft.rs | 227 +++++++++++++++++- crates/photonlayer-core/src/propagate.rs | 11 +- .../tests/propagation_speedup.rs | 156 +++++++++++- 3 files changed, 380 insertions(+), 14 deletions(-) diff --git a/crates/photonlayer-core/src/fft.rs b/crates/photonlayer-core/src/fft.rs index 8771c4052..b5559b282 100644 --- a/crates/photonlayer-core/src/fft.rs +++ b/crates/photonlayer-core/src/fft.rs @@ -14,14 +14,62 @@ pub fn is_pow2(n: usize) -> bool { n != 0 && (n & (n - 1)) == 0 } +/// Precomputed twiddle factors for a length-`n` FFT of a fixed direction. +/// +/// Holds `tw[j] = exp(sign · 2π · j / n)` for `j in 0..n/2`. The stage-`len` +/// butterfly twiddle for index `k` is `tw[k * (n / len)]`, so every factor is +/// read straight from the table by index — never accumulated with repeated +/// complex multiplies. This both removes the per-butterfly `w *= wlen` cost and +/// eliminates the f32 drift that accumulation injects (a determinism *gain*: +/// the angles are computed once at full `cos/sin` precision). +#[derive(Clone)] +pub struct TwiddleTable { + n: usize, + inverse: bool, + tw: Vec, +} + +impl TwiddleTable { + /// Build the table for a length-`n` (power-of-two) transform. + /// + /// # Panics + /// Panics if `n` is not a power of two. + pub fn new(n: usize, inverse: bool) -> Self { + assert!(is_pow2(n), "FFT length must be a power of two, got {n}"); + let sign = if inverse { 1.0 } else { -1.0 }; + let half = n / 2; // 0 when n == 1; table is unused at that size. + let mut tw = Vec::with_capacity(half); + let scale = sign * 2.0 * PI / n as f32; + for j in 0..half { + // Index the angle directly: no `w *= wlen` accumulation, no drift. + tw.push(Complex::from_phase(j as f32 * scale)); + } + Self { n, inverse, tw } + } +} + /// In-place 1D FFT. `inverse = true` computes the inverse transform and /// applies the `1/N` normalization so that `ifft(fft(x)) == x`. /// +/// Builds a one-shot [`TwiddleTable`]; callers transforming many equal-length +/// rows/columns should build the table once and use [`fft_1d_with`]. +/// /// # Panics /// Panics if `data.len()` is not a power of two. pub fn fft_1d(data: &mut [Complex], inverse: bool) { + let table = TwiddleTable::new(data.len().max(1), inverse); + fft_1d_with(data, &table); +} + +/// In-place 1D FFT using a precomputed [`TwiddleTable`] (must match the buffer +/// length and direction). +/// +/// # Panics +/// Panics if `data.len()` is not a power of two or does not match `table.n`. +pub fn fft_1d_with(data: &mut [Complex], table: &TwiddleTable) { let n = data.len(); assert!(is_pow2(n), "FFT length must be a power of two, got {n}"); + assert_eq!(n, table.n, "twiddle table length mismatch"); if n == 1 { return; } @@ -40,29 +88,27 @@ pub fn fft_1d(data: &mut [Complex], inverse: bool) { } } - // Danielson–Lanczos butterflies. - let sign = if inverse { 1.0 } else { -1.0 }; + // Danielson–Lanczos butterflies. Index the twiddle table by stride instead + // of accumulating `w *= wlen` — same math, no per-stage drift. let mut len = 2; while len <= n { - let ang = sign * 2.0 * PI / len as f32; - let wlen = Complex::from_phase(ang); let half = len / 2; + let stride = n / len; // table[k * stride] == exp(sign · 2π · k / len) let mut i = 0; while i < n { - let mut w = Complex::ONE; for k in 0..half { + let w = table.tw[k * stride]; let u = data[i + k]; let v = data[i + k + half] * w; data[i + k] = u + v; data[i + k + half] = u - v; - w = w * wlen; } i += len; } len <<= 1; } - if inverse { + if table.inverse { let inv = 1.0 / n as f32; for c in data.iter_mut() { *c = c.scale(inv); @@ -78,25 +124,53 @@ pub fn fft_2d(data: &mut [Complex], width: usize, height: usize, inverse: bool) assert_eq!(data.len(), width * height, "buffer size mismatch"); assert!(is_pow2(width) && is_pow2(height), "dims must be power of two"); - // Rows. + // Build each dimension's twiddle table once and reuse it across every + // row / column transform (OPT-B) — angles are computed a single time. + let row_tw = TwiddleTable::new(width, inverse); for r in 0..height { let row = &mut data[r * width..(r + 1) * width]; - fft_1d(row, inverse); + fft_1d_with(row, &row_tw); } // Columns (gather/scatter to keep the 1D kernel contiguous). + let col_tw = TwiddleTable::new(height, inverse); let mut col = vec![Complex::ZERO; height]; for c in 0..width { for r in 0..height { col[r] = data[r * width + c]; } - fft_1d(&mut col, inverse); + fft_1d_with(&mut col, &col_tw); for r in 0..height { data[r * width + c] = col[r]; } } } +/// Checkerboard premultiply: negate every sample at an odd `(row + col)`. +/// +/// By the DFT shift theorem, modulating the input by `(-1)^(x+y)` shifts the +/// transform output by `(N/2, M/2)` — i.e. forward-FFT of the premultiplied +/// buffer equals `fftshift_2d` of the forward-FFT of the original. This lets a +/// Fraunhofer path do `premult → fft_2d` instead of `fft_2d → fftshift_2d`, +/// avoiding the full-buffer allocation + quadrant copy in [`fftshift_2d`]. +/// +/// The negation is exact (`{-re, -im}`), so the substitution is bit-identical +/// to the fft-then-fftshift sequence on every platform. +pub fn checkerboard_premultiply(data: &mut [Complex], width: usize, height: usize) { + debug_assert_eq!(data.len(), width * height, "buffer size mismatch"); + for row in 0..height { + // First column negated when the row index is odd; flips every column. + let mut neg = row & 1 == 1; + let base = row * width; + for c in &mut data[base..base + width] { + if neg { + *c = Complex::ZERO - *c; + } + neg = !neg; + } + } +} + /// 2D fftshift: swaps quadrants so the zero-frequency component moves to the /// center. `width` and `height` must be even (always true for power-of-two). pub fn fftshift_2d(data: &mut [Complex], width: usize, height: usize) { @@ -140,6 +214,139 @@ mod tests { } } + #[test] + fn checkerboard_premult_equals_fft_then_fftshift() { + // OPT-A correctness gate: `premult → fft` must be ELEMENT-FOR-ELEMENT + // identical to `fft → fftshift` (shift theorem, exact ±1 negation). + for &(w, h) in &[(8usize, 8usize), (16, 4), (4, 16), (32, 32), (2, 2)] { + let src: Vec = (0..w * h) + .map(|i| Complex::new((i % 7) as f32 - 3.0, (i % 5) as f32 - 2.0)) + .collect(); + + // Old path: forward FFT, then quadrant fftshift. + let mut old = src.clone(); + fft_2d(&mut old, w, h, false); + fftshift_2d(&mut old, w, h); + + // New path: checkerboard premultiply, then forward FFT. + let mut new = src.clone(); + checkerboard_premultiply(&mut new, w, h); + fft_2d(&mut new, w, h, false); + + assert_eq!(new, old, "checkerboard path differs at {w}x{h}"); + } + } + + #[test] + fn checkerboard_is_exact_pm_one() { + // Negation must be exact ±1.0 (true negate), not a multiply by -1.0f32 + // that could differ; applying it twice restores the original bits. + let src: Vec = (0..16) + .map(|i| Complex::new(i as f32 * 0.123 - 1.0, i as f32 * -0.071)) + .collect(); + let mut x = src.clone(); + checkerboard_premultiply(&mut x, 4, 4); + checkerboard_premultiply(&mut x, 4, 4); + assert_eq!(x, src, "double checkerboard must be identity (bit-exact)"); + } + + /// Reference forward DFT in f64 (no FFT factorization, no f32 accumulation) + /// — the ground truth OPT-B's twiddle tables are measured against. + fn dft_1d_ref_f64(x: &[Complex]) -> Vec<(f64, f64)> { + let n = x.len(); + let mut out = vec![(0.0f64, 0.0f64); n]; + for (k, slot) in out.iter_mut().enumerate() { + let (mut re, mut im) = (0.0f64, 0.0f64); + for (j, c) in x.iter().enumerate() { + let ang = -2.0 * std::f64::consts::PI * (k * j) as f64 / n as f64; + let (s, co) = ang.sin_cos(); + re += c.re as f64 * co - c.im as f64 * s; + im += c.re as f64 * s + c.im as f64 * co; + } + *slot = (re, im); + } + out + } + + #[test] + fn fft_1d_is_deterministic_bitexact() { + // OPT-B determinism gate: identical input -> identical output bytes. + let src: Vec = (0..64) + .map(|i| Complex::new((i as f32 * 0.37).sin(), (i as f32 * 0.11).cos())) + .collect(); + let mut a = src.clone(); + let mut b = src.clone(); + fft_1d(&mut a, false); + fft_1d(&mut b, false); + assert_eq!(a, b, "FFT must be bit-for-bit reproducible across runs"); + } + + #[test] + fn twiddle_table_error_does_not_increase() { + // OPT-B accuracy gate: indexing a precomputed table must not worsen + // max-abs error vs an f64 reference DFT — drift removal should help. + let n = 256; + let src: Vec = (0..n) + .map(|i| Complex::new((i as f32 * 0.21).sin(), (i as f32 * 0.05).cos())) + .collect(); + let reference = dft_1d_ref_f64(&src); + + // New (table-indexed) path. + let mut new = src.clone(); + fft_1d(&mut new, false); + let err_new = new + .iter() + .zip(&reference) + .map(|(c, &(re, im))| ((c.re as f64 - re).abs()).max((c.im as f64 - im).abs())) + .fold(0.0f64, f64::max); + + // Old (accumulated `w *= wlen`) path, recomputed here for comparison. + let mut old = src.clone(); + let nn = old.len(); + { + let mut jj = 0usize; + for i in 1..nn { + let mut bit = nn >> 1; + while jj & bit != 0 { + jj ^= bit; + bit >>= 1; + } + jj ^= bit; + if i < jj { + old.swap(i, jj); + } + } + let mut len = 2; + while len <= nn { + let wlen = Complex::from_phase(-2.0 * PI / len as f32); + let half = len / 2; + let mut i = 0; + while i < nn { + let mut w = Complex::ONE; + for k in 0..half { + let u = old[i + k]; + let v = old[i + k + half] * w; + old[i + k] = u + v; + old[i + k + half] = u - v; + w = w * wlen; + } + i += len; + } + len <<= 1; + } + } + let err_old = old + .iter() + .zip(&reference) + .map(|(c, &(re, im))| ((c.re as f64 - re).abs()).max((c.im as f64 - im).abs())) + .fold(0.0f64, f64::max); + + assert!( + err_new <= err_old, + "table FFT error {err_new:e} must not exceed accumulated-twiddle error {err_old:e}" + ); + } + #[test] fn roundtrip_2d() { let (w, h) = (8, 4); diff --git a/crates/photonlayer-core/src/propagate.rs b/crates/photonlayer-core/src/propagate.rs index 319ddf787..73dba14e2 100644 --- a/crates/photonlayer-core/src/propagate.rs +++ b/crates/photonlayer-core/src/propagate.rs @@ -8,7 +8,7 @@ use crate::complex::Complex; use crate::config::{OpticalConfig, PropagationMode}; use crate::error::{PhotonError, Result}; -use crate::fft::{fft_2d, fftshift_2d, is_pow2}; +use crate::fft::{checkerboard_premultiply, fft_2d, is_pow2}; use crate::field::OpticalField; use core::f32::consts::PI; @@ -47,8 +47,11 @@ pub fn propagate(field: &OpticalField, config: &OpticalConfig) -> Result Result { let (w, h) = (field.width, field.height); let mut data = field.data.clone(); + // fftshift(FFT(x)) == FFT((-1)^(x+y) · x): premultiply by a ±1 checkerboard + // before the transform instead of shifting quadrants after it. Exact ±1.0 + // negation -> bit-identical to `fft_2d` + `fftshift_2d`, but no shift alloc. + checkerboard_premultiply(&mut data, w, h); fft_2d(&mut data, w, h, false); - fftshift_2d(&mut data, w, h); // Normalize so total power stays in a sane range for downstream metrics. let norm = 1.0 / (w as f32 * h as f32).sqrt(); for c in &mut data { @@ -181,8 +184,10 @@ impl Propagator { } match &self.kind { PropKind::Fraunhofer => { + // OPT-A: ±1 checkerboard premultiply folds the post-FFT fftshift + // into the input (shift theorem) — bit-identical, no shift alloc. + checkerboard_premultiply(data, w, h); fft_2d(data, w, h, false); - fftshift_2d(data, w, h); let norm = 1.0 / (w as f32 * h as f32).sqrt(); for c in data.iter_mut() { *c = c.scale(norm); diff --git a/crates/photonlayer-core/tests/propagation_speedup.rs b/crates/photonlayer-core/tests/propagation_speedup.rs index 0771df92b..b01976e03 100644 --- a/crates/photonlayer-core/tests/propagation_speedup.rs +++ b/crates/photonlayer-core/tests/propagation_speedup.rs @@ -7,9 +7,12 @@ use std::time::Instant; -use photonlayer_core::config::OpticalConfig; +use photonlayer_core::complex::Complex; +use photonlayer_core::config::{OpticalConfig, PropagationMode}; +use photonlayer_core::fft::{fftshift_2d, is_pow2}; use photonlayer_core::field::{InputImage, OpticalField}; use photonlayer_core::propagate::{propagate, Propagator}; +use std::f32::consts::PI; const N: usize = 64; // grid (learn-loop regime where H-recompute is a large fraction) const ITERS: usize = 3000; @@ -86,3 +89,154 @@ fn cached_propagator_is_faster() { "cached+in-place propagator must be >= 1.5x the naive path; got {speedup:.2}x" ); } + +// --------------------------------------------------------------------------- +// OPT-A + OPT-B benchmark: the new Fraunhofer path (±1 checkerboard premultiply +// that folds away `fftshift`, plus a table-indexed FFT that replaces the +// per-butterfly `w *= wlen` accumulation) vs a self-contained reimplementation +// of the OLD path (accumulated-twiddle 2D FFT, then `fftshift_2d`). The old +// path is rebuilt locally so the "before" number is real, not assumed. +// --------------------------------------------------------------------------- + +/// Old 1D FFT: accumulates `w *= wlen` per stage (the pre-OPT-B behavior). +fn old_fft_1d(data: &mut [Complex], inverse: bool) { + let n = data.len(); + assert!(is_pow2(n)); + if n == 1 { + return; + } + let mut j = 0usize; + for i in 1..n { + let mut bit = n >> 1; + while j & bit != 0 { + j ^= bit; + bit >>= 1; + } + j ^= bit; + if i < j { + data.swap(i, j); + } + } + let sign = if inverse { 1.0 } else { -1.0 }; + let mut len = 2; + while len <= n { + let wlen = Complex::from_phase(sign * 2.0 * PI / len as f32); + let half = len / 2; + let mut i = 0; + while i < n { + let mut w = Complex::ONE; + for k in 0..half { + let u = data[i + k]; + let v = data[i + k + half] * w; + data[i + k] = u + v; + data[i + k + half] = u - v; + w = w * wlen; + } + i += len; + } + len <<= 1; + } + if inverse { + let inv = 1.0 / n as f32; + for c in data.iter_mut() { + *c = c.scale(inv); + } + } +} + +/// Old 2D FFT: rebuilds `wlen` per row and per column (no shared table). +fn old_fft_2d(data: &mut [Complex], width: usize, height: usize, inverse: bool) { + for r in 0..height { + old_fft_1d(&mut data[r * width..(r + 1) * width], inverse); + } + let mut col = vec![Complex::ZERO; height]; + for c in 0..width { + for r in 0..height { + col[r] = data[r * width + c]; + } + old_fft_1d(&mut col, inverse); + for r in 0..height { + data[r * width + c] = col[r]; + } + } +} + +/// Old Fraunhofer: `old_fft_2d` then `fftshift_2d` then normalize. +fn old_fraunhofer_into(data: &mut [Complex], w: usize, h: usize) { + old_fft_2d(data, w, h, false); + fftshift_2d(data, w, h); + let norm = 1.0 / (w as f32 * h as f32).sqrt(); + for c in data.iter_mut() { + *c = c.scale(norm); + } +} + +#[test] +#[ignore = "timing benchmark — run with --release --ignored"] +fn fraunhofer_optab_is_faster() { + let field = test_field(N); + let mut config = OpticalConfig::demo(N, N); + config.propagation = PropagationMode::Fraunhofer; + let prop = Propagator::new(N, N, &config).unwrap(); + + // Correctness gate (always meaningful): the new in-place Fraunhofer path is + // bit-for-bit identical to the locally-rebuilt OLD fft+fftshift path? NO — + // OPT-B deliberately changes bits (drift removed). So assert they agree to a + // tight f32 tolerance, and assert the new path is internally deterministic. + let mut new_buf = field.data.clone(); + prop.propagate_into(&mut new_buf).unwrap(); + let mut new_buf2 = field.data.clone(); + prop.propagate_into(&mut new_buf2).unwrap(); + assert_eq!(new_buf, new_buf2, "new Fraunhofer path must be deterministic"); + + let mut old_buf = field.data.clone(); + old_fraunhofer_into(&mut old_buf, N, N); + let max_diff = new_buf + .iter() + .zip(&old_buf) + .map(|(a, b)| (a.re - b.re).abs().max((a.im - b.im).abs())) + .fold(0.0f32, f32::max); + assert!( + max_diff < 1e-3, + "OPT-B should only shift bits within f32 noise vs old path; got {max_diff:e}" + ); + + // Warm up. + for _ in 0..64 { + let mut b = field.data.clone(); + prop.propagate_into(&mut b).unwrap(); + } + + // Old path timing. + let t = Instant::now(); + let mut sink = 0.0f32; + let mut scratch = vec![Complex::ZERO; N * N]; + for _ in 0..ITERS { + scratch.copy_from_slice(&field.data); + old_fraunhofer_into(&mut scratch, N, N); + sink += scratch[0].re; + } + let old = t.elapsed().as_secs_f64(); + + // New path timing (OPT-A checkerboard + OPT-B twiddle table, in-place). + let t = Instant::now(); + for _ in 0..ITERS { + scratch.copy_from_slice(&field.data); + prop.propagate_into(&mut scratch).unwrap(); + sink += scratch[0].re; + } + let new = t.elapsed().as_secs_f64(); + std::hint::black_box(sink); + + let speedup = old / new; + eprintln!( + "fraunhofer OPT-A+B {N}x{N} x{ITERS}: old(fft+fftshift,accum-twiddle)={:.1}ms \ + new(checkerboard+table)={:.1}ms speedup={speedup:.2}x max_diff_vs_old={max_diff:e}", + old * 1e3, + new * 1e3 + ); + assert!( + speedup >= 1.0, + "OPT-A+B Fraunhofer path must not be slower than the old path; got {speedup:.2}x" + ); +}