diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5c7966cc..a600636f 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,11 +167,11 @@ 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) == 43, 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: >- diff --git a/docs/components/xfit.md b/docs/components/xfit.md index 181a2589..f7b9950c 100644 --- a/docs/components/xfit.md +++ b/docs/components/xfit.md @@ -66,6 +66,26 @@ 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 +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. + +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..cab6e5b6 --- /dev/null +++ b/src/cuphoton/xfit/_cupy_fused.py @@ -0,0 +1,310 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""CuPy kernels for Gaussian residuals and analytic Jacobians.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from ._types import BackendArray, FitMode + +_BLOCK_SIZE = 128 + +# 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; + +__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; + + __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 = 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 = 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 = rn_mul(amplitude, unit); + if (derivatives) { + derivative[0] = unit; + 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; + } +}; + +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 = 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( + 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] = rn_mul( + rn_sub(prediction, images[source]), weights[source]); +} + + +extern "C" __global__ void materialized_jacobian( + const T* parameters, const long long* indices, + 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 < 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 parameter = 0; parameter < 8; ++parameter) { + derivative[parameter] = rn_mul(derivative[parameter], weight); + output[(fit * 8 + parameter) * observations + observation] + = derivative[parameter]; + } +} +""" + + +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._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=true"), + ) + self._residual_kernel = self._module.get_function("residual") + self._jacobian_kernel = self._module.get_function( + "materialized_jacobian" + ) + + 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 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] + jacobian = self._cp.empty( + (batch, 8, self._observations), dtype=self._dtype, order="C" + ) + if batch: + blocks = ( + batch * self._observations + _BLOCK_SIZE - 1 + ) // _BLOCK_SIZE + self._jacobian_kernel( + (blocks,), + (_BLOCK_SIZE,), + ( + parameters, + indices, + self._weights, + jacobian, + np.int64(batch), + self._height, + self._width, + np.int64(self._observations), + ), + ) + 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 8e6fa971..48d22f1f 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,14 @@ 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 analytic-Jacobian + kernels with ``backend="cupy"``. Normal equations retain the default CuPy + contractions. 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 +1105,7 @@ def fit_dipoles( mode=mode, resolved=resolved, config=config, + fusion=fusion, ) return _materialize_dipole_fit_result(backend_result) @@ -1076,6 +1120,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 +1129,10 @@ 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 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 diagnostic masks to the host for control flow. Parameters, residuals, @@ -1098,6 +1147,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 +1184,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..fab105fc 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 analytic Jacobians; " + "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..9fb1245a --- /dev/null +++ b/tests/xfit/test_cupy_fusion.py @@ -0,0 +1,513 @@ +# 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", "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'"): + 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) + ) + + +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 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..cff96385 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": True} if context.config.xfit.fusion else {}), } stage_calls.append("xfit") return _fake_xfit_result(cp, 2) @@ -1761,6 +1763,119 @@ 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( + "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, + backend: str, + 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, backend=backend, 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, + **({"backend": backend} if backend != "cupy" else {}), + **({"fusion": True} if fusion else {}), + } + ] + + def test_loaded_model_contract_matches_checkpoint_contract() -> None: with pytest.raises(ValueError, match="triplet"): _validate_loaded_model_contract( 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)