Thomas Ales
11ff3b62e2
At pad 40 the old path materialised a ~9 GB padded spectrum per row,
which collapsed the chunk planner to one worker and one rfft call with
workers=1 — synthesis ran single-threaded, ~1 hour per angle on real
files.
The padded spectrum is never materialised now. Each block of 512
waveforms gets a coarse rfft at next_fast_len(2*spf); every coarse bin
within 0.7 of its row's max (plus the DC-adjacent window, which coarse
DC suppression would otherwise blind) is refined onto the exact n_fft
grid by a small complex gemm. The selected bin is bit-identical to the
full padded argmax — enforced by test_zoom_identity, a 25-seed fuzz
test over adversarial spectra, and a clean golden-hash diff against the
pre-rewrite baseline across pads {1,2,4,8,40}, masked/unmasked, bg
on/off, int8/int16, and both backends.
Blocks fan out over a persistent thread pool; pyFFTW runs through
per-thread FFTW_MEASURE builder plans with wisdom persisted to
~/.cache/sras-viewer, and threadpoolctl clamps BLAS under the pool.
compute_rf_image(exact=True) (or SRAS_FFT_EXACT=1) keeps the reference
padded path for audits.
tools/bench_fft.py measures: pad 40, 16 cores, 8192x2500 synthetic —
exact serial 717 wf/s -> zoom pool 25100 wf/s (35x, pyFFTW backend;
19x scipy), every variant verified equal to the reference.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
107 lines
4.1 KiB
Python
107 lines
4.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Benchmark the FFT peak-search path: exact vs zoom, serial vs pooled.
|
|
|
|
Reports wall time, waveforms/s, CPU utilization (utime+stime over wall, in
|
|
cores), and verifies every variant against the exact reference image.
|
|
|
|
Usage:
|
|
python tools/bench_fft.py # synthetic, pads 1/8/40
|
|
python tools/bench_fft.py --pads 40 --spf 2500 --rows 8 --frames 1024
|
|
python tools/bench_fft.py --real /path/big.sras --real-rows 32 --pads 40
|
|
"""
|
|
|
|
import argparse
|
|
import resource
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
|
|
import sras_compute as compute # noqa: E402
|
|
from sras_compute import compute_rf_image, set_fft_backend # noqa: E402
|
|
from sras_format import SrasFile # noqa: E402
|
|
import tools.make_test_sras as gen # noqa: E402
|
|
from tools.check_equivalence import row_slice # noqa: E402
|
|
|
|
|
|
def _timed(fn):
|
|
r0 = resource.getrusage(resource.RUSAGE_SELF)
|
|
t0 = time.perf_counter()
|
|
out = fn()
|
|
wall = time.perf_counter() - t0
|
|
r1 = resource.getrusage(resource.RUSAGE_SELF)
|
|
cpu = (r1.ru_utime - r0.ru_utime) + (r1.ru_stime - r0.ru_stime)
|
|
return out, wall, cpu / max(wall, 1e-9)
|
|
|
|
|
|
def bench(sras, pads, backends):
|
|
n_wf = sum(int(sras.n_rows[a]) * int(sras.n_frames[a])
|
|
for a in range(sras.n_angles))
|
|
spf = sras.samples_per_frame
|
|
print(f"{n_wf} waveforms x {spf} samples, {sras.n_angles} angle(s)")
|
|
print(f"{'pad':>4} {'backend':>8} {'variant':>16} {'wall':>9} "
|
|
f"{'wf/s':>10} {'util':>6} match")
|
|
|
|
for pad in pads:
|
|
n_fft = spf * pad if pad > 1 else None
|
|
for backend in backends:
|
|
set_fft_backend(backend)
|
|
|
|
def run(**kw):
|
|
imgs = [compute_rf_image(sras, a, dc_threshold_mv=None,
|
|
apply_bg_sub=True, n_fft=n_fft, **kw)
|
|
for a in range(sras.n_angles)]
|
|
return np.concatenate([i.ravel() for i in imgs])
|
|
|
|
ref, wall, util = _timed(lambda: run(exact=True))
|
|
rows = [("exact(serial)", ref, wall, util, True)]
|
|
for label, kw in (("zoom(serial)", dict(max_workers=1)),
|
|
("zoom(pool)", {})):
|
|
img, wall, util = _timed(lambda: run(**kw))
|
|
rows.append((label, img, wall, util, bool(np.array_equal(img, ref))))
|
|
for label, img, wall, util, ok in rows:
|
|
print(f"{pad:>4} {backend:>8} {label:>16} {wall:>8.2f}s "
|
|
f"{n_wf / wall:>10.0f} {util:>5.1f}x "
|
|
f"{'OK' if ok else 'MISMATCH'}")
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser(description=__doc__)
|
|
p.add_argument("--pads", default="1,8,40",
|
|
help="comma-separated pad factors (default 1,8,40)")
|
|
p.add_argument("--spf", type=int, default=2500)
|
|
p.add_argument("--rows", type=int, default=8)
|
|
p.add_argument("--frames", type=int, default=1024)
|
|
p.add_argument("--backends", default=None,
|
|
help="comma-separated (default: numpy,pyfftw if available)")
|
|
p.add_argument("--real", help="path to a real .sras file")
|
|
p.add_argument("--real-rows", type=int, default=32,
|
|
help="rows of angle 0 to use from the real file")
|
|
args = p.parse_args()
|
|
|
|
pads = [int(x) for x in args.pads.split(",")]
|
|
if args.backends:
|
|
backends = args.backends.split(",")
|
|
else:
|
|
backends = ["numpy"] + (["pyfftw"] if compute.PYFFTW_AVAILABLE else [])
|
|
|
|
if args.real:
|
|
sras = row_slice(SrasFile(args.real), 0, args.real_rows)
|
|
sras.data = [sras.data[0]]
|
|
sras.n_angles = 1
|
|
bench(sras, pads, backends)
|
|
else:
|
|
with tempfile.TemporaryDirectory(prefix="sras_bench_") as tmp:
|
|
path = Path(tmp) / "bench.sras"
|
|
gen.write(path, n_angles=1, seed=0, samples_per_frame=args.spf,
|
|
geometry=[(args.rows, args.frames)])
|
|
bench(SrasFile(str(path)), pads, backends)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|