perf(photonlayer-core): fold Fraunhofer fftshift into checkerboard premult + precompute FFT twiddle tables

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 <ruv@ruv.net>
This commit is contained in:
ruv 2026-06-18 00:34:20 -04:00
parent 00e401ac79
commit cbcd0eb2eb
3 changed files with 380 additions and 14 deletions

View file

@ -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<Complex>,
}
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) {
}
}
// DanielsonLanczos butterflies.
let sign = if inverse { 1.0 } else { -1.0 };
// DanielsonLanczos 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<Complex> = (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<Complex> = (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<Complex> = (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<Complex> = (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);

View file

@ -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<Optical
fn fraunhofer(field: &OpticalField) -> Result<OpticalField> {
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);

View file

@ -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"
);
}