From 34b198e50f798309e037a787d703f6ee61a9b03d Mon Sep 17 00:00:00 2001 From: Jacob Taylor Date: Wed, 7 Oct 2026 08:29:35 -1000 Subject: [PATCH 1/2] Fuse the GPU frame chain into fewer kernels, bit-identically Small GPU frames were launch-bound: about twenty CuPy calls per frame, three of them uploads of a scalar input. The GPU path now runs the photo and total expectations as one ElementwiseKernel and the whole readout (clip, defects, noise terms, gain, offsets, common mode, rint, ADC clip, uint32) as another, keeps scalar inputs and scalar distribution parameters on the host (or a cached device copy), samples Poisson counts directly in the working precision, and reuses the chain's total as the truth. Noise is drawn by the same calls in the same order and the fused kernels repeat the same float operations with FMA contraction off, so seeded frames and truth are unchanged; a GPU test checks this bit for bit against the separate operations. Both paths skip whole-frame identity arithmetic by Python scalars. Co-Authored-By: Claude Opus 5.5 --- src/getframes/_cuda.py | 224 +++++++++++++++++++++++++++++++ src/getframes/backend.py | 46 ++++++- src/getframes/camera.py | 8 +- src/getframes/noise.py | 278 +++++++++++++++++++++++++++------------ tests/test_gpu.py | 82 ++++++++++++ 5 files changed, 552 insertions(+), 86 deletions(-) create mode 100644 src/getframes/_cuda.py 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) + ) From a2ee67907f6185c8715660b2ce0519a0c1a10114 Mon Sep 17 00:00:00 2001 From: Jacob Taylor Date: Wed, 7 Oct 2026 08:29:35 -1000 Subject: [PATCH 2/2] Document the fused GPU path and its measured speedup Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 24 +++++++++++++++ CHANGELOG.md | 30 ++++++++++++++++++ benchmarks/device-results-neoverse-n1.md | 6 ++++ docs/guides/gpu.md | 39 ++++++++++++++++++++++++ 4 files changed, 99 insertions(+) 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