From 0aa98315b18351300daf193770c0aba3fb555a59 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 29 Sep 2026 22:13:44 -0700 Subject: [PATCH 1/7] Fuse Gaussian xFit evaluation with public cuTile APIs Add opt-in weighted residual and analytic-Jacobian kernels with fit-owned buffers. Preserve CuPy width powers, normal contractions, and the existing solver and diagnostics. Wire fusion through Python, the CLI, and device pipeline configuration. Cover scientific behavior, concurrent use, teardown, and optional dependencies, and exercise the public cuTile path in GPU CI. Signed-off-by: Trent Nelson --- .github/workflows/ci.yml | 64 ++ docs/components/xfit.md | 14 +- examples/xfit/benchmark_gaussian.py | 13 +- src/cuphoton/xfit/_cutile_fused.py | 488 +++++++++++++ src/cuphoton/xfit/api.py | 89 ++- src/cuphoton/xfit/commands.py | 21 + src/cuphoton/xfit/executor.py | 16 +- src/cuphoton/xscan/device_pipeline.py | 35 +- .../xscan/pipeline_benchmark/runner.py | 6 +- src/cuphoton/xscan/xfit_features.py | 4 +- tests/core/test_cli_contract.py | 4 +- tests/test_public_api.py | 2 + tests/xfit/test_commands.py | 57 +- tests/xfit/test_cutile_fused.py | 649 ++++++++++++++++++ tests/xfit/test_executor.py | 59 +- tests/xfit/test_models.py | 201 +++++- tests/xscan/test_device_pipeline.py | 60 ++ tests/xscan/test_pipeline_benchmark.py | 16 +- tests/xscan/test_pipeline_benchmark_cli.py | 78 +++ tests/xscan/test_pipeline_benchmark_stages.py | 80 +++ 20 files changed, 1915 insertions(+), 41 deletions(-) create mode 100644 src/cuphoton/xfit/_cutile_fused.py create mode 100644 tests/xfit/test_cutile_fused.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5c7966cc..be718256 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -196,6 +196,66 @@ jobs: run: >- uv run --no-sync python examples/run_quickstarts.py --require-gpu --output-dir "${RUNNER_TEMP}/cuphoton-quickstart" + - name: Run cuTile xFit GPU coverage + run: | + # CPU Torch supports feature imports without its CUDA 13.0 pin. + cutile_env="${RUNNER_TEMP}/cuphoton-cutile" + uv venv --python 3.12 "$cutile_env" + uv pip install --python "$cutile_env/bin/python" \ + --index https://download.pytorch.org/whl/cpu 'torch==2.13.0+cpu' + uv pip install --python "$cutile_env/bin/python" --editable '.[cutile]' \ + 'pytest==9.1.1' 'cupy-cuda13x[ctk]==14.1.1' \ + 'cuda-tile[tileiras]==1.6.0' \ + 'cuda-toolkit[nvcc,nvvm,tileiras]==13.2.2' 'nvidia-nvjitlink==13.2.86' + uv pip freeze --python "$cutile_env/bin/python" \ + > "${RUNNER_TEMP}/cutile-packages.txt" + "$cutile_env/bin/python" - <<'PYCOMPILER' > "${RUNNER_TEMP}/cutile-compiler.txt" 2>&1 + from pathlib import Path + import subprocess + import sysconfig + + compiler = Path(sysconfig.get_path("purelib")) / "nvidia/cu13/bin" + for name in ("tileiras", "nvcc", "ptxas"): + executable = compiler / name + print(executable, flush=True) + subprocess.run([str(executable), "--version"], check=True) + PYCOMPILER + cat "${RUNNER_TEMP}/cutile-compiler.txt" + cutile_status=0 + "$cutile_env/bin/python" -m pytest -q -s -rs \ + --junitxml="${RUNNER_TEMP}/cutile-gpu.xml" \ + tests/xfit/test_cutile_fused.py \ + tests/xfit/test_models.py::test_cutile_gaussian_fit_matches_cupy \ + tests/xfit/test_models.py::test_cutile_gaussian_normal_equations_match_noisy_float32 \ + tests/xfit/test_models.py::test_cutile_weighting_and_compacted_indices_match_cupy \ + tests/xfit/test_models.py::test_cutile_preserves_rank_and_budget_failure_semantics \ + tests/xfit/test_models.py::test_cutile_device_results_survive_worker_reuse \ + > "${RUNNER_TEMP}/cutile-gpu.log" 2>&1 || cutile_status=$? + cat "${RUNNER_TEMP}/cutile-gpu.log" + CUTILE_PYTEST_STATUS="$cutile_status" "$cutile_env/bin/python" - <<'PY' + import os + from pathlib import Path + import re + import xml.etree.ElementTree as ET + + root = Path(os.environ["RUNNER_TEMP"]) + output = (root / "cutile-gpu.log").read_text() + lifecycle_error = re.compile( + r"CUDA_ERROR_|cudaError|CUDADriverError|CUDARuntimeError|CUDA error|CUDA driver error|" + r"CUDA runtime error|PytestUnraisableExceptionWarning|unraisable exception|" + r"Exception ignored in:|Exception ignored while|moduleUnload|Module\.__dealloc__", + re.IGNORECASE, + ) + if lifecycle_error.search(output): + raise SystemExit("cuTile tests emitted CUDA or interpreter teardown errors") + if int(os.environ["CUTILE_PYTEST_STATUS"]) != 0: + raise SystemExit("cuTile pytest process failed") + report = ET.parse(root / "cutile-gpu.xml") + cases = report.findall(".//testcase") + if len(cases) != 48 or any(report.findall(".//" + tag) for tag in ("skipped", "failure", "error")): + raise SystemExit("Require all 48 selected cuTile tests without skips, failures or errors") + print("All 48 cuTile tests passed without skips or teardown errors") + PY - name: Upload test evidence if: always() uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 @@ -204,6 +264,10 @@ jobs: path: | ${{ runner.temp }}/xfit-gpu.xml ${{ runner.temp }}/xfit-native-gpu.xml + ${{ runner.temp }}/cutile-gpu.xml + ${{ runner.temp }}/cutile-gpu.log + ${{ runner.temp }}/cutile-packages.txt + ${{ runner.temp }}/cutile-compiler.txt ${{ runner.temp }}/cuphoton-quickstart/**/summary.json if-no-files-found: error retention-days: 14 diff --git a/docs/components/xfit.md b/docs/components/xfit.md index 181a2589..0dbfae3d 100644 --- a/docs/components/xfit.md +++ b/docs/components/xfit.md @@ -66,14 +66,24 @@ 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. -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 +With the default `fusion=False`, 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 3.12 through 3.14. Final rank and covariance diagnostics use the analytic Jacobian and a singular-value factorization. Sampled-stamp fits stay on the NumPy or CuPy backends. The Tile backend requires analytic derivatives and rejects finite-difference fitting. +Pass `fusion=True` to `fit_dipoles`, or add `--fusion` to +`cuphoton xfit fit-dipoles --backend cutile`, to fuse Gaussian residual and +analytic-Jacobian evaluation, including weighting. Width powers and +normal-equation contractions use CuPy to preserve the reference arithmetic. +Each fit owns its workspace and uses the caller's current +CuPy stream. The default remains +`fusion=False`; CLI artifacts record the selected option in the effective +configuration. + The `cutile` extra installs cuda-tile 1.6 or newer and CuPy. Kernel compilation also requires a compatible CUDA toolkit. Follow the shared [cuTile compiler setup](../getting-started.md#optional-runtimes) to configure diff --git a/examples/xfit/benchmark_gaussian.py b/examples/xfit/benchmark_gaussian.py index 3a0b2189..4e2f9cc9 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: @@ -214,6 +218,9 @@ def main() -> None: default=180, help="Per-round wait, including cold compilation (default: 180).", ) + parser.add_argument( + "--fusion", action="store_true", help="Use fused cutile evaluation." + ) args = parser.parse_args() if not math.isfinite(args.timeout_seconds) or args.timeout_seconds <= 0: parser.error("--timeout-seconds must be positive and finite") @@ -228,6 +235,8 @@ def main() -> None: if args.compare_native and args.backend: parser.error("--compare-native cannot be combined with --backend") backends = args.backend or ["cupy", "cutile"] + if args.fusion and backends != ["cutile"]: + parser.error("--fusion requires --backend cutile") 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,6 +283,7 @@ def main() -> None: context.Value("i", 0), context.Barrier(args.workers), any(backend != "numpy" for backend in treatments), + args.fusion, ), **executor_options, ) as pool: @@ -345,6 +355,7 @@ def main() -> None: json.dumps( { "backend": treatment, + "fusion": args.fusion, "execution": args.execution, "workers": args.workers, "batch": batch, diff --git a/src/cuphoton/xfit/_cutile_fused.py b/src/cuphoton/xfit/_cutile_fused.py new file mode 100644 index 00000000..78f2ec90 --- /dev/null +++ b/src/cuphoton/xfit/_cutile_fused.py @@ -0,0 +1,488 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Fused Gaussian evaluation using public ``cuda.tile`` operations.""" + +import math +from functools import cache +from typing import TYPE_CHECKING, Any + +from ._cutile import _load_cutile +from ._types import BackendArray, FitMode + +_PARAMETER_COUNT = 8 +_EVALUATION_TILE = 128 + + +@cache +def _evaluation_kernel() -> Any: + if TYPE_CHECKING: + import cuda.tile as ct + else: + _, ct = _load_cutile() + + @ct.kernel + def kernel( + parameters, + width_cubes, + images, + weights, + indices, + output, + height: ct.Constant[int], + width: ct.Constant[int], + plane_count: ct.Constant[int], + compute_jacobian: ct.Constant[bool], + ): + block = ct.bid(0) + observation_tile = ct.bid(1) + weight_row = ct.reshape(ct.load(indices, (block,), (1,)), ()) + parameter_tile = ct.load( + parameters, + index=(block, 0), + shape=(1, _PARAMETER_COUNT), + ) + amplitude = ct.extract(parameter_tile, (0, 0), (1, 1)) + sigma_x = ct.exp(ct.extract(parameter_tile, (0, 1), (1, 1))) + sigma_y = ct.exp(ct.extract(parameter_tile, (0, 2), (1, 1))) + theta = ct.extract(parameter_tile, (0, 3), (1, 1)) + x_positive = ct.extract(parameter_tile, (0, 4), (1, 1)) + y_positive = ct.extract(parameter_tile, (0, 5), (1, 1)) + x_negative = ct.extract(parameter_tile, (0, 6), (1, 1)) + y_negative = ct.extract(parameter_tile, (0, 7), (1, 1)) + # Match invalid physical widths after log-sigma conversion. + sigma_x = ct.where( + (sigma_x > 0) & (sigma_x < math.inf), sigma_x, math.nan + ) + sigma_y = ct.where( + (sigma_y > 0) & (sigma_y < math.inf), sigma_y, math.nan + ) + if compute_jacobian: + cubes = ct.load(width_cubes, (block, 0), (1, 2)) + sigma_x_cube = ct.extract(cubes, (0, 0), (1, 1)) + sigma_y_cube = ct.extract(cubes, (0, 1), (1, 1)) + else: + sigma_x_cube = 1.0 + sigma_y_cube = 1.0 + cosine = ct.cos(theta) + sine = ct.sin(theta) + inverse_x_squared = 1.0 / (sigma_x * sigma_x) + inverse_y_squared = 1.0 / (sigma_y * sigma_y) + pixel_count = height * width + + # Two-term reductions preserve the separately rounded products in + # the CuPy model. Direct multiply/add expressions may contract into + # an FMA and change the trajectory of poorly conditioned fits. + observation = observation_tile * _EVALUATION_TILE + ct.arange( + _EVALUATION_TILE, dtype=ct.int32 + ) + plane = observation // pixel_count + pixel = observation - plane * pixel_count + y = pixel // width + x = pixel - y * width + x_coordinate = x.astype(parameters.dtype) - (width - 1) / 2 + y_coordinate = y.astype(parameters.dtype) - (height - 1) / 2 + + x_hat_positive = x_coordinate - x_positive + y_hat_positive = y_coordinate - y_positive + x_rotated_positive = ct.sum( + ct.cat((cosine * x_hat_positive, sine * y_hat_positive), axis=0), + axis=0, + keepdims=True, + ) + y_rotated_positive = ct.sum( + ct.cat( + (cosine * y_hat_positive, -(sine * x_hat_positive)), axis=0 + ), + axis=0, + keepdims=True, + ) + unit_positive = ct.exp( + -0.5 + * ( + ct.sum( + ct.cat( + ( + x_rotated_positive + * x_rotated_positive + * inverse_x_squared, + y_rotated_positive + * y_rotated_positive + * inverse_y_squared, + ), + axis=0, + ), + axis=0, + keepdims=True, + ) + ) + ) + star_positive = amplitude * unit_positive + width_x_positive = ( + star_positive + * x_rotated_positive + * x_rotated_positive + / sigma_x_cube + ) + width_y_positive = ( + star_positive + * y_rotated_positive + * y_rotated_positive + / sigma_y_cube + ) + derivatives_positive = ( + unit_positive, + width_x_positive, + width_y_positive, + -star_positive + * x_rotated_positive + * y_rotated_positive + * (inverse_x_squared - inverse_y_squared), + star_positive + * ( + ct.sum( + ct.cat( + ( + cosine * x_rotated_positive * inverse_x_squared, + -(sine * y_rotated_positive * inverse_y_squared), + ), + axis=0, + ), + axis=0, + keepdims=True, + ) + ), + star_positive + * ( + ct.sum( + ct.cat( + ( + sine * x_rotated_positive * inverse_x_squared, + cosine * y_rotated_positive * inverse_y_squared, + ), + axis=0, + ), + axis=0, + keepdims=True, + ) + ), + ) + + x_hat_negative = x_coordinate - x_negative + y_hat_negative = y_coordinate - y_negative + x_rotated_negative = ct.sum( + ct.cat((cosine * x_hat_negative, sine * y_hat_negative), axis=0), + axis=0, + keepdims=True, + ) + y_rotated_negative = ct.sum( + ct.cat( + (cosine * y_hat_negative, -(sine * x_hat_negative)), axis=0 + ), + axis=0, + keepdims=True, + ) + unit_negative = ct.exp( + -0.5 + * ( + ct.sum( + ct.cat( + ( + x_rotated_negative + * x_rotated_negative + * inverse_x_squared, + y_rotated_negative + * y_rotated_negative + * inverse_y_squared, + ), + axis=0, + ), + axis=0, + keepdims=True, + ) + ) + ) + star_negative = amplitude * unit_negative + width_x_negative = ( + star_negative + * x_rotated_negative + * x_rotated_negative + / sigma_x_cube + ) + width_y_negative = ( + star_negative + * y_rotated_negative + * y_rotated_negative + / sigma_y_cube + ) + derivatives_negative = ( + unit_negative, + width_x_negative, + width_y_negative, + -star_negative + * x_rotated_negative + * y_rotated_negative + * (inverse_x_squared - inverse_y_squared), + star_negative + * ( + ct.sum( + ct.cat( + ( + cosine * x_rotated_negative * inverse_x_squared, + -(sine * y_rotated_negative * inverse_y_squared), + ), + axis=0, + ), + axis=0, + keepdims=True, + ) + ), + star_negative + * ( + ct.sum( + ct.cat( + ( + sine * x_rotated_negative * inverse_x_squared, + cosine * y_rotated_negative * inverse_y_squared, + ), + axis=0, + ), + axis=0, + keepdims=True, + ) + ), + ) + + if plane_count == 1: + derivative_0 = derivatives_positive[0] - derivatives_negative[0] + derivative_1 = derivatives_positive[1] - derivatives_negative[1] + derivative_2 = derivatives_positive[2] - derivatives_negative[2] + derivative_3 = ct.sum( + ct.cat( + (derivatives_positive[3], -(derivatives_negative[3])), + axis=0, + ), + axis=0, + keepdims=True, + ) + derivative_4 = derivatives_positive[4] + derivative_5 = derivatives_positive[5] + derivative_6 = -derivatives_negative[4] + derivative_7 = -derivatives_negative[5] + else: + zero = ct.zeros((_EVALUATION_TILE,), dtype=parameters.dtype) + derivative_0 = ct.where( + plane == 0, + derivatives_positive[0] - derivatives_negative[0], + ct.where( + plane == 1, + derivatives_positive[0], + derivatives_negative[0], + ), + ) + derivative_1 = ct.where( + plane == 0, + derivatives_positive[1] - derivatives_negative[1], + ct.where( + plane == 1, + derivatives_positive[1], + derivatives_negative[1], + ), + ) + derivative_2 = ct.where( + plane == 0, + derivatives_positive[2] - derivatives_negative[2], + ct.where( + plane == 1, + derivatives_positive[2], + derivatives_negative[2], + ), + ) + derivative_3 = ct.where( + plane == 0, + ct.sum( + ct.cat( + (derivatives_positive[3], -(derivatives_negative[3])), + axis=0, + ), + axis=0, + keepdims=True, + ), + ct.where( + plane == 1, + derivatives_positive[3], + derivatives_negative[3], + ), + ) + derivative_4 = ct.where(plane < 2, derivatives_positive[4], zero) + derivative_5 = ct.where(plane < 2, derivatives_positive[5], zero) + derivative_6 = ct.where( + plane == 0, + -derivatives_negative[4], + ct.where(plane == 2, derivatives_negative[4], zero), + ) + derivative_7 = ct.where( + plane == 0, + -derivatives_negative[5], + ct.where(plane == 2, derivatives_negative[5], zero), + ) + + # Preserve the model's physical-width derivative and chain + # rule order, including overflow/underflow of sigma**3. + derivative_1 *= sigma_x + derivative_2 *= sigma_y + + weight = ct.load( + weights, + index=(weight_row, observation_tile), + shape=(1, _EVALUATION_TILE), + padding_mode=ct.PaddingMode.ZERO, + ) + jacobian_01 = ct.cat( + ( + ct.reshape(derivative_0, (1, _EVALUATION_TILE)), + ct.reshape(derivative_1, (1, _EVALUATION_TILE)), + ), + axis=0, + ) + jacobian_23 = ct.cat( + ( + ct.reshape(derivative_2, (1, _EVALUATION_TILE)), + ct.reshape(derivative_3, (1, _EVALUATION_TILE)), + ), + axis=0, + ) + jacobian_45 = ct.cat( + ( + ct.reshape(derivative_4, (1, _EVALUATION_TILE)), + ct.reshape(derivative_5, (1, _EVALUATION_TILE)), + ), + axis=0, + ) + jacobian_67 = ct.cat( + ( + ct.reshape(derivative_6, (1, _EVALUATION_TILE)), + ct.reshape(derivative_7, (1, _EVALUATION_TILE)), + ), + axis=0, + ) + jacobian_03 = ct.cat((jacobian_01, jacobian_23), axis=0) + jacobian_47 = ct.cat((jacobian_45, jacobian_67), axis=0) + jacobian = ct.cat((jacobian_03, jacobian_47), axis=0) + jacobian *= weight + + if compute_jacobian: + ct.store(output, (block, observation_tile), jacobian) + else: + prediction = ct.sum( + ct.cat((star_positive, -star_negative), axis=0), + axis=0, + keepdims=True, + ) + if plane_count == 3: + prediction = ct.where( + plane == 0, + prediction, + ct.where(plane == 1, star_positive, star_negative), + ) + image = ct.load( + images, + (weight_row, observation_tile), + (1, _EVALUATION_TILE), + padding_mode=ct.PaddingMode.ZERO, + ) + ct.store( + output, + (block, observation_tile), + (prediction - image) * weight, + ) + + return kernel + + +class GaussianWorkspace: + """Buffers owned by one fit invocation and its current CuPy stream. + + Jacobian outputs are borrowed until the next Jacobian call. The solver + copies them into its own storage before reuse. Residuals are separately + owned because accepted residuals must survive trial evaluations. + """ + + def __init__( + self, + images: BackendArray, + weights: BackendArray, + *, + image_shape: tuple[int, int], + mode: FitMode, + ) -> None: + self._cp, self._ct = _load_cutile() + self._stream = self._cp.cuda.get_current_stream() + self._height, self._width = image_shape + self._planes = 1 if mode == "difference" else 3 + batch = int(images.shape[0]) + self._observations = self._planes * self._height * self._width + self._images = self._cp.ascontiguousarray(images).reshape( + batch, self._observations + ) + self._weights = self._cp.ascontiguousarray(weights).reshape( + batch, self._observations + ) + self._jacobians = self._cp.empty( + (batch * 8, self._observations), dtype=images.dtype + ) + self._width_cubes = self._cp.empty((batch, 2), dtype=images.dtype) + self._kernel = _evaluation_kernel() + + def _evaluate( + self, + parameters: BackendArray, + indices: BackendArray, + output: BackendArray, + *, + compute_jacobian: bool, + ) -> None: + width_cubes = self._width_cubes[: parameters.shape[0]] + if compute_jacobian: + # Match the model's power rounding. A one-ULP difference here + # can alter the trajectory of an ill-conditioned fit. + self._cp.exp(parameters[:, 1:3], out=width_cubes) + self._cp.power(width_cubes, 3, out=width_cubes) + self._ct.launch( + self._stream, + ( + parameters.shape[0], + math.ceil(self._observations / _EVALUATION_TILE), + 1, + ), + self._kernel, + ( + parameters, + width_cubes, + self._images, + self._weights, + indices, + output, + self._height, + self._width, + self._planes, + compute_jacobian, + ), + ) + + def residual( + self, parameters: BackendArray, *, indices: BackendArray + ) -> BackendArray: + output = self._cp.empty( + (parameters.shape[0], self._observations), dtype=parameters.dtype + ) + self._evaluate(parameters, indices, output, compute_jacobian=False) + return output + + def jacobian( + self, parameters: BackendArray, *, indices: BackendArray + ) -> BackendArray: + output = self._jacobians[: parameters.shape[0] * 8] + self._evaluate(parameters, indices, output, compute_jacobian=True) + return output.reshape(parameters.shape[0], 8, self._observations) diff --git a/src/cuphoton/xfit/api.py b/src/cuphoton/xfit/api.py index 8e6fa971..cda65e93 100644 --- a/src/cuphoton/xfit/api.py +++ b/src/cuphoton/xfit/api.py @@ -131,15 +131,19 @@ class DeviceDipoleFitResult: init=False, default="levenberg-marquardt", ) - backend: Literal["cupy", "native"] = "cupy" + backend: Literal["cupy", "cutile", "native"] = "cupy" result_location: Literal["device"] = field( init=False, default="device", ) def __post_init__(self) -> None: - if self.backend not in {"cupy", "native"}: - raise ValueError("device fit backend must be cupy or native") + if self.backend not in {"cupy", "cutile", "native"}: + raise ValueError( + "device fit backend must be cupy, cutile or native" + ) + if self.backend == "cutile" and self.model != "gaussian": + raise ValueError("backend='cutile' supports only Gaussian fits") if isinstance(self.device_id, bool) or not isinstance( self.device_id, (int, np.integer) ): @@ -615,6 +619,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``.""" @@ -747,9 +752,22 @@ def evaluate_model(parameters: BackendArray) -> BackendArray: ) return selected_model.evaluate(parameters, mode=mode) + workspace = None + if fusion: + from ._cutile_fused import GaussianWorkspace + + workspace = GaussianWorkspace( + images_array.reshape(batch, -1), + weights.reshape(batch, -1), + image_shape=(height, width), + mode=mode, + ) + def residual( parameters: BackendArray, *, indices: BackendArray ) -> BackendArray: + if workspace is not None: + return workspace.residual(parameters, indices=indices) prediction = evaluate_model(model_parameters(parameters)) value = (prediction - images_array[indices]) * weights[indices] return value.reshape(parameters.shape[0], -1) @@ -760,6 +778,8 @@ def residual( def jacobian( parameters: BackendArray, *, indices: BackendArray ) -> BackendArray: + if workspace is not None: + return workspace.jacobian(parameters, indices=indices) assert isinstance(selected_model, GaussianDipoleModel) physical = model_parameters(parameters) value = selected_model._jacobian_positive_unchecked( @@ -773,7 +793,7 @@ def jacobian( jacobian_function = jacobian normal_equations = None - if gaussian_model and resolved.name == "cutile": + if gaussian_model and resolved.name == "cutile" and workspace is None: from ._cutile import gaussian_normal_equations def normal_equations( @@ -1018,6 +1038,25 @@ def _validate_execution_options( ) +def _validate_fusion( + fusion: bool, + *, + backend: BackendRequest, + gaussian_model: bool, + use_finite_difference: bool, +) -> None: + if not isinstance(fusion, bool): + raise TypeError("fusion must be a boolean") + if not fusion: + return + if backend != "cutile": + raise ValueError("fusion requires backend='cutile'") + if not gaussian_model: + raise ValueError("fusion supports only the Gaussian model") + if use_finite_difference: + raise ValueError("fusion does not support finite-difference fitting") + + def fit_dipoles( images: ArrayLike, *, @@ -1028,6 +1067,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 +1089,25 @@ 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 + evaluation with ``backend="cutile"``. Normal equations retain CuPy's + contractions and the solver's convergence and diagnostics controls. """ _validate_execution_options(model, backend, config) + _validate_fusion( + fusion, + backend=backend, + gaussian_model=( + model == "gaussian" + if isinstance(model, str) + else isinstance(model, GaussianDipoleModel) + ), + use_finite_difference=( + config is not None and config.use_finite_difference + ), + ) resolved = resolve_backend(backend) backend_result = _fit_dipoles_backend( images, @@ -1062,6 +1118,7 @@ def fit_dipoles( mode=mode, resolved=resolved, config=config, + fusion=fusion, ) return _materialize_dipole_fit_result(backend_result) @@ -1075,11 +1132,14 @@ def fit_dipoles_device( variance: ArrayLike | None = None, mode: FitMode = "difference", config: LMConfig | None = None, - backend: Literal["cupy", "native"] = "cupy", + backend: Literal["cupy", "cutile", "native"] = "cupy", + fusion: bool = False, ) -> DeviceDipoleFitResult: """Fit dipoles while keeping every result array on the active GPU. - This experimental seam uses CuPy-owned buffers. The optional native + This experimental seam uses CuPy-owned buffers. The cuTile backend + supports Gaussian analytic fits; ``fusion=True`` selects fused residual + and analytic-Jacobian evaluation. The optional native backend runs Gaussian LM iterations in CUDA C++. NumPy inputs are copied to the active device. Existing CuPy inputs must already reside on that device and are checked before any input conversion. Input arrays are @@ -1098,8 +1158,20 @@ def fit_dipoles_device( """ _validate_execution_options(model, backend, config) - if backend not in {"cupy", "native"}: - raise ValueError("device fit backend must be cupy or native") + _validate_fusion( + fusion, + backend=backend, + gaussian_model=( + model == "gaussian" + if isinstance(model, str) + else isinstance(model, GaussianDipoleModel) + ), + use_finite_difference=( + config is not None and config.use_finite_difference + ), + ) + if backend not in {"cupy", "cutile", "native"}: + raise ValueError("device fit backend must be cupy, cutile or native") if backend == "native": from ._native import load_extension @@ -1134,6 +1206,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..fda3e8ff 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 = None chunk_size = None class ChunkSizeArg(PositiveIntegerInvariant): @@ -266,6 +267,12 @@ class ComputeDtypeArg(ComputeDtypeInvariant): _mandatory = False _default = "input" + class FusionArg(BoolInvariant): + _arg = "--fusion" + _help = "Fuse Gaussian evaluation with the cutile backend." + _mandatory = False + _default = False + class StampEvaluationArg(StampEvaluationInvariant): _arg = "--stamp-evaluation" _help = ( @@ -330,6 +337,17 @@ def _model_specification( ) def run(self) -> None: + from .api import _validate_fusion + + assert self.backend is not None + assert self.fusion is not None + self._call( + _validate_fusion, + self.fusion, + backend=self.backend, + gaussian_model=self.model == "gaussian", + 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 +380,7 @@ def run(self) -> None: "g_tol", "max_evaluations", "use_finite_difference", + "fusion", "fits_reader", ) } @@ -439,6 +458,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 +471,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..32b7436f 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, } @@ -101,6 +107,12 @@ def _plan_xfit_chunks( settings = json_mapping( {**_DEFAULTS, **fit_options}, field="xFit options" ) + _validate_fusion( + settings["fusion"], + backend=settings["backend"], + gaussian_model=settings["model"] == "gaussian", + use_finite_difference=settings["use_finite_difference"], + ) if settings["compute_dtype"] not in {"input", "float32", "float64"}: raise ValueError("unsupported xFit compute dtype") if settings["fits_reader"] not in {"auto", "astropy", "xdr"}: @@ -281,6 +293,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 +456,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 diff --git a/src/cuphoton/xscan/device_pipeline.py b/src/cuphoton/xscan/device_pipeline.py index 37754c1b..04b0247e 100644 --- a/src/cuphoton/xscan/device_pipeline.py +++ b/src/cuphoton/xscan/device_pipeline.py @@ -650,17 +650,28 @@ class DeviceXFitPipelineConfig: damping_decrease: float = 0.3 finite_difference_step: float | None = None use_finite_difference: bool = False - backend: Literal["cupy", "native"] = "cupy" + backend: Literal["cupy", "cutile", "native"] = "cupy" + fusion: bool = False def __post_init__(self) -> None: if self.model != "gaussian": raise ValueError("xfit model must be 'gaussian'") if self.mode != "difference": raise ValueError("xfit mode must be 'difference'") - if self.backend not in ("cupy", "native"): - raise ValueError("xfit backend must be 'cupy' or 'native'") + if self.backend not in ("cupy", "cutile", "native"): + raise ValueError( + "xfit backend must be 'cupy', 'cutile' 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 != "cutile": + raise ValueError("xfit fusion requires backend='cutile'") + if self.backend == "cutile" and self.use_finite_difference: + raise ValueError( + "xfit backend='cutile' does not support finite differences" + ) if self.backend == "native" and self.use_finite_difference: raise ValueError( "native xfit does not support finite differences" @@ -715,6 +726,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 +736,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 +747,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 @@ -740,9 +756,10 @@ def from_payload( ) -> DeviceXFitPipelineConfig: """Restore exact LM settings from JSON-compatible values.""" + expected = frozenset(value.name for value in fields(cls)) values = json_mapping(payload, field="device pipeline xfit config") values.setdefault("backend", "cupy") - expected = frozenset(value.name for value in fields(cls)) + values.setdefault("fusion", False) values = _require_exact_fields( values, expected=expected, @@ -1079,9 +1096,10 @@ def _device_pipeline_evidence_layout_contract( "mode": "difference", "dtype": "float64", } - if xfit["backend"] not in ("cupy", "native"): + if xfit["backend"] not in ("cupy", "cutile", "native"): raise ValueError( - "scientific evidence xfit.backend must be 'cupy' or 'native'" + "scientific evidence xfit.backend must be 'cupy', 'cutile' " + "or 'native'" ) for name, expected in expected_xfit.items(): if xfit[name] != expected: @@ -2158,6 +2176,7 @@ def _pack_scientific_evidence( candidate_count: int, kernel_shape: tuple[int, int], flux_conserve: bool, + xfit_backend: Literal["cupy", "cutile", "native"] = "cupy", ) -> _PackedScientificEvidence: """Pack all parity evidence on the producer stream as float64.""" @@ -2183,7 +2202,8 @@ def _pack_scientific_evidence( str(xfit_result.mode), str(xfit_result.dtype), ) - if xfit_contract[0] not in ("cupy", "native") or xfit_contract[1:] != ( + if xfit_contract != ( + xfit_backend, "levenberg-marquardt", "gaussian", "difference", @@ -3352,6 +3372,7 @@ def run_device_pipeline_item( candidate_count=len(item.candidates), kernel_shape=context.config.xpois.kernel_shape, flux_conserve=context.config.xpois.flux_conserve, + xfit_backend=context.config.xfit.backend, ), ) evidence_layout = _append_prediction_evidence_layout( diff --git a/src/cuphoton/xscan/pipeline_benchmark/runner.py b/src/cuphoton/xscan/pipeline_benchmark/runner.py index 543bc7b2..d0f2cc3f 100644 --- a/src/cuphoton/xscan/pipeline_benchmark/runner.py +++ b/src/cuphoton/xscan/pipeline_benchmark/runner.py @@ -383,6 +383,8 @@ def audit(args: SimpleNamespace) -> dict[str, Any]: if metadata["xpois"][f"xpois_{name}"] != value: raise ValueError(f"xPois {name} differs") evidence = result["scientific_evidence"] + if metadata["xfit"].get("backend") != evidence["xfit"]["backend"]: + raise ValueError("xFit backend differs") if ( metadata["xpois"]["basis_terms"] != evidence["xpois"]["basis_terms"] @@ -427,6 +429,8 @@ def provenance() -> dict[str, Any]: "xfit/api.py", "xfit/backend.py", "xfit/_native.py", + "xfit/_cutile.py", + "xfit/_cutile_fused.py", "xfit/models.py", "xfit/solver.py", "xscan/xfit_features.py", @@ -437,7 +441,7 @@ def provenance() -> dict[str, Any]: cupy_packages = importlib.metadata.packages_distributions().get( "cupy", [] ) - for name in ("cuphoton", "numpy", "torch", *cupy_packages): + for name in ("cuphoton", "numpy", "torch", "cuda-tile", *cupy_packages): try: packages[name] = importlib.metadata.version(name) except importlib.metadata.PackageNotFoundError: diff --git a/src/cuphoton/xscan/xfit_features.py b/src/cuphoton/xscan/xfit_features.py index aaa5d0bd..fd379d57 100644 --- a/src/cuphoton/xscan/xfit_features.py +++ b/src/cuphoton/xscan/xfit_features.py @@ -1292,10 +1292,12 @@ def _validate_device_feature_source( if result.solver != "levenberg-marquardt": raise ValueError("xFit device result solver is unsupported") if ( - result.backend not in {"cupy", "native"} + result.backend not in {"cupy", "cutile", "native"} or result.result_location != "device" ): raise ValueError("xFit result must be a CuPy device result") + if result.backend == "cutile" and result.model != "gaussian": + raise ValueError("cuTile xFit features require the Gaussian model") parameter_names = _validate_live_feature_source(result) if int(result.device_id) != active_device_id: raise ValueError( 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_commands.py b/tests/xfit/test_commands.py index f975d85c..e03dc38a 100644 --- a/tests/xfit/test_commands.py +++ b/tests/xfit/test_commands.py @@ -477,10 +477,12 @@ def test_artifacts_retain_archive_hash_captured_at_load_time( assert loaded_sha256 != file_sha256(input_path) +@pytest.mark.parametrize("fusion", [False, True]) def test_fit_dipoles_command_writes_complete_artifact_set( monkeypatch, tmp_path, capsys, + fusion, ) -> None: import cuphoton.xfit as xfit @@ -510,6 +512,7 @@ def fake_fit_dipoles(received_images, **kwargs): assert received_images.dtype == np.float64 assert kwargs["model"] == "gaussian" assert isinstance(kwargs["config"], FakeConfig) + assert kwargs["fusion"] is fusion return _fake_result(received_images) monkeypatch.setattr(xfit, "LMConfig", FakeConfig, raising=False) @@ -531,7 +534,8 @@ def fake_fit_dipoles(received_images, **kwargs): "--model", "gaussian", "--backend", - "numpy", + "cutile" if fusion else "numpy", + *(["--fusion"] if fusion else []), "--compute-dtype", "float64", ], @@ -557,6 +561,56 @@ def fake_fit_dipoles(received_images, **kwargs): "requested": "float64", "resolved": "float64", } + assert effective["fusion"] is fusion + assert effective["backend"] == ("cutile" if fusion else "numpy") + + +@pytest.mark.parametrize( + ("backend", "model", "extra_args", "message"), + [ + ("cupy", "gaussian", [], "requires backend='cutile'"), + ("cutile", "stamp", [], "only the Gaussian model"), + ( + "cutile", + "gaussian", + ["--use-finite-difference"], + "finite-difference fitting", + ), + ], +) +def test_fusion_command_rejects_invalid_settings_before_loading_input( + tmp_path, monkeypatch, capsys, backend, model, extra_args, message +) -> None: + from cuphoton.xfit.commands import FitDipolesCommand + + input_path = tmp_path / "input.npz" + _write_gaussian_input(input_path) + monkeypatch.setattr( + FitDipolesCommand, + "_load_dataset", + lambda *args, **kwargs: pytest.fail( + "invalid fusion settings read the input" + ), + ) + rc = run_component( + "xfit", + [ + "fit-dipoles", + "--input", + str(input_path), + "--output-dir", + str(tmp_path / "fit"), + "--backend", + backend, + "--model", + model, + "--fusion", + *extra_args, + ], + ) + assert rc == 1 + assert message in capsys.readouterr().err + assert not (tmp_path / "fit").exists() def test_fit_dipoles_command_runs_real_synthetic_gaussian_fit( @@ -754,6 +808,7 @@ def test_command_help_exposes_stable_safe_input_options(capsys) -> None: assert "vignetted" in captured.out assert "finite-volume" in captured.out assert "--use-finite-difference" in captured.out + assert "--fusion" in captured.out @pytest.mark.parametrize("runtime", ["dragon", "mpi"]) diff --git a/tests/xfit/test_cutile_fused.py b/tests/xfit/test_cutile_fused.py new file mode 100644 index 00000000..9803a041 --- /dev/null +++ b/tests/xfit/test_cutile_fused.py @@ -0,0 +1,649 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import re +import subprocess +import sys +from concurrent.futures import ThreadPoolExecutor +from dataclasses import fields +from threading import Barrier + +import numpy as np +import pytest + +import cuphoton.xfit.api as xfit_api +from cuphoton.xfit import GaussianDipoleModel, LMConfig, fit_dipoles + + +@pytest.fixture +def cutile_gpu(): + cp = pytest.importorskip("cupy") + pytest.importorskip("cuda.tile") + try: + count = cp.cuda.runtime.getDeviceCount() + except Exception as exc: + pytest.skip(f"CUDA runtime is unavailable: {exc}") + if count < 1: + pytest.skip("CUDA device is unavailable") + return cp + + +def _case(mode, dtype): + model = GaussianDipoleModel( + (11, 15), split_weights=(1.0, 0.25, 1.75), dtype=dtype + ) + truth = 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, + ) + images = np.asarray(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 + ) + # The first row finishes before the other rows, exercising original + # candidate indices after the active batch has been compacted. + initial[0] = truth[0] + mask = np.ones(images.shape, dtype=bool) + mask.reshape(3, -1)[0, ::17] = False + mask.reshape(3, -1)[1, 3::19] = False + mask.reshape(3, -1)[2, 5::23] = False + variance = np.linspace(0.5, 2.0, images.size, dtype=dtype).reshape( + images.shape + ) + images[~mask] = np.nan + variance[~mask] = np.nan + return dict( + images=images, + model=model, + initial=initial, + mode=mode, + mask=mask, + variance=variance, + ) + + +def _assert_fit_equal(actual, expected, dtype): + tolerance = 3e-5 if dtype is np.float32 else 5e-12 + for name in ( + "status", + "converged", + "evaluations", + "valid_pixel_count", + "degrees_of_freedom", + "uncertainty_valid", + ): + np.testing.assert_array_equal( + getattr(actual, name), getattr(expected, name), err_msg=name + ) + assert actual.uncertainty_reason == expected.uncertainty_reason + for name in ( + "parameters", + "residual_norm", + "chi_square", + "valid_pixel_fraction", + "null_chi_square", + "delta_chi_square", + "fractional_null_improvement", + "reduced_chi_square", + "covariance", + "standard_errors", + "residuals", + ): + value = getattr(actual, name) + reference = getattr(expected, name) + np.testing.assert_array_equal( + np.isnan(value), np.isnan(reference), err_msg=name + ) + np.testing.assert_allclose( + value, + reference, + rtol=tolerance, + atol=tolerance, + equal_nan=True, + err_msg=name, + ) + + +def _fit_and_rank(cp, monkeypatch, **kwargs): + # Rank is a solver diagnostic; the public result reports its effect on + # covariance through uncertainty_reason instead of exposing the number. + solve = xfit_api.batched_levenberg_marquardt + ranks = [] + + def record_rank(*args, **options): + result = solve(*args, **options) + ranks.append(cp.asnumpy(result.rank).copy()) + return result + + with monkeypatch.context() as patch: + patch.setattr(xfit_api, "batched_levenberg_marquardt", record_rank) + result = fit_dipoles(**kwargs) + assert len(ranks) == 1 + return result, ranks[0] + + +@pytest.mark.parametrize("mode", ["difference", "split"]) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_fused_fit_preserves_weighted_compacted_results( + cutile_gpu, monkeypatch, mode, dtype +): + case = _case(mode, dtype) + fused, fused_rank = _fit_and_rank( + cutile_gpu, monkeypatch, **case, backend="cutile", fusion=True + ) + assert fused.evaluations[0] == 1 + assert np.all(fused.evaluations[1:] > 1) + assert fused.converged.all() + assert fused.uncertainty_valid.all() + for backend in ("cupy", "cutile"): + reference, rank = _fit_and_rank( + cutile_gpu, monkeypatch, **case, backend=backend + ) + _assert_fit_equal(fused, reference, dtype) + np.testing.assert_array_equal(fused_rank, rank) + + +@pytest.mark.parametrize("scenario", ["rank", "budget", "rejected-step"]) +def test_fused_fit_preserves_termination_and_failed_uncertainty( + cutile_gpu, monkeypatch, scenario +): + model = GaussianDipoleModel((11, 15), dtype=np.float64) + truth = np.asarray([[5.0, 1.3, 1.8, 0.25, -2.0, 0.5, 2.2, -0.7]]) + initial = truth.copy() + config = LMConfig() + if scenario == "rank": + truth[:, 0] = 0 + initial = truth.copy() + elif scenario == "budget": + initial[:, 0] += 0.2 + config = LMConfig(max_evaluations=1) + else: + initial = np.asarray([[1.0, 0.5, 0.7, 0.0, 0.0, 0.0, 0.5, 0.0]]) + config = LMConfig(max_evaluations=2, initial_damping=1e-8) + case = dict( + images=model.evaluate(truth), + model=model, + initial=initial, + config=config, + ) + fused, rank = _fit_and_rank( + cutile_gpu, monkeypatch, **case, backend="cutile", fusion=True + ) + unfused, unfused_rank = _fit_and_rank( + cutile_gpu, monkeypatch, **case, backend="cutile" + ) + _assert_fit_equal(fused, unfused, np.float64) + np.testing.assert_array_equal(rank, unfused_rank) + # Analytic final diagnostics do not consume residual evaluations, even + # when the normal-equation path needs to refresh its Jacobian. + reference, reference_rank = _fit_and_rank( + cutile_gpu, monkeypatch, **case, backend="cupy" + ) + _assert_fit_equal(fused, reference, np.float64) + np.testing.assert_array_equal(rank, reference_rank) + assert not fused.uncertainty_valid.any() + assert np.isnan(fused.covariance).all() + if scenario == "rank": + assert np.all(rank < 8) + assert fused.uncertainty_reason == ("rank-deficient Jacobian",) + elif scenario == "budget": + np.testing.assert_array_equal(fused.evaluations, [1]) + np.testing.assert_array_equal(rank, [-1]) + else: + np.testing.assert_array_equal(fused.evaluations, [2]) + np.testing.assert_array_equal(fused.status, ["max_evaluations"]) + # A rejected trial must preserve the accepted parameters/residuals. + np.testing.assert_allclose( + fused.parameters, initial, rtol=5e-12, atol=5e-12 + ) + + +@pytest.mark.parametrize("mode", ["difference", "split"]) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_fused_kernel_outputs_match_noisy_strided_inputs( + cutile_gpu, mode, dtype +): + from cuphoton.xfit._cutile_fused import GaussianWorkspace + + cp = cutile_gpu + rng = np.random.default_rng(1219) + shape = (8, 10) + physical = cp.asarray(_case(mode, dtype)["initial"]) + model = GaussianDipoleModel(shape, backend="cupy", dtype=dtype) + prediction = model.evaluate(physical, mode=mode).reshape(3, -1) + images = prediction + cp.asarray( + rng.normal(scale=0.2, size=prediction.shape), dtype=dtype + ) + weights = cp.asarray( + rng.uniform(0.25, 2.0, size=prediction.shape), dtype=dtype + ) + weights[:, ::7] = 0 + # Both inputs have non-unit observation strides; the active rows also + # arrive in a different order from their original image/weight storage. + images = cp.repeat(images, 2, axis=1)[:, ::2] + weights = cp.repeat(weights, 2, axis=1)[:, ::2] + assert not images.flags.c_contiguous + assert not weights.flags.c_contiguous + indices = cp.asarray([2, 0], dtype=cp.int64) + active = physical[indices].copy() + active[:, 0] *= 1.1 + parameters = active.copy() + parameters[:, 1:3] = cp.log(parameters[:, 1:3]) + expected_residuals = ( + model.evaluate(active, mode=mode).reshape(2, -1) - images[indices] + ) * weights[indices] + jacobian = model.jacobian(active, mode=mode).reshape(2, 8, -1) + jacobian[:, 1:3] *= active[:, 1:3, None] + jacobian *= weights[indices, None, :] + workspace = GaussianWorkspace( + images, weights, image_shape=shape, mode=mode + ) + residuals = workspace.residual(parameters, indices=indices) + actual_jacobian = workspace.jacobian(parameters, indices=indices) + tolerance = 3e-5 if dtype is np.float32 else 5e-12 + for actual, expected in ( + (residuals, expected_residuals), + (actual_jacobian, jacobian), + ): + np.testing.assert_allclose( + cp.asnumpy(actual), + cp.asnumpy(expected), + rtol=tolerance, + atol=tolerance, + ) + + +def test_fused_jacobian_preserves_cancellation_sensitive_hessian(cutile_gpu): + from cuphoton.xfit._cutile_fused import GaussianWorkspace + + cp = cutile_gpu + shape = (51, 51) + parameters = cp.asarray( + [[2000.0, 1.5, 2.25, 0.4, -2.0, 0.0, 2.0, 0.0]], + dtype=cp.float64, + ) + parameters[:, 1:3] = cp.log(parameters[:, 1:3]) + physical = parameters.copy() + physical[:, 1:3] = cp.exp(physical[:, 1:3]) + model = GaussianDipoleModel(shape, backend="cupy", dtype=np.float64) + expected_jacobian = model.jacobian(physical).reshape(1, 8, -1) + expected_jacobian[:, 1:3] *= physical[:, 1:3, None] + images = cp.zeros((1, shape[0] * shape[1]), dtype=cp.float64) + workspace = GaussianWorkspace( + images, cp.ones_like(images), image_shape=shape, mode="difference" + ) + jacobian = workspace.jacobian( + parameters, indices=cp.asarray([0], dtype=cp.int64) + ) + # Large symmetric stars make cancellation in Hessian cross terms + # sensitive to the separately rounded products in the analytic model. + expected_hessian = cp.einsum( + "knm,kpm->knp", expected_jacobian, expected_jacobian + ) + hessian = cp.einsum("knm,kpm->knp", jacobian, jacobian) + for actual, expected in ( + (jacobian, expected_jacobian), + (hessian, expected_hessian), + ): + np.testing.assert_allclose( + cp.asnumpy(actual), cp.asnumpy(expected), rtol=5e-12, atol=5e-12 + ) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_fused_invalid_sigma_trials_keep_accepted_residuals( + cutile_gpu, dtype +): + from cuphoton.xfit._cutile_fused import GaussianWorkspace + + cp = cutile_gpu + model = GaussianDipoleModel((7, 9), dtype=dtype) + physical = np.asarray([[5, 1.3, 1.8, 0.25, -2, 0.5, 2.2, -0.7]], dtype) + image = cp.asarray(model.evaluate(physical)).reshape(1, -1) + workspace = GaussianWorkspace( + image, cp.ones_like(image), image_shape=(7, 9), mode="difference" + ) + parameters = cp.asarray(physical) + parameters[:, 1:3] = cp.log(parameters[:, 1:3]) + indices = cp.asarray([0], dtype=cp.int64) + accepted = workspace.residual(parameters, indices=indices) + before = cp.asnumpy(accepted) + for log_sigma in (-1000.0, 1000.0): + trial = parameters.copy() + trial[:, 1:3] = log_sigma + invalid = workspace.residual(trial, indices=indices) + assert cp.isnan(invalid).all().item() + np.testing.assert_array_equal(cp.asnumpy(accepted), before) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("fusion", [False, True]) +def test_cutile_large_width_keeps_physical_derivative_order( + cutile_gpu, monkeypatch, dtype, fusion +): + from cuphoton.xfit._cutile import gaussian_normal_equations + from cuphoton.xfit._cutile_fused import GaussianWorkspace + + cp = cutile_gpu + width = 1e13 if dtype is np.float32 else 1e103 + initial = np.asarray([[1, width, 1, 0, -width, 0, width, 0]], dtype) + model = GaussianDipoleModel((5, 7), split_weights=(1, 1, 1), dtype=dtype) + images = np.asarray(model.evaluate(initial, mode="split")) * dtype(1.01) + mask = np.zeros_like(images, dtype=bool) + mask[..., 2, :] = True + parameters = cp.asarray(initial) + parameters[:, 1:3] = cp.log(parameters[:, 1:3]) + physical = parameters.copy() + physical[:, 1:3] = cp.exp(physical[:, 1:3]) + gpu_model = model.to_backend("cupy") + weights = cp.asarray(mask, dtype=dtype).reshape(1, -1) + data = cp.asarray(images).reshape(1, -1) + expected_residual = ( + gpu_model.evaluate(physical, mode="split").reshape(1, -1) - data + ) * weights + jacobian = gpu_model.jacobian(physical, mode="split").reshape(1, 8, -1) + jacobian[:, 1:3] *= physical[:, 1:3, None] + jacobian *= weights[:, None, :] + # The physical sigma**3 denominator overflows, while sigma**2 and + # every weighted Jacobian entry remain finite. Simplifying the chain + # rule to star*x*x/sigma**2 would introduce a spurious width direction. + assert cp.isfinite(jacobian).all().item() + assert (jacobian[:, 1] == 0).all().item() + workspace = GaussianWorkspace( + data, weights, image_shape=(5, 7), mode="split" + ) + indices = cp.asarray([0], dtype=cp.int64) + if fusion: + actual_jacobian = workspace.jacobian(parameters, indices=indices) + tolerance = 3e-5 if dtype is np.float32 else 5e-12 + np.testing.assert_allclose( + cp.asnumpy(actual_jacobian), + cp.asnumpy(jacobian), + rtol=tolerance, + atol=tolerance, + ) + assert cp.isfinite(actual_jacobian).all().item() + assert (actual_jacobian[:, 1] == 0).all().item() + else: + gradient, hessian = gaussian_normal_equations( + parameters, + expected_residual, + weights, + image_shape=(5, 7), + mode="split", + ) + assert cp.isfinite(gradient).all().item() + assert cp.isfinite(hessian).all().item() + assert (gradient[:, 1] == 0).all().item() + assert (hessian[:, 1, :] == 0).all().item() + case = dict( + images=images, + model=model, + initial=initial, + mask=mask, + mode="split", + config=LMConfig(max_evaluations=2, f_tol=0, x_tol=0, g_tol=0), + ) + reference, reference_rank = _fit_and_rank( + cp, monkeypatch, **case, backend="cupy" + ) + fused, rank = _fit_and_rank( + cp, monkeypatch, **case, backend="cutile", fusion=fusion + ) + np.testing.assert_array_equal(reference.status, ["max_evaluations"]) + np.testing.assert_array_equal(reference.evaluations, [2]) + _assert_fit_equal(fused, reference, dtype) + np.testing.assert_array_equal(rank, reference_rank) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("masked", [False, True]) +@pytest.mark.parametrize("fusion", [False, True]) +def test_cutile_tiny_width_preserves_invalid_derivative( + cutile_gpu, monkeypatch, dtype, masked, fusion +): + from cuphoton.xfit._cutile import gaussian_normal_equations + from cuphoton.xfit._cutile_fused import GaussianWorkspace + + cp = cutile_gpu + width = 1e-16 if dtype is np.float32 else 1e-110 + initial = np.asarray([[1, width, 1, 0, -1, 0, 1, 0]], dtype) + images = np.full((1, 2, 2), np.nan if masked else 0, dtype=dtype) + mask = np.zeros_like(images, dtype=bool) if masked else None + parameters = cp.asarray(initial) + parameters[:, 1:3] = cp.log(parameters[:, 1:3]) + # Every model pixel is zero, but the physical width derivative is 0/0 + # because sigma**3 underflows. Zero mask weights preserve that NaN. + weights = cp.full((1, 4), 0 if masked else 1, dtype=dtype) + workspace = GaussianWorkspace( + cp.zeros((1, 4), dtype=dtype), + weights, + image_shape=(2, 2), + mode="difference", + ) + indices = cp.asarray([0], dtype=cp.int64) + residual = workspace.residual(parameters, indices=indices) + assert (residual == 0).all().item() + if fusion: + jacobian = workspace.jacobian(parameters, indices=indices) + physical = parameters.copy() + physical[:, 1:3] = cp.exp(physical[:, 1:3]) + model = GaussianDipoleModel((2, 2), backend="cupy", dtype=dtype) + expected_jacobian = model.jacobian(physical).reshape(1, 8, 4) + expected_jacobian[:, 1:3] *= physical[:, 1:3, None] + expected_jacobian *= weights[:, None, :] + assert not cp.isfinite(jacobian).all().item() + np.testing.assert_array_equal( + cp.asnumpy(cp.isnan(jacobian)), + cp.asnumpy(cp.isnan(expected_jacobian)), + ) + tolerance = 3e-5 if dtype is np.float32 else 5e-12 + np.testing.assert_allclose( + cp.asnumpy(jacobian), + cp.asnumpy(expected_jacobian), + rtol=tolerance, + atol=tolerance, + equal_nan=True, + ) + else: + gradient, hessian = gaussian_normal_equations( + parameters, + residual, + weights, + image_shape=(2, 2), + mode="difference", + ) + assert not cp.isfinite(gradient).all().item() + assert not cp.isfinite(hessian).all().item() + case = dict( + images=images, + model="gaussian", + initial=initial, + mask=mask, + config=LMConfig(max_evaluations=2), + ) + reference, reference_rank = _fit_and_rank( + cp, monkeypatch, **case, backend="cupy" + ) + fused, rank = _fit_and_rank( + cp, monkeypatch, **case, backend="cutile", fusion=fusion + ) + np.testing.assert_array_equal(reference.status, ["invalid_residual"]) + np.testing.assert_array_equal(reference.evaluations, [1]) + _assert_fit_equal(fused, reference, dtype) + np.testing.assert_array_equal(rank, reference_rank) + + +@pytest.mark.parametrize( + ("dtype", "amplitude"), [(np.float32, 1e25), (np.float64, 1e200)] +) +@pytest.mark.parametrize("nonzero_residual", [False, True]) +@pytest.mark.parametrize("fusion", [False, True]) +def test_cutile_preserves_finite_jacobian_overflow_results( + cutile_gpu, monkeypatch, dtype, amplitude, nonzero_residual, fusion +): + cp = cutile_gpu + model = GaussianDipoleModel((5, 7), backend="cupy", dtype=dtype) + # Coincident stars keep the model finite and exactly zero. Only the + # middle row's finite position derivatives overflow their products. + 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: + images[1, 1, 2] = amplitude + jacobian = model.jacobian(cp.asarray(initial)).reshape(3, 8, -1) + assert cp.isfinite(jacobian).all().item() + 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], + ) + case = dict(images=images, model=model, initial=initial) + reference, reference_rank = _fit_and_rank( + cp, monkeypatch, **case, backend="cupy" + ) + actual, rank = _fit_and_rank( + cp, monkeypatch, **case, backend="cutile", fusion=fusion + ) + assert reference.status.tolist() == [ + "converged_g_tol", + "singular" if nonzero_residual else "converged_g_tol", + "converged_g_tol", + ] + np.testing.assert_array_equal(reference.evaluations, [1, 1, 1]) + np.testing.assert_array_equal(actual.parameters, reference.parameters) + _assert_fit_equal(actual, reference, dtype) + np.testing.assert_array_equal(rank, reference_rank) + + +def test_fused_concurrent_workers_keep_independent_results(cutile_gpu): + cp = cutile_gpu + cases = [_case("split", np.float64) for _ in range(2)] + cases[1]["images"] = cases[1]["images"] * 1.25 + cases[1]["initial"] = cases[1]["initial"].copy() + cases[1]["initial"][:, 0] *= 1.25 + # Warm all specializations on the default stream before testing worker + # ownership; compilation concurrency is a separate concern. + with cp.cuda.Stream.null: + expected = [fit_dipoles(**case, backend="cupy") for case in cases] + for case in cases: + fit_dipoles(**case, backend="cutile", fusion=True) + cp.cuda.Stream.null.synchronize() + barrier = Barrier(2) + + def worker(index): + pool = cp.cuda.MemoryPool() + stream = cp.cuda.Stream(non_blocking=True) + retained = [] + with cp.cuda.Device(0), stream, cp.cuda.using_allocator(pool.malloc): + barrier.wait(timeout=30) + for _ in range(3): + result = fit_dipoles( + **cases[index], backend="cutile", fusion=True + ) + snapshots = { + field.name: value.copy() + for field in fields(result) + if isinstance( + value := getattr(result, field.name), np.ndarray + ) + } + retained.append((result, snapshots)) + stream.synchronize() + pool.free_all_blocks() + for result, snapshots in retained: + _assert_fit_equal(result, expected[index], np.float64) + for name, snapshot in snapshots.items(): + np.testing.assert_array_equal(getattr(result, name), snapshot) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(worker, index) for index in range(2)] + for future in futures: + future.result(timeout=180) + + +def test_fused_cli_matches_unfused_artifacts(cutile_gpu, tmp_path): + import pyarrow.parquet as pq + import yaml + + case = _case("difference", np.float64) + input_path = tmp_path / "synthetic.npz" + np.savez( + input_path, + candidate_id=np.asarray(["a", "b", "c"]), + **{ + name: case[name] + for name in ("images", "initial", "mask", "variance") + }, + ) + outputs = [] + for fusion in (False, True): + output = tmp_path / ("fused" if fusion else "unfused") + command = [ + sys.executable, + "-m", + "cuphoton", + "xfit", + "fit-dipoles", + "--input", + str(input_path), + "--output-dir", + str(output), + "--model", + "gaussian", + "--backend", + "cutile", + ] + if fusion: + command.append("--fusion") + completed = subprocess.run( + command, text=True, capture_output=True, timeout=240, check=False + ) + assert completed.returncode == 0, completed.stdout + completed.stderr + assert not re.search( + r"CUDA_ERROR_|Exception ignored (?:in:|while)", + completed.stdout + completed.stderr, + ), completed.stdout + completed.stderr + effective = yaml.safe_load( + (output / "effective-config.yaml").read_text() + ) + assert effective["fusion"] is fusion + outputs.append(output) + reference, fused = [pq.read_table(p / "fits.parquet") for p in outputs] + for name in ("candidate_id", "status", "evaluations", "converged"): + assert reference[name].to_pylist() == fused[name].to_pylist() + for name in case["model"].parameter_names: + np.testing.assert_allclose( + fused[name].to_numpy(), + reference[name].to_numpy(), + rtol=5e-12, + atol=5e-12, + ) + with ( + np.load(outputs[0] / "fit-arrays.npz", allow_pickle=False) as before, + np.load(outputs[1] / "fit-arrays.npz", allow_pickle=False) as after, + ): + for name in ("covariance", "residuals"): + np.testing.assert_array_equal( + np.isnan(after[name]), np.isnan(before[name]) + ) + np.testing.assert_allclose( + after[name], + before[name], + rtol=5e-12, + atol=5e-12, + equal_nan=True, + ) diff --git a/tests/xfit/test_executor.py b/tests/xfit/test_executor.py index 6f058f81..71967854 100644 --- a/tests/xfit/test_executor.py +++ b/tests/xfit/test_executor.py @@ -13,6 +13,7 @@ import numpy as np import pyarrow.parquet as pq import pytest +import yaml from cuphoton.core.artifacts import file_sha256 from cuphoton.xfit import ( @@ -103,8 +104,9 @@ def test_chunk_size_is_part_of_configuration_identity(tmp_path): assert first["configuration_sha256"] != second["configuration_sha256"] +@pytest.mark.parametrize("fusion", [False, True]) def test_finalizer_retains_planning_input_and_still_checks_content( - tmp_path, monkeypatch + tmp_path, monkeypatch, fusion ): path = _input(tmp_path) loads = [] @@ -115,8 +117,11 @@ def load(*args, **kwargs): return original_load(*args, **kwargs) def cpu_fit(*args, **kwargs): - result = fit_dipoles(*args, **{**kwargs, "backend": "numpy"}) - return replace(result, backend="cupy", device="cuda:0") + assert kwargs["fusion"] is fusion + result = fit_dipoles( + *args, **{**kwargs, "backend": "numpy", "fusion": False} + ) + return replace(result, backend=kwargs["backend"], device="cuda:0") monkeypatch.setattr(executor, "load_xfit_dataset", load) monkeypatch.setattr(executor, "fit_dipoles", cpu_fit) @@ -124,7 +129,11 @@ def cpu_fit(*args, **kwargs): spec = executor.prepare_xfit_workload( input_path=path, chunk_size=2, - fit_options={"backend": "cupy", "max_evaluations": 2}, + fit_options={ + "backend": "cutile" if fusion else "cupy", + "fusion": fusion, + "max_evaluations": 2, + }, ) worker = spec.worker_factory(spec.options_payload) for index in range(2): @@ -134,6 +143,10 @@ def cpu_fit(*args, **kwargs): worker.run_item(item, round_dir / "items" / item.item_id) records.append({"item_id": item.item_id, "status": "success"}) assert spec.finalize_round(round_dir, records)["candidate_count"] == 3 + effective = yaml.safe_load( + (round_dir / "scientific" / "effective-config.yaml").read_text() + ) + assert effective["fusion"] is fusion assert len(loads) == 2 # One planning load and one worker load. before = path.stat() changed = bytearray(path.read_bytes()) @@ -146,18 +159,46 @@ def cpu_fit(*args, **kwargs): worker.close() -def test_preflight_accepts_analytic_gaussian_cutile(tmp_path): +@pytest.mark.parametrize("fusion", [False, True]) +def test_preflight_accepts_analytic_gaussian_cutile(tmp_path, fusion): spec = executor.prepare_xfit_workload( - input_path=_input(tmp_path), fit_options={"backend": "cutile"} + input_path=_input(tmp_path), + fit_options={"backend": "cutile", "fusion": fusion}, ) assert spec.backend == "cutile" + assert spec.options_payload["fit_options"]["fusion"] is fusion + + +@pytest.mark.parametrize( + ("backend", "fusion", "message"), + [ + ("cupy", True, "requires backend='cutile'"), + ("cutile", "yes", "must be a boolean"), + ], +) +def test_preflight_rejects_invalid_fusion_before_input_reads( + tmp_path, monkeypatch, backend, fusion, message +): + monkeypatch.setattr( + executor, + "load_xfit_dataset", + lambda *args, **kwargs: pytest.fail( + "invalid fusion settings read the input" + ), + ) + with pytest.raises((ValueError, TypeError), match=message): + executor.prepare_xfit_workload( + input_path=tmp_path / "missing.npz", + fit_options={"backend": backend, "fusion": fusion}, + ) @pytest.mark.parametrize( "runtime,rank", [("dragon", 0), ("mpi", 0), ("mpi", 1)] ) +@pytest.mark.parametrize("fusion", [False, True]) def test_cli_dispatches_collective_preflight_and_preserves_options( - tmp_path, monkeypatch, capsys, runtime, rank + tmp_path, monkeypatch, capsys, runtime, rank, fusion ): from types import SimpleNamespace @@ -197,7 +238,8 @@ def run(**kwargs): "--model", "gaussian", "--backend", - "cupy", + "cutile" if fusion else "cupy", + *(["--fusion"] if fusion else []), "--executor", runtime, "--chunk-size", @@ -222,6 +264,7 @@ def run(**kwargs): assert kwargs["benchmark"].measure_rounds == 2 assert len(spec.items) == 2 assert spec.options_payload["fit_options"]["max_evaluations"] == 10 + assert spec.options_payload["fit_options"]["fusion"] is fusion assert json.loads(capsys.readouterr().out)["status"] == "success" diff --git a/tests/xfit/test_models.py b/tests/xfit/test_models.py index 13adbc3d..3c5c7b59 100644 --- a/tests/xfit/test_models.py +++ b/tests/xfit/test_models.py @@ -920,14 +920,14 @@ def bookkeeping_only_solver_as_numpy(value): with pytest.raises(FrozenInstanceError): device.device_id = device.device_id + 1 # type: ignore[misc] - for name, value in ( - ("schema", "other"), - ("solver", "other"), - ("backend", "numpy"), - ("result_location", "host"), - ): - with pytest.raises(ValueError): - replace(device, **{name: value}) + for name in ("schema", "solver", "result_location"): + # Python 3.14 changed replace()'s init=False error to TypeError. + with pytest.raises( + (ValueError, TypeError), match=rf"field {name}.*init=False" + ): + replace(device, **{name: "other"}) + with pytest.raises(ValueError, match="backend"): + replace(device, backend="numpy") with pytest.raises(ValueError, match="CUDA device"): replace(device, device_id=device.device_id + 1) with pytest.raises(TypeError, match=device.dtype): @@ -1097,6 +1097,40 @@ def test_fit_dipoles_device_rejects_stamp_model_dtype_conversion( fit_dipoles_device(images, model=model) +@pytest.mark.parametrize( + ("backend", "model_name", "finite_difference", "fusion", "message"), + [ + ("auto", "gaussian", False, False, "device fit backend"), + ("numpy", "gaussian", False, False, "device fit backend"), + ("cupy", "gaussian", False, True, "requires backend='cutile'"), + ("cutile", "stamp", False, False, "only the Gaussian model"), + ("cutile", "gaussian", True, False, "finite-difference fitting"), + ("cutile", "gaussian", False, 1, "must be a boolean"), + ], +) +def test_device_options_reject_invalid_settings_before_cuda( + monkeypatch, backend, model_name, finite_difference, fusion, message +) -> None: + model = ( + StampDipoleModel(np.ones((5, 7)), image_shape=(5, 7)) + if model_name == "stamp" + else "gaussian" + ) + monkeypatch.setattr( + xfit_api, + "_load_cupy", + lambda: pytest.fail("invalid device options loaded CUDA"), + ) + with pytest.raises((TypeError, ValueError), match=message): + fit_dipoles_device( + np.zeros((1, 5, 7)), + model=model, + backend=backend, + fusion=fusion, + config=LMConfig(use_finite_difference=finite_difference), + ) + + def test_fit_dipoles_device_reports_no_visible_device( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -1222,6 +1256,40 @@ def test_cupy_stamp_fit_matches_numpy_when_cuda_is_available() -> None: assert np.allclose(gpu.parameters, cpu.parameters, rtol=2e-8, atol=2e-8) +@pytest.mark.parametrize( + ("backend", "model_name", "finite_difference", "fusion", "message"), + [ + ("numpy", "gaussian", False, True, "requires backend='cutile'"), + ("cupy", "gaussian", False, True, "requires backend='cutile'"), + ("auto", "gaussian", False, True, "requires backend='cutile'"), + ("cutile", "stamp", False, True, "only the Gaussian model"), + ("cutile", "gaussian", True, True, "finite-difference fitting"), + ("cutile", "gaussian", False, "yes", "must be a boolean"), + ], +) +def test_fusion_rejects_unsupported_settings_before_backend_resolution( + monkeypatch, backend, model_name, finite_difference, fusion, message +) -> None: + model = ( + StampDipoleModel(np.ones((5, 7)), image_shape=(5, 7)) + if model_name == "stamp" + else "gaussian" + ) + monkeypatch.setattr( + xfit_api, + "resolve_backend", + lambda _: pytest.fail("invalid fusion settings resolved a backend"), + ) + with pytest.raises((TypeError, ValueError), match=message): + fit_dipoles( + np.zeros((1, 5, 7)), + model=model, + backend=backend, + fusion=fusion, + config=LMConfig(use_finite_difference=finite_difference), + ) + + def test_cutile_rejects_finite_difference_provenance() -> None: model = GaussianDipoleModel((7, 9), dtype=np.float64) truth = _gaussian_truth()[:1] @@ -1543,3 +1611,120 @@ def test_cutile_preserves_rank_and_budget_failure_semantics() -> None: assert np.array_equal(cutile_budget.status, cupy_budget.status) assert np.array_equal(cutile_budget.evaluations, cupy_budget.evaluations) assert cutile_budget.uncertainty_reason == cupy_budget.uncertainty_reason + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("fusion", [False, True]) +def test_cutile_device_results_survive_worker_reuse(dtype, fusion) -> None: + from concurrent.futures import ThreadPoolExecutor + from dataclasses import fields + from threading import Barrier + + cp = _require_cutile() + pytest.importorskip("torch") + + from cuphoton.xscan.xfit_features import ( + transform_xfit_result_features_device, + ) + + device_id = int(cp.cuda.Device().id) + model = GaussianDipoleModel((11, 15), dtype=dtype) + cases = [] + for scale in (1.0, 1.5): + truth = _gaussian_truth(dtype) + truth[:, 0] *= scale + images = model.evaluate(truth, mode="difference") + 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] + cases.append( + dict( + images=images, + model=model, + initial=initial, + variance=np.full_like(images, 0.5), + mode="difference", + ) + ) + + def snapshot(result): + return { + field.name: cp.asnumpy(value) + for field in fields(result) + if isinstance(value := getattr(result, field.name), cp.ndarray) + } + + def features(result): + transformed = transform_xfit_result_features_device( + result, image_shape=(11, 15), variance_present=True + ) + assert transformed.backend == "cupy" + return transformed + + expected = [] + with cp.cuda.Stream.null: + for case in cases: + reference = fit_dipoles_device(**case) + expected.append( + (snapshot(reference), cp.asnumpy(features(reference).values)) + ) + fit_dipoles_device(**case, backend="cutile", fusion=fusion) + cp.cuda.Stream.null.synchronize() + barrier = Barrier(2) + tolerance = 3e-5 if dtype is np.float32 else 5e-12 + + def worker(index): + with cp.cuda.Device(device_id): + pool = cp.cuda.MemoryPool() + stream = cp.cuda.Stream(non_blocking=True) + retained = [] + with stream, cp.cuda.using_allocator(pool.malloc): + barrier.wait(timeout=30) + for iteration in range(2): + case_index = (index + iteration) % len(cases) + result = fit_dipoles_device( + **cases[case_index], backend="cutile", fusion=fusion + ) + assert result.backend == "cutile" + assert result.device_id == device_id + transformed = features(result) + arrays = snapshot(result) + reference, reference_features = expected[case_index] + assert arrays.keys() == reference.keys() + for name, actual in arrays.items(): + target = reference[name] + if actual.dtype.kind in "fc": + np.testing.assert_array_equal( + np.isnan(actual), np.isnan(target) + ) + np.testing.assert_allclose( + actual, target, rtol=tolerance, atol=tolerance + ) + else: + np.testing.assert_array_equal(actual, target) + feature_values = cp.asnumpy(transformed.values) + # Canonical xScan features have float32 storage even + # when the fit was computed in float64. + np.testing.assert_allclose( + feature_values, + reference_features, + rtol=3e-5, + atol=3e-5, + ) + retained.append( + (result, arrays, transformed, feature_values) + ) + stream.synchronize() + pool.free_all_blocks() + for result, arrays, transformed, feature_values in retained: + for name, actual in snapshot(result).items(): + np.testing.assert_array_equal(actual, arrays[name]) + np.testing.assert_array_equal( + cp.asnumpy(transformed.values), feature_values + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(worker, index) for index in range(2)] + for future in futures: + future.result(timeout=180) diff --git a/tests/xscan/test_device_pipeline.py b/tests/xscan/test_device_pipeline.py index a4f41f7a..40314b00 100644 --- a/tests/xscan/test_device_pipeline.py +++ b/tests/xscan/test_device_pipeline.py @@ -795,6 +795,20 @@ def test_scientific_evidence_config_binds_xpois_metadata_and_layout( ) +@pytest.mark.parametrize("backend", ["cupy", "cutile", "native"]) +def test_scientific_evidence_binds_selected_xfit_backend( + tmp_path: Path, backend: str +) -> None: + config = _evidence_config(tmp_path) + config = replace(config, xfit=replace(config.xfit, backend=backend)) + evidence = _scientific_evidence_payload() + evidence["xfit"]["backend"] = backend + decode_device_pipeline_evidence(evidence, config=config) + evidence["xfit"]["backend"] = "cutile" if backend == "cupy" else "cupy" + with pytest.raises(ValueError, match="xFit backend does not match"): + decode_device_pipeline_evidence(evidence, config=config) + + def test_scientific_evidence_config_binds_flux_basis_order_and_reference( tmp_path: Path, ) -> None: @@ -1366,7 +1380,9 @@ def fit(images: Any, **kwargs: Any) -> Any: assert kwargs == { "model": "gaussian", "mode": "difference", + "backend": "cupy", "config": context.solver_config, + "fusion": False, } stage_calls.append("xfit") return _fake_xfit_result(cp, 2) @@ -1748,6 +1764,50 @@ def test_changed_input_fails_before_device_ingress( _read_item_inputs(item) +def test_xfit_backend_defaults_preserve_legacy_payload_and_hash( + tmp_path: Path, +) -> None: + positional = DeviceXFitPipelineConfig("gaussian", "difference", 1e-6) + assert positional.f_tol == 1e-6 + assert positional.backend == "cupy" + assert positional.fusion is False + config = _evidence_config(tmp_path) + legacy = config.to_payload() + assert "backend" not in legacy["xfit"] + assert "fusion" not in legacy["xfit"] + explicit = config.to_payload() + explicit["xfit"].update(backend="cupy", fusion=False) + restored = DevicePipelineConfig.from_payload(explicit) + assert restored.to_payload() == legacy + assert restored.configuration_sha256 == config.configuration_sha256 + + fused = replace( + config, xfit=replace(config.xfit, backend="cutile", fusion=True) + ) + assert DevicePipelineConfig.from_payload(fused.to_payload()) == fused + assert fused.configuration_sha256 != config.configuration_sha256 + assert fused.xfit.solver_payload() == config.xfit.solver_payload() + + +@pytest.mark.parametrize( + ("options", "message"), + [ + ({"backend": "numpy"}, "xfit backend"), + ({"fusion": 1}, "fusion must be boolean"), + ({"fusion": True}, "fusion requires backend='cutile'"), + ( + {"backend": "cutile", "use_finite_difference": True}, + "does not support finite differences", + ), + ], +) +def test_xfit_config_rejects_unsupported_backend_options( + options: dict[str, Any], message: str +) -> None: + with pytest.raises((TypeError, ValueError), match=message): + DeviceXFitPipelineConfig(**options) + + def test_xfit_config_rejects_values_the_solver_cannot_use() -> None: with pytest.raises(ValueError, match="model"): DeviceXFitPipelineConfig(model="stamp") # type: ignore[arg-type] diff --git a/tests/xscan/test_pipeline_benchmark.py b/tests/xscan/test_pipeline_benchmark.py index fdd0d949..150449ee 100644 --- a/tests/xscan/test_pipeline_benchmark.py +++ b/tests/xscan/test_pipeline_benchmark.py @@ -307,7 +307,11 @@ def test_compare_cpu_orchestration_and_audit( ] metadata = { "xpois": {"xpois_chi2": 1.0, "basis_terms": []}, - "xfit": {"parameter_names": ["amplitude"], "feature_names": ["flux"]}, + "xfit": { + "parameter_names": ["amplitude"], + "feature_names": ["flux"], + "backend": config.xfit.backend, + }, "xscan": {"predictions": predictions}, } # Exercise the real file contracts, stage summaries, round bookkeeping and @@ -352,6 +356,7 @@ def child(command, log, *, timeout): "xpois": {"chi2": 1.0}, "scientific_evidence": { "xpois": {"basis_terms": []}, + "xfit": {"backend": config.xfit.backend}, "parameter_names": ["amplitude"], "feature_names": ["flux"], }, @@ -413,6 +418,15 @@ def child(command, log, *, timeout): assert row["batch_seconds_without_extra_hashing"] == pytest.approx( row["batch_seconds"] - row["extra_hashing_seconds"] ) + stage_path = args.output / "staged/round-001/xfit/summary.json" + original_stage = stage_path.read_text() + changed_stage = json.loads(original_stage) + changed_stage["items"][0]["metadata"]["backend"] = "cutile" + benchmark.write_json(stage_path, changed_stage) + with pytest.raises(ValueError, match="xFit backend differs"): + benchmark.audit(args) + stage_path.write_text(original_stage) + # Candidate identity and exact scientific values both gate acceptance. path = args.output / "pipeline/round-001/item-0000.json" result = json.loads(path.read_text()) diff --git a/tests/xscan/test_pipeline_benchmark_cli.py b/tests/xscan/test_pipeline_benchmark_cli.py index ed424daf..f758f0e8 100644 --- a/tests/xscan/test_pipeline_benchmark_cli.py +++ b/tests/xscan/test_pipeline_benchmark_cli.py @@ -4,6 +4,7 @@ """CPU-only checks for the benchmark's shared CLI adapter.""" +import json from pathlib import Path import pytest @@ -114,6 +115,83 @@ def test_benchmark_child_options(monkeypatch): ) +def test_benchmark_config_selects_cutile_fusion(tmp_path, monkeypatch): + from cuphoton.xscan.device_pipeline import ( + DevicePipelineCandidate, + DevicePipelineConfig, + DevicePipelineItem, + DeviceXFitPipelineConfig, + DeviceXPOISPipelineConfig, + NpyArrayDescriptor, + ) + from cuphoton.xscan.pipeline_benchmark import stages + + config = DevicePipelineConfig( + device="cuda:0", + checkpoint_dir=str(tmp_path), + checkpoint_sha256="a" * 64, + feature_schema_path=str(tmp_path / "schema.json"), + feature_schema_sha256="b" * 64, + stamp_shape=(17, 17), + decision_threshold=0.5, + xpois=DeviceXPOISPipelineConfig( + kernel_shape=(3, 3), basis_sigmas=(0.8,), basis_degrees=(0,) + ), + xfit=DeviceXFitPipelineConfig( + backend="cutile", fusion=True, max_evaluations=13 + ), + ) + image = NpyArrayDescriptor( + path=str(tmp_path / "image.npy"), + sha256="c" * 64, + shape=(64, 64), + dtype="float64", + ) + item = DevicePipelineItem( + item_id="pair-0", + reference=image, + target=image, + candidates=(DevicePipelineCandidate("candidate-0", 32, 32, 0),), + ) + config_path, items_path = ( + tmp_path / "config.json", + tmp_path / "items.json", + ) + config_path.write_text(json.dumps(config.to_payload())) + items_path.write_text(json.dumps([item.to_payload()])) + received = [] + monkeypatch.setattr( + stages, "run_stage", lambda *args: received.append(args) + ) + output = tmp_path / "result" + + assert ( + run_component( + "xscan", + [ + "benchmark-pipeline", + "--config", + str(config_path), + "--items", + str(items_path), + "--stage", + "xfit", + "--output", + str(output), + ], + ) + == 0 + ) + ((stage, actual_config, actual_items, actual_output),) = received + assert stage == "xfit" + assert actual_config == config + assert actual_config.xfit.backend == "cutile" + assert actual_config.xfit.fusion is True + assert actual_config.xfit.max_evaluations == 13 + assert actual_items == [item] + assert actual_output == output + + @pytest.mark.parametrize("timeout", ["0", "-1", "nan", "inf"]) def test_benchmark_rejects_invalid_timeout(monkeypatch, timeout): received = [] diff --git a/tests/xscan/test_pipeline_benchmark_stages.py b/tests/xscan/test_pipeline_benchmark_stages.py index bff69100..1cfa9919 100644 --- a/tests/xscan/test_pipeline_benchmark_stages.py +++ b/tests/xscan/test_pipeline_benchmark_stages.py @@ -177,6 +177,86 @@ def test_synchronized_timer_waits_for_completion(monkeypatch) -> None: assert timings == {"compute": 3.0} +@pytest.mark.parametrize( + "backend,fusion", + [("cupy", False), ("cutile", False), ("cutile", True), ("native", False)], +) +def test_xfit_stage_forwards_backend_and_fusion( + tmp_path, monkeypatch, backend, fusion +) -> None: + from cuphoton import xfit + + config = SimpleNamespace( + stamp_shape=(17, 17), + xfit=stages.pipeline.DeviceXFitPipelineConfig( + backend=backend, fusion=fusion, max_evaluations=13, g_tol=1.0e-7 + ), + ) + difference = np.zeros((1, 17, 17), dtype=np.float64) + monkeypatch.setattr(stages, "_read_array", lambda *args: difference) + result = SimpleNamespace( + **{name: np.ones(1) for name in stages.XFIT_FIELDS}, + parameter_names=("amplitude",), + model="gaussian", + mode="difference", + backend=backend, + solver="lm", + dtype="float64", + ) + fit_calls = [] + + def fit(images, **kwargs): + fit_calls.append((images, kwargs)) + return result + + monkeypatch.setattr(xfit, "fit_dipoles_device", fit) + feature_values = np.zeros((1, 2), dtype=np.float32) + + def features(actual, **kwargs): + assert actual is result + assert kwargs == { + "image_shape": config.stamp_shape, + "variance_present": False, + } + return SimpleNamespace( + values=feature_values, feature_names=("first", "second") + ) + + monkeypatch.setitem( + sys.modules, + "cuphoton.xscan.xfit_features", + SimpleNamespace(transform_xfit_result_features_device=features), + ) + fake_cupy = SimpleNamespace( + asarray=np.asarray, + asnumpy=np.asarray, + cuda=SimpleNamespace( + Device=lambda: SimpleNamespace(synchronize=lambda: None) + ), + ) + arrays, metadata, _ = stages._run_item( + "xfit", + SimpleNamespace(), + config, + {"cp": fake_cupy, "torch": None}, + tmp_path, + {"xpois": {"items": [{"artifacts": {"difference": {}}}]}}, + 0, + ) + + ((images, kwargs),) = fit_calls + assert images is difference + assert kwargs == { + "model": "gaussian", + "mode": "difference", + **({"backend": backend} if backend != "cupy" else {}), + **({"fusion": True} if fusion else {}), + "config": xfit.LMConfig(max_evaluations=13, g_tol=1.0e-7), + } + assert metadata["backend"] == backend + assert arrays["xfit.features"] is feature_values + + @pytest.mark.parametrize("stage", stages.STAGES) def test_setup_imports_only_its_numerical_runtime( tmp_path, monkeypatch, stage From 2f5002bcbc054842f28f3944db05782fc85002c8 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Fri, 9 Oct 2026 12:13:46 -0700 Subject: [PATCH 2/7] Wrap fused cuTile evaluation failures consistently Signed-off-by: Trent Nelson --- src/cuphoton/xfit/_cutile_fused.py | 49 +++++++++++++++++------------- 1 file changed, 28 insertions(+), 21 deletions(-) diff --git a/src/cuphoton/xfit/_cutile_fused.py b/src/cuphoton/xfit/_cutile_fused.py index 78f2ec90..1d9ba4ce 100644 --- a/src/cuphoton/xfit/_cutile_fused.py +++ b/src/cuphoton/xfit/_cutile_fused.py @@ -449,27 +449,34 @@ def _evaluate( # can alter the trajectory of an ill-conditioned fit. self._cp.exp(parameters[:, 1:3], out=width_cubes) self._cp.power(width_cubes, 3, out=width_cubes) - self._ct.launch( - self._stream, - ( - parameters.shape[0], - math.ceil(self._observations / _EVALUATION_TILE), - 1, - ), - self._kernel, - ( - parameters, - width_cubes, - self._images, - self._weights, - indices, - output, - self._height, - self._width, - self._planes, - compute_jacobian, - ), - ) + try: + self._ct.launch( + self._stream, + ( + parameters.shape[0], + math.ceil(self._observations / _EVALUATION_TILE), + 1, + ), + self._kernel, + ( + parameters, + width_cubes, + self._images, + self._weights, + indices, + output, + self._height, + self._width, + self._planes, + compute_jacobian, + ), + ) + except Exception as exc: + raise RuntimeError( + "the xFit fused cuda.tile evaluation kernel failed to " + "compile or launch; check that cuda-tile matches the CUDA " + "runtime" + ) from exc def residual( self, parameters: BackendArray, *, indices: BackendArray From b8c1794b815bd5660abf6f87a59e934c564c68f8 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Fri, 9 Oct 2026 13:49:52 -0700 Subject: [PATCH 3/7] 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) From 1e30b19e15691614b0027990e636877e18176c3e Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Mon, 19 Oct 2026 09:29:33 -0700 Subject: [PATCH 4/7] Stabilize FP32 normal equations and damped solves Accumulate built-in FP32 normal equations with FP64 intermediates, then round them back before convergence checks. Solve FP32 damped systems in FP64 and return FP32 steps on both backends. Preserve model, state, covariance and rank precision, and retain GPU subnormals across workspace conversions. Cover the public variance discrepancy with nearby noise, masked inputs and conversion boundaries. Signed-off-by: Trent Nelson --- src/cuphoton/xfit/solver.py | 87 +++++++++++++++-- tests/xfit/test_models.py | 189 +++++++++++++++++++++++++++++++++++- 2 files changed, 264 insertions(+), 12 deletions(-) diff --git a/src/cuphoton/xfit/solver.py b/src/cuphoton/xfit/solver.py index c73c4e81..53a1d22f 100644 --- a/src/cuphoton/xfit/solver.py +++ b/src/cuphoton/xfit/solver.py @@ -8,6 +8,7 @@ from dataclasses import dataclass from enum import IntEnum +from functools import lru_cache from types import ModuleType from typing import Any, Protocol @@ -378,6 +379,55 @@ def _factorize_jacobians( return left_vectors, singular_values +@lru_cache(maxsize=2) +def _fp32_workspace_cast_kernel(ap: ModuleType, to_float32: bool) -> Any: + # CuPy appends -ftz=true even after a caller requests --ftz=false. + # Explicit PTX without .ftz preserves subnormals in both directions. + if to_float32: + return ap.ElementwiseKernel( + "float64 x", + "float32 y", + 'asm("cvt.rn.f32.f64 %0, %1;" : "=f"(y) : "d"(x));', + "xfit_workspace_float32", + ) + return ap.ElementwiseKernel( + "float32 x", + "float64 y", + 'asm("cvt.f64.f32 %0, %1;" : "=d"(y) : "f"(x));', + "xfit_workspace_float64", + ) + + +def _cast_fp32_workspace( + values: BackendArray, dtype: Any, ap: ModuleType +) -> BackendArray: + """Convert FP32 solver workspaces without flushing subnormal values.""" + + if ap is np: + return values.astype(dtype) + return _fp32_workspace_cast_kernel(ap, dtype == np.float32)(values) + + +def _normal_equations_from_jacobian( + jacobian: BackendArray, residual: BackendArray, ap: ModuleType +) -> tuple[BackendArray, BackendArray]: + """Accumulate FP32 products in FP64, then retain the input precision.""" + + widen = jacobian.dtype == np.float32 and residual.dtype == np.float32 + normal_jacobian = ( + _cast_fp32_workspace(jacobian, np.float64, ap) if widen else jacobian + ) + normal_residual = ( + _cast_fp32_workspace(residual, np.float64, ap) if widen else residual + ) + gradient = ap.einsum("knm,km->kn", normal_jacobian, normal_residual) + hessian = ap.einsum("knm,kpm->knp", normal_jacobian, normal_jacobian) + if widen: + gradient = _cast_fp32_workspace(gradient, np.float32, ap) + hessian = _cast_fp32_workspace(hessian, np.float32, ap) + return gradient, hessian + + def batched_levenberg_marquardt( problem: BatchedLeastSquaresProblem, x0: BackendArray, @@ -390,6 +440,12 @@ def batched_levenberg_marquardt( interpreted as parameter-major, with shape ``(K, N, M)``. When it is absent, or finite differences are requested, the solver allocates the perturbation workspace in the same backend as ``x0``. + + Built-in FP32 normal equations accumulate in FP64 and are rounded back + before convergence checks. FP32 damped systems likewise use FP64 solve + intermediates and return FP32 steps. Model evaluations, solver state, + covariance and rank thresholds retain their input precision. A supplied + normal-equations callback controls its own accumulation precision. """ if not hasattr(x0, "ndim") or x0.ndim != 2: @@ -512,13 +568,13 @@ def finite_rows(value: BackendArray) -> BackendArray: fallback_jacobian = fallback_jacobian[valid_fallback] gradient = gradient.copy() hessian = hessian.copy() - gradient[recovered_positions] = ap.einsum( - "knm,km->kn", + ( + gradient[recovered_positions], + hessian[recovered_positions], + ) = _normal_equations_from_jacobian( fallback_jacobian, residual_active[recovered_positions], - ) - hessian[recovered_positions] = ap.einsum( - "knm,kpm->knp", fallback_jacobian, fallback_jacobian + ap, ) valid_normal_equations[recovered_positions] = True jacobians[recovered_indices] = fallback_jacobian @@ -566,8 +622,9 @@ def finite_rows(value: BackendArray) -> BackendArray: jacobian = jacobian[valid_jacobian] jacobians[active_indices] = jacobian jacobian_current[active_indices] = True - gradient = ap.einsum("knm,km->kn", jacobian, residual_active) - hessian = ap.einsum("knm,kpm->knp", jacobian, jacobian) + gradient, hessian = _normal_equations_from_jacobian( + jacobian, residual_active, ap + ) diagonal = ap.maximum(ap.diagonal(hessian, axis1=1, axis2=2), eps) scaled_gradient = ap.max(ap.abs(gradient) / ap.sqrt(diagonal), axis=1) residual_norm = ap.linalg.norm(residual_active, axis=1) @@ -592,7 +649,21 @@ def finite_rows(value: BackendArray) -> BackendArray: * diagonal[:, None, :] ) try: - step = ap.linalg.solve(system, -gradient[..., None])[..., 0] + if system.dtype == np.float32 and gradient.dtype == np.float32: + # NumPy already solves FP32 inputs with FP64 intermediates. + # Match that policy on the GPU before rounding the step. + step = _cast_fp32_workspace( + ap.linalg.solve( + _cast_fp32_workspace(system, np.float64, ap), + -_cast_fp32_workspace( + gradient[..., None], np.float64, ap + ), + )[..., 0], + np.float32, + ap, + ) + else: + step = ap.linalg.solve(system, -gradient[..., None])[..., 0] except Exception: status[active_indices] = int(LMStatus.SINGULAR) continue diff --git a/tests/xfit/test_models.py b/tests/xfit/test_models.py index 3c5c7b59..3fb8569a 100644 --- a/tests/xfit/test_models.py +++ b/tests/xfit/test_models.py @@ -645,16 +645,24 @@ def test_vignetted_bilinear_preserves_visible_requested_flux() -> None: assert vignetted.to_backend("numpy").evaluation == "bilinear-vignetted" -def test_rank_deficient_fit_reports_invalid_public_uncertainty() -> None: - model = GaussianDipoleModel((7, 9), dtype=np.float64) - initial = np.asarray([[0.0, 1.0, 1.0, 0.0, -1.0, 0.0, 1.0, 0.0]]) +@pytest.mark.parametrize("backend", ["numpy", "cupy"]) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_rank_deficient_fit_reports_invalid_public_uncertainty( + backend, dtype +) -> None: + if backend == "cupy": + _require_cupy_device() + model = GaussianDipoleModel((7, 9), dtype=dtype) + initial = np.asarray( + [[0.0, 1.0, 1.0, 0.0, -1.0, 0.0, 1.0, 0.0]], dtype=dtype + ) images = model.evaluate(initial) result = fit_dipoles( images, model=model, initial=initial, - backend="numpy", + backend=backend, ) assert result.converged.all() @@ -662,6 +670,7 @@ def test_rank_deficient_fit_reports_invalid_public_uncertainty() -> None: assert result.uncertainty_reason == ("rank-deficient Jacobian",) assert np.isnan(result.covariance).all() assert np.isnan(result.standard_errors).all() + assert result.covariance.dtype == dtype def test_stamp_multi_mode_basis_requires_explicit_weights() -> None: @@ -1728,3 +1737,175 @@ def worker(index): futures = [executor.submit(worker, index) for index in range(2)] for future in futures: future.result(timeout=180) + + +@pytest.mark.parametrize( + "variation", ["original", "nearby-noise", "masked-nonfinite"] +) +def test_fp32_public_covariance_matches_with_variance( + variation, +) -> None: + _require_cupy_device() + truth = np.asarray( + [ + [4.0, 1.2, 1.6, 0.13, -1.0, 0.2, 1.0, -0.2], + [3.1, 1.4, 1.1, -0.17, -0.8, -0.3, 0.9, 0.35], + ], + dtype=np.float32, + ) + model = GaussianDipoleModel((7, 9), dtype=np.float32) + images = model.evaluate(truth) + images += ( + np.random.default_rng(20260929) + .normal(0, 0.004, images.shape) + .astype(np.float32) + ) + if variation == "nearby-noise": + images += ( + np.random.default_rng(20261013) + .normal(0, 1e-4, images.shape) + .astype(np.float32) + ) + weights = np.linspace(0.7, 1.6, images.size, dtype=np.float32).reshape( + images.shape + ) + weights.reshape(2, -1)[:, ::11] = 0 + mask = weights != 0 + images[~mask] = np.nan + variance = 1 / np.where(mask, weights, 1) ** 2 + if variation == "masked-nonfinite": + variance[~mask] = np.inf + initial = truth + np.asarray( + [0.1, 0.04, -0.04, 0.015, 0.05, -0.04, -0.05, 0.04], + dtype=np.float32, + ) + originals = [value.copy() for value in (images, initial, variance, mask)] + results = [] + for backend in ("numpy", "cupy"): + results.append( + fit_dipoles( + images, + model=model, + initial=initial, + mask=mask, + variance=variance, + config=LMConfig(max_evaluations=80), + backend=backend, + ) + ) + with pytest.raises(ValueError, match="finite on included pixels"): + fit_dipoles(images, model=model, initial=initial, backend=backend) + cpu, gpu = results + assert cpu.uncertainty_valid.all() and gpu.uncertainty_valid.all() + assert gpu.parameters.dtype == gpu.covariance.dtype == np.float32 + for field, expected in vars(cpu).items(): + if field in {"backend", "device"}: + continue + actual = getattr(gpu, field) + if isinstance(expected, np.ndarray): + assert actual.dtype == expected.dtype + if isinstance(expected, np.ndarray) and expected.dtype.kind == "f": + np.testing.assert_allclose( + actual, expected, rtol=3e-5, atol=3e-5, equal_nan=True + ) + else: + np.testing.assert_array_equal(actual, expected) + for actual, expected in zip( + (images, initial, variance, mask), originals, strict=True + ): + assert actual.tobytes() == expected.tobytes() + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_fp32_workspace_cast_preserves_subnormals_on_device(dtype) -> None: + cp = _require_cupy_device() + values = np.asarray( + [ + 0, + 0x80000000, + 1, + 0x80000001, + 16384, + 0x80004000, + 0x7FFFFF, + 0x807FFFFF, + 0x800000, + 0x80800000, + 0x3F800000, + 0xBF800000, + 0x7F7FFFFF, + 0xFF7FFFFF, + 0x7F800000, + 0xFF800000, + 0x7FC00000, + ], + dtype=np.uint32, + ).view(np.float32) + if dtype == np.float32: + tiny = float(np.finfo(np.float32).smallest_subnormal) + values = np.concatenate( + ( + values.astype(np.float64), + np.asarray( + [ + tiny * 0.5, + tiny * 1.5, + tiny * 2.5, + -tiny * 0.5, + -tiny * 1.5, + -tiny * 2.5, + ] + ), + ) + ) + expected = values.astype(dtype) + stream = cp.cuda.Stream(non_blocking=True) + with stream: + device = cp.asarray(np.repeat(values, 2))[::2] + actual_device = xfit_solver._cast_fp32_workspace(device, dtype, cp) + stream.synchronize() + actual = cp.asnumpy(actual_device) + assert not device.flags.c_contiguous + 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)) + finite = np.isfinite(expected) + bits = np.uint32 if dtype == np.float32 else np.uint64 + np.testing.assert_array_equal( + actual[finite].view(bits), expected[finite].view(bits) + ) + assert cp.asnumpy(device).tobytes() == values.tobytes() + + +@pytest.mark.parametrize( + ("jacobian_scale", "residual_scale"), + [(2.0**-65, 2.0**-70), (2.0**-135, 2.0**100)], +) +def test_fp32_normal_equations_preserve_subnormal_contributions( + jacobian_scale, residual_scale +) -> None: + cp = _require_cupy_device() + jacobian = np.full((2, 2, 16), jacobian_scale, dtype=np.float32) + residual = np.full((2, 16), residual_scale, dtype=np.float32) + expected_gradient = np.full( + (2, 2), 16 * jacobian_scale * residual_scale, dtype=np.float32 + ) + expected_hessian = np.full( + (2, 2, 2), 16 * jacobian_scale**2, dtype=np.float32 + ) + for ap in (np, cp): + gradient, hessian = xfit_solver._normal_equations_from_jacobian( + ap.asarray(jacobian), ap.asarray(residual), ap + ) + if ap is cp: + gradient, hessian = cp.asnumpy(gradient), cp.asnumpy(hessian) + for actual, expected in ( + (gradient, expected_gradient), + (hessian, expected_hessian), + ): + assert actual.dtype == np.float32 + np.testing.assert_array_equal( + actual.view(np.uint32), expected.view(np.uint32) + ) From 56693d6865244ed1adeb7b76d5c298b82f893ddf Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Mon, 19 Oct 2026 09:29:57 -0700 Subject: [PATCH 5/7] Require predicted progress to satisfy xFit function tolerance A small accepted improvement can conceal a large predicted reduction. Require both reductions to be small and the local model not to underpredict by more than a factor of two. Retain regressions for premature and legitimate function stops. Signed-off-by: Trent Nelson --- src/cuphoton/xfit/solver.py | 22 +++++++++++++-- tests/xfit/test_solver.py | 54 +++++++++++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+), 2 deletions(-) diff --git a/src/cuphoton/xfit/solver.py b/src/cuphoton/xfit/solver.py index 53a1d22f..13c2fa75 100644 --- a/src/cuphoton/xfit/solver.py +++ b/src/cuphoton/xfit/solver.py @@ -667,6 +667,14 @@ def finite_rows(value: BackendArray) -> BackendArray: except Exception: status[active_indices] = int(LMStatus.SINGULAR) continue + # For (H + damping * D) step = -gradient, the quadratic model + # predicts -g.T step - step.T H step / 2. Use the equivalent + # O(n) expression with the damping that produced this step. + predicted_reduction = 0.5 * ap.einsum( + "kn,kn->k", + step, + damping[active_indices, None] * diagonal * step - gradient, + ) valid_step = finite_rows(step) status[active_indices[~valid_step]] = int(LMStatus.SINGULAR) active_indices = active_indices[valid_step] @@ -675,6 +683,7 @@ def finite_rows(value: BackendArray) -> BackendArray: x_active = x_active[valid_step] residual_active = residual_active[valid_step] step = step[valid_step] + predicted_reduction = predicted_reduction[valid_step] can_evaluate = evaluations[active_indices] < settings.max_evaluations status[active_indices[~can_evaluate]] = int(LMStatus.MAX_EVALUATIONS) @@ -684,6 +693,7 @@ def finite_rows(value: BackendArray) -> BackendArray: x_active = x_active[can_evaluate] residual_active = residual_active[can_evaluate] step = step[can_evaluate] + predicted_reduction = predicted_reduction[can_evaluate] trial_x = x_active + step trial_residual = _call_residual(problem, trial_x, active_indices) @@ -709,8 +719,16 @@ def finite_rows(value: BackendArray) -> BackendArray: damping[accepted_indices] * settings.damping_decrease, ) rejected_steps[accepted_indices] = 0 - f_converged = reduction[accepted] <= ( - settings.f_tol * ap.maximum(1.0, cost[accepted]) + # As in MINPACK, both actual and predicted progress must be + # small, without the model underpredicting by more than 2x. + prediction = predicted_reduction[accepted] + f_threshold = settings.f_tol * ap.maximum(1.0, cost[accepted]) + f_converged = ( + (reduction[accepted] <= f_threshold) + & ap.isfinite(prediction) + & (prediction > 0) + & (prediction <= f_threshold) + & (0.5 * reduction[accepted] <= prediction) ) x_converged = ap.linalg.norm(accepted_step, axis=1) <= ( settings.x_tol * (settings.x_tol + 1.0) diff --git a/tests/xfit/test_solver.py b/tests/xfit/test_solver.py index 079bf479..01ec7435 100644 --- a/tests/xfit/test_solver.py +++ b/tests/xfit/test_solver.py @@ -71,6 +71,60 @@ def jacobian(x, *, indices, out=None): assert np.isfinite(result.covariance).all() +def test_f_tol_continues_after_poorly_predicted_nonlinear_progress() -> None: + coefficient = 1.001 - 1.0e-8 + + def residual(x, *, indices): + del indices + return 1.0 - x + coefficient * x * x + + def jacobian(x, *, indices): + del indices + return (-1.0 + 2.0 * coefficient * x)[:, :, None] + + # The first damped step is almost one: actual cost improvement is about + # 1e-8, although the linear model predicts about 0.5. Small actual + # progress alone would stop far from this positive parabola's minimum. + problem = BatchedLeastSquaresProblem(residual, jacobian) + initial = np.zeros((1, 1)) + limited = batched_levenberg_marquardt( + problem, initial, config=LMConfig(max_evaluations=2) + ) + assert limited.status.tolist() == [LMStatus.MAX_EVALUATIONS] + assert limited.evaluations.tolist() == [2] + np.testing.assert_allclose(limited.parameters, 1.0 / 1.001) + + result = batched_levenberg_marquardt(problem, initial) + + assert result.converged.all() + assert result.evaluations[0] > 2 + np.testing.assert_allclose( + result.parameters, 1.0 / (2.0 * coefficient), atol=5.0e-5, rtol=0 + ) + np.testing.assert_allclose( + result.residuals, 1.0 - 1.0 / (4.0 * coefficient), atol=1.0e-8 + ) + + +def test_f_tol_still_accepts_small_well_predicted_progress() -> None: + def residual(x, *, indices): + del indices + return np.concatenate((x, np.ones_like(x)), axis=1) + + def jacobian(x, *, indices): + del indices + return np.broadcast_to(np.asarray([[[1.0, 0.0]]]), (len(x), 1, 2)) + + result = batched_levenberg_marquardt( + BatchedLeastSquaresProblem(residual, jacobian), + np.asarray([[1.0e-5]]), + ) + + assert result.status.tolist() == [LMStatus.CONVERGED_F_TOL] + assert result.evaluations.tolist() == [2] + np.testing.assert_allclose(result.parameters, 1.0e-5 / 1001.0) + + def test_active_batch_is_compacted_after_independent_termination() -> None: target = np.asarray([[0.0], [2.0], [-3.0]]) seen_indices: list[tuple[int, ...]] = [] From d5a3e04c17d68e451ade708bc185822acb6a2f41 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Mon, 19 Oct 2026 09:30:22 -0700 Subject: [PATCH 6/7] Improve xFit covariance accuracy without changing rank Equilibrate the full-rank covariance factorization to reduce error when parameter scales differ substantially. Preserve the original unscaled rank decision and isolate covariance failures by row. Cover the analytic ill-scaled case, unchanged rank eligibility and row-level failure behavior. Signed-off-by: Trent Nelson --- src/cuphoton/xfit/solver.py | 59 +++++++++++++++++++++---------- tests/xfit/test_solver.py | 70 ++++++++++++++++++++++++++++++++++--- 2 files changed, 105 insertions(+), 24 deletions(-) diff --git a/src/cuphoton/xfit/solver.py b/src/cuphoton/xfit/solver.py index 13c2fa75..e08ab527 100644 --- a/src/cuphoton/xfit/solver.py +++ b/src/cuphoton/xfit/solver.py @@ -357,8 +357,8 @@ def _final_jacobian( def _factorize_jacobians( jacobians: BackendArray, ap: ModuleType -) -> tuple[BackendArray, BackendArray]: - """Return left singular vectors and values for covariance diagnostics.""" +) -> tuple[BackendArray, BackendArray, BackendArray]: + """Return unscaled rank factors and the matrix used to compute them.""" batch, parameters, observations = jacobians.shape module_name = type(jacobians).__module__.split(".", maxsplit=1)[0] @@ -372,11 +372,26 @@ def _factorize_jacobians( left_vectors, singular_values, _ = ap.linalg.svd( reduced, full_matrices=False ) - return left_vectors, singular_values + return left_vectors, singular_values, reduced left_vectors, singular_values, _ = ap.linalg.svd( jacobians, full_matrices=False ) - return left_vectors, singular_values + return left_vectors, singular_values, jacobians + + +def _covariance_from_factor( + factor: BackendArray, ap: ModuleType +) -> BackendArray: + """Equilibrate a full-rank parameter-major factor for covariance only.""" + + scale = ap.max(ap.abs(factor), axis=-1) + left_vectors, singular_values, _ = ap.linalg.svd( + factor / scale[..., None], full_matrices=False + ) + inverse_factor = (left_vectors / singular_values[..., None, :]) / scale[ + ..., None + ] + return inverse_factor @ inverse_factor.swapaxes(-2, -1) @lru_cache(maxsize=2) @@ -908,7 +923,7 @@ def finite_rows(value: BackendArray) -> BackendArray: hessians[diagnostic_indices] = ap.einsum( "knm,kpm->knp", diagnostic_jacobians, diagnostic_jacobians ) - left_vectors, singular_values = _factorize_jacobians( + _, singular_values, covariance_factor = _factorize_jacobians( diagnostic_jacobians, ap ) tolerance = max(n, m) * eps * ap.max(singular_values, axis=1) @@ -917,14 +932,23 @@ def finite_rows(value: BackendArray) -> BackendArray: full_rank = rank_values == n full_rank_indices = diagnostic_indices[full_rank] if full_rank_indices.shape[0] > 0: - full_rank_vectors = left_vectors[full_rank] - full_rank_singular_values = singular_values[full_rank] - inverse_squared = 1.0 / ( - full_rank_singular_values * full_rank_singular_values - ) - covariance[full_rank_indices] = ( - full_rank_vectors * inverse_squared[:, None, :] - ) @ full_rank_vectors.transpose(0, 2, 1) + full_rank_factor = covariance_factor[full_rank] + try: + covariance[full_rank_indices] = _covariance_from_factor( + full_rank_factor, ap + ) + except Exception: + # A covariance-only failure must not replace successful + # unscaled rank decisions with a different factorization. + for index in range(full_rank_indices.shape[0]): + try: + covariance[full_rank_indices[index]] = ( + _covariance_from_factor( + full_rank_factor[index], ap + ) + ) + except Exception: + continue except Exception: # Preserve row-level failure isolation if a backend cannot perform # one of the batched linear-algebra operations. @@ -933,18 +957,15 @@ def finite_rows(value: BackendArray) -> BackendArray: jacobian = jacobians[row] gradients[row] = jacobian @ residuals[row] hessians[row] = jacobian @ jacobian.T - left_vectors, singular_values, _ = ap.linalg.svd( + _, singular_values, _ = ap.linalg.svd( jacobian, full_matrices=False ) tolerance = max(n, m) * eps * ap.max(singular_values) rank[row] = ap.sum(singular_values > tolerance) if _int(rank[row]) == n: - inverse_squared = 1.0 / ( - singular_values * singular_values + covariance[row] = _covariance_from_factor( + jacobian, ap ) - covariance[row] = ( - left_vectors * inverse_squared[None, :] - ) @ left_vectors.T except Exception: continue diff --git a/tests/xfit/test_solver.py b/tests/xfit/test_solver.py index 01ec7435..2b6abd25 100644 --- a/tests/xfit/test_solver.py +++ b/tests/xfit/test_solver.py @@ -570,7 +570,48 @@ def jacobian(x, *, indices): ) -def test_final_diagnostics_factorize_batch_once(monkeypatch) -> None: +def test_ill_scaled_covariance_preserves_unscaled_rank_policy() -> None: + orthogonal = ( + np.asarray( + [[1, 1, 1, 1], [1, -1, 1, -1], [1, 1, -1, -1], [1, -1, -1, 1]], + dtype=np.float64, + ) + / 2 + ) + singular_values = np.asarray([1, 0.5, 0.25, 2.0**-8]) + scales = np.asarray([2.0**-24, 1, 32, 2.0**-8]) + factor = scales[:, None] * orthogonal * singular_values[None, :] + matrix = np.tile(factor, (1, 3)) + matrices = np.stack((matrix, matrix.copy())) + # The second row remains numerically rank deficient under the existing + # unscaled threshold, even though equilibration can recover its rank. + matrices[1, 0] *= 2.0**-36 + expected = ( + ((orthogonal / singular_values**2) @ orthogonal.T) + / scales[:, None] + / scales[None, :] + / 3 + ) + + def residual(x, *, indices): + return np.einsum("kn,knm->km", x, matrices[indices]) + + def jacobian(x, *, indices): + del x + return matrices[indices] + + result = batched_levenberg_marquardt( + BatchedLeastSquaresProblem(residual, jacobian), np.zeros((2, 4)) + ) + + assert result.rank.tolist() == [4, 3] + np.testing.assert_allclose( + result.covariance[0], expected, rtol=2.0e-8, atol=2.0e-8 + ) + assert np.isnan(result.covariance[1]).all() + + +def test_final_diagnostics_factorize_each_stage_as_batch(monkeypatch) -> None: target = np.asarray( [[1.0, -2.0], [-3.0, 4.0], [5.0, -6.0]], dtype=np.float64 ) @@ -601,7 +642,7 @@ def recording_svd(value, *args, **kwargs): target.copy(), ) - assert seen_shapes == [(3, 2, 3)] + assert seen_shapes == [(3, 2, 3), (3, 2, 3)] assert result.converged.all() assert result.rank.tolist() == [2, 2, 2] assert np.isfinite(result.covariance).all() @@ -637,8 +678,10 @@ def jacobian(x, *, indices): assert np.isnan(result.gradient[1]).all() +@pytest.mark.parametrize("failing_stage", ["rank", "covariance"]) def test_final_diagnostics_fall_back_to_rows_when_batched_svd_fails( monkeypatch, + failing_stage, ) -> None: target = np.asarray( [[1.0, -2.0], [-3.0, 4.0], [5.0, -6.0]], dtype=np.float64 @@ -665,23 +708,40 @@ def jacobian(x, *, indices): def row_only_svd(value, *args, **kwargs): seen_shapes.append(value.shape) - if value.ndim == 3: + if value.ndim == 3 and ( + failing_stage == "rank" or len(seen_shapes) == 2 + ): raise RuntimeError("batched svd is unavailable") + if failing_stage == "covariance" and len(seen_shapes) == 4: + raise RuntimeError("row 1 covariance factorization failed") return original_svd(value, *args, **kwargs) monkeypatch.setattr(np.linalg, "svd", row_only_svd) result = batched_levenberg_marquardt(problem, target.copy()) - assert seen_shapes == [(3, 2, 3), (2, 3), (2, 3), (2, 3)] + if failing_stage == "rank": + assert seen_shapes == [(3, 2, 3), *[(2, 3)] * 6] + else: + # The successful batched rank factorization is not repeated per row. + assert seen_shapes == [(3, 2, 3), (3, 2, 3), *[(2, 3)] * 3] assert result.rank.tolist() == [2, 2, 2] assert np.array_equal(result.status, expected.status) - for field in ("gradient", "gn_hessian", "covariance"): + for field in ("gradient", "gn_hessian"): assert np.allclose( getattr(result, field), getattr(expected, field), rtol=1.0e-12, atol=1.0e-12, ) + if failing_stage == "covariance": + expected.covariance[1] = np.nan + np.testing.assert_allclose( + result.covariance, + expected.covariance, + rtol=1.0e-12, + atol=1.0e-12, + equal_nan=True, + ) @pytest.mark.parametrize("dtype", [np.float32, np.float64]) From e64418a75d80208011177ed1b420fada60937d95 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Mon, 19 Oct 2026 09:31:07 -0700 Subject: [PATCH 7/7] Exercise FP32 xFit solver regressions in GPU CI Run the variance comparison and subnormal workspace cases in the existing GPU parity job, with skips treated as failures. Signed-off-by: Trent Nelson --- .github/workflows/ci.yml | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index be718256..f40bd2a4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -157,6 +157,9 @@ 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_models.py::test_fp32_public_covariance_matches_with_variance + tests/xfit/test_models.py::test_fp32_workspace_cast_preserves_subnormals_on_device + tests/xfit/test_models.py::test_fp32_normal_equations_preserve_subnormal_contributions - name: Require all selected tests to execute run: | uv run --no-sync python - <<'PY' @@ -166,11 +169,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) == 13, 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("All 13 selected tests passed without skips") PY - name: Run native xFit GPU tests run: >-