Fuse the GPU frame chain: up to 2.6x on small frames, bit-identical output

AGENTS.md

@@ -36,6 +36,7 @@ removals. A JOSS paper + citation remain a post-2.0 follow-up.
 | `config.py` | `CameraConfig` (frozen dataclass of detector params) and `SensorType` enum. Pure data + validation + temperature scaling. No randomness. |
 | `noise.py` | The physics. Pure functions: `CameraConfig` + exposure + temperature + seeded `Generator` → electrons/ADU. This is where noise models and the opt-in reusable `DetectorWorkspace` live. |
 | `backend.py` | Optional NumPy/CuPy array and RNG boundary, explicit host conversion, and backend convolution. NumPy is the reference/default. Owns the stack's `device` vocabulary (`"cpu"`, `"gpu"`, `"gpu:N"`, `"auto"`; parsed by `_parse_device`, resolved by `get_backend`) and `precision` vocabulary (`resolve_precision`: `"single"`/`"double"`, aliases `"float32"`/`"float64"`). A GPU `ArrayBackend` carries a `device_id`; `activate()` makes it current. Kept local rather than `aocore.Backend`: getframes draws per-frame noise from a device-native (cuRAND) stream, which `aocore.Backend.random` (host-side `Generator`) does not provide. |
+| `_cuda.py` | Private fused CuPy `ElementwiseKernel`s for the launch-bound GPU hot path: the photo + total expectation (`fused_expectation`, called from `simulate_frame`) and the whole readout (`fused_readout`, called from `digitize`). Each mirrors the separate CuPy operations it replaces op for op (same order, working precision, Python scalars rounded to it, `--fmad=false`), so seeded GPU frames are bit-identical to the unfused path, which remains the fallback for any operand it cannot take exactly. Imported only on the GPU backend; CuPy is imported lazily. |
 | `frame.py` | `Frame` container: a NumPy array (ADU) plus metadata; array-like; optional FITS export. |
 | `camera.py` | `Camera`, the main user-facing object. Orchestrates config + scene + noise into `Frame`s. Holds the RNG and high-level methods (`dark_frame`, `dark_series`, reset-correlated `nondestructive_series`, `correlated_double_sample`, `expose`, `observe`, `*_series`, `master_*`). Reset-correlated readout has one core, the private `_ramp_reads`, which walks a ramp on an arbitrary per-read interval pattern; `nondestructive_series` drives it uniformly and `correlated_double_sample` drives it as pedestal-then-signal. Put new ramp-readout modes there rather than in a second loop — the interval-scaled bias/settling/avalanche terms are easy to get subtly wrong twice. Every public method that touches device arrays is decorated with `backend._on_device`, which runs it (each step, for generator methods) with the camera's CUDA device current; decorate new ones too, or `device="gpu:N"` silently allocates on the wrong card. |
 | `calibrate.py` | Master-frame builders (`combine`) and `calibrate` reduction — the raw → reduced → truth loop (phase 1.1). |
@@ -210,6 +211,29 @@ make the smallest change that fits the existing patterns; and if a change spans
 layers, respect the one-way data flow. When something is genuinely done and
 verified, say so plainly; when a step was skipped or a test fails, say *that*.
 
