From a4e1735a9b421532d1ab2e3a100d2139ba60652e Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 29 Sep 2026 15:28:00 -0700 Subject: [PATCH 1/5] Fuse CuPy Gaussian residuals and normal equations Add opt-in kernels that combine Gaussian evaluation, weighting and normal-equation accumulation while retaining the existing LM solver. Expose fusion through the fit API, CLI and device pipeline settings. Signed-off-by: Trent Nelson --- .github/workflows/ci.yml | 3 +- docs/components/xfit.md | 19 + examples/xfit/benchmark_gaussian.py | 34 +- src/cuphoton/xfit/_cupy_fused.py | 348 +++++++++++++ src/cuphoton/xfit/api.py | 51 +- src/cuphoton/xfit/commands.py | 29 +- src/cuphoton/xfit/executor.py | 18 +- src/cuphoton/xscan/device_pipeline.py | 13 + .../xscan/pipeline_benchmark/runner.py | 1 + tests/core/test_cli_contract.py | 4 +- tests/test_public_api.py | 2 + tests/xfit/test_cupy_fusion.py | 469 ++++++++++++++++++ tests/xscan/test_device_pipeline.py | 107 ++++ 13 files changed, 1090 insertions(+), 8 deletions(-) create mode 100644 src/cuphoton/xfit/_cupy_fused.py create mode 100644 tests/xfit/test_cupy_fusion.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5c7966cc..1b828127 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -157,6 +157,7 @@ jobs: tests/xfit/test_models.py::test_cupy_finite_difference_theta_fix_matches_numpy tests/xfit/test_models.py::test_fit_dipoles_device_matches_portable_result_without_bulk_d2h tests/xfit/test_solver.py::test_cupy_tall_jacobian_diagnostics_match_numpy + tests/xfit/test_cupy_fusion.py - name: Require all selected tests to execute run: | uv run --no-sync python - <<'PY' @@ -166,7 +167,7 @@ jobs: report = ET.parse(Path(os.environ["RUNNER_TEMP"]) / "xfit-gpu.xml") cases = report.findall(".//testcase") - assert len(cases) == 6, len(cases) + assert len(cases) == 41, len(cases) assert not report.findall(".//skipped"), "GPU tests must not skip" assert not report.findall(".//failure") assert not report.findall(".//error") diff --git a/docs/components/xfit.md b/docs/components/xfit.md index 181a2589..913906bd 100644 --- a/docs/components/xfit.md +++ b/docs/components/xfit.md @@ -66,6 +66,25 @@ applies one mask or variance plane to every channel of each candidate. If the batch size is three, `(3, y, x)` keeps that per-candidate meaning; use the explicit `(1, 3, y, x)` shape for per-plane values. +With `backend="cupy"`, `fusion=True` evaluates Gaussian residuals and weighted +normal equations in CuPy CUDA kernels. It supports float32 and float64, +difference and split modes, masks, and variances. The solver uses analytic +derivatives; final uncertainty diagnostics use the full analytic Jacobian. +Each fit owns its scratch buffers and runs on the current CuPy stream. +Select it in the CLI with `--backend cupy --fusion`; the default is +`fusion=False`. Device pipeline manifests select it with `xfit.fusion: true`. +Sampled-stamp and finite-difference fits use the default path. + +The public synthetic benchmark records the first call separately from every +warmed sample, including input transfer and result materialization: + +```bash +uv run --locked --extra gpu python examples/xfit/benchmark_gaussian.py \ + --backend cupy --fusion --batch 16 --height 51 --width 51 --repeat 5 +``` + +Run the same command without `--fusion` for the default CuPy implementation. + An explicit `backend="cutile"` uses one `cuda.tile` CTA per Gaussian fit to form its weighted 8-by-8 normal equations directly. Select this backend explicitly and install the `cuphoton[cutile]` extra on Linux with Python diff --git a/examples/xfit/benchmark_gaussian.py b/examples/xfit/benchmark_gaussian.py index 3a0b2189..b07104aa 100644 --- a/examples/xfit/benchmark_gaussian.py +++ b/examples/xfit/benchmark_gaussian.py @@ -59,7 +59,9 @@ def _fixture( return images, initial -def _initialize(images, initial, mode, workers, counter, barrier, gpu): +def _initialize( + images, initial, mode, workers, counter, barrier, gpu, fusion +): with counter.get_lock(): index = counter.value counter.value += 1 @@ -72,6 +74,7 @@ def _initialize(images, initial, mode, workers, counter, barrier, gpu): _worker.initial = initial[start:stop].copy() _worker.variance = np.ones_like(_worker.images) _worker.mode = mode + _worker.fusion = fusion _worker.barrier = barrier _worker.stream = None _worker.pool = None @@ -113,6 +116,7 @@ def _run(treatment): variance=_worker.variance, mode=_worker.mode, backend=backend, + fusion=_worker.fusion, ) seconds = time.perf_counter() - start if result.backend != backend: @@ -198,6 +202,7 @@ def main() -> None: help="Restrict all workers together to this many inherited CPUs.", ) parser.add_argument("--batch", type=int, action="append") + parser.add_argument("--fusion", action="store_true") parser.add_argument( "--dtype", choices=("float32", "float64"), default="float64" ) @@ -227,7 +232,13 @@ def main() -> None: parser.error("workers and image dimensions must be positive") if args.compare_native and args.backend: parser.error("--compare-native cannot be combined with --backend") - backends = args.backend or ["cupy", "cutile"] + backends = args.backend or ( + ["cupy"] if args.fusion else ["cupy", "cutile"] + ) + if args.fusion and ( + args.compare_native or any(backend != "cupy" for backend in backends) + ): + parser.error("--fusion requires --backend cupy") batches = args.batch or [1, 16, 256, 4096] if any(batch < args.workers for batch in batches): parser.error("every batch must contain at least --workers candidates") @@ -274,17 +285,33 @@ def main() -> None: context.Value("i", 0), context.Barrier(args.workers), any(backend != "numpy" for backend in treatments), + args.fusion, ), **executor_options, ) as pool: _round(pool, None, args.workers, args.timeout_seconds) startup = time.perf_counter() - start warmups, samples, last = {}, {t: [] for t in treatments}, {} + cold = {} try: expected, reference_seconds = _round( pool, reference, args.workers, args.timeout_seconds ) for treatment in treatments: + if treatment == reference: + cold[treatment] = 1e3 * reference_seconds + else: + results, seconds = _round( + pool, + treatment, + args.workers, + args.timeout_seconds, + ) + cold[treatment] = 1e3 * seconds + for actual, baseline in zip( + results, expected, strict=True + ): + _validate(actual[1], baseline[1], dtype) start = time.perf_counter() for _ in range(args.warmup): results, _ = _round( @@ -345,6 +372,9 @@ def main() -> None: json.dumps( { "backend": treatment, + "fusion": args.fusion, + "cold_milliseconds": cold[treatment], + "milliseconds_samples": wall_ms, "execution": args.execution, "workers": args.workers, "batch": batch, diff --git a/src/cuphoton/xfit/_cupy_fused.py b/src/cuphoton/xfit/_cupy_fused.py new file mode 100644 index 00000000..6c2cfe1b --- /dev/null +++ b/src/cuphoton/xfit/_cupy_fused.py @@ -0,0 +1,348 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""CuPy kernels for Gaussian residuals and analytic normal equations.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from ._types import BackendArray, FitMode + +_BLOCK_SIZE = 128 +_NORMAL_COUNT = 44 # Eight gradient and 36 upper-triangular Hessian entries. + +# Keep arithmetic in the input dtype and preserve the model's operation order, +# including physical-sigma derivatives followed by the log-sigma chain rule. +# Disable FMA to avoid contracting separate CuPy operations. +_CUDA_SOURCE = r""" +typedef SCALAR_TYPE T; +const int block_size = 128; +const int normal_count = 44; + +struct Gaussian { + T amplitude, sigma_x, sigma_y, inverse_x2, inverse_y2, cosine, sine; + + __device__ Gaussian(const T* parameters) { + amplitude = parameters[0]; + sigma_x = exp(parameters[1]); + sigma_y = exp(parameters[2]); + const T valid_x = isfinite(sigma_x) && sigma_x > T(0) + ? sigma_x : T(nan("")); + const T valid_y = isfinite(sigma_y) && sigma_y > T(0) + ? sigma_y : T(nan("")); + inverse_x2 = T(1) / (valid_x * valid_x); + inverse_y2 = T(1) / (valid_y * valid_y); + cosine = cos(parameters[3]); + sine = sin(parameters[3]); + } + + template + __device__ T star(T x, T y, T* derivative) const { + const T rotated_x = cosine * x + sine * y; + const T rotated_y = cosine * y - sine * x; + const T exponent = T(-0.5) * ( + rotated_x * rotated_x * inverse_x2 + + rotated_y * rotated_y * inverse_y2); + const T unit = exp(exponent); + const T value = amplitude * unit; + if (derivatives) { + derivative[0] = unit; + derivative[1] = value * rotated_x * rotated_x + / pow(sigma_x, T(3)); + derivative[2] = value * rotated_y * rotated_y + / pow(sigma_y, T(3)); + derivative[3] = -value * rotated_x * rotated_y + * (inverse_x2 - inverse_y2); + derivative[4] = value * ( + cosine * rotated_x * inverse_x2 + - sine * rotated_y * inverse_y2); + derivative[5] = value * ( + sine * rotated_x * inverse_x2 + + cosine * rotated_y * inverse_y2); + } + return value; + } +}; + +extern "C" __global__ void residual( + const T* parameters, const long long* indices, + const T* images, const T* weights, T* output, + long long batch, long long height, long long width, + long long observations) { + const long long offset = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (offset >= batch * observations) return; + const long long fit = offset / observations; + const long long observation = offset % observations; + const long long pixels = height * width; + const int plane = observation / pixels; + const long long pixel = observation % pixels; + const T x = T(pixel % width) - T(width - 1) / T(2); + const T y = T(pixel / width) - T(height - 1) / T(2); + const T* p = parameters + fit * 8; + const Gaussian model(p); + const T positive = model.star(x - p[4], y - p[5], nullptr); + const T negative = model.star(x - p[6], y - p[7], nullptr); + const T prediction = plane == 0 ? positive - negative + : (plane == 1 ? positive : negative); + const long long source = indices[fit] * observations + observation; + output[offset] = (prediction - images[source]) * weights[source]; +} + +__device__ T warp_sum(T value) { + #pragma unroll + for (int offset = 16; offset > 0; offset /= 2) + value += __shfl_down_sync(0xffffffff, value, offset); + return value; +} + +extern "C" __global__ void normal_partials( + const T* parameters, const long long* indices, + const T* residuals, const T* weights, T* partials, + long long height, long long width, long long observations, + long long tile_count) { + const long long fit = blockIdx.x / tile_count; + const long long tile = blockIdx.x % tile_count; + const long long observation = tile * block_size + threadIdx.x; + const int lane = threadIdx.x % 32; + const int warp = threadIdx.x / 32; + __shared__ T warp_partials[block_size / 32][normal_count]; + T derivative[8] = {}; + T residual_value = T(0); + + if (observation < observations) { + const long long pixels = height * width; + const int plane = observation / pixels; + const long long pixel = observation % pixels; + const T x = T(pixel % width) - T(width - 1) / T(2); + const T y = T(pixel / width) - T(height - 1) / T(2); + const T* p = parameters + fit * 8; + const Gaussian model(p); + T positive[6], negative[6]; + model.star(x - p[4], y - p[5], positive); + model.star(x - p[6], y - p[7], negative); + #pragma unroll + for (int parameter = 0; parameter < 4; ++parameter) + derivative[parameter] = plane == 0 + ? positive[parameter] - negative[parameter] + : (plane == 1 ? positive[parameter] : negative[parameter]); + derivative[4] = plane < 2 ? positive[4] : T(0); + derivative[5] = plane < 2 ? positive[5] : T(0); + derivative[6] = plane == 0 ? -negative[4] + : (plane == 2 ? negative[4] : T(0)); + derivative[7] = plane == 0 ? -negative[5] + : (plane == 2 ? negative[5] : T(0)); + derivative[1] *= model.sigma_x; + derivative[2] *= model.sigma_y; + const T weight = weights[indices[fit] * observations + observation]; + #pragma unroll + for (int parameter = 0; parameter < 8; ++parameter) + derivative[parameter] *= weight; + residual_value = residuals[fit * observations + observation]; + } + + #pragma unroll + for (int parameter = 0; parameter < 8; ++parameter) { + const T total = warp_sum(derivative[parameter] * residual_value); + if (lane == 0) warp_partials[warp][parameter] = total; + } + int component = 8; + #pragma unroll + for (int row = 0; row < 8; ++row) { + #pragma unroll + for (int column = row; column < 8; ++column) { + const T total = warp_sum(derivative[row] * derivative[column]); + if (lane == 0) warp_partials[warp][component] = total; + ++component; + } + } + __syncthreads(); + if (threadIdx.x < normal_count) { + T total = warp_partials[0][threadIdx.x]; + #pragma unroll + for (int source_warp = 1; source_warp < block_size / 32; + ++source_warp) + total += warp_partials[source_warp][threadIdx.x]; + partials[(fit * tile_count + tile) * normal_count + threadIdx.x] + = total; + } +} + +extern "C" __global__ void normal_finish( + const T* partials, T* gradients, T* hessians, long long tile_count) { + const long long fit = blockIdx.x; + const int component = threadIdx.x; + if (component >= normal_count) return; + T total = T(0); + for (long long tile = 0; tile < tile_count; ++tile) + total += partials[ + (fit * tile_count + tile) * normal_count + component]; + if (component < 8) { + gradients[fit * 8 + component] = total; + } else { + int column = component - 8; + int row = 0; + while (column >= 8 - row) { + column -= 8 - row; + ++row; + } + column += row; + hessians[fit * 64 + row * 8 + column] = total; + hessians[fit * 64 + column * 8 + row] = total; + } +} +""" + + +class GaussianFusedProblem: + """Own CuPy Gaussian callbacks and input storage for a single fit call. + + The caller supplies sanitized images and precomputed multiplicative + weights. Parameters use log sigma, and indices refer to original input + rows. Callbacks return fresh arrays because the LM solver retains earlier + residuals while evaluating trial steps. CuPy is supplied by the resolved + backend; importing this module does not import CuPy or initialize CUDA. + """ + + def __init__( + self, + cp: Any, + images: BackendArray, + weights: BackendArray, + image_shape: tuple[int, int], + mode: FitMode, + ) -> None: + if mode not in ("difference", "split"): + raise ValueError("mode must be 'difference' or 'split'") + height, width = image_shape + if height < 1 or width < 1: + raise ValueError("image_shape dimensions must be positive") + self._dtype = np.dtype(images.dtype) + if self._dtype not in (np.dtype(np.float32), np.dtype(np.float64)): + raise TypeError( + "fused Gaussian kernels require float32 or float64" + ) + expected_shape = ( + (height, width) if mode == "difference" else (3, height, width) + ) + if ( + images.shape[1:] != expected_shape + or weights.shape != images.shape + ): + raise ValueError( + "images and weights must match the Gaussian mode" + ) + self._cp = cp + self._height = np.int64(height) + self._width = np.int64(width) + self._observations = int(np.prod(expected_shape)) + self._tiles = (self._observations + _BLOCK_SIZE - 1) // _BLOCK_SIZE + self._partial_capacity = 0 + self._partials: Any = None + self._images = cp.ascontiguousarray(images) + self._weights = cp.ascontiguousarray(weights, dtype=self._dtype) + scalar_type = ( + "float" if self._dtype == np.dtype(np.float32) else "double" + ) + self._module = cp.RawModule( + code=_CUDA_SOURCE.replace("SCALAR_TYPE", scalar_type), + options=("--std=c++11", "--fmad=false"), + ) + self._residual_kernel = self._module.get_function("residual") + self._normal_kernel = self._module.get_function("normal_partials") + self._finish_kernel = self._module.get_function("normal_finish") + + def _inputs( + self, parameters: BackendArray, indices: BackendArray + ) -> tuple[Any, Any]: + batch = parameters.shape[0] + if parameters.shape != (batch, 8): + raise ValueError("Gaussian parameters must have shape (batch, 8)") + if indices.shape != (batch,): + raise ValueError("indices must contain one original row per fit") + return ( + self._cp.ascontiguousarray(parameters, dtype=self._dtype), + self._cp.ascontiguousarray(indices, dtype=np.int64), + ) + + def residual( + self, parameters: BackendArray, *, indices: BackendArray + ) -> BackendArray: + """Evaluate weighted residuals with one kernel launch.""" + + parameters, indices = self._inputs(parameters, indices) + batch = parameters.shape[0] + output = self._cp.empty( + (batch, self._observations), dtype=self._dtype + ) + if batch: + blocks = ( + batch * self._observations + _BLOCK_SIZE - 1 + ) // _BLOCK_SIZE + self._residual_kernel( + (blocks,), + (_BLOCK_SIZE,), + ( + parameters, + indices, + self._images, + self._weights, + output, + np.int64(batch), + self._height, + self._width, + np.int64(self._observations), + ), + ) + return output + + def normal_equations( + self, + parameters: BackendArray, + residuals: BackendArray, + *, + indices: BackendArray, + ) -> tuple[BackendArray, BackendArray]: + """Accumulate gradient and Hessian without a global Jacobian array.""" + + parameters, indices = self._inputs(parameters, indices) + batch = parameters.shape[0] + if residuals.shape != (batch, self._observations): + raise ValueError("residuals must match the active observations") + residuals = self._cp.ascontiguousarray(residuals, dtype=self._dtype) + gradient = self._cp.empty((batch, 8), dtype=self._dtype) + hessian = self._cp.empty((batch, 8, 8), dtype=self._dtype) + if batch: + tiles = self._tiles + if batch > self._partial_capacity: + self._partials = self._cp.empty( + (batch, tiles, _NORMAL_COUNT), dtype=self._dtype + ) + self._partial_capacity = batch + # Calls belong to one fit and execute on its current stream. The + # previous reduction finishes before this scratch is overwritten. + partials = self._partials + self._normal_kernel( + (batch * tiles,), + (_BLOCK_SIZE,), + ( + parameters, + indices, + residuals, + self._weights, + partials, + self._height, + self._width, + np.int64(self._observations), + np.int64(tiles), + ), + ) + self._finish_kernel( + (batch,), + (64,), + (partials, gradient, hessian, np.int64(tiles)), + ) + return gradient, hessian diff --git a/src/cuphoton/xfit/api.py b/src/cuphoton/xfit/api.py index 8e6fa971..78fbb999 100644 --- a/src/cuphoton/xfit/api.py +++ b/src/cuphoton/xfit/api.py @@ -605,6 +605,27 @@ def _validate_cupy_stamp_model_dtype( ) +def _validate_fusion( + fusion: bool, + *, + backend: BackendRequest, + model: ModelName | GaussianDipoleModel | StampDipoleModel, + config: LMConfig | None, +) -> None: + if not isinstance(fusion, bool): + raise TypeError("fusion must be boolean") + if not fusion: + return + if backend != "cupy": + raise ValueError("fusion requires backend='cupy'") + if isinstance(model, StampDipoleModel) or ( + isinstance(model, str) and model != "gaussian" + ): + raise ValueError("fusion supports only the Gaussian model") + if config is not None and config.use_finite_difference: + raise ValueError("fusion does not support finite-difference fitting") + + def _fit_dipoles_backend( images: ArrayLike, *, @@ -615,6 +636,7 @@ def _fit_dipoles_backend( mode: FitMode = "difference", resolved: Backend, config: LMConfig | None = None, + fusion: bool = False, ) -> _BackendDipoleFitResult: """Fit dipoles while retaining all result arrays in ``resolved``.""" @@ -791,8 +813,23 @@ def normal_equations( mode=mode, ) + if fusion: + from ._cupy_fused import GaussianFusedProblem + + fused = GaussianFusedProblem( + ap, + images_array, + weights, + image_shape=(height, width), + mode=mode, + ) + residual_function = fused.residual + normal_equations = fused.normal_equations + else: + residual_function = residual + problem = BatchedLeastSquaresProblem( - residual, + residual_function, jacobian_function, normal_equations=normal_equations, ) @@ -1028,6 +1065,7 @@ def fit_dipoles( mode: FitMode = "difference", backend: BackendRequest = "auto", config: LMConfig | None = None, + fusion: bool = False, ) -> DipoleFitResult: """Fit batched Gaussian or sampled-stamp dipoles. @@ -1049,9 +1087,13 @@ def fit_dipoles( treats all three planes as independent; when difference equals positive minus negative, the reported degrees of freedom and uncertainties are not statistically calibrated for that dependence. + + ``fusion=True`` selects fused Gaussian residual and normal-equation + kernels with ``backend="cupy"``. It requires analytic derivatives. """ _validate_execution_options(model, backend, config) + _validate_fusion(fusion, backend=backend, model=model, config=config) resolved = resolve_backend(backend) backend_result = _fit_dipoles_backend( images, @@ -1062,6 +1104,7 @@ def fit_dipoles( mode=mode, resolved=resolved, config=config, + fusion=fusion, ) return _materialize_dipole_fit_result(backend_result) @@ -1076,6 +1119,7 @@ def fit_dipoles_device( mode: FitMode = "difference", config: LMConfig | None = None, backend: Literal["cupy", "native"] = "cupy", + fusion: bool = False, ) -> DeviceDipoleFitResult: """Fit dipoles while keeping every result array on the active GPU. @@ -1084,6 +1128,9 @@ def fit_dipoles_device( to the active device. Existing CuPy inputs must already reside on that device and are checked before any input conversion. Input arrays are borrowed and are not mutated. + ``fusion=True`` selects fused CuPy Gaussian residual and normal-equation + kernels; native fits, sampled stamps and finite differences require + ``fusion=False``. The solver may copy per-fit status, evaluation counts and boolean diagnostic masks to the host for control flow. Parameters, residuals, @@ -1098,6 +1145,7 @@ def fit_dipoles_device( """ _validate_execution_options(model, backend, config) + _validate_fusion(fusion, backend=backend, model=model, config=config) if backend not in {"cupy", "native"}: raise ValueError("device fit backend must be cupy or native") if backend == "native": @@ -1134,6 +1182,7 @@ def fit_dipoles_device( mode=mode, resolved=resolve_backend(backend), config=config, + fusion=fusion, ) return DeviceDipoleFitResult( parameters=result.parameters, diff --git a/src/cuphoton/xfit/commands.py b/src/cuphoton/xfit/commands.py index 8d19d3c8..e37f2030 100644 --- a/src/cuphoton/xfit/commands.py +++ b/src/cuphoton/xfit/commands.py @@ -236,6 +236,7 @@ class FitDipolesCommand(ExecutorOptions, _ValidatedDatasetCommand): g_tol = None max_evaluations = None use_finite_difference = None + fusion: bool = False chunk_size = None class ChunkSizeArg(PositiveIntegerInvariant): @@ -258,6 +259,15 @@ class BackendArg(BackendInvariant): _mandatory = False _default = "auto" + class FusionArg(BoolInvariant): + _arg = "--fusion" + _help = ( + "Fuse Gaussian residuals and normal equations; " + "requires --backend cupy." + ) + _mandatory = False + _default = False + class ComputeDtypeArg(ComputeDtypeInvariant): _arg = "--compute-dtype" _help = ( @@ -330,6 +340,20 @@ def _model_specification( ) def run(self) -> None: + from . import LMConfig + from .api import _validate_fusion + + assert self.backend is not None + assert self.model is not None + self._call( + _validate_fusion, + self.fusion, + backend=self.backend, + model=self.model, + config=LMConfig( + use_finite_difference=bool(self.use_finite_difference) + ), + ) executor_options = self.executor_options() if self.executor != "local": from cuphoton.core.executors import run_workload @@ -362,6 +386,7 @@ def run(self) -> None: "g_tol", "max_evaluations", "use_finite_difference", + "fusion", "fits_reader", ) } @@ -389,7 +414,7 @@ def run(self) -> None: "--chunk-size requires --executor dragon or mpi" ) - from . import LMConfig, fit_dipoles + from . import fit_dipoles from .api import _validate_execution_options assert self.output_dir is not None @@ -439,6 +464,7 @@ def run(self) -> None: mode=self.mode, backend=self.backend, config=config, + fusion=self.fusion, ) dtype = np.dtype(result.dtype) default_tolerance = float(np.sqrt(np.finfo(dtype).eps)) @@ -451,6 +477,7 @@ def run(self) -> None: "model": self.model, "mode": self.mode, "backend": self.backend, + "fusion": self.fusion, **( {"fits_reader": self.fits_reader} if dataset.input_sources diff --git a/src/cuphoton/xfit/executor.py b/src/cuphoton/xfit/executor.py index cbff8f8b..2461c52c 100644 --- a/src/cuphoton/xfit/executor.py +++ b/src/cuphoton/xfit/executor.py @@ -26,7 +26,12 @@ ) from cuphoton.core.fits_options import normalize_xdr_options -from .api import DipoleFitResult, _floating_dtype, fit_dipoles +from .api import ( + DipoleFitResult, + _floating_dtype, + _validate_fusion, + fit_dipoles, +) from .io import ( XFitDataset, _broadcast_image_auxiliary, @@ -62,6 +67,7 @@ "g_tol": None, "max_evaluations": None, "use_finite_difference": False, + "fusion": False, "fits_reader": "auto", "xdr_options": None, } @@ -281,6 +287,7 @@ def run_item(self, item: WorkItem, item_dir: Path) -> Mapping[str, Any]: mode=self.settings["mode"], backend=self.settings["backend"], config=self.config, + fusion=self.settings["fusion"], ) solve_seconds = time.perf_counter() - solve_start _check_input_stat(self.options) @@ -443,6 +450,7 @@ def _effective_config( "model": settings["model"], "mode": settings["mode"], "backend": settings["backend"], + "fusion": settings["fusion"], **( {"fits_reader": settings["fits_reader"]} if dataset.input_sources @@ -613,6 +621,14 @@ def prepare_xfit_workload( settings = {**_DEFAULTS, **fit_options} backend = settings["backend"] + _validate_fusion( + settings["fusion"], + backend=backend, + model=settings["model"], + config=LMConfig( + use_finite_difference=settings["use_finite_difference"] + ), + ) if backend not in {"cupy", "cutile", "native"}: raise ValueError( "distributed xFit requires backend cupy, cutile, or native" diff --git a/src/cuphoton/xscan/device_pipeline.py b/src/cuphoton/xscan/device_pipeline.py index 37754c1b..77973e33 100644 --- a/src/cuphoton/xscan/device_pipeline.py +++ b/src/cuphoton/xscan/device_pipeline.py @@ -651,6 +651,7 @@ class DeviceXFitPipelineConfig: finite_difference_step: float | None = None use_finite_difference: bool = False backend: Literal["cupy", "native"] = "cupy" + fusion: bool = False def __post_init__(self) -> None: if self.model != "gaussian": @@ -661,6 +662,12 @@ def __post_init__(self) -> None: raise ValueError("xfit backend must be 'cupy' or 'native'") if not isinstance(self.use_finite_difference, bool): raise TypeError("xfit use_finite_difference must be boolean") + if not isinstance(self.fusion, bool): + raise TypeError("xfit fusion must be boolean") + if self.fusion and self.backend != "cupy": + raise ValueError("xfit fusion requires backend='cupy'") + if self.fusion and self.use_finite_difference: + raise ValueError("xfit fusion requires analytic derivatives") if self.backend == "native" and self.use_finite_difference: raise ValueError( "native xfit does not support finite differences" @@ -715,6 +722,8 @@ def to_payload(self) -> dict[str, Any]: payload = asdict(self) if self.backend == "cupy": del payload["backend"] + if not self.fusion: + del payload["fusion"] return payload def execution_payload(self) -> dict[str, Any]: @@ -723,6 +732,8 @@ def execution_payload(self) -> dict[str, Any]: payload: dict[str, Any] = {} if self.backend != "cupy": payload["backend"] = self.backend + if self.fusion: + payload["fusion"] = True return payload def solver_payload(self) -> dict[str, Any]: @@ -732,6 +743,7 @@ def solver_payload(self) -> dict[str, Any]: del payload["model"] del payload["mode"] payload.pop("backend", None) + payload.pop("fusion", None) return payload @classmethod @@ -742,6 +754,7 @@ def from_payload( values = json_mapping(payload, field="device pipeline xfit config") values.setdefault("backend", "cupy") + values.setdefault("fusion", False) expected = frozenset(value.name for value in fields(cls)) values = _require_exact_fields( values, diff --git a/src/cuphoton/xscan/pipeline_benchmark/runner.py b/src/cuphoton/xscan/pipeline_benchmark/runner.py index 543bc7b2..b5ab8948 100644 --- a/src/cuphoton/xscan/pipeline_benchmark/runner.py +++ b/src/cuphoton/xscan/pipeline_benchmark/runner.py @@ -428,6 +428,7 @@ def provenance() -> dict[str, Any]: "xfit/backend.py", "xfit/_native.py", "xfit/models.py", + "xfit/_cupy_fused.py", "xfit/solver.py", "xscan/xfit_features.py", "xscan/training.py", diff --git a/tests/core/test_cli_contract.py b/tests/core/test_cli_contract.py index ad4fdec1..ca2493aa 100644 --- a/tests/core/test_cli_contract.py +++ b/tests/core/test_cli_contract.py @@ -71,7 +71,7 @@ def test_public_command_surface_counts_are_exact() -> None: assert per_group == [ ("xdr", 1, 0, 17, 0), - ("xfit", 3, 3, 38, 1), + ("xfit", 3, 3, 39, 1), ("xpois", 7, 7, 162, 1), ("xscan", 44, 44, 229, 1), ("xrep", 6, 6, 123, 1), @@ -81,7 +81,7 @@ def test_public_command_surface_counts_are_exact() -> None: assert sum(item[1] for item in per_group) == 94 assert sum(item[1] + item[4] for item in per_group) == 99 assert sum(item[2] for item in per_group) == 91 - assert sum(item[3] for item in per_group) == 961 + assert sum(item[3] for item in per_group) == 962 def test_public_registry_order_and_component_derivations() -> None: diff --git a/tests/test_public_api.py b/tests/test_public_api.py index 6787184b..9d54d87e 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -123,6 +123,7 @@ def test_xfit_curated_exports_and_fit_signature() -> None: "mode", "backend", "config", + "fusion", ] assert list(inspect.signature(xfit.fit_dipoles_device).parameters) == [ "images", @@ -133,6 +134,7 @@ def test_xfit_curated_exports_and_fit_signature() -> None: "mode", "config", "backend", + "fusion", ] assert [field.name for field in fields(xfit.DipoleFitResult)] == [ "parameters", diff --git a/tests/xfit/test_cupy_fusion.py b/tests/xfit/test_cupy_fusion.py new file mode 100644 index 00000000..a928563a --- /dev/null +++ b/tests/xfit/test_cupy_fusion.py @@ -0,0 +1,469 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from dataclasses import fields + +import numpy as np +import pytest +import yaml + +import cuphoton.xfit.api as xfit_api +from cuphoton.core.cli import run_component +from cuphoton.xfit import ( + GaussianDipoleModel, + LMConfig, + StampDipoleModel, + fit_dipoles, + fit_dipoles_device, +) + + +def _require_cupy(): + cp = pytest.importorskip("cupy") + try: + count = cp.cuda.runtime.getDeviceCount() + except Exception as exc: + pytest.skip(f"CuPy CUDA runtime is not usable: {exc}") + if count < 1: + pytest.skip("test requires a visible CUDA device") + return cp + + +def _truth(dtype=np.float64): + return np.asarray( + [ + [5.0, 1.3, 1.8, 0.25, -2.0, 0.5, 2.2, -0.7], + [3.0, 1.6, 1.1, -0.3, -1.0, -1.0, 2.0, 1.0], + [4.2, 1.4, 2.0, 0.1, -1.7, 0.8, 1.9, -0.9], + ], + dtype=dtype, + ) + + +def _assert_numeric(actual, expected, dtype): + # Match the existing same-backend Gaussian normal-equation tolerances. + tolerance = 3e-5 if np.dtype(dtype) == np.dtype(np.float32) else 5e-12 + assert actual.shape == expected.shape + assert actual.dtype == expected.dtype + np.testing.assert_array_equal(np.isnan(actual), np.isnan(expected)) + np.testing.assert_array_equal(np.isposinf(actual), np.isposinf(expected)) + np.testing.assert_array_equal(np.isneginf(actual), np.isneginf(expected)) + np.testing.assert_allclose( + actual, expected, rtol=tolerance, atol=tolerance, equal_nan=True + ) + + +def _assert_result(actual, expected): + for field in fields(expected): + observed = getattr(actual, field.name) + reference = getattr(expected, field.name) + if isinstance(reference, np.ndarray): + if reference.dtype.kind == "f": + _assert_numeric(observed, reference, expected.dtype) + else: + np.testing.assert_array_equal( + observed, reference, err_msg=field.name + ) + else: + assert observed == reference, field.name + + +def _no_cuda(*args, **kwargs): + pytest.fail("unsupported fusion settings reached CUDA/backend discovery") + + +@pytest.mark.parametrize("backend", ["numpy", "auto", "cutile"]) +def test_fusion_rejects_unsupported_backend_before_cuda(monkeypatch, backend): + monkeypatch.setattr(xfit_api, "resolve_backend", _no_cuda) + with pytest.raises(ValueError, match="fusion requires backend='cupy'"): + fit_dipoles( + np.zeros((1, 5, 7)), + model="gaussian", + backend=backend, + fusion=True, + ) + + +@pytest.mark.parametrize("device", [False, True]) +@pytest.mark.parametrize("case", ["stamp", "finite_difference", "nonboolean"]) +def test_fusion_rejects_unsupported_configuration_before_cuda( + monkeypatch, device, case +): + model = "gaussian" + config = None + fusion = True + exception = ValueError + match = "finite-difference" + if case == "stamp": + model = StampDipoleModel(np.ones((3, 3)), image_shape=(5, 7)) + match = "Gaussian model" + elif case == "finite_difference": + config = LMConfig(use_finite_difference=True) + else: + fusion = 1 + exception = TypeError + match = "fusion must be boolean" + monkeypatch.setattr(xfit_api, "resolve_backend", _no_cuda) + monkeypatch.setattr(xfit_api, "_load_cupy", _no_cuda) + kwargs = {} if device else {"backend": "cupy"} + with pytest.raises(exception, match=match): + (fit_dipoles_device if device else fit_dipoles)( + np.zeros((1, 5, 7)), + model=model, + config=config, + fusion=fusion, + **kwargs, + ) + + +@pytest.mark.parametrize( + "arguments,message", + [ + (["--backend", "numpy"], "fusion requires backend='cupy'"), + (["--backend", "cupy", "--model", "stamp"], "Gaussian model"), + ( + ["--backend", "cupy", "--use-finite-difference"], + "finite-difference", + ), + ], +) +def test_fusion_cli_rejects_options_before_loading_or_backend_probe( + tmp_path, capsys, monkeypatch, arguments, message +): + import cuphoton.xfit.backend as backend + from cuphoton.xfit.commands import FitDipolesCommand + + monkeypatch.setattr(backend, "resolve_backend", _no_cuda) + monkeypatch.setattr(xfit_api, "resolve_backend", _no_cuda) + monkeypatch.setattr(FitDipolesCommand, "_load_dataset", _no_cuda) + input_path = tmp_path / "unread-manifest.json" + input_path.write_text("{}") + rc = run_component( + "xfit", + [ + "fit-dipoles", + "--input", + str(input_path), + "--output-dir", + str(tmp_path / "output"), + "--model", + "gaussian", + "--fusion", + *arguments, + ], + ) + captured = capsys.readouterr() + assert rc == 1 + assert message in captured.err + assert not (tmp_path / "output").exists() + + +def _reference_callbacks(cp, parameters, images, weights, indices, mode): + physical = parameters.copy() + physical[:, 1:3] = cp.exp(parameters[:, 1:3]) + model = GaussianDipoleModel( + images.shape[-2:], backend="cupy", dtype=images.dtype + ) + prediction = model._evaluate_positive_unchecked(physical, mode=mode) + residual = ((prediction - images[indices]) * weights[indices]).reshape( + parameters.shape[0], -1 + ) + jacobian = model._jacobian_positive_unchecked(physical, mode=mode) + chain_shape = (parameters.shape[0], 2) + (1,) * (jacobian.ndim - 2) + jacobian[:, 1:3] *= physical[:, 1:3].reshape(chain_shape) + jacobian *= weights[indices, None, ...] + return residual, jacobian.reshape(parameters.shape[0], 8, -1) + + +@pytest.mark.parametrize("shape", [(11, 15), (8, 10)]) +@pytest.mark.parametrize("mode", ["difference", "split"]) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_cupy_fusion_callbacks_preserve_weighting_indices_and_layout( + shape, mode, dtype +): + cp = _require_cupy() + from cuphoton.xfit._cupy_fused import GaussianFusedProblem + + rng = np.random.default_rng(1219) + truth = _truth(dtype) + truth[2, 0] *= -1 + model = GaussianDipoleModel(shape, dtype=dtype) + prediction = model.evaluate(truth, mode=mode) + images = cp.asarray( + prediction + rng.normal(scale=0.2, size=prediction.shape), dtype=dtype + )[..., ::-1] + weights = cp.asarray( + rng.uniform(0.25, 2.0, size=prediction.shape), dtype=dtype + )[..., ::-1] + weights[..., ::3, ::2] = 0 + indices = cp.asarray([2, 0], dtype=cp.int64) + parameters = cp.asarray(truth)[indices] + parameters[:, 1:3] = cp.log(parameters[:, 1:3]) + problem = GaussianFusedProblem( + cp, images, weights, image_shape=shape, mode=mode + ) + expected, jacobian = _reference_callbacks( + cp, parameters, images, weights, indices, mode + ) + actual = problem.residual(parameters, indices=indices) + _assert_numeric(cp.asnumpy(actual), cp.asnumpy(expected), dtype) + + # A supplied residual must be consumed, not silently recomputed. + residual = expected + dtype(0.03) + gradient, hessian = problem.normal_equations( + parameters, residual, indices=indices + ) + expected_gradient = cp.einsum("knm,km->kn", jacobian, residual) + expected_hessian = cp.einsum("knm,kpm->knp", jacobian, jacobian) + _assert_numeric( + cp.asnumpy(gradient), cp.asnumpy(expected_gradient), dtype + ) + _assert_numeric(cp.asnumpy(hessian), cp.asnumpy(expected_hessian), dtype) + np.testing.assert_array_equal( + cp.asnumpy(hessian), cp.asnumpy(hessian).swapaxes(1, 2) + ) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_cupy_fusion_extreme_widths_preserve_nonfinite_results(dtype): + cp = _require_cupy() + from cuphoton.xfit._cupy_fused import GaussianFusedProblem + + parameters = cp.asarray(_truth(dtype)) + parameters[:, 3] = 0 + parameters[:, 4:8] = cp.asarray([-1, 0, 1, 0], dtype=dtype) + small = 1e-16 if dtype is np.float32 else 1e-120 + parameters[:, 1:3] = cp.asarray( + [[-1000, -1000], [1000, 1000], [np.log(small), np.log(small)]], + dtype=dtype, + ) + images = cp.zeros((3, 3, 5, 7), dtype=dtype) + weights = cp.ones_like(images) + weights[0] = 0 + indices = cp.arange(3, dtype=cp.int64) + problem = GaussianFusedProblem( + cp, images, weights, image_shape=(5, 7), mode="split" + ) + residual, jacobian = _reference_callbacks( + cp, parameters, images, weights, indices, "split" + ) + actual = problem.residual(parameters, indices=indices) + _assert_numeric(cp.asnumpy(actual), cp.asnumpy(residual), dtype) + gradient, hessian = problem.normal_equations( + parameters, residual, indices=indices + ) + _assert_numeric( + cp.asnumpy(gradient), + cp.asnumpy(cp.einsum("knm,km->kn", jacobian, residual)), + dtype, + ) + _assert_numeric( + cp.asnumpy(hessian), + cp.asnumpy(cp.einsum("knm,kpm->knp", jacobian, jacobian)), + dtype, + ) + + +@pytest.mark.parametrize("mode", ["difference", "split"]) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_cupy_fusion_full_results_preserve_compacted_rows(mode, dtype): + _require_cupy() + truth = _truth(dtype) + model = GaussianDipoleModel( + (11, 15), split_weights=(1.0, 0.25, 1.75), dtype=dtype + ) + images = model.evaluate(truth, mode=mode) + initial = truth + np.asarray( + [0.2, 0.1, -0.1, 0.03, 0.1, -0.1, -0.1, 0.1], dtype=dtype + ) + initial[0] = truth[0] + mask = np.ones(images.shape, dtype=bool) + mask[..., 0, :3] = False + images[..., 0, :3] = np.nan + variance = np.linspace(0.5, 1.5, images.size, dtype=dtype).reshape( + images.shape + ) + variance[~mask] = np.nan + kwargs = dict( + model=model, + initial=initial, + mask=mask, + variance=variance, + mode=mode, + backend="cupy", + ) + expected = fit_dipoles(images, **kwargs) + actual = fit_dipoles(images, fusion=True, **kwargs) + assert expected.evaluations[0] == 1 + assert np.all(expected.evaluations[1:] > 1) + assert expected.uncertainty_valid.all() + assert np.max(np.abs(expected.covariance)) > 1e-3 + _assert_result(actual, expected) + + +@pytest.mark.parametrize("case", ["masked", "rank_deficient", "exhausted"]) +def test_cupy_fusion_preserves_degenerate_and_budget_results(case): + _require_cupy() + truth = _truth()[:1] + model = GaussianDipoleModel((5, 7), split_weights=(1.0, 0.0, 0.5)) + if case == "rank_deficient": + truth[:, 0] = 0 + images = model.evaluate(truth, mode="split") + mask = np.ones(images.shape, dtype=bool) + if case == "masked": + mask[:] = False + images[:] = np.nan + initial = truth.copy() + config = None + if case == "exhausted": + initial[:, 0] += 0.2 + config = LMConfig(max_evaluations=1) + kwargs = dict( + model=model, + initial=initial, + mask=mask, + mode="split", + config=config, + backend="cupy", + ) + expected = fit_dipoles(images, **kwargs) + actual = fit_dipoles(images, fusion=True, **kwargs) + assert not actual.uncertainty_valid.any() + _assert_result(actual, expected) + + +@pytest.mark.parametrize( + ("dtype", "amplitude"), [(np.float32, 1e25), (np.float64, 1e200)] +) +@pytest.mark.parametrize("nonzero_residual", [False, True]) +def test_cupy_fusion_preserves_finite_jacobian_overflow_results( + dtype, amplitude, nonzero_residual +): + cp = _require_cupy() + model = GaussianDipoleModel((5, 7), backend="cupy", dtype=dtype) + # Coincident stars give an exactly zero difference while retaining finite, + # large position derivatives. Only the middle row overflows J.T @ J. + initial = np.asarray( + [ + [5, 1, 1, 0, 0, 0, 0, 0], + [amplitude, 1, 1, 0, 0, 0, 0, 0], + [3, 1, 1, 0, 0, 0, 0, 0], + ], + dtype=dtype, + ) + images = np.zeros((3, 5, 7), dtype=dtype) + if nonzero_residual: + # A noncentral pixel also overflows J.T @ r, exercising the solver + # after its gradient stopping check rather than only at zero residual. + images[1, 1, 2] = amplitude + jacobian = model.jacobian(cp.asarray(initial)).reshape(3, 8, -1) + assert bool(cp.isfinite(jacobian).all()) + hessian = cp.einsum("knm,kpm->knp", jacobian, jacobian) + np.testing.assert_array_equal( + cp.asnumpy(cp.isfinite(hessian).all(axis=(1, 2))), + [True, False, True], + ) + + kwargs = dict(model=model, initial=initial, backend="cupy") + expected = fit_dipoles(images, **kwargs) + actual = fit_dipoles(images, fusion=True, **kwargs) + assert expected.status.tolist() == [ + "converged_g_tol", + "singular" if nonzero_residual else "converged_g_tol", + "converged_g_tol", + ] + np.testing.assert_array_equal(expected.evaluations, [1, 1, 1]) + np.testing.assert_array_equal(actual.parameters, expected.parameters) + _assert_result(actual, expected) + + +def test_cupy_fusion_repeated_calls_and_independent_worker_streams(): + cp = _require_cupy() + device = cp.cuda.runtime.getDevice() + + def run_worker(offset): + with cp.cuda.Device(device), cp.cuda.Stream(non_blocking=True): + saved = [] + for shift in (offset, offset + 0.4): + truth = _truth()[:2].copy() + truth[:, 4:8] += shift + model = GaussianDipoleModel((9, 13)) + images = cp.asarray(model.evaluate(truth)) + initial = cp.asarray(truth + 0.025) + expected = fit_dipoles( + images, model=model, initial=initial, backend="cupy" + ) + actual = fit_dipoles_device( + images, model=model, initial=initial, fusion=True + ) + snapshot = cp.asnumpy(actual.parameters) + _assert_numeric(snapshot, expected.parameters, np.float64) + saved.append((actual, snapshot)) + for actual, snapshot in saved: + np.testing.assert_array_equal( + cp.asnumpy(actual.parameters), snapshot + ) + cp.cuda.get_current_stream().synchronize() + return saved[-1][1] + + with ThreadPoolExecutor(max_workers=2) as executor: + results = list(executor.map(run_worker, [0.0, 0.3])) + assert not np.array_equal(results[0], results[1]) + + +def test_cupy_fusion_cli_runs_selected_mode_and_persists_outputs( + tmp_path, capsys +): + _require_cupy() + truth = _truth()[:2] + model = GaussianDipoleModel((9, 13)) + images = model.evaluate(truth) + initial = truth + 0.025 + input_path = tmp_path / "input.npz" + np.savez( + input_path, + candidate_id=np.asarray(["first", "second"]), + images=images, + initial=initial, + variance=np.ones_like(images), + ) + output = tmp_path / "fused" + rc = run_component( + "xfit", + [ + "fit-dipoles", + "--input", + str(input_path), + "--output-dir", + str(output), + "--model", + "gaussian", + "--backend", + "cupy", + "--fusion", + ], + ) + captured = capsys.readouterr() + assert rc == 0, captured.err + configuration = yaml.safe_load( + (output / "effective-config.yaml").read_text() + ) + assert configuration["fusion"] is True + assert configuration["backend"] == "cupy" + expected = fit_dipoles( + images, + model=model, + initial=initial, + variance=np.ones_like(images), + backend="cupy", + ) + with np.load(output / "fit-arrays.npz", allow_pickle=False) as result: + np.testing.assert_array_equal( + result["candidate_id"], ["first", "second"] + ) + _assert_numeric(result["residuals"], expected.residuals, np.float64) + _assert_numeric(result["covariance"], expected.covariance, np.float64) diff --git a/tests/xscan/test_device_pipeline.py b/tests/xscan/test_device_pipeline.py index a4f41f7a..19e175ba 100644 --- a/tests/xscan/test_device_pipeline.py +++ b/tests/xscan/test_device_pipeline.py @@ -14,6 +14,7 @@ import subprocess import sys import weakref +from contextlib import nullcontext from dataclasses import replace from pathlib import Path from types import SimpleNamespace @@ -1367,6 +1368,7 @@ def fit(images: Any, **kwargs: Any) -> Any: "model": "gaussian", "mode": "difference", "config": context.solver_config, + "fusion": context.config.xfit.fusion, } stage_calls.append("xfit") return _fake_xfit_result(cp, 2) @@ -1761,6 +1763,111 @@ def test_xfit_config_rejects_values_the_solver_cannot_use() -> None: DeviceXFitPipelineConfig(finite_difference_step=0.0) +def test_xfit_fusion_configuration_roundtrips_and_defaults_old_payloads( + tmp_path: Path, +) -> None: + checkpoint_dir = tmp_path / "checkpoint" + checkpoint_dir.mkdir() + (checkpoint_dir / "checkpoint.pt").write_bytes(b"checkpoint") + ordinary = _config(checkpoint_dir) + fused = replace(ordinary, xfit=replace(ordinary.xfit, fusion=True)) + assert ordinary.xfit.fusion is False + assert fused.configuration_sha256 != ordinary.configuration_sha256 + assert DevicePipelineConfig.from_payload(fused.to_payload()) == fused + + legacy = ordinary.to_payload() + assert "fusion" not in legacy["xfit"] + restored = DevicePipelineConfig.from_payload(legacy) + assert restored == ordinary + assert restored.xfit.fusion is False + + +def test_xfit_fusion_validation_and_solver_settings_are_independent() -> None: + from cuphoton.xfit import LMConfig + + with pytest.raises(TypeError, match="fusion must be boolean"): + DeviceXFitPipelineConfig(fusion=1) # type: ignore[arg-type] + with pytest.raises(ValueError, match="analytic derivatives"): + DeviceXFitPipelineConfig(fusion=True, use_finite_difference=True) + config = DeviceXFitPipelineConfig(fusion=True, max_evaluations=9) + assert "fusion" not in config.solver_payload() + assert LMConfig(**config.solver_payload()).max_evaluations == 9 + + +@pytest.mark.parametrize("fusion", [False, True]) +def test_xfit_fusion_option_reaches_device_pipeline_callback_without_cuda( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, fusion: bool +) -> None: + import cuphoton.xscan.device_pipeline as pipeline + from cuphoton.xfit import LMConfig + + checkpoint_dir = tmp_path / "checkpoint" + checkpoint_dir.mkdir() + (checkpoint_dir / "checkpoint.pt").write_bytes(b"checkpoint") + ordinary = _config(checkpoint_dir) + config = replace(ordinary, xfit=replace(ordinary.xfit, fusion=fusion)) + path = tmp_path / "image.npy" + np.save(path, np.ones((31, 31), dtype=np.float64)) + item = _item(path, path) + + class Stream(nullcontext): + def synchronize(self) -> None: + pass + + cp = SimpleNamespace( + asarray=np.asarray, + empty=np.empty, + float32=np.float32, + float64=np.float64, + cuda=SimpleNamespace(runtime=SimpleNamespace(getDevice=lambda: 0)), + ) + context = DeviceWorkerContext( + config=config, + cp=cp, + torch=SimpleNamespace(cuda=SimpleNamespace(current_device=lambda: 0)), + model=None, + performance=None, + feature_variance_present=False, + device="cuda:0", + producer_stream=Stream(), + consumer_stream=Stream(), + solver_config=LMConfig(**config.xfit.solver_payload()), + load_seconds=0.0, + ) + calls = [] + + def fit(images: Any, **kwargs: Any) -> Any: + calls.append(kwargs) + np.testing.assert_array_equal(images, np.zeros((2, 11, 11))) + raise RuntimeError("reached xFit callback") + + functions = _PipelineFunctions( + gaussian_basis_component=lambda **values: values, + solve_constant_kernel_device=lambda reference, target, *_a, **_k: ( + SimpleNamespace(residual=target - reference) + ), + fit_dipoles_device=fit, + transform_xfit_result_features_device=lambda *_a, **_k: pytest.fail( + "unexpected feature stage" + ), + cupy_to_torch=lambda *_a, **_k: pytest.fail("unexpected bridge"), + predict_tensors=lambda *_a, **_k: pytest.fail("unexpected inference"), + ) + monkeypatch.setattr( + pipeline, "_load_pipeline_functions", lambda: functions + ) + with pytest.raises(RuntimeError, match="reached xFit callback"): + run_device_pipeline_item(item, context) + assert calls == [ + { + "model": "gaussian", + "mode": "difference", + "config": context.solver_config, + "fusion": fusion, + } + ] + + def test_loaded_model_contract_matches_checkpoint_contract() -> None: with pytest.raises(ValueError, match="triplet"): _validate_loaded_model_contract( From 8cb08e34df7324841ab2fed0e0357dca9b601cf8 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 29 Sep 2026 17:31:47 -0700 Subject: [PATCH 2/5] Preserve CuPy arithmetic in fused Gaussian callbacks Use explicit rounding at pointwise operation boundaries while retaining the default library math and CuPy normal-equation contractions. This preserves sensitive Gaussian fits and covers cancelling derivatives. Signed-off-by: Trent Nelson --- docs/components/xfit.md | 9 +- src/cuphoton/xfit/_cupy_fused.py | 289 +++++++++++++------------------ src/cuphoton/xfit/api.py | 10 +- src/cuphoton/xfit/commands.py | 2 +- tests/xfit/test_cupy_fusion.py | 44 +++++ 5 files changed, 179 insertions(+), 175 deletions(-) diff --git a/docs/components/xfit.md b/docs/components/xfit.md index 913906bd..f7b9950c 100644 --- a/docs/components/xfit.md +++ b/docs/components/xfit.md @@ -67,10 +67,11 @@ batch size is three, `(3, y, x)` keeps that per-candidate meaning; use the explicit `(1, 3, y, x)` shape for per-plane values. With `backend="cupy"`, `fusion=True` evaluates Gaussian residuals and weighted -normal equations in CuPy CUDA kernels. It supports float32 and float64, -difference and split modes, masks, and variances. The solver uses analytic -derivatives; final uncertainty diagnostics use the full analytic Jacobian. -Each fit owns its scratch buffers and runs on the current CuPy stream. +analytic Jacobians in fused CuPy CUDA kernels. The Jacobian is materialized +for the same CuPy normal-equation contractions used by the default solver. +It supports float32 and float64, difference and split modes, masks, and +variances. Final uncertainty diagnostics use the full analytic Jacobian. +Each fit owns its buffers and runs on the current CuPy stream. Select it in the CLI with `--backend cupy --fusion`; the default is `fusion=False`. Device pipeline manifests select it with `xfit.fusion: true`. Sampled-stamp and finite-difference fits use the default path. diff --git a/src/cuphoton/xfit/_cupy_fused.py b/src/cuphoton/xfit/_cupy_fused.py index 6c2cfe1b..6406eed0 100644 --- a/src/cuphoton/xfit/_cupy_fused.py +++ b/src/cuphoton/xfit/_cupy_fused.py @@ -2,7 +2,7 @@ # # SPDX-License-Identifier: Apache-2.0 -"""CuPy kernels for Gaussian residuals and analytic normal equations.""" +"""CuPy kernels for Gaussian residuals and analytic Jacobians.""" from __future__ import annotations @@ -13,15 +13,23 @@ from ._types import BackendArray, FitMode _BLOCK_SIZE = 128 -_NORMAL_COUNT = 44 # Eight gradient and 36 upper-triangular Hessian entries. -# Keep arithmetic in the input dtype and preserve the model's operation order, -# including physical-sigma derivatives followed by the log-sigma chain rule. -# Disable FMA to avoid contracting separate CuPy operations. +# Match CuPy's separate ufunc rounding, including the physical-sigma +# derivatives before the log-sigma chain rule. Explicit RN intrinsics preserve +# these boundaries while allowing normal FMA use inside libdevice functions: +# disabling FMA globally can change pow results and ill-conditioned fits. _CUDA_SOURCE = r""" typedef SCALAR_TYPE T; -const int block_size = 128; -const int normal_count = 44; + +__device__ float rn_add(float a, float b) { return __fadd_rn(a, b); } +__device__ double rn_add(double a, double b) { return __dadd_rn(a, b); } +__device__ float rn_sub(float a, float b) { return __fsub_rn(a, b); } +__device__ double rn_sub(double a, double b) { return __dsub_rn(a, b); } +__device__ float rn_mul(float a, float b) { return __fmul_rn(a, b); } +__device__ double rn_mul(double a, double b) { return __dmul_rn(a, b); } +__device__ float rn_div(float a, float b) { return __fdiv_rn(a, b); } +__device__ double rn_div(double a, double b) { return __ddiv_rn(a, b); } + struct Gaussian { T amplitude, sigma_x, sigma_y, inverse_x2, inverse_y2, cosine, sine; @@ -34,35 +42,38 @@ ? sigma_x : T(nan("")); const T valid_y = isfinite(sigma_y) && sigma_y > T(0) ? sigma_y : T(nan("")); - inverse_x2 = T(1) / (valid_x * valid_x); - inverse_y2 = T(1) / (valid_y * valid_y); + inverse_x2 = rn_div(T(1), rn_mul(valid_x, valid_x)); + inverse_y2 = rn_div(T(1), rn_mul(valid_y, valid_y)); cosine = cos(parameters[3]); sine = sin(parameters[3]); } template __device__ T star(T x, T y, T* derivative) const { - const T rotated_x = cosine * x + sine * y; - const T rotated_y = cosine * y - sine * x; - const T exponent = T(-0.5) * ( - rotated_x * rotated_x * inverse_x2 - + rotated_y * rotated_y * inverse_y2); + const T rotated_x = rn_add(rn_mul(cosine, x), rn_mul(sine, y)); + const T rotated_y = rn_sub(rn_mul(cosine, y), rn_mul(sine, x)); + const T exponent = rn_mul(T(-0.5), rn_add( + rn_mul(rn_mul(rotated_x, rotated_x), inverse_x2), + rn_mul(rn_mul(rotated_y, rotated_y), inverse_y2))); const T unit = exp(exponent); - const T value = amplitude * unit; + const T value = rn_mul(amplitude, unit); if (derivatives) { derivative[0] = unit; - derivative[1] = value * rotated_x * rotated_x - / pow(sigma_x, T(3)); - derivative[2] = value * rotated_y * rotated_y - / pow(sigma_y, T(3)); - derivative[3] = -value * rotated_x * rotated_y - * (inverse_x2 - inverse_y2); - derivative[4] = value * ( - cosine * rotated_x * inverse_x2 - - sine * rotated_y * inverse_y2); - derivative[5] = value * ( - sine * rotated_x * inverse_x2 - + cosine * rotated_y * inverse_y2); + derivative[1] = rn_div( + rn_mul(rn_mul(value, rotated_x), rotated_x), + pow(sigma_x, T(3))); + derivative[2] = rn_div( + rn_mul(rn_mul(value, rotated_y), rotated_y), + pow(sigma_y, T(3))); + derivative[3] = rn_mul( + rn_mul(rn_mul(-value, rotated_x), rotated_y), + rn_sub(inverse_x2, inverse_y2)); + derivative[4] = rn_mul(value, rn_sub( + rn_mul(rn_mul(cosine, rotated_x), inverse_x2), + rn_mul(rn_mul(sine, rotated_y), inverse_y2))); + derivative[5] = rn_mul(value, rn_add( + rn_mul(rn_mul(sine, rotated_x), inverse_x2), + rn_mul(rn_mul(cosine, rotated_y), inverse_y2))); } return value; } @@ -80,118 +91,60 @@ const long long pixels = height * width; const int plane = observation / pixels; const long long pixel = observation % pixels; - const T x = T(pixel % width) - T(width - 1) / T(2); - const T y = T(pixel / width) - T(height - 1) / T(2); + const T x = rn_sub(T(pixel % width), rn_div(T(width - 1), T(2))); + const T y = rn_sub(T(pixel / width), rn_div(T(height - 1), T(2))); const T* p = parameters + fit * 8; const Gaussian model(p); - const T positive = model.star(x - p[4], y - p[5], nullptr); - const T negative = model.star(x - p[6], y - p[7], nullptr); - const T prediction = plane == 0 ? positive - negative + const T positive = model.star( + rn_sub(x, p[4]), rn_sub(y, p[5]), nullptr); + const T negative = model.star( + rn_sub(x, p[6]), rn_sub(y, p[7]), nullptr); + const T prediction = plane == 0 ? rn_sub(positive, negative) : (plane == 1 ? positive : negative); const long long source = indices[fit] * observations + observation; - output[offset] = (prediction - images[source]) * weights[source]; + output[offset] = rn_mul( + rn_sub(prediction, images[source]), weights[source]); } -__device__ T warp_sum(T value) { - #pragma unroll - for (int offset = 16; offset > 0; offset /= 2) - value += __shfl_down_sync(0xffffffff, value, offset); - return value; -} -extern "C" __global__ void normal_partials( +extern "C" __global__ void materialized_jacobian( const T* parameters, const long long* indices, - const T* residuals, const T* weights, T* partials, - long long height, long long width, long long observations, - long long tile_count) { - const long long fit = blockIdx.x / tile_count; - const long long tile = blockIdx.x % tile_count; - const long long observation = tile * block_size + threadIdx.x; - const int lane = threadIdx.x % 32; - const int warp = threadIdx.x / 32; - __shared__ T warp_partials[block_size / 32][normal_count]; - T derivative[8] = {}; - T residual_value = T(0); - - if (observation < observations) { - const long long pixels = height * width; - const int plane = observation / pixels; - const long long pixel = observation % pixels; - const T x = T(pixel % width) - T(width - 1) / T(2); - const T y = T(pixel / width) - T(height - 1) / T(2); - const T* p = parameters + fit * 8; - const Gaussian model(p); - T positive[6], negative[6]; - model.star(x - p[4], y - p[5], positive); - model.star(x - p[6], y - p[7], negative); - #pragma unroll - for (int parameter = 0; parameter < 4; ++parameter) - derivative[parameter] = plane == 0 - ? positive[parameter] - negative[parameter] - : (plane == 1 ? positive[parameter] : negative[parameter]); - derivative[4] = plane < 2 ? positive[4] : T(0); - derivative[5] = plane < 2 ? positive[5] : T(0); - derivative[6] = plane == 0 ? -negative[4] - : (plane == 2 ? negative[4] : T(0)); - derivative[7] = plane == 0 ? -negative[5] - : (plane == 2 ? negative[5] : T(0)); - derivative[1] *= model.sigma_x; - derivative[2] *= model.sigma_y; - const T weight = weights[indices[fit] * observations + observation]; - #pragma unroll - for (int parameter = 0; parameter < 8; ++parameter) - derivative[parameter] *= weight; - residual_value = residuals[fit * observations + observation]; - } - + const T* weights, T* output, + long long batch, long long height, long long width, + long long observations) { + const long long offset = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (offset >= batch * observations) return; + const long long fit = offset / observations; + const long long observation = offset % observations; + const long long pixels = height * width; + const int plane = observation / pixels; + const long long pixel = observation % pixels; + const T x = rn_sub(T(pixel % width), rn_div(T(width - 1), T(2))); + const T y = rn_sub(T(pixel / width), rn_div(T(height - 1), T(2))); + const T* p = parameters + fit * 8; + const Gaussian model(p); + T positive[6], negative[6], derivative[8]; + model.star(rn_sub(x, p[4]), rn_sub(y, p[5]), positive); + model.star(rn_sub(x, p[6]), rn_sub(y, p[7]), negative); #pragma unroll - for (int parameter = 0; parameter < 8; ++parameter) { - const T total = warp_sum(derivative[parameter] * residual_value); - if (lane == 0) warp_partials[warp][parameter] = total; - } - int component = 8; + for (int parameter = 0; parameter < 4; ++parameter) + derivative[parameter] = plane == 0 + ? rn_sub(positive[parameter], negative[parameter]) + : (plane == 1 ? positive[parameter] : negative[parameter]); + derivative[4] = plane < 2 ? positive[4] : T(0); + derivative[5] = plane < 2 ? positive[5] : T(0); + derivative[6] = plane == 0 ? -negative[4] + : (plane == 2 ? negative[4] : T(0)); + derivative[7] = plane == 0 ? -negative[5] + : (plane == 2 ? negative[5] : T(0)); + derivative[1] = rn_mul(derivative[1], model.sigma_x); + derivative[2] = rn_mul(derivative[2], model.sigma_y); + const T weight = weights[indices[fit] * observations + observation]; #pragma unroll - for (int row = 0; row < 8; ++row) { - #pragma unroll - for (int column = row; column < 8; ++column) { - const T total = warp_sum(derivative[row] * derivative[column]); - if (lane == 0) warp_partials[warp][component] = total; - ++component; - } - } - __syncthreads(); - if (threadIdx.x < normal_count) { - T total = warp_partials[0][threadIdx.x]; - #pragma unroll - for (int source_warp = 1; source_warp < block_size / 32; - ++source_warp) - total += warp_partials[source_warp][threadIdx.x]; - partials[(fit * tile_count + tile) * normal_count + threadIdx.x] - = total; - } -} - -extern "C" __global__ void normal_finish( - const T* partials, T* gradients, T* hessians, long long tile_count) { - const long long fit = blockIdx.x; - const int component = threadIdx.x; - if (component >= normal_count) return; - T total = T(0); - for (long long tile = 0; tile < tile_count; ++tile) - total += partials[ - (fit * tile_count + tile) * normal_count + component]; - if (component < 8) { - gradients[fit * 8 + component] = total; - } else { - int column = component - 8; - int row = 0; - while (column >= 8 - row) { - column -= 8 - row; - ++row; - } - column += row; - hessians[fit * 64 + row * 8 + column] = total; - hessians[fit * 64 + column * 8 + row] = total; + for (int parameter = 0; parameter < 8; ++parameter) { + derivative[parameter] = rn_mul(derivative[parameter], weight); + output[(fit * 8 + parameter) * observations + observation] + = derivative[parameter]; } } """ @@ -239,9 +192,6 @@ def __init__( self._height = np.int64(height) self._width = np.int64(width) self._observations = int(np.prod(expected_shape)) - self._tiles = (self._observations + _BLOCK_SIZE - 1) // _BLOCK_SIZE - self._partial_capacity = 0 - self._partials: Any = None self._images = cp.ascontiguousarray(images) self._weights = cp.ascontiguousarray(weights, dtype=self._dtype) scalar_type = ( @@ -249,11 +199,12 @@ def __init__( ) self._module = cp.RawModule( code=_CUDA_SOURCE.replace("SCALAR_TYPE", scalar_type), - options=("--std=c++11", "--fmad=false"), + options=("--std=c++11", "--fmad=true"), ) self._residual_kernel = self._module.get_function("residual") - self._normal_kernel = self._module.get_function("normal_partials") - self._finish_kernel = self._module.get_function("normal_finish") + self._jacobian_kernel = self._module.get_function( + "materialized_jacobian" + ) def _inputs( self, parameters: BackendArray, indices: BackendArray @@ -299,50 +250,56 @@ def residual( ) return output - def normal_equations( - self, - parameters: BackendArray, - residuals: BackendArray, - *, - indices: BackendArray, - ) -> tuple[BackendArray, BackendArray]: - """Accumulate gradient and Hessian without a global Jacobian array.""" + def weighted_log_jacobian( + self, parameters: BackendArray, *, indices: BackendArray + ) -> BackendArray: + """Materialize a weighted log-sigma Jacobian in baseline layout.""" parameters, indices = self._inputs(parameters, indices) batch = parameters.shape[0] - if residuals.shape != (batch, self._observations): - raise ValueError("residuals must match the active observations") - residuals = self._cp.ascontiguousarray(residuals, dtype=self._dtype) - gradient = self._cp.empty((batch, 8), dtype=self._dtype) - hessian = self._cp.empty((batch, 8, 8), dtype=self._dtype) + jacobian = self._cp.empty( + (batch, 8, self._observations), dtype=self._dtype, order="C" + ) if batch: - tiles = self._tiles - if batch > self._partial_capacity: - self._partials = self._cp.empty( - (batch, tiles, _NORMAL_COUNT), dtype=self._dtype - ) - self._partial_capacity = batch - # Calls belong to one fit and execute on its current stream. The - # previous reduction finishes before this scratch is overwritten. - partials = self._partials - self._normal_kernel( - (batch * tiles,), + blocks = ( + batch * self._observations + _BLOCK_SIZE - 1 + ) // _BLOCK_SIZE + self._jacobian_kernel( + (blocks,), (_BLOCK_SIZE,), ( parameters, indices, - residuals, self._weights, - partials, + jacobian, + np.int64(batch), self._height, self._width, np.int64(self._observations), - np.int64(tiles), ), ) - self._finish_kernel( - (batch,), - (64,), - (partials, gradient, hessian, np.int64(tiles)), + return jacobian + + def normal_equations( + self, + parameters: BackendArray, + residuals: BackendArray, + *, + indices: BackendArray, + ) -> tuple[BackendArray, BackendArray]: + """Form normal equations with the baseline CuPy contractions.""" + + batch = parameters.shape[0] + if residuals.shape != (batch, self._observations): + raise ValueError("residuals must match the active observations") + jacobian = self.weighted_log_jacobian(parameters, indices=indices) + residuals = self._cp.ascontiguousarray(residuals, dtype=self._dtype) + if not batch: + return ( + self._cp.empty((0, 8), dtype=self._dtype), + self._cp.empty((0, 8, 8), dtype=self._dtype), ) + # Keep the solver's contraction order for ill-conditioned fits. + gradient = self._cp.einsum("knm,km->kn", jacobian, residuals) + hessian = self._cp.einsum("knm,kpm->knp", jacobian, jacobian) return gradient, hessian diff --git a/src/cuphoton/xfit/api.py b/src/cuphoton/xfit/api.py index 78fbb999..48d22f1f 100644 --- a/src/cuphoton/xfit/api.py +++ b/src/cuphoton/xfit/api.py @@ -1088,8 +1088,9 @@ def fit_dipoles( minus negative, the reported degrees of freedom and uncertainties are not statistically calibrated for that dependence. - ``fusion=True`` selects fused Gaussian residual and normal-equation - kernels with ``backend="cupy"``. It requires analytic derivatives. + ``fusion=True`` selects fused Gaussian residual and analytic-Jacobian + kernels with ``backend="cupy"``. Normal equations retain the default CuPy + contractions. It requires analytic derivatives. """ _validate_execution_options(model, backend, config) @@ -1128,8 +1129,9 @@ def fit_dipoles_device( to the active device. Existing CuPy inputs must already reside on that device and are checked before any input conversion. Input arrays are borrowed and are not mutated. - ``fusion=True`` selects fused CuPy Gaussian residual and normal-equation - kernels; native fits, sampled stamps and finite differences require + ``fusion=True`` selects fused CuPy Gaussian residual and analytic-Jacobian + kernels while retaining the default CuPy normal-equation contractions. + Native fits, sampled stamps and finite differences require ``fusion=False``. The solver may copy per-fit status, evaluation counts and boolean diff --git a/src/cuphoton/xfit/commands.py b/src/cuphoton/xfit/commands.py index e37f2030..fab105fc 100644 --- a/src/cuphoton/xfit/commands.py +++ b/src/cuphoton/xfit/commands.py @@ -262,7 +262,7 @@ class BackendArg(BackendInvariant): class FusionArg(BoolInvariant): _arg = "--fusion" _help = ( - "Fuse Gaussian residuals and normal equations; " + "Fuse Gaussian residuals and analytic Jacobians; " "requires --backend cupy." ) _mandatory = False diff --git a/tests/xfit/test_cupy_fusion.py b/tests/xfit/test_cupy_fusion.py index a928563a..14363a06 100644 --- a/tests/xfit/test_cupy_fusion.py +++ b/tests/xfit/test_cupy_fusion.py @@ -228,6 +228,50 @@ def test_cupy_fusion_callbacks_preserve_weighting_indices_and_layout( ) +def test_cupy_fusion_preserves_cancelling_hessian_terms(): + cp = _require_cupy() + from cuphoton.xfit._cupy_fused import GaussianFusedProblem + + shape = (51, 51) + parameters = cp.asarray( + [[2**20, 8.0, 4.0, 0.0, -0.25, 0.0, 0.25, 0.0]], + dtype=cp.float64, + ) + parameters[:, 1:3] = cp.log(parameters[:, 1:3]) + images = cp.zeros((1, *shape), dtype=cp.float64) + coordinate = cp.arange(51, dtype=cp.float64) - 25 + radius_squared = coordinate[:, None] ** 2 + coordinate[None, :] ** 2 + weights = (1 + radius_squared[None, ...] / 4096).copy() + indices = cp.asarray([0], dtype=cp.int64) + _, jacobian = _reference_callbacks( + cp, parameters, images, weights, indices, "difference" + ) + residual = cp.ones((1, 51 * 51), dtype=cp.float64) + expected_gradient = cp.einsum("knm,km->kn", jacobian, residual) + expected_hessian = cp.einsum("knm,kpm->knp", jacobian, jacobian) + + # Symmetry makes these bright-lobe derivative products cancel. A small + # relative error in each partial sum can change the near-zero Hessian. + absolute_sum = float( + cp.sum(cp.abs(jacobian[0, 4] * jacobian[0, 5])).item() + ) + assert absolute_sum > 1e8 + assert abs(float(expected_hessian[0, 4, 5].item())) < 1e-12 * absolute_sum + + problem = GaussianFusedProblem( + cp, images, weights, image_shape=shape, mode="difference" + ) + gradient, hessian = problem.normal_equations( + parameters, residual, indices=indices + ) + _assert_numeric( + cp.asnumpy(gradient), cp.asnumpy(expected_gradient), np.float64 + ) + _assert_numeric( + cp.asnumpy(hessian), cp.asnumpy(expected_hessian), np.float64 + ) + + @pytest.mark.parametrize("dtype", [np.float32, np.float64]) def test_cupy_fusion_extreme_widths_preserve_nonfinite_results(dtype): cp = _require_cupy() From 151fc645a764c8b6b8843071aa1531dc420b4d1b Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 30 Sep 2026 04:11:51 -0700 Subject: [PATCH 3/5] Correct the xFit GPU test execution count Signed-off-by: Trent Nelson --- .github/workflows/ci.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1b828127..2ffda3bd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -167,11 +167,11 @@ jobs: report = ET.parse(Path(os.environ["RUNNER_TEMP"]) / "xfit-gpu.xml") cases = report.findall(".//testcase") - assert len(cases) == 41, len(cases) + assert len(cases) == 42, len(cases) assert not report.findall(".//skipped"), "GPU tests must not skip" assert not report.findall(".//failure") assert not report.findall(".//error") - print("All 6 selected tests passed without skips") + print(f"All {len(cases)} selected tests passed without skips") PY - name: Run native xFit GPU tests run: >- From b72bc0affa7fa16f69f4a6cc8317a7303c4d683b Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Fri, 9 Oct 2026 12:21:51 -0700 Subject: [PATCH 4/5] Preserve native dispatch alongside CuPy fusion Signed-off-by: Trent Nelson --- .github/workflows/ci.yml | 2 +- src/cuphoton/xfit/_cupy_fused.py | 11 ++++++++--- tests/xfit/test_cupy_fusion.py | 2 +- tests/xscan/test_device_pipeline.py | 18 +++++++++++++----- 4 files changed, 23 insertions(+), 10 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2ffda3bd..a600636f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -167,7 +167,7 @@ jobs: report = ET.parse(Path(os.environ["RUNNER_TEMP"]) / "xfit-gpu.xml") cases = report.findall(".//testcase") - assert len(cases) == 42, len(cases) + assert len(cases) == 43, len(cases) assert not report.findall(".//skipped"), "GPU tests must not skip" assert not report.findall(".//failure") assert not report.findall(".//error") diff --git a/src/cuphoton/xfit/_cupy_fused.py b/src/cuphoton/xfit/_cupy_fused.py index 6406eed0..cab6e5b6 100644 --- a/src/cuphoton/xfit/_cupy_fused.py +++ b/src/cuphoton/xfit/_cupy_fused.py @@ -85,7 +85,9 @@ long long batch, long long height, long long width, long long observations) { const long long offset = (long long)blockIdx.x * blockDim.x + threadIdx.x; - if (offset >= batch * observations) return; + if (offset >= batch * observations) { + return; + } const long long fit = offset / observations; const long long observation = offset % observations; const long long pixels = height * width; @@ -113,7 +115,9 @@ long long batch, long long height, long long width, long long observations) { const long long offset = (long long)blockIdx.x * blockDim.x + threadIdx.x; - if (offset >= batch * observations) return; + if (offset >= batch * observations) { + return; + } const long long fit = offset / observations; const long long observation = offset % observations; const long long pixels = height * width; @@ -127,10 +131,11 @@ model.star(rn_sub(x, p[4]), rn_sub(y, p[5]), positive); model.star(rn_sub(x, p[6]), rn_sub(y, p[7]), negative); #pragma unroll - for (int parameter = 0; parameter < 4; ++parameter) + for (int parameter = 0; parameter < 4; ++parameter) { derivative[parameter] = plane == 0 ? rn_sub(positive[parameter], negative[parameter]) : (plane == 1 ? positive[parameter] : negative[parameter]); + } derivative[4] = plane < 2 ? positive[4] : T(0); derivative[5] = plane < 2 ? positive[5] : T(0); derivative[6] = plane == 0 ? -negative[4] diff --git a/tests/xfit/test_cupy_fusion.py b/tests/xfit/test_cupy_fusion.py index 14363a06..9fb1245a 100644 --- a/tests/xfit/test_cupy_fusion.py +++ b/tests/xfit/test_cupy_fusion.py @@ -76,7 +76,7 @@ def _no_cuda(*args, **kwargs): pytest.fail("unsupported fusion settings reached CUDA/backend discovery") -@pytest.mark.parametrize("backend", ["numpy", "auto", "cutile"]) +@pytest.mark.parametrize("backend", ["numpy", "auto", "cutile", "native"]) def test_fusion_rejects_unsupported_backend_before_cuda(monkeypatch, backend): monkeypatch.setattr(xfit_api, "resolve_backend", _no_cuda) with pytest.raises(ValueError, match="fusion requires backend='cupy'"): diff --git a/tests/xscan/test_device_pipeline.py b/tests/xscan/test_device_pipeline.py index 19e175ba..cff96385 100644 --- a/tests/xscan/test_device_pipeline.py +++ b/tests/xscan/test_device_pipeline.py @@ -1368,7 +1368,7 @@ def fit(images: Any, **kwargs: Any) -> Any: "model": "gaussian", "mode": "difference", "config": context.solver_config, - "fusion": context.config.xfit.fusion, + **({"fusion": True} if context.config.xfit.fusion else {}), } stage_calls.append("xfit") return _fake_xfit_result(cp, 2) @@ -1794,9 +1794,14 @@ def test_xfit_fusion_validation_and_solver_settings_are_independent() -> None: assert LMConfig(**config.solver_payload()).max_evaluations == 9 -@pytest.mark.parametrize("fusion", [False, True]) +@pytest.mark.parametrize( + "backend,fusion", [("cupy", False), ("cupy", True), ("native", False)] +) def test_xfit_fusion_option_reaches_device_pipeline_callback_without_cuda( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch, fusion: bool + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + backend: str, + fusion: bool, ) -> None: import cuphoton.xscan.device_pipeline as pipeline from cuphoton.xfit import LMConfig @@ -1805,7 +1810,9 @@ def test_xfit_fusion_option_reaches_device_pipeline_callback_without_cuda( checkpoint_dir.mkdir() (checkpoint_dir / "checkpoint.pt").write_bytes(b"checkpoint") ordinary = _config(checkpoint_dir) - config = replace(ordinary, xfit=replace(ordinary.xfit, fusion=fusion)) + config = replace( + ordinary, xfit=replace(ordinary.xfit, backend=backend, fusion=fusion) + ) path = tmp_path / "image.npy" np.save(path, np.ones((31, 31), dtype=np.float64)) item = _item(path, path) @@ -1863,7 +1870,8 @@ def fit(images: Any, **kwargs: Any) -> Any: "model": "gaussian", "mode": "difference", "config": context.solver_config, - "fusion": fusion, + **({"backend": backend} if backend != "cupy" else {}), + **({"fusion": True} if fusion else {}), } ] From 7a4a32b08af1c212ade9eb92892f56b2479f14cc Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Fri, 9 Oct 2026 13:49:53 -0700 Subject: [PATCH 5/5] Keep the legacy pipeline fixture free of fusion controls Signed-off-by: Trent Nelson --- tests/xscan/test_native_pipeline_config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/xscan/test_native_pipeline_config.py b/tests/xscan/test_native_pipeline_config.py index cdecc97b..44b0b2c1 100644 --- a/tests/xscan/test_native_pipeline_config.py +++ b/tests/xscan/test_native_pipeline_config.py @@ -42,6 +42,7 @@ def test_default_execution_preserves_legacy_payload_and_hash(tmp_path): legacy = config.to_payload() legacy_xfit = asdict(config.xfit) del legacy_xfit["backend"] + del legacy_xfit["fusion"] legacy["xfit"] = legacy_xfit restored = DevicePipelineConfig.from_payload(legacy)