Add a benchmark harness and record the performance baseline

benchmarks/BASELINE.md

@@ -0,0 +1,101 @@
+# Performance baseline
+
+Recorded before any optimization work, as the reference every later change is
+measured against. Nothing here has been optimized yet; these are the numbers as
+the library stands.
+
+**Hardware:** NVIDIA H200 NVL, 139.8 GB, CUDA 12.9, CuPy 14.0.1, NumPy 2.4.3,
+Python 3.13. Measured at PtyLab 0.2.8.
+
+**Reproduce:** `python benchmarks/bench_engines.py --markdown`
+
+Numbers move a few percent run to run. Treat anything under ~5% as noise.
+
+## Regimes
+
+Which optimization pays depends on the field size, so it is worth naming the
+regimes rather than talking about "the" performance of the library:
+
+| | field size | bound by | where the headroom is |
+|---|---|---|---|
+| A | ≤ ~12 MB | host dispatch / kernel launch | removing syncs, fusion, CUDA graphs |
+| B | ≥ ~50 MB | HBM bandwidth + cuFFT | very little |
+| C | OPR | orthogonalization time and memory | the linear algebra |
+
+`field MB` = `nlambda × nosm × npsm × nslice × Np² × 8 bytes`.
+
+## Engine baseline
+
+3 timed iterations after a warm-up iteration, GPU backend.
+
+| config | reg | field MB | ms/iter | µs/pos | peak GB |
+|---|---|---|---|---|---|
+| simu-like 128, 100fr | A | 0.1 | 58.6 | 586 | 0.03 |
+| USAF-like 364, 102fr | A | 1.0 | 71.4 | 700 | 0.23 |
+| mixed npsm=4 364 | A | 4.0 | 72.5 | 711 | 0.26 |
+| Brain-like poly=7 182 | A | 1.8 | 74.9 | 749 | 0.06 |
+| multislice nslice=4 364 | A | 4.0 | 150.6 | 1477 | 0.26 |
+| heavy 7×2×4 364 | B | 56.6 | 90.0 | 2249 | 1.02 |
+| **OPR 364, 202fr, 4 modes** | C | 4.0 | **505.4** | **2502** | **2.85** |
+
+**The regime-A rows are flat at 586-749 µs/position across a 40x range of field
+size** (0.1 → 4.0 MB). That flatness is the finding: the per-position cost barely
+depends on how much data is being processed, which means the GPU is not the
+thing being waited on. The cost is host-side dispatch.
+
+OPR is 7x the cost of the comparable non-OPR row at the same field size
+(505.4 vs 72.5 ms/iter at 4.0 MB) and 11x the memory.
+
+## OPR iteration breakdown (364 px, 202 frames, 4 OPR modes, subspace 4)
+
+| stage | ms/iter | share |
+|---|---|---|
+| `orthogonalizeIncoherentModes` (`OPR.py:149`) | 199.9 | **39.6%** |
+| `orthogonalizeProbeStack` (`OPR.py:204`) | 133.8 | 26.5% |
+| position loop and everything else | 170.8 | 33.9% |
+
+**Two thirds of an OPR iteration is linear algebra, not ptychography.**
+`orthogonalizeIncoherentModes` — a Python loop running one small SVD per frame —
+is the larger of the two.
+
+This matters for small cards. `probe_stack` is
+`(1, 1, nModes, 1, Np, Np, nFrames)` complex64, so it grows linearly in frames
+and modes: 0.80 GB at 364/202/4, 10.43 GB at 512/890/6, before the transient the
+SVD itself needs on top.
+
+## Negative results — measured, do not re-investigate
+
+**FFT-size padding is not worth it.** Np=364 (`2²·7·13`) and Np=182 (`2·7·13`)
+look like awkward cuFFT sizes, but measure at 2.6 µs/MB against 1.9 µs/MB for
+N=1024 — only 1.4x off best-case, not the several-fold penalty a genuinely bad
+size would show. Padding 364→384 measured **0.93x, i.e. slower**; 364→512 was
+0.65x. Padding 182→192 does give 1.16x, but costs 11% more memory everywhere.
+
+**Batching the existing SVD buys nothing.** Replacing the per-frame Python loop
+in `orthogonalizeIncoherentModes` with a single batched `cp.linalg.svd` measured
+1.02x at 364/202/4 and 1.03x at 512/890/6. The cost is inside the SVD itself,
+not in per-call launch overhead, so batching the same algorithm does not help.
+Only changing the algorithm does.
+
+## Environment note
+
+`cupy.linalg.eigh` **does not work in this environment**: `cupyx.cusolver` fails
+to import with `libcusolver.so.11: cannot open shared object file`, while
+`cupy.linalg.svd` works normally. This affects any user code calling `eigh`, and
+it rules out eigendecomposition-based approaches here until the CUDA
+installation is repaired.
+
+## Test-suite baseline
+
+The safety net these benchmarks are measured against.
+
+| | before | now (0.2.8) |
+|---|---|---|
+| passing | 14 | 61 |
+| skipped | 19 | 19 |
+| engine configurations with a numerical golden | 0 | 9 |
+| propagators with an output golden | 0 | 8 |
+| runtime | 1.6 s | 1.9 s |
+
+Verified that the goldens reproduce bit-exactly and that a 0.1% change to
+`betaObject` is caught.