+## Performance notes (GPU hot path)
+
+- Small GPU frames are launch-bound (each CuPy call costs 10–20 µs of host time on a
+  slow core), so per-frame work counts CuPy calls, not FLOPs. Large frames, and the
+  CPU path, are bound by the Poisson/Gamma samplers; their draws come from one
+  sequential stream, so they cannot be threaded or reordered without changing
+  seeded output.
+- **Change the arithmetic in two places.** `photo_signal_map`, the
+  `photo + dark + extra` total in `simulate_frame`, and `digitize`'s readout each
+  have a fused twin in `_cuda.py`. Changing the order, precision, or terms of one
+  without the other silently changes seeded GPU frames;
+  `tests/test_gpu.py::test_gpu_fused_kernels_are_bit_identical_to_separate_operations`
+  catches it on a CUDA machine (ordinary CI has no GPU). A new readout term goes
+  into both paths.
+- Keep per-frame scalars on the host: use `noise._working_operand` rather than
+  `backend.asarray` for a value that may be a scalar, and pass scalar distribution
+  parameters through the `_CuPyGenerator` wrapper (it reuses a cached device copy),
+  so a GPU frame does not upload a fresh 0-d array per call.
+- Identity arithmetic on a whole frame (`+ 0`, `/ 1`, `* 1` by a *Python* scalar) is
+  skipped via `noise._is_scalar_value`; it cannot change any ADU (at most the sign
+  of a zero, which rounding erases). Keep the test to Python scalars: reading a 0-d
+  device array to compare it would synchronise the GPU every frame.
+
 ## Things to avoid
 
 - Don't add heavy core runtime dependencies. Core runtime deps are `numpy`, `scipy`,

CHANGELOG.md

@@ -12,6 +12,36 @@ to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
   device table on an Ampere Neoverse-N1 host (16 pinned cores) with an RTX
   4060 and an RTX A400.
 
