diff --git a/AGENTS.md b/AGENTS.md index f8c471b..d91d3b8 100644 --- a/AGENTS.md +++ b/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`, diff --git a/CHANGELOG.md b/CHANGELOG.md index fd69de7..4057e01 100644 --- a/CHANGELOG.md +++ b/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`; diff --git a/benchmarks/device-results-neoverse-n1.md b/benchmarks/device-results-neoverse-n1.md index 0217ae2..b1abda7 100644 --- a/benchmarks/device-results-neoverse-n1.md +++ b/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). diff --git a/docs/guides/gpu.md b/docs/guides/gpu.md index dddf9db..279d003 100644 --- a/docs/guides/gpu.md +++ b/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 diff --git a/src/getframes/_cuda.py b/src/getframes/_cuda.py new file mode 100644 index 0000000..7ab9946 --- /dev/null +++ b/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 diff --git a/src/getframes/backend.py b/src/getframes/backend.py index 6608601..d056e79 100644 --- a/src/getframes/backend.py +++ b/src/getframes/backend.py @@ -50,14 +50,53 @@ def _cupy_seed(seed: Any) -> int | None: class _CuPyGenerator: """Expose the NumPy Generator spellings over a fast CuPy RandomState.""" + # Bound on cached device copies of scalar distribution parameters. A camera + # uses a handful (EM-gain scale, CIC rate); the bound only matters for a + # caller sweeping a parameter through one generator. + _MAX_DEVICE_SCALARS = 64 + def __init__(self, generator: Any, xp: Any, float_dtype: Any) -> None: self._generator = generator self._xp = xp self._float_dtype = float_dtype + self._device_scalars: dict[tuple[int, type, Any], Any] = {} def __getattr__(self, name: str) -> Any: return getattr(self._generator, name) + def _device_parameter(self, value: Any) -> Any: + """Return a host scalar parameter as the cached 0-d device array CuPy would build. + + CuPy's ``poisson`` and ``gamma`` call ``cupy.asarray`` on their + parameters, which uploads a Python scalar to a new device array on every + draw. Passing the identical device array (same value, same dtype) leaves + the samples unchanged; ``cupy.asarray`` returns an existing device array + as is. + """ + if not isinstance(value, (int, float)) or isinstance(value, bool): + return value + key = (int(self._xp.cuda.runtime.getDevice()), type(value), value) + cached = self._device_scalars.get(key) + if cached is None: + if len(self._device_scalars) >= self._MAX_DEVICE_SCALARS: + self._device_scalars.clear() + cached = self._xp.asarray(value) + self._device_scalars[key] = cached + return cached + + def poisson(self, lam: Any = 1.0, size: Any = None) -> Any: + """Draw Poisson counts (``int64``) without re-uploading a scalar rate.""" + return self._generator.poisson(lam=self._device_parameter(lam), size=size) + + def _poisson_as(self, lam: Any) -> Any: + """Poisson counts for an array rate, stored directly in the rate's float dtype. + + CuPy's sampler converts each ``int64`` count to the output dtype in the + same kernel, which is exactly ``poisson(lam).astype(lam.dtype)`` without + the ``int64`` intermediate or a second kernel. + """ + return self._generator.poisson(lam=lam, dtype=lam.dtype) + def normal(self, loc: Any = 0.0, scale: Any = 1.0, size: Any = None) -> Any: """Draw a scaled normal variate directly in the working precision.""" return self._generator.normal(loc=loc, scale=scale, size=size, dtype=self._float_dtype) @@ -73,7 +112,12 @@ def lognormal(self, mean: Any = 0.0, sigma: Any = 1.0, size: Any = None) -> Any: def gamma(self, shape: Any, scale: Any = 1.0, size: Any = None) -> Any: """Draw Gamma variates directly in the detector working precision.""" - return self._generator.gamma(shape=shape, scale=scale, size=size, dtype=self._float_dtype) + return self._generator.gamma( + shape=self._device_parameter(shape), + scale=self._device_parameter(scale), + size=size, + dtype=self._float_dtype, + ) def integers(self, low: Any, high: Any = None, size: Any = None) -> Any: """NumPy-Generator spelling for CuPy RandomState's ``randint``.""" diff --git a/src/getframes/camera.py b/src/getframes/camera.py index 35f0dda..8293168 100644 --- a/src/getframes/camera.py +++ b/src/getframes/camera.py @@ -903,6 +903,9 @@ def _expose_into( self.sensor_resolution[1] // binning, ) detector_out = workspace._buffer("full_adu", self._backend, detector_shape, np.uint32) + # Without a workspace or ROI the chain's own ``photo + dark + extra`` sum is + # exactly the truth total, so it is reused rather than recomputed. + mean_total: list[Any] = [] result = noise.simulate_frame( self.config, self._full_detector_input(photon_rate, "photon_rate", workspace=workspace), @@ -928,6 +931,7 @@ def _expose_into( _preserve_truth=include_truth, _output_slices=output_slices, _out_validated=out is not None or direct_roi_output, + _mean_total_out=mean_total if include_truth and self.roi is None else None, ) cropped_adu = result.adu if direct_roi_output else self._crop_to_roi(result.adu, binning) if self.roi is not None and out is not None: @@ -945,7 +949,9 @@ def _expose_into( ) truth = ( FrameTruth( - mean_electrons=result.mean_photoelectrons + mean_electrons=mean_total[0] + if mean_total + else result.mean_photoelectrons + result.mean_dark_electrons + self._binned_extra(extra_electrons, binning), mean_photoelectrons=result.mean_photoelectrons, diff --git a/src/getframes/noise.py b/src/getframes/noise.py index 0e409b6..de72b18 100644 --- a/src/getframes/noise.py +++ b/src/getframes/noise.py @@ -39,9 +39,11 @@ import numpy as np from aocore import block_sum -from .backend import ArrayBackend, _working_dtype, get_backend +from .backend import ArrayBackend, _CuPyGenerator, _working_dtype, get_backend if TYPE_CHECKING: + from collections.abc import Iterable, Iterator + from numpy.typing import DTypeLike, NDArray from .config import CameraConfig @@ -155,6 +157,32 @@ def _buffer( return buffer +def _working_operand(value: Any, float_dtype: DTypeLike, backend: ArrayBackend) -> Any: + """``backend.asarray(value, dtype=float_dtype)``, keeping a host scalar on the host. + + A Python or NumPy scalar (or 0-d NumPy array) becomes a NumPy scalar of the + working dtype. It rounds exactly as the 0-d array would and combines with + backend arrays the same way, but on the GPU it avoids allocating and + uploading a 0-d device array on every frame. + """ + if isinstance(value, (int, float, np.integer, np.floating)) or ( + isinstance(value, np.ndarray) and value.ndim == 0 + ): + return np.asarray(value, dtype=float_dtype)[()] + return backend.asarray(value, dtype=float_dtype) + + +def _is_scalar_value(value: Any, target: float) -> bool: + """Whether ``value`` is a plain Python number equal to ``target``. + + Used to skip identity arithmetic (``x * 1``, ``x / 1``, ``x + 0``) on whole + frames; skipping it changes at most the sign of a zero, which rounding to + integer ADU erases. Only Python scalars qualify, since comparing a 0-d device + array would synchronise the GPU. + """ + return isinstance(value, (int, float)) and not isinstance(value, bool) and value == target + + def _validate_output_buffer( out: Any, backend: ArrayBackend, @@ -545,8 +573,8 @@ def photo_signal_map( xp = resolved.xp height, width = config.resolution qe = config.quantum_efficiency if quantum_efficiency is None else quantum_efficiency - rate = resolved.asarray(photon_rate, dtype=float_dtype) - background = resolved.asarray(background_photon_rate, dtype=float_dtype) + rate = _working_operand(photon_rate, float_dtype, resolved) + background = _working_operand(background_photon_rate, float_dtype, resolved) if rate.ndim not in (0, 2) or background.ndim not in (0, 2): raise ValueError("photon_rate/background must be a scalar or a 2-D array.") @@ -1063,15 +1091,12 @@ def digitize( readout_full_well_e = ( config.output_full_well_e if config.output_full_well_e is not None else config.full_well_e ) - xp.clip(signal, 0.0, readout_full_well_e, out=signal) # Dead pixels/columns: a fixed defect map that collects no charge (they still # carry read/reset noise and the bias pedestal, so they read as dark defects). defects = ( fixed_patterns.defect_mask if fixed_patterns is not None else _defect_mask(config, resolved) ) - if defects is not None: - signal[defects] = 0.0 def normal_noise(sigma: Any) -> Any: if resolved.is_cpu and signal.dtype == np.dtype(np.float32): @@ -1080,53 +1105,64 @@ def normal_noise(sigma: Any) -> Any: return draw return rng.normal(0.0, sigma, size=signal.shape) - # An ordinary exposure draws a fresh kTC uncertainty. A nondestructive ramp - # passes one cached draw back on every read so reset noise remains common to - # the ramp and cancels under correlated double sampling. - if reset_noise_e is None: - if config.reset_noise_e > 0: - signal += normal_noise(config.reset_noise_e) - else: - signal += reset_noise_e - - # Some eAPD stacks carry an additional per-read component that scales with - # avalanche multiplication rather than remaining fixed at the output - # amplifier. Keep it separate so ``read_noise_e`` retains its usual output - # reference and high-gain data can identify the two terms independently. - avalanche_noise = ( - config.avalanche_input_noise_e - if avalanche_input_noise_e is None - else avalanche_input_noise_e - ) - if avalanche_noise > 0: - gain_scale = config.em_gain * ( - config.em_gain / config.avalanche_input_noise_reference_gain - ) ** (config.avalanche_input_noise_gain_exponent - 1.0) - signal += normal_noise(avalanche_noise * gain_scale) - - # Read noise in electrons, added at the amplifier. The per-pixel RMS is a - # fixed property of the sensor (see :func:`_read_noise_sigma_map`), so only the - # Gaussian draw itself is per-frame. - # - # ``read_noise_correlated_fraction`` splits that draw in two. The correlated - # part is passed in by a nondestructive ramp, which holds one draw for the - # whole ramp so that differencing two reads removes it; the independent part - # is redrawn every read and survives the difference. An ordinary exposure has - # nothing to correlate against and draws the full RMS. - if config.read_noise_e > 0: - sigma_map = ( - fixed_patterns.read_noise_sigma - if fixed_patterns is not None - else _read_noise_sigma_map(config, resolved, float_dtype=signal.dtype) - ) - correlated_weight = ( - config.read_noise_correlated_fraction if correlated_read_noise_e is not None else 0.0 - ) - if correlated_weight > 0.0 and correlated_read_noise_e is not None: - signal += correlated_read_noise_e - signal += normal_noise(sigma_map * np.sqrt(1.0 - correlated_weight)) + def electron_noise() -> Iterator[Any]: + """Additive electron-domain noise terms, drawn lazily in stream order.""" + # An ordinary exposure draws a fresh kTC uncertainty. A nondestructive ramp + # passes one cached draw back on every read so reset noise remains common to + # the ramp and cancels under correlated double sampling. + if reset_noise_e is None: + if config.reset_noise_e > 0: + yield normal_noise(config.reset_noise_e) else: - signal += normal_noise(sigma_map) + yield reset_noise_e + + # Some eAPD stacks carry an additional per-read component that scales with + # avalanche multiplication rather than remaining fixed at the output + # amplifier. Keep it separate so ``read_noise_e`` retains its usual output + # reference and high-gain data can identify the two terms independently. + avalanche_noise = ( + config.avalanche_input_noise_e + if avalanche_input_noise_e is None + else avalanche_input_noise_e + ) + if avalanche_noise > 0: + gain_scale = config.em_gain * ( + config.em_gain / config.avalanche_input_noise_reference_gain + ) ** (config.avalanche_input_noise_gain_exponent - 1.0) + yield normal_noise(avalanche_noise * gain_scale) + + # Read noise in electrons, added at the amplifier. The per-pixel RMS is a + # fixed property of the sensor (see :func:`_read_noise_sigma_map`), so only + # the Gaussian draw itself is per-frame. + # + # ``read_noise_correlated_fraction`` splits that draw in two. The correlated + # part is passed in by a nondestructive ramp, which holds one draw for the + # whole ramp so that differencing two reads removes it; the independent part + # is redrawn every read and survives the difference. An ordinary exposure + # has nothing to correlate against and draws the full RMS. + if config.read_noise_e > 0: + sigma_map = ( + fixed_patterns.read_noise_sigma + if fixed_patterns is not None + else _read_noise_sigma_map(config, resolved, float_dtype=signal.dtype) + ) + correlated_weight = ( + config.read_noise_correlated_fraction + if correlated_read_noise_e is not None + else 0.0 + ) + if correlated_weight > 0.0 and correlated_read_noise_e is not None: + yield correlated_read_noise_e + yield normal_noise(sigma_map * np.sqrt(1.0 - correlated_weight)) + else: + yield normal_noise(sigma_map) + + def common_mode() -> Any: + if common_mode_adu is not None: + return common_mode_adu + if config.readout_common_mode_noise_adu > 0: + return rng.normal(0.0, config.readout_common_mode_noise_adu) + return 0.0 if fixed_patterns is None: gain_map, amp_offset = _amplifier_maps(config, resolved, float_dtype=signal.dtype) @@ -1135,15 +1171,53 @@ def normal_noise(sigma: Any) -> Any: gain_map = fixed_patterns.amplifier_gain amp_offset = fixed_patterns.amplifier_offset bias_structure = fixed_patterns.bias_structure - signal /= gain_map - signal += config.bias_offset_adu - signal += amp_offset - signal += bias_structure - if common_mode_adu is None: - if config.readout_common_mode_noise_adu > 0: - signal += rng.normal(0.0, config.readout_common_mode_noise_adu) - else: - signal += common_mode_adu + + noise_terms: Iterable[Any] = electron_noise() + common: Any | None = None + if not resolved.is_cpu: + # On the GPU the readout is launch-bound for WFS-sized frames, so draw the + # noise now (same stream order) and run the whole readout as one kernel. + noise_terms = list(noise_terms) + common = common_mode() + from . import _cuda + + fused = _cuda.fused_readout( + signal, + backend=resolved, + full_well_e=readout_full_well_e, + defects=defects, + noise_terms=noise_terms, + gain=gain_map, + bias_offset_adu=config.bias_offset_adu, + amplifier_offset_adu=amp_offset, + bias_structure_adu=bias_structure, + common_mode_adu=common, + max_adu=config.max_adu, + out=out, + output_slices=_output_slices, + out_validated=_out_validated, + ) + if fused is not None: + return fused + + # The same steps one array operation at a time (the NumPy reference, and the + # GPU fallback for operands the fused kernel does not take). Adding a Python + # zero or dividing by a Python one is skipped: it cannot change any ADU. + xp.clip(signal, 0.0, readout_full_well_e, out=signal) + if defects is not None: + signal[defects] = 0.0 + for term in noise_terms: + if not _is_scalar_value(term, 0.0): + signal += term + if not _is_scalar_value(gain_map, 1.0): + signal /= gain_map + for offset in (config.bias_offset_adu, amp_offset, bias_structure): + if not _is_scalar_value(offset, 0.0): + signal += offset + if common is None: + common = common_mode() + if not _is_scalar_value(common, 0.0): + signal += common xp.rint(signal, out=signal) xp.clip(signal, 0, config.max_adu, out=signal) if out is None: @@ -1182,7 +1256,11 @@ def frame_electrons( (``float64`` exact, or ``float32`` for the memory-light fast path). """ resolved = backend or get_backend() - electrons = rng.poisson(mean_electrons).astype(mean_electrons.dtype) + if isinstance(rng, _CuPyGenerator): + # The same draw, converted to the working dtype inside the sampler kernel. + electrons = rng._poisson_as(mean_electrons) + else: + electrons = rng.poisson(mean_electrons).astype(mean_electrons.dtype) if config.clock_induced_charge_e > 0: electrons += rng.poisson(config.clock_induced_charge_e, size=electrons.shape) @@ -1224,7 +1302,8 @@ def frame_electrons( float_dtype=electrons.dtype, ).avalanche_gain_multiplier ) - electrons *= gain_multiplier + if not _is_scalar_value(gain_multiplier, 1.0): # x * 1.0 == x: skip the pass + electrons *= gain_multiplier return electrons @@ -1253,6 +1332,7 @@ def simulate_frame( _preserve_truth: bool = True, _output_slices: tuple[slice, slice] | None = None, _out_validated: bool = False, + _mean_total_out: list[Any] | None = None, ) -> SimulationResult: """Simulate one frame end-to-end, returning ADU and the noise-free truth. @@ -1366,17 +1446,6 @@ def simulate_frame( photo_out = None if workspace is not None and not _preserve_truth: photo_out = workspace._buffer("mean_photo", resolved, config.resolution, float_dtype) - mean_photo = photo_signal_map( - config, - photon_rate, - exposure_s, - background_photon_rate, - quantum_efficiency, - float_dtype, - backend=resolved, - fixed_patterns=fixed_patterns, - out=photo_out, - ) mean_dark = ( dark_signal_map( config, @@ -1389,17 +1458,58 @@ def simulate_frame( if _dark_signal is None else _dark_signal ) - if workspace is None: - mean_total = mean_photo + mean_dark - elif not _preserve_truth: - # The photo expectation is private scratch when truth is disabled, so it - # can become the total in place with no allocation or extra copy kernel. - mean_total = mean_photo - mean_total += mean_dark + expectation = None + if not resolved.is_cpu and fixed_patterns is not None: + # Launch-bound GPU frames: one kernel for the photo and total expectations. + from . import _cuda + + expectation = _cuda.fused_expectation( + _working_operand(photon_rate, float_dtype, resolved), + _working_operand(background_photon_rate, float_dtype, resolved), + exposure_scale=exposure_s + * (config.quantum_efficiency if quantum_efficiency is None else quantum_efficiency), + prnu=fixed_patterns.prnu_multiplier if config.prnu > 0 else 1.0, + dark=mean_dark, + extra=_working_operand(extra_electrons, float_dtype, resolved), + shape=config.resolution, + dtype=float_dtype, + photo_out=photo_out, + total_out=( + workspace._buffer("mean_total", resolved, config.resolution, float_dtype) + if workspace is not None and photo_out is None + else None + ), + ) + if expectation is not None: + mean_photo, mean_total = expectation else: - mean_total = workspace._buffer("mean_total", resolved, config.resolution, float_dtype) - resolved.xp.add(mean_photo, mean_dark, out=mean_total) - mean_total += resolved.asarray(extra_electrons, dtype=float_dtype) + mean_photo = photo_signal_map( + config, + photon_rate, + exposure_s, + background_photon_rate, + quantum_efficiency, + float_dtype, + backend=resolved, + fixed_patterns=fixed_patterns, + out=photo_out, + ) + if workspace is None: + mean_total = mean_photo + mean_dark + elif not _preserve_truth: + # The photo expectation is private scratch when truth is disabled, so it + # can become the total in place with no allocation or extra copy kernel. + mean_total = mean_photo + mean_total += mean_dark + else: + mean_total = workspace._buffer("mean_total", resolved, config.resolution, float_dtype) + resolved.xp.add(mean_photo, mean_dark, out=mean_total) + mean_total += _working_operand(extra_electrons, float_dtype, resolved) + if _mean_total_out is not None and workspace is None and binning == 1: + # ``mean_total`` is a fresh array here and the chain below only reads it, + # so the caller may keep it as the ``photo + dark + extra`` truth instead + # of summing the same arrays again in the same order. + _mean_total_out.append(mean_total) if binning > 1: height, width = config.resolution diff --git a/tests/test_gpu.py b/tests/test_gpu.py index d0ead1f..d004546 100644 --- a/tests/test_gpu.py +++ b/tests/test_gpu.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import Any + import numpy as np import pytest @@ -279,3 +281,83 @@ def test_gpu_number_beyond_the_device_count_is_rejected() -> None: count = cupy.cuda.runtime.getDeviceCount() with pytest.raises(ValueError, match=f"CuPy sees {count} device"): gf.Camera.from_preset("generic_cmos", device=f"gpu:{count}") + + +_STRUCTURED_DETECTOR = { + "prnu": 0.02, + "dark_current_nonuniformity": 0.05, + "reset_noise_e": 1.5, + "read_noise_nonuniformity": 0.2, + "readout_common_mode_noise_adu": 1.3, + "amplifier_layout": (2, 2), + "amp_gain_nonuniformity": 0.02, + "amp_offset_spread_adu": 3.0, + "bias_structure_amplitude_adu": 2.0, + "dead_pixel_fraction": 0.02, + "bad_column_fraction": 0.05, +} + + +@pytest.mark.parametrize("precision", ["single", "double"]) +@pytest.mark.parametrize( + ("preset", "changes"), + [ + ("generic_cmos", {}), + ("generic_cmos", _STRUCTURED_DETECTOR), + ("generic_cmos", {**_STRUCTURED_DETECTOR, "roi": (4, 2, 20, 16)}), + ("generic_emccd", {"avalanche_gain_nonuniformity": 0.05}), + ("generic_eapd", {"avalanche_input_noise_e": 0.4, "read_noise_correlated_fraction": 0.5}), + ], +) +def test_gpu_fused_kernels_are_bit_identical_to_separate_operations( + monkeypatch: pytest.MonkeyPatch, preset: str, changes: dict[str, object], precision: str +) -> None: + cupy = _cupy() + from getframes import _cuda + + config = gf.load_preset(preset).replace(resolution=(24, 32), **changes) + camera = gf.Camera(config, device="gpu", precision=precision, seed=3) + dtype = camera._float_dtype + rate = cupy.asarray(np.random.default_rng(0).uniform(0.0, 2e6, camera.resolution), dtype=dtype) + fused_calls = {"readout": 0, "expectation": 0} + + def counted(name: str, function: Any) -> Any: + def wrapper(*args: Any, **kwargs: Any) -> Any: + result = function(*args, **kwargs) + fused_calls[name] += result is not None + return result + + return wrapper + + def frames() -> list[gf.Frame]: + workspace = noise.DetectorWorkspace() + out = cupy.empty(camera.resolution, dtype=cupy.uint32) + series = [ + camera.expose(rate, 1e-3, seed=4), + camera.expose(rate, 1e-3, seed=5, background=12.5, extra_electrons=3.0), + camera.expose(2e5, 1e-3, seed=6), + camera.expose(rate, 1e-3, seed=7, workspace=workspace), + camera.expose(rate, 1e-3, seed=8, workspace=workspace, out=out, include_truth=False), + camera.dark_frame(0.5, seed=9), + ] + if not config.nonlinearity: + series.append(camera.correlated_double_sample(rate, 1e-3, seed=10)) + return series + + monkeypatch.setattr(_cuda, "fused_readout", counted("readout", _cuda.fused_readout)) + monkeypatch.setattr(_cuda, "fused_expectation", counted("expectation", _cuda.fused_expectation)) + fused = frames() + assert fused_calls["readout"] > 0 and fused_calls["expectation"] > 0 + + monkeypatch.setattr(_cuda, "fused_readout", lambda *args, **kwargs: None) + monkeypatch.setattr(_cuda, "fused_expectation", lambda *args, **kwargs: None) + reference = frames() + + bits = cupy.uint32 if np.dtype(dtype) == np.float32 else cupy.uint64 + for got, want in zip(fused, reference, strict=True): + assert cupy.array_equal(got.data, want.data) + if want.truth is not None: + for name in ("mean_electrons", "mean_photoelectrons"): + assert cupy.array_equal( + getattr(got.truth, name).view(bits), getattr(want.truth, name).view(bits) + )