From 37d1c68f77d18367a3561108f84838a96c2805fe Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 29 Sep 2026 22:13:44 -0700 Subject: [PATCH 1/8] 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 | 202 +++++- 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, 1913 insertions(+), 44 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 c36f4e9d..2fbbc824 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -202,6 +202,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 @@ -210,6 +270,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 3ab26e61..29740bcc 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 d9fb8bfc..260837f1 100644 --- a/src/cuphoton/xfit/api.py +++ b/src/cuphoton/xfit/api.py @@ -151,7 +151,7 @@ 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", @@ -159,8 +159,12 @@ class DeviceDipoleFitResult: covariance_factor: Any | None = None 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) ): @@ -647,6 +651,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``.""" @@ -779,9 +784,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) @@ -792,6 +810,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( @@ -805,7 +825,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( @@ -1068,6 +1088,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, *, @@ -1078,6 +1117,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. @@ -1099,9 +1139,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, @@ -1112,6 +1168,7 @@ def fit_dipoles( mode=mode, resolved=resolved, config=config, + fusion=fusion, ) return _materialize_dipole_fit_result(backend_result) @@ -1125,11 +1182,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 @@ -1148,8 +1208,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 @@ -1184,6 +1256,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 9de2a43e..c870be7f 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 = ( @@ -338,6 +345,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 @@ -370,6 +388,7 @@ def run(self) -> None: "g_tol", "max_evaluations", "use_finite_difference", + "fusion", "fits_reader", ) } @@ -447,6 +466,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)) @@ -459,6 +479,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 87d710ab..9441b358 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, @@ -64,6 +69,7 @@ "g_tol": None, "max_evaluations": None, "use_finite_difference": False, + "fusion": False, "fits_reader": "auto", "xdr_options": None, } @@ -103,6 +109,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"}: @@ -283,6 +295,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) @@ -445,6 +458,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 7b9f61cd..c8f36c9a 100644 --- a/src/cuphoton/xscan/device_pipeline.py +++ b/src/cuphoton/xscan/device_pipeline.py @@ -651,17 +651,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" @@ -716,6 +727,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]: @@ -724,6 +737,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]: @@ -733,6 +748,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 @@ -741,9 +757,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, @@ -1080,9 +1097,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: @@ -2159,6 +2177,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.""" @@ -2184,7 +2203,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", @@ -3353,6 +3373,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 4596ede0..04d9e1d0 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), @@ -82,7 +82,7 @@ def test_public_command_surface_counts_are_exact() -> None: assert sum(item[1] for item in per_group) == 99 assert sum(item[1] + item[4] for item in per_group) == 105 assert sum(item[2] for item in per_group) == 92 - assert sum(item[3] for item in per_group) == 985 + assert sum(item[3] for item in per_group) == 986 def test_public_registry_order_and_component_derivations() -> None: diff --git a/tests/test_public_api.py b/tests/test_public_api.py index b1b073c0..b3b0405b 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -213,6 +213,7 @@ def test_xfit_curated_exports_and_fit_signature() -> None: "mode", "backend", "config", + "fusion", ] assert list(inspect.signature(xfit.fit_dipoles_device).parameters) == [ "images", @@ -223,6 +224,7 @@ def test_xfit_curated_exports_and_fit_signature() -> None: "mode", "config", "backend", + "fusion", ] assert list(inspect.signature(xfit.fit_single_sources).parameters) == [ "images", 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 a76fefa5..4ea6ff84 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 141ece4a..647e2cb2 100644 --- a/tests/xfit/test_models.py +++ b/tests/xfit/test_models.py @@ -929,19 +929,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"), - ("result_location", "host"), - ): + for name in ("schema", "solver", "result_location"): + # Python 3.14 changed replace()'s init=False error to TypeError. with pytest.raises( - (TypeError, ValueError), match="declared with init=False" + (ValueError, TypeError), match=rf"field {name}.*init=False" ): - replace(device, **{name: value}) - with pytest.raises( - ValueError, match="device fit backend must be cupy or native" - ): - replace(device, backend="numpy") # type: ignore[arg-type] + 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): @@ -1111,6 +1106,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: @@ -1236,6 +1265,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] @@ -1813,3 +1876,120 @@ def normal_equations(x, residuals, *, indices): 1 / (2 * float(coefficient) ** 2), rtol=1e-6, ) + + +@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 6ab4fa83..6dd0835e 100644 --- a/tests/xscan/test_device_pipeline.py +++ b/tests/xscan/test_device_pipeline.py @@ -873,6 +873,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: @@ -1507,7 +1521,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) @@ -1892,6 +1908,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 310cd4d1041eb8a5c54cf1e4d94876b784c334af Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Fri, 9 Oct 2026 12:13:46 -0700 Subject: [PATCH 2/8] 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 4f2f55f9a750b85c715a2d822408e057d56f6784 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Fri, 9 Oct 2026 13:49:52 -0700 Subject: [PATCH 3/8] 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 4580ed9f163b91916fe3f0f597c589973b8b47c6 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Sat, 24 Oct 2026 10:20:02 -0700 Subject: [PATCH 4/8] Check pipeline defaults in GPU CI Match the device pipeline fixture to the default keyword omission. Run the device-resident pipeline regression in the GPU lane so stream and handoff expectations are exercised. Signed-off-by: Trent Nelson --- .github/workflows/ci.yml | 5 +++-- tests/xscan/test_device_pipeline.py | 2 -- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2fbbc824..be990e70 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -149,7 +149,7 @@ jobs: assert float((tensor * tensor).sum().item()) == 1240.0 print("CuPy and PyTorch CUDA kernels and DLPack passed") PY - - name: Run xFit GPU parity tests + - name: Run xFit GPU parity and device pipeline tests run: >- uv run --no-sync pytest -q -rs --junitxml="${RUNNER_TEMP}/xfit-gpu.xml" @@ -163,6 +163,7 @@ jobs: 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 tests/xfit/test_models.py::test_fp32_solver_preserves_subnormal_curvature + tests/xscan/test_device_pipeline.py::test_persistent_context_pipeline_is_device_resident_and_deterministic - name: Require all selected tests to execute run: | uv run --no-sync python - <<'PY' @@ -172,7 +173,7 @@ jobs: report = ET.parse(Path(os.environ["RUNNER_TEMP"]) / "xfit-gpu.xml") cases = report.findall(".//testcase") - assert len(cases) == 18, len(cases) + assert len(cases) == 19, len(cases) assert not report.findall(".//skipped"), "GPU tests must not skip" assert not report.findall(".//failure") assert not report.findall(".//error") diff --git a/tests/xscan/test_device_pipeline.py b/tests/xscan/test_device_pipeline.py index 6dd0835e..f0690cb6 100644 --- a/tests/xscan/test_device_pipeline.py +++ b/tests/xscan/test_device_pipeline.py @@ -1521,9 +1521,7 @@ 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) From e2928520f66251f1848aec34e459e416fb4fcfba Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Sat, 24 Oct 2026 11:33:00 -0700 Subject: [PATCH 5/8] Exercise real cuTile classifier pipeline in GPU CI Run subtraction, fitting, feature extraction and classification for both cuTile modes against the CuPy pipeline. Preserve repeated-call, transfer and scientific-result checks at their existing tolerances. Use the locked CUDA Torch runtime with the separate cuTile compiler and require both pipeline cases to execute without skips. Signed-off-by: Trent Nelson --- .github/workflows/ci.yml | 56 ++++++++++ tests/xscan/test_cutile_pipeline.py | 165 ++++++++++++++++++++++++++++ 2 files changed, 221 insertions(+) create mode 100644 tests/xscan/test_cutile_pipeline.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index be990e70..211c9aac 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -263,6 +263,59 @@ jobs: 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: Run real cuTile classifier pipeline + run: | + # Keep CUDA Torch's runtime pins; reuse the separate compiler tools. + uv pip install --python .venv/bin/python --no-deps 'cuda-tile==1.6.0' + CUTILE_COMPILER_ROOT="$( + "${RUNNER_TEMP}/cuphoton-cutile/bin/python" -c \ + 'import sysconfig; print(sysconfig.get_path("purelib") + "/nvidia/cu13")' + )" + export CUTILE_COMPILER_ROOT + export CUDA_HOME="$CUTILE_COMPILER_ROOT" + export PATH="$CUTILE_COMPILER_ROOT/bin:$PATH" + uv pip freeze --python .venv/bin/python \ + > "${RUNNER_TEMP}/cutile-pipeline-packages.txt" + .venv/bin/python - <<'PY' + import os + from pathlib import Path + from cuda.tile._compile import _find_compiler_bin + + compiler = _find_compiler_bin() + expected = Path(os.environ["CUTILE_COMPILER_ROOT"]) / "bin/tileiras" + assert Path(compiler.path).resolve() == expected.resolve() + print(f"Pipeline tileiras: {compiler.path}") + PY + pipeline_status=0 + .venv/bin/python -m pytest -q -s -rs \ + --junitxml="${RUNNER_TEMP}/cutile-pipeline-gpu.xml" \ + tests/xscan/test_cutile_pipeline.py \ + > "${RUNNER_TEMP}/cutile-pipeline-gpu.log" 2>&1 || pipeline_status=$? + cat "${RUNNER_TEMP}/cutile-pipeline-gpu.log" + CUTILE_PIPELINE_STATUS="$pipeline_status" .venv/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-pipeline-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 pipeline emitted CUDA or interpreter teardown errors") + if int(os.environ["CUTILE_PIPELINE_STATUS"]) != 0: + raise SystemExit("cuTile pipeline pytest process failed") + report = ET.parse(root / "cutile-pipeline-gpu.xml") + cases = report.findall(".//testcase") + if len(cases) != 2 or any(report.findall(".//" + tag) for tag in ("skipped", "failure", "error")): + raise SystemExit("Require both real cuTile pipeline cases without skips, failures or errors") + print("Both real cuTile pipeline cases passed without skips or teardown errors") + PY - name: Upload test evidence if: always() uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 @@ -275,6 +328,9 @@ jobs: ${{ runner.temp }}/cutile-gpu.log ${{ runner.temp }}/cutile-packages.txt ${{ runner.temp }}/cutile-compiler.txt + ${{ runner.temp }}/cutile-pipeline-gpu.xml + ${{ runner.temp }}/cutile-pipeline-gpu.log + ${{ runner.temp }}/cutile-pipeline-packages.txt ${{ runner.temp }}/cuphoton-quickstart/**/summary.json if-no-files-found: error retention-days: 14 diff --git a/tests/xscan/test_cutile_pipeline.py b/tests/xscan/test_cutile_pipeline.py new file mode 100644 index 00000000..8aabe9f7 --- /dev/null +++ b/tests/xscan/test_cutile_pipeline.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import numpy as np +import pytest + +import cuphoton.xscan.device_pipeline as pipeline +from cuphoton.xscan.pipeline_benchmark import stages + + +@pytest.mark.parametrize("fusion", [False, True], ids=["unfused", "fused"]) +def test_cutile_pipeline_cli_runs_real_stages_and_matches_cupy( + tmp_path, capsys, fusion +): + import json + from dataclasses import replace + + from cuphoton.core.cli import run_component + from cuphoton.xscan.pipeline_benchmark.fixture import prepare_fixture + from cuphoton.xscan.pipeline_benchmark.runner import read_inputs + + cp = pytest.importorskip("cupy") + torch = pytest.importorskip("torch") + pytest.importorskip("cuda.tile") + try: + count = cp.cuda.runtime.getDeviceCount() + except cp.cuda.runtime.CUDARuntimeError as exc: + if exc.status != 100: + raise + pytest.skip("CUDA device is unavailable") + if count < 1 or not torch.cuda.is_available(): + pytest.skip("CUDA device is unavailable") + + config_path, items_path = prepare_fixture( + tmp_path / "inputs", + images=1, + image_size=48, + candidates=2, + stamp_size=9, + seed=2026, + device="cuda:0", + ) + config, items = read_inputs(config_path, items_path) + assert "backend" not in config.to_payload()["xfit"] + exact_fields = { + "xfit.status_codes", + "xfit.converged", + "xfit.evaluations", + "xfit.valid_pixel_count", + "xfit.valid_pixel_fraction", + "xfit.null_chi_square", + "xfit.degrees_of_freedom", + "xfit.uncertainty_valid", + "xfit.uncertainty_reason_codes", + } + reference = reference_payload = None + for backend in ("cupy", "cutile"): + selected = replace( + config, + xfit=replace( + config.xfit, + backend=backend, + fusion=fusion if backend == "cutile" else False, + ), + ) + selected_path = tmp_path / f"{backend}.json" + selected_path.write_text(json.dumps(selected.to_payload())) + output = tmp_path / backend + rc = run_component( + "xscan", + [ + "benchmark-pipeline", + "--stage", + "pipeline", + "--config", + str(selected_path), + "--items", + str(items_path), + "--output", + str(output), + "--warmup", + "1", + "--repeat", + "1", + ], + ) + captured = capsys.readouterr() + assert rc == 0, captured.err + summary = json.loads((output / "summary.json").read_text()) + assert [row["completed_items"] for row in summary["rounds"]] == [1, 1] + previous = None + for index in range(2): + result = json.loads( + (output / f"round-{index:03d}" / "item-0000.json").read_text() + ) + assert ( + result["configuration_sha256"] + == selected.configuration_sha256 + ) + evidence = result["scientific_evidence"] + assert evidence["xfit"]["backend"] == backend + arrays = pipeline.decode_device_pipeline_evidence( + evidence, config=selected + ) + assert set(arrays) == set(stages.SCIENCE_KEYS) + assert len(result["predictions"]) == len(items[0].candidates) == 2 + assert result["transfers"]["pipeline_full_array_d2h_bytes"] == 0 + assert result["transfers"]["pipeline_compact_h2d_bytes"] == 0 + assert result["transfers"]["terminal_d2h_calls"] == 1 + if previous is not None: + for name, values in arrays.items(): + np.testing.assert_array_equal(values, previous[name]) + previous = arrays + if reference is None: + reference, reference_payload = arrays, result + for key in ( + "item_id", + "input_sha256", + "checkpoint_sha256", + "feature_schema_sha256", + "xpois", + "transfers", + ): + assert result[key] == reference_payload[key] + for actual, expected in zip( + result["predictions"], + reference_payload["predictions"], + strict=True, + ): + assert actual.keys() == expected.keys() + for key in actual: + if key in {"logit", "probability"}: + np.testing.assert_allclose( + actual[key], expected[key], rtol=3e-5, atol=3e-5 + ) + else: + assert actual[key] == expected[key] + for name, actual in arrays.items(): + expected = reference[name] + assert ( + actual.shape == expected.shape + and actual.dtype == expected.dtype + ) + if name.startswith("xpois.") or name in exact_fields: + np.testing.assert_array_equal(actual, expected) + continue + for mask in (np.isnan, np.isposinf, np.isneginf): + np.testing.assert_array_equal( + mask(actual), mask(expected) + ) + tolerance = ( + 3e-5 + if name == "xfit.features" or name.startswith("xscan.") + else 5e-12 + ) + np.testing.assert_allclose( + actual, + expected, + rtol=tolerance, + atol=tolerance, + equal_nan=True, + ) From 6615de90a81aa098e224bf29f96287a89bfa8647 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Sat, 24 Oct 2026 13:43:03 -0700 Subject: [PATCH 6/8] Check cuTile factor ownership without fixing its basis Validate the covariance represented by each factor and keep the existing retained-buffer checks across device streams. Signed-off-by: Trent Nelson --- tests/xfit/test_models.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/tests/xfit/test_models.py b/tests/xfit/test_models.py index 647e2cb2..fb48fbee 100644 --- a/tests/xfit/test_models.py +++ b/tests/xfit/test_models.py @@ -1914,12 +1914,23 @@ def test_cutile_device_results_survive_worker_reuse(dtype, fusion) -> None: ) def snapshot(result): - return { + arrays = { field.name: cp.asnumpy(value) for field in fields(result) if isinstance(value := getattr(result, field.name), cp.ndarray) } + factor = arrays["covariance_factor"] + factor_tolerance = 3e-5 if dtype is np.float32 else 5e-12 + np.testing.assert_allclose( + factor @ factor.swapaxes(-2, -1), + arrays["covariance"], + rtol=factor_tolerance, + atol=factor_tolerance, + equal_nan=True, + ) + return arrays + def features(result): transformed = transform_xfit_result_features_device( result, image_shape=(11, 15), variance_present=True @@ -1958,6 +1969,9 @@ def worker(index): reference, reference_features = expected[case_index] assert arrays.keys() == reference.keys() for name, actual in arrays.items(): + # Factor columns have no unique cross-backend basis. + if name == "covariance_factor": + continue target = reference[name] if actual.dtype.kind in "fc": np.testing.assert_array_equal( From 4f16f6eb370500e10c0798f947be73a0697be355 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Sat, 24 Oct 2026 15:21:38 -0700 Subject: [PATCH 7/8] Describe the expanded GPU pipeline selection Signed-off-by: Trent Nelson --- CONTRIBUTING.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index c3509cbb..38f7d9f3 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -210,7 +210,7 @@ checks in separate fresh containers. Image pulls have bounded retries; installation and test failures fail the job without retrying. The GPU job uses one L40G with Python 3.12 and the locked CUDA 13 dependencies. -It checks CuPy/PyTorch execution, runs 18 xFit GPU numerical cases, and runs +It checks CuPy/PyTorch execution, runs 19 xFit GPU and device pipeline cases, and runs all five synthetic quickstarts with `--require-gpu`. Missing CUDA support or skipped parity cases fail the job. JUnit results and quickstart summaries are uploaded with the tested commit SHA. GPU jobs run one at a time; a new From 7d6423189f07466c3ca0e8f3f21c240d4f6d66e9 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Sat, 24 Oct 2026 15:26:43 -0700 Subject: [PATCH 8/8] Preserve cuTile model and benchmark contracts Reject custom Gaussian subclasses before fused execution. Compare fused benchmarks against unfused CuPy and keep threaded checks on the active device. Signed-off-by: Trent Nelson --- examples/xfit/benchmark_gaussian.py | 10 +++++--- src/cuphoton/xfit/api.py | 4 +-- tests/xfit/test_cutile_fused.py | 38 +++++++++++++++-------------- tests/xfit/test_models.py | 28 +++++++++++++++++++++ 4 files changed, 56 insertions(+), 24 deletions(-) diff --git a/examples/xfit/benchmark_gaussian.py b/examples/xfit/benchmark_gaussian.py index 4e2f9cc9..92e5fef7 100644 --- a/examples/xfit/benchmark_gaussian.py +++ b/examples/xfit/benchmark_gaussian.py @@ -116,7 +116,7 @@ def _run(treatment): variance=_worker.variance, mode=_worker.mode, backend=backend, - fusion=_worker.fusion, + fusion=_worker.fusion and backend == "cutile", ) seconds = time.perf_counter() - start if result.backend != backend: @@ -254,9 +254,11 @@ def main() -> None: os.sched_setaffinity(0, allowed[: args.cpus]) allowed = sorted(os.sched_getaffinity(0)) treatments = ["cupy", "native"] if args.compare_native else backends - # Native comparisons always validate against CuPy, including when only - # native timings were requested. Reference work is outside measurements. - reference = "cupy" if "native" in treatments else treatments[0] + # Fused and native comparisons validate against independent CuPy fits. + # Reference work is outside measurements. + reference = ( + "cupy" if args.fusion or "native" in treatments else treatments[0] + ) dtype = np.dtype(args.dtype) context = multiprocessing.get_context("spawn") for batch in batches: diff --git a/src/cuphoton/xfit/api.py b/src/cuphoton/xfit/api.py index 260837f1..3fe46d3f 100644 --- a/src/cuphoton/xfit/api.py +++ b/src/cuphoton/xfit/api.py @@ -1152,7 +1152,7 @@ def fit_dipoles( gaussian_model=( model == "gaussian" if isinstance(model, str) - else isinstance(model, GaussianDipoleModel) + else type(model) is GaussianDipoleModel ), use_finite_difference=( config is not None and config.use_finite_difference @@ -1214,7 +1214,7 @@ def fit_dipoles_device( gaussian_model=( model == "gaussian" if isinstance(model, str) - else isinstance(model, GaussianDipoleModel) + else type(model) is GaussianDipoleModel ), use_finite_difference=( config is not None and config.use_finite_difference diff --git a/tests/xfit/test_cutile_fused.py b/tests/xfit/test_cutile_fused.py index 9803a041..0bc2165a 100644 --- a/tests/xfit/test_cutile_fused.py +++ b/tests/xfit/test_cutile_fused.py @@ -531,13 +531,14 @@ def test_cutile_preserves_finite_jacobian_overflow_results( def test_fused_concurrent_workers_keep_independent_results(cutile_gpu): cp = cutile_gpu + device_id = cp.cuda.Device().id 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: + with cp.cuda.Device(device_id), 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) @@ -545,25 +546,26 @@ def test_fused_concurrent_workers_keep_independent_results(cutile_gpu): 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 + with cp.cuda.Device(device_id): + pool = cp.cuda.MemoryPool() + stream = cp.cuda.Stream(non_blocking=True) + with stream, cp.cuda.using_allocator(pool.malloc): + barrier.wait(timeout=30) + for _ in range(3): + result = fit_dipoles( + **cases[index], backend="cutile", fusion=True ) - } - retained.append((result, snapshots)) - stream.synchronize() - pool.free_all_blocks() + 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(): diff --git a/tests/xfit/test_models.py b/tests/xfit/test_models.py index fb48fbee..5c49edd7 100644 --- a/tests/xfit/test_models.py +++ b/tests/xfit/test_models.py @@ -1299,6 +1299,34 @@ def test_fusion_rejects_unsupported_settings_before_backend_resolution( ) +@pytest.mark.parametrize("device", [False, True]) +def test_fused_cutile_rejects_custom_gaussian_before_cuda( + monkeypatch, device +) -> None: + class CustomGaussian(GaussianDipoleModel): + def to_backend(self, *args, **kwargs): + pytest.fail("custom model conversion must not run") + + monkeypatch.setattr( + xfit_api, + "resolve_backend", + lambda _: pytest.fail("custom fused model resolved a backend"), + ) + monkeypatch.setattr( + xfit_api, + "_load_cupy", + lambda: pytest.fail("custom fused model loaded CUDA"), + ) + fit = fit_dipoles_device if device else fit_dipoles + with pytest.raises(ValueError, match="only the Gaussian model"): + fit( + np.zeros((1, 5, 7)), + model=CustomGaussian((5, 7)), + backend="cutile", + fusion=True, + ) + + def test_cutile_rejects_finite_difference_provenance() -> None: model = GaussianDipoleModel((7, 9), dtype=np.float64) truth = _gaussian_truth()[:1]