+### Performance
+
+- **Small GPU frames run up to 2.6x faster, with bit-identical seeded output.**
+  WFS-sized frames were launch-bound: the host spent longer issuing about twenty
+  small CuPy kernels per frame than the GPU spent running them. The GPU path now
+  computes the photo and total expectations in one fused kernel and the whole
+  readout (full-well clip, defects, reset/avalanche/read noise, gain, bias
+  pedestal and structure, common mode, rounding, ADC saturation, `uint32`
+  conversion) in another; keeps scalar inputs (`background`, `extra_electrons`,
+  the EM-gain scale, the CIC rate) on the host instead of uploading them every
+  frame; samples Poisson counts directly in the working precision; and reuses
+  the chain's `photo + dark + extra` sum as `FrameTruth.mean_electrons`. The
+  noise is drawn by the same calls in the same order and the fused kernels
+  repeat the same floating-point operations (FMA contraction off), so every
+  seeded GPU frame and truth array is unchanged; a new GPU test compares them
+  bit for bit against the separate operations. `bench_devices.py` on an Ampere
+  Neoverse-N1 host (12 pinned cores), frames/s, 2.4.0 → now, median of three
+  interleaved runs: RTX 4060 — Pyramid 80x80 2,422 → 6,418 (2.65x),
+  Shack-Hartmann 160x160 2,424 → 6,340 (2.62x), OCAM2K 240x240 1,650 → 2,733
+  (1.66x), SAPHIRA 256x320 1,809 → 2,072 (1.15x), 1024x1024 240 → 247 (1.03x);
+  RTX A400 — Pyramid 2,403 → 6,315 (2.63x), the larger frames 1.02–1.06x.
+  Larger frames are bound by CuPy's double-precision Poisson and Gamma
+  samplers. See the [GPU guide](docs/guides/gpu.md#small-frames-fused-kernels).
+- The NumPy path skips whole-frame identity arithmetic (adding a zero offset,
+  dividing by a unit gain, multiplying by a unit avalanche-gain map) and reuses
+  the expectation sum as the truth: 1–5% faster (1024x1024: 13.4 → 14.1 frames/s
+  on the same host), with identical seeded output. It remains bound by NumPy's
+  Poisson sampler (about 80% of a 1024x1024 frame), whose single sequential
+  stream cannot be threaded without changing seeded frames.
+
 ### Fixed
 
 - `benchmarks/bench_devices.py` reported the CPU of Arm hosts as `aarch64`;

benchmarks/device-results-neoverse-n1.md

@@ -31,3 +31,9 @@ Reading this:
   1.1x the RTX A400 at 80x80, 3.3x on the OCAM2K and 5.6x at 1024x1024.
 - **16 Neoverse-N1 cores reach a steady 0.40–0.44x a 16-core Ryzen 9
   9950X3D** across every workflow.
+
+These are getframes 2.4.0 numbers. From 2.5.0 the GPU fuses the per-frame
+chain into far fewer kernels, which lifts the launch-bound Pyramid 80x80
+frame on this host to about 6,300–6,400 frames/s on either card, and
+Shack-Hartmann 160x160 to about 6,300 on the RTX 4060 (2.6x); see
+[Small frames: fused kernels](../docs/guides/gpu.md#small-frames-fused-kernels).

docs/guides/gpu.md

@@ -122,6 +122,45 @@ and [raw JSON](https://github.com/jacotay7/getframes/blob/main/benchmarks/device
 An Arm data point (Ampere Neoverse-N1, 16 pinned cores, RTX 4060 and RTX A400)
 is in [device-results-neoverse-n1.md](https://github.com/jacotay7/getframes/blob/main/benchmarks/device-results-neoverse-n1.md).
 
+### Small frames: fused kernels
+
+A wavefront-sensor frame (80x80 to 240x240 pixels) is limited by the host
+issuing kernels, not by the GPU running them: on a slow host core each CuPy call
+costs 10–20 µs, and the 2.4.0 chain issued about twenty per frame, three of them
+uploads of a scalar input. Since 2.5.0 the GPU path:
+
+- computes the photo and total expectations in one kernel, and the whole
+  readout (full-well clip, defects, reset/avalanche/read noise, gain, bias
+  pedestal and structure, common mode, rounding, ADC saturation and the
+  `uint32` conversion) in another;
+- keeps scalar inputs (`background`, `extra_electrons`, the EM-gain scale, the
+  CIC rate) on the host instead of uploading them every frame;
+- samples Poisson counts directly in the working precision, and reuses the
+  chain's `photo + dark + extra` sum as the truth instead of adding it again.
+
+The noise is still drawn by the same CuPy calls in the same order, and the fused
+kernels perform the same floating-point operations in the same order and
+precision with FMA contraction disabled, so seeded frames are bit-identical to
+2.4.0 (`tests/test_gpu.py` checks this against the separate operations).
+Measured on an Ampere Neoverse-N1 host (12 pinned cores), median of three
+interleaved `bench_devices.py` runs per cell, frames/s:
+
+| Workflow | Native shape | RTX 4060 2.4.0 | RTX 4060 2.5.0 | RTX A400 2.4.0 | RTX A400 2.5.0 |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| Pyramid WFS CMOS | 80x80 | 2,422 | 6,418 (2.65x) | 2,403 | 6,315 (2.63x) |
+| Shack-Hartmann WFS CMOS | 160x160 | 2,424 | 6,340 (2.62x) | 1,698 | 1,794 (1.06x) |
+| OCAM2K EMCCD | 240x240 | 1,650 | 2,733 (1.66x) | 517 | 531 (1.03x) |
+| SAPHIRA eAPD | 256x320 | 1,809 | 2,072 (1.15x) | 384 | 393 (1.02x) |
+| Large science CMOS | 1024x1024 | 240 | 247 (1.03x) | 42 | 45 (1.06x) |
+
+Once the host overhead is gone a frame is bound by CuPy's Poisson and Gamma
+samplers, which work in double precision: consumer GPUs run that at a small
+fraction of their single-precision rate (1/64 on the RTX 4060), so larger frames
+and the A400 gain little. The NumPy path is bound the same way by NumPy's
+Poisson sampler (about 80% of a 1024x1024 frame); its draws come from one
+sequential stream, so they cannot be split across threads without changing
+seeded output.
+
 An additional owner-isolation benchmark covers structured-detector digitization.
 On this repository's Quadro P620, native float32 amplifier and bias maps reduced
 the 2048x2048 digitization median from 9.196 ms to 6.654 ms (1.382x); the local

src/getframes/_cuda.py

@@ -0,0 +1,224 @@
+# SPDX-License-Identifier: MIT
+"""Fused CuPy kernels for the GPU detector hot path.
+
+A WFS-sized frame (80x80 to 240x240 pixels) is launch-bound on the GPU: the
+host spends longer issuing a dozen small elementwise kernels than the device
+spends running them. These kernels fold a run of those operations into one
+launch. Each kernel performs the *same* floating-point operations, in the same
+order and the same working precision, as the separate CuPy calls it replaces,
+and is compiled with FMA contraction disabled, so seeded GPU frames stay
+bit-identical to the unfused path.
+
+Imported only by the GPU backend; CuPy itself is imported lazily.
+"""
+
+from __future__ import annotations
+
+from functools import cache
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from .backend import ArrayBackend
+
+# Additive electron-domain noise terms (reset, avalanche input, correlated read,
+# read) that :func:`getframes.noise.digitize` can produce for one read.
+_MAX_NOISE_TERMS = 4
+
+# ``clip`` mirrors CuPy's own ``cupy_clip`` ufunc expression exactly, and
+# ``adu = v`` is the same float-to-uint32 conversion as ``astype(np.uint32)``.
+_READOUT_OPERATION = """
+const T zero = 0;
+T v = s;
+v = zero > full_well ? full_well : (v < zero ? zero : (v > full_well ? full_well : v));
+if (dead) v = zero;
+v = v + noise0;
+v = v + noise1;
+v = v + noise2;
+v = v + noise3;
+v = v / gain;
+v = v + bias;
+v = v + amplifier_offset;
+v = v + structure;
+v = v + common_mode;
+v = rint(v);
+v = zero > max_adu ? max_adu : (v < zero ? zero : (v > max_adu ? max_adu : v));
+adu = v;
+"""
+
+
+@cache
+def _readout_kernel() -> Any:
+    import cupy
+
+    return cupy.ElementwiseKernel(
+        "T s, T full_well, bool dead, T noise0, T noise1, T noise2, T noise3, T gain, T bias, "
+        "T amplifier_offset, T structure, T common_mode, T max_adu",
+        "uint32 adu",
+        _READOUT_OPERATION,
+        "getframes_fused_readout",
+        options=("--fmad=false",),
+    )
+
+
+# The photo expectation ``(rate + background) * exposure_scale * prnu`` (as in
+# :func:`getframes.noise.photo_signal_map`) and the total ``photo + dark + extra``
+# (as in :func:`getframes.noise.simulate_frame`), each product and sum rounded in
+# the same order as the separate kernels.
+_EXPECTATION_OPERATION = """
+T p = rate + background;
+p = p * exposure_scale;
+p = p * prnu;
+T t = p + dark;
+t = t + extra;
+"""
+
+
+@cache
+def _expectation_kernel(keep_photo: bool) -> Any:
+    import cupy
+
+    outputs = "T photo, T total" if keep_photo else "T total"
+    store = "photo = p;\ntotal = t;\n" if keep_photo else "total = t;\n"
+    return cupy.ElementwiseKernel(
+        "T rate, T background, T exposure_scale, T prnu, T dark, T extra",
+        outputs,
+        _EXPECTATION_OPERATION + store,
+        "getframes_fused_expectation",
+        options=("--fmad=false",),
+    )
+
+
+class _NotFusable(Exception):
+    """An operand the fused kernel cannot take with identical semantics."""
+
+
+def fused_readout(
+    signal: Any,
+    *,
+    backend: ArrayBackend,
+    full_well_e: float,
+    defects: Any | None,
+    noise_terms: list[Any],
+    gain: Any,
+    bias_offset_adu: Any,
+    amplifier_offset_adu: Any,
+    bias_structure_adu: Any,
+    common_mode_adu: Any,
+    max_adu: int,
+    out: Any | None,
+    output_slices: tuple[slice, slice] | None,
+    out_validated: bool,
+) -> Any | None:
+    """Run :func:`getframes.noise.digitize`'s readout arithmetic as one kernel.
+
+    The noise terms must already be drawn. Returns the ``uint32`` frame, or
+    ``None`` when an operand cannot be fused without changing its arithmetic
+    (for example an array in a different dtype, which CuPy would promote); the
+    caller then falls back to the separate operations. Python scalars are
+    rounded to the working dtype, exactly as CuPy rounds a Python scalar
+    combined with an array of that dtype.
+    """
+    import cupy
+
+    from .noise import _validate_output_buffer
+
+    dtype = signal.dtype
+    if dtype not in (np.dtype(np.float32), np.dtype(np.float64)) or signal.ndim != 2:
+        return None
+    if len(noise_terms) > _MAX_NOISE_TERMS:
+        return None
+    slices = output_slices if out is not None else None
+    shape = tuple(signal.shape)
+
+    def operand(value: Any, value_dtype: np.dtype[Any] = dtype) -> Any:
+        if isinstance(value, cupy.ndarray):
+            if value.dtype != value_dtype or value.ndim not in (0, 2):
+                raise _NotFusable
+            if value.ndim == 2:
+                if tuple(value.shape) != shape:
+                    raise _NotFusable
+                return value if slices is None else value[slices]
+            return value
+        if isinstance(value, (int, float)) and not isinstance(value, bool):
+            return value_dtype.type(value)
+        raise _NotFusable
+
+    try:
+        dead = False if defects is None else operand(defects, np.dtype(np.bool_))
+        # Missing noise terms add zero after the real ones; that can only change
+        # the sign of a zero, which the integer conversion erases.
+        noise = [operand(term) for term in noise_terms]
+        noise += [dtype.type(0)] * (_MAX_NOISE_TERMS - len(noise))
+        args = (
+            signal if slices is None else signal[slices],
+            operand(full_well_e),
+            dead,
+            *noise,
+            operand(gain),
+            operand(bias_offset_adu),
+            operand(amplifier_offset_adu),
+            operand(bias_structure_adu),
+            operand(common_mode_adu),
+            operand(max_adu),
+        )
+    except _NotFusable:
+        return None
+
+    kernel = _readout_kernel()
+    if out is None:
+        return kernel(*args)
+    if not out_validated:
+        rows, columns = args[0].shape
+        _validate_output_buffer(out, backend, (int(rows), int(columns)))
+    kernel(*args, out)
+    return out
+
+
+def fused_expectation(
+    rate: Any,
+    background: Any,
+    *,
+    exposure_scale: Any,
+    prnu: Any,
+    dark: Any,
+    extra: Any,
+    shape: tuple[int, int],
+    dtype: Any,
+    photo_out: Any | None,
+    total_out: Any | None,
+) -> tuple[Any, Any] | None:
+    """The photo and total expectations of one exposure in one kernel.
+
+    ``rate``, ``background`` and ``extra`` are working-dtype operands (see
+    ``noise._working_operand``); ``exposure_scale`` is ``exposure_s * qe`` and
+    ``prnu`` the fixed PRNU map, or ``1.0`` without PRNU. With ``photo_out`` (truth
+    disabled) the total overwrites that private buffer, as the unfused path does,
+    and both returned arrays are it. Returns ``None`` for an operand the unfused
+    path should handle, including every input it would reject with an error.
+    """
+    import cupy
+
+    dtype = np.dtype(dtype)
+
+    def operand(value: Any) -> Any:
+        if isinstance(value, (cupy.ndarray, np.generic)):
+            if value.dtype != dtype or (value.ndim != 0 and tuple(value.shape) != shape):
+                raise _NotFusable
+            return value
+        if isinstance(value, (int, float)) and not isinstance(value, bool):
+            return dtype.type(value)
+        raise _NotFusable
+
+    try:
+        args = [operand(value) for value in (rate, background, exposure_scale, prnu, dark, extra)]
+    except _NotFusable:
+        return None
+    if photo_out is not None:
+        _expectation_kernel(False)(*args, photo_out)
+        return photo_out, photo_out
+    photo = cupy.empty(shape, dtype=dtype)
+    total = cupy.empty(shape, dtype=dtype) if total_out is None else total_out
+    _expectation_kernel(True)(*args, photo, total)
+    return photo, total