benchmarks/bench_engines.py

@@ -0,0 +1,247 @@
+"""Benchmark harness for PtyLab reconstruction engines.
+
+Reports wall time *and* peak GPU memory. Both are acceptance criteria: this
+group runs on cards as small as 32 GB, so an optimization that trades memory
+for speed is a regression, not an improvement.
+
+Datasets are synthesized in-process at a requested size, so the harness runs
+anywhere without the (gitignored) files in ``example_data/``.
+
+Usage::
+
+    python benchmarks/bench_engines.py                 # default sweep
+    python benchmarks/bench_engines.py --quick         # one config per regime
+    python benchmarks/bench_engines.py --cpu           # CPU baseline too
+    python benchmarks/bench_engines.py --markdown      # BASELINE.md table
+
+Regimes (see BASELINE.md for the measurements behind these):
+  A  field <= ~12 MB   launch/dispatch bound   -- graphs & fusion pay
+  B  field >= ~50 MB   HBM bandwidth + cuFFT   -- little headroom
+  C  OPR               orthogonalization       -- two thirds linear algebra
+"""
+
+import argparse
+import gc
+import logging
+import os
+import sys
+import time
+from pathlib import Path
+
+# Progress bars cost host time in the loop being measured (mPIE draws a bar per
+# scan position); disable before PtyLab imports tqdm.
+os.environ.setdefault("TQDM_DISABLE", "1")
+
+import numpy as np
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+import matplotlib
+
+matplotlib.use("Agg")  # never open a window; showReconstruction is disabled anyway
+
+from PtyLab import Engines  # noqa: E402
+from PtyLab.ExperimentalData.ExperimentalData import ExperimentalData  # noqa: E402
+from PtyLab.Monitor.Monitor import DummyMonitor  # noqa: E402
+from PtyLab.Params.Params import Params  # noqa: E402
+from PtyLab.Reconstruction.Reconstruction import Reconstruction  # noqa: E402
+
+try:
+    import cupy as cp
+
+    HAS_GPU = cp.cuda.is_available()
+except Exception:  # pragma: no cover
+    cp = None
+    HAS_GPU = False
+
+
+# name, engine, propagator, (nlambda,nosm,npsm,nslice), Nd, n_frames, regime
+CONFIGS = [
+    ("simu-like  128, 100fr",       "mPIE", "Fraunhofer", (1, 1, 1, 1), 128, 100, "A"),
+    ("USAF-like  364, 102fr",       "mPIE", "Fraunhofer", (1, 1, 1, 1), 364, 102, "A"),
+    ("mixed npsm=4  364",           "mPIE", "Fraunhofer", (1, 1, 4, 1), 364, 102, "A"),
+    ("Brain-like poly=7  182",      "ePIE", "polychromeASP", (7, 1, 1, 1), 182, 100, "A"),
+    ("multislice nslice=4  364",    "e3PIE", "Fraunhofer", (1, 1, 1, 4), 364, 102, "A"),
+    ("heavy 7x2x4  364",            "mPIE", "Fraunhofer", (7, 2, 4, 1), 364, 40, "B"),
+    ("OPR 364, 202fr, 4 modes",     "OPR", "Fraunhofer", (1, 1, 4, 1), 364, 202, "C"),
+]
+
+QUICK = {"simu-like  128, 100fr", "heavy 7x2x4  364", "OPR 364, 202fr, 4 modes"}
+
+
+def synth_dataset(path, nd, n_frames, seed=7):
+    """Write a deterministic CPM dataset of the requested size."""
+    import h5py
+
+    rng = np.random.default_rng(seed)
+    grid = int(np.ceil(np.sqrt(n_frames)))
+    step = 3e-6
+    coords = (np.arange(grid) - (grid - 1) / 2) * step
+    yy, xx = np.meshgrid(coords, coords, indexing="ij")
+    encoder = np.stack([yy.ravel(), xx.ravel()], axis=1)[:n_frames]
+
+    ptychogram = rng.random((n_frames, nd, nd)).astype(np.float32)
+
+    with h5py.File(path, "w") as hf:
+        hf.create_dataset("ptychogram", data=ptychogram, dtype="f")
+        hf.create_dataset("encoder", data=encoder, dtype="f")
+        hf.create_dataset("dxd", data=np.array(75e-6))
+        hf.create_dataset("zo", data=np.array(0.05))
+        hf.create_dataset("wavelength", data=np.array(632.8e-9))
+        hf.create_dataset("entrancePupilDiameter", data=np.array(400e-6))
+    return path
+
+
+def build(path, config, gpu):
+    _name, engine_name, propagator, modes, _nd, _nfr, _regime = config
+    nlambda, nosm, npsm, nslice = modes
+
+    data = ExperimentalData(str(path), operationMode="CPM")
+    params = Params()
+    params.gpuSwitch = gpu
+    params.propagatorType = propagator
+    params.positionOrder = "sequential"
+
+    reconstruction = Reconstruction(data, params)
+    reconstruction.nlambda = nlambda
+    reconstruction.nosm = nosm
+    reconstruction.npsm = npsm
+    reconstruction.nslice = nslice
+    if nlambda > 1:
+        base = float(np.atleast_1d(reconstruction.wavelength)[0])
+        reconstruction.spectralDensity = base * np.linspace(0.98, 1.02, nlambda)
+    if nslice > 1:
+        reconstruction.dz = 1e-4
+        reconstruction.refrIndex = 1.0
+
+    np.random.seed(0)
+    reconstruction.initializeObjectProbe()
+
+    monitor = DummyMonitor()
+    engine = getattr(Engines, engine_name)(reconstruction, data, params, monitor)
+    if engine_name == "OPR":
+        params.OPR_modes = np.arange(npsm)
+        params.OPR_subspace = min(4, _nfr)
+        engine.OPR_modes = params.OPR_modes
+        engine.n_subspace = params.OPR_subspace
+    return data, reconstruction, params, engine
+
+
+def field_mb(config):
+    _n, _e, _p, (nl, nos, nps, nsl), nd, _f, _r = config
+    return nl * nos * nps * nsl * nd * nd * 8 / 2**20
+
+
+def reset_memory():
+    if HAS_GPU:
+        cp.get_default_memory_pool().free_all_blocks()
+        cp.get_default_memory_pool().set_limit(size=0)  # no cap; just reset accounting
+    gc.collect()
+
+
+def peak_gpu_gb():
+    if not HAS_GPU:
+        return float("nan")
+    return cp.get_default_memory_pool().total_bytes() / 2**30
+
+
+def sync():
+    if HAS_GPU:
+        cp.cuda.runtime.deviceSynchronize()
+
+
+def run_one(path, config, gpu, iters, warmup_iters=1):
+    """Time ``iters`` iterations after a warm-up, returning (s/iter, peak GB)."""
+    _name, _engine_name, _prop, _modes, _nd, n_frames, _regime = config
+
+    data, reconstruction, params, engine = build(path, config, gpu)
+
+    # Warm up: JIT, cuFFT plans, lru_cache'd transfer functions, first-touch
+    # allocations. Without this the first iteration dominates the measurement.
+    engine.numIterations = warmup_iters
+    engine.reconstruct()
+    sync()
+
+    if HAS_GPU:
+        cp.get_default_memory_pool().free_all_blocks()
+    before = peak_gpu_gb()
+
+    engine.numIterations = iters
+    sync()
+    t0 = time.perf_counter()
+    engine.reconstruct()
+    sync()
+    elapsed = time.perf_counter() - t0
+
+    peak = max(peak_gpu_gb(), before)
+    del data, reconstruction, params, engine
+    reset_memory()
+    return elapsed / iters, peak, n_frames
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--quick", action="store_true", help="one config per regime")
+    ap.add_argument("--cpu", action="store_true", help="also benchmark the CPU path")
+    ap.add_argument("--iters", type=int, default=3, help="timed iterations")
+    ap.add_argument("--markdown", action="store_true", help="emit a markdown table")
+    ap.add_argument("--only", type=str, default=None, help="substring filter")
+    args = ap.parse_args()
+
+    logging.disable(logging.WARNING)
+
+    configs = [c for c in CONFIGS if not args.quick or c[0] in QUICK]
+    if args.only:
+        configs = [c for c in configs if args.only.lower() in c[0].lower()]
+
+    backends = [("gpu", True)] if HAS_GPU else []
+    if args.cpu or not HAS_GPU:
+        backends.append(("cpu", False))
+    if not backends:
+        print("no backend available")
+        return
+
+    if HAS_GPU:
+        dev = cp.cuda.Device()
+        _free, total = dev.mem_info
+        name = cp.cuda.runtime.getDeviceProperties(dev.id)["name"].decode()
+        print(f"# GPU: {name}, {total / 2**30:.1f} GB total, CuPy {cp.__version__}")
+    print(f"# numpy {np.__version__}, {args.iters} timed iterations after warm-up\n")
+
+    header = ["config", "reg", "field MB", "backend", "ms/iter", "us/pos", "peak GB"]
+    if args.markdown:
+        print("| " + " | ".join(header) + " |")
+        print("|" + "|".join("---" for _ in header) + "|")
+    else:
+        print(f"{'config':<28s} {'reg':>3s} {'fieldMB':>8s} {'backend':>7s} "
+              f"{'ms/iter':>9s} {'us/pos':>8s} {'peakGB':>8s}")
+        print("-" * 80)
+
+    import tempfile
+
+    tmpdir = Path(tempfile.mkdtemp(prefix="ptylab_bench_"))
+    for config in configs:
+        name, engine_name, _p, _m, nd, n_frames, regime = config
+        path = synth_dataset(tmpdir / f"{nd}_{n_frames}.hdf5", nd, n_frames)
+        for backend_name, gpu in backends:
+            if engine_name == "OPR" and not gpu:
+                continue  # OPR has no CPU path (OPR.py calls cp.* directly)
+            try:
+                s_per_iter, peak, nfr = run_one(path, config, gpu, args.iters)
+                row = [name, regime, f"{field_mb(config):.1f}", backend_name,
+                       f"{s_per_iter * 1e3:.1f}",
+                       f"{s_per_iter / nfr * 1e6:.0f}",
+                       f"{peak:.2f}"]
+            except Exception as exc:  # keep going; report the failure in-band
+                row = [name, regime, f"{field_mb(config):.1f}", backend_name,
+                       f"FAILED: {type(exc).__name__}", "-", "-"]
+            if args.markdown:
+                print("| " + " | ".join(row) + " |")
+            else:
+                print(f"{row[0]:<28s} {row[1]:>3s} {row[2]:>8s} {row[3]:>7s} "
+                      f"{row[4]:>9s} {row[5]:>8s} {row[6]:>8s}")
+            sys.stdout.flush()
+
+
+if __name__ == "__main__":
+    main()