diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5c7966cc..f40bd2a4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -157,6 +157,9 @@ jobs: tests/xfit/test_models.py::test_cupy_finite_difference_theta_fix_matches_numpy tests/xfit/test_models.py::test_fit_dipoles_device_matches_portable_result_without_bulk_d2h tests/xfit/test_solver.py::test_cupy_tall_jacobian_diagnostics_match_numpy + tests/xfit/test_models.py::test_fp32_public_covariance_matches_with_variance + tests/xfit/test_models.py::test_fp32_workspace_cast_preserves_subnormals_on_device + tests/xfit/test_models.py::test_fp32_normal_equations_preserve_subnormal_contributions - name: Require all selected tests to execute run: | uv run --no-sync python - <<'PY' @@ -166,11 +169,11 @@ jobs: report = ET.parse(Path(os.environ["RUNNER_TEMP"]) / "xfit-gpu.xml") cases = report.findall(".//testcase") - assert len(cases) == 6, len(cases) + assert len(cases) == 13, len(cases) assert not report.findall(".//skipped"), "GPU tests must not skip" assert not report.findall(".//failure") assert not report.findall(".//error") - print("All 6 selected tests passed without skips") + print("All 13 selected tests passed without skips") PY - name: Run native xFit GPU tests run: >- @@ -196,6 +199,66 @@ jobs: run: >- uv run --no-sync python examples/run_quickstarts.py --require-gpu --output-dir "${RUNNER_TEMP}/cuphoton-quickstart" + - name: Run cuTile xFit GPU coverage + run: | + # CPU Torch supports feature imports without its CUDA 13.0 pin. + cutile_env="${RUNNER_TEMP}/cuphoton-cutile" + uv venv --python 3.12 "$cutile_env" + uv pip install --python "$cutile_env/bin/python" \ + --index https://download.pytorch.org/whl/cpu 'torch==2.13.0+cpu' + uv pip install --python "$cutile_env/bin/python" --editable '.[cutile]' \ + 'pytest==9.1.1' 'cupy-cuda13x[ctk]==14.1.1' \ + 'cuda-tile[tileiras]==1.6.0' \ + 'cuda-toolkit[nvcc,nvvm,tileiras]==13.2.2' 'nvidia-nvjitlink==13.2.86' + uv pip freeze --python "$cutile_env/bin/python" \ + > "${RUNNER_TEMP}/cutile-packages.txt" + "$cutile_env/bin/python" - <<'PYCOMPILER' > "${RUNNER_TEMP}/cutile-compiler.txt" 2>&1 + from pathlib import Path + import subprocess + import sysconfig + + compiler = Path(sysconfig.get_path("purelib")) / "nvidia/cu13/bin" + for name in ("tileiras", "nvcc", "ptxas"): + executable = compiler / name + print(executable, flush=True) + subprocess.run([str(executable), "--version"], check=True) + PYCOMPILER + cat "${RUNNER_TEMP}/cutile-compiler.txt" + cutile_status=0 + "$cutile_env/bin/python" -m pytest -q -s -rs \ + --junitxml="${RUNNER_TEMP}/cutile-gpu.xml" \ + tests/xfit/test_cutile_fused.py \ + tests/xfit/test_models.py::test_cutile_gaussian_fit_matches_cupy \ + tests/xfit/test_models.py::test_cutile_gaussian_normal_equations_match_noisy_float32 \ + tests/xfit/test_models.py::test_cutile_weighting_and_compacted_indices_match_cupy \ + tests/xfit/test_models.py::test_cutile_preserves_rank_and_budget_failure_semantics \ + tests/xfit/test_models.py::test_cutile_device_results_survive_worker_reuse \ + > "${RUNNER_TEMP}/cutile-gpu.log" 2>&1 || cutile_status=$? + cat "${RUNNER_TEMP}/cutile-gpu.log" + CUTILE_PYTEST_STATUS="$cutile_status" "$cutile_env/bin/python" - <<'PY' + import os + from pathlib import Path + import re + import xml.etree.ElementTree as ET + + root = Path(os.environ["RUNNER_TEMP"]) + output = (root / "cutile-gpu.log").read_text() + lifecycle_error = re.compile( + r"CUDA_ERROR_|cudaError|CUDADriverError|CUDARuntimeError|CUDA error|CUDA driver error|" + r"CUDA runtime error|PytestUnraisableExceptionWarning|unraisable exception|" + r"Exception ignored in:|Exception ignored while|moduleUnload|Module\.__dealloc__", + re.IGNORECASE, + ) + if lifecycle_error.search(output): + raise SystemExit("cuTile tests emitted CUDA or interpreter teardown errors") + if int(os.environ["CUTILE_PYTEST_STATUS"]) != 0: + raise SystemExit("cuTile pytest process failed") + report = ET.parse(root / "cutile-gpu.xml") + cases = report.findall(".//testcase") + if len(cases) != 48 or any(report.findall(".//" + tag) for tag in ("skipped", "failure", "error")): + raise SystemExit("Require all 48 selected cuTile tests without skips, failures or errors") + print("All 48 cuTile tests passed without skips or teardown errors") + PY - name: Upload test evidence if: always() uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 @@ -204,6 +267,10 @@ jobs: path: | ${{ runner.temp }}/xfit-gpu.xml ${{ runner.temp }}/xfit-native-gpu.xml + ${{ runner.temp }}/cutile-gpu.xml + ${{ runner.temp }}/cutile-gpu.log + ${{ runner.temp }}/cutile-packages.txt + ${{ runner.temp }}/cutile-compiler.txt ${{ runner.temp }}/cuphoton-quickstart/**/summary.json if-no-files-found: error retention-days: 14 diff --git a/docs/components/xfit.md b/docs/components/xfit.md index 181a2589..0dbfae3d 100644 --- a/docs/components/xfit.md +++ b/docs/components/xfit.md @@ -66,14 +66,24 @@ applies one mask or variance plane to every channel of each candidate. If the batch size is three, `(3, y, x)` keeps that per-candidate meaning; use the explicit `(1, 3, y, x)` shape for per-plane values. -An explicit `backend="cutile"` uses one `cuda.tile` CTA per Gaussian fit to -form its weighted 8-by-8 normal equations directly. Select this backend +With the default `fusion=False`, an explicit `backend="cutile"` uses one +`cuda.tile` CTA per Gaussian fit to form its weighted 8-by-8 normal equations +directly. Select this backend explicitly and install the `cuphoton[cutile]` extra on Linux with Python 3.12 through 3.14. Final rank and covariance diagnostics use the analytic Jacobian and a singular-value factorization. Sampled-stamp fits stay on the NumPy or CuPy backends. The Tile backend requires analytic derivatives and rejects finite-difference fitting. +Pass `fusion=True` to `fit_dipoles`, or add `--fusion` to +`cuphoton xfit fit-dipoles --backend cutile`, to fuse Gaussian residual and +analytic-Jacobian evaluation, including weighting. Width powers and +normal-equation contractions use CuPy to preserve the reference arithmetic. +Each fit owns its workspace and uses the caller's current +CuPy stream. The default remains +`fusion=False`; CLI artifacts record the selected option in the effective +configuration. + The `cutile` extra installs cuda-tile 1.6 or newer and CuPy. Kernel compilation also requires a compatible CUDA toolkit. Follow the shared [cuTile compiler setup](../getting-started.md#optional-runtimes) to configure diff --git a/examples/xfit/benchmark_gaussian.py b/examples/xfit/benchmark_gaussian.py index 3a0b2189..4e2f9cc9 100644 --- a/examples/xfit/benchmark_gaussian.py +++ b/examples/xfit/benchmark_gaussian.py @@ -59,7 +59,9 @@ def _fixture( return images, initial -def _initialize(images, initial, mode, workers, counter, barrier, gpu): +def _initialize( + images, initial, mode, workers, counter, barrier, gpu, fusion +): with counter.get_lock(): index = counter.value counter.value += 1 @@ -72,6 +74,7 @@ def _initialize(images, initial, mode, workers, counter, barrier, gpu): _worker.initial = initial[start:stop].copy() _worker.variance = np.ones_like(_worker.images) _worker.mode = mode + _worker.fusion = fusion _worker.barrier = barrier _worker.stream = None _worker.pool = None @@ -113,6 +116,7 @@ def _run(treatment): variance=_worker.variance, mode=_worker.mode, backend=backend, + fusion=_worker.fusion, ) seconds = time.perf_counter() - start if result.backend != backend: @@ -214,6 +218,9 @@ def main() -> None: default=180, help="Per-round wait, including cold compilation (default: 180).", ) + parser.add_argument( + "--fusion", action="store_true", help="Use fused cutile evaluation." + ) args = parser.parse_args() if not math.isfinite(args.timeout_seconds) or args.timeout_seconds <= 0: parser.error("--timeout-seconds must be positive and finite") @@ -228,6 +235,8 @@ def main() -> None: if args.compare_native and args.backend: parser.error("--compare-native cannot be combined with --backend") backends = args.backend or ["cupy", "cutile"] + if args.fusion and backends != ["cutile"]: + parser.error("--fusion requires --backend cutile") batches = args.batch or [1, 16, 256, 4096] if any(batch < args.workers for batch in batches): parser.error("every batch must contain at least --workers candidates") @@ -274,6 +283,7 @@ def main() -> None: context.Value("i", 0), context.Barrier(args.workers), any(backend != "numpy" for backend in treatments), + args.fusion, ), **executor_options, ) as pool: @@ -345,6 +355,7 @@ def main() -> None: json.dumps( { "backend": treatment, + "fusion": args.fusion, "execution": args.execution, "workers": args.workers, "batch": batch, diff --git a/src/cuphoton/xfit/_cutile_fused.py b/src/cuphoton/xfit/_cutile_fused.py new file mode 100644 index 00000000..1d9ba4ce --- /dev/null +++ b/src/cuphoton/xfit/_cutile_fused.py @@ -0,0 +1,495 @@ +# 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) + 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 + ) -> BackendArray: + output = self._cp.empty( + (parameters.shape[0], self._observations), dtype=parameters.dtype + ) + self._evaluate(parameters, indices, output, compute_jacobian=False) + return output + + def jacobian( + self, parameters: BackendArray, *, indices: BackendArray + ) -> BackendArray: + output = self._jacobians[: parameters.shape[0] * 8] + self._evaluate(parameters, indices, output, compute_jacobian=True) + return output.reshape(parameters.shape[0], 8, self._observations) diff --git a/src/cuphoton/xfit/api.py b/src/cuphoton/xfit/api.py index 8e6fa971..cda65e93 100644 --- a/src/cuphoton/xfit/api.py +++ b/src/cuphoton/xfit/api.py @@ -131,15 +131,19 @@ class DeviceDipoleFitResult: init=False, default="levenberg-marquardt", ) - backend: Literal["cupy", "native"] = "cupy" + backend: Literal["cupy", "cutile", "native"] = "cupy" result_location: Literal["device"] = field( init=False, default="device", ) def __post_init__(self) -> None: - if self.backend not in {"cupy", "native"}: - raise ValueError("device fit backend must be cupy or native") + if self.backend not in {"cupy", "cutile", "native"}: + raise ValueError( + "device fit backend must be cupy, cutile or native" + ) + if self.backend == "cutile" and self.model != "gaussian": + raise ValueError("backend='cutile' supports only Gaussian fits") if isinstance(self.device_id, bool) or not isinstance( self.device_id, (int, np.integer) ): @@ -615,6 +619,7 @@ def _fit_dipoles_backend( mode: FitMode = "difference", resolved: Backend, config: LMConfig | None = None, + fusion: bool = False, ) -> _BackendDipoleFitResult: """Fit dipoles while retaining all result arrays in ``resolved``.""" @@ -747,9 +752,22 @@ def evaluate_model(parameters: BackendArray) -> BackendArray: ) return selected_model.evaluate(parameters, mode=mode) + workspace = None + if fusion: + from ._cutile_fused import GaussianWorkspace + + workspace = GaussianWorkspace( + images_array.reshape(batch, -1), + weights.reshape(batch, -1), + image_shape=(height, width), + mode=mode, + ) + def residual( parameters: BackendArray, *, indices: BackendArray ) -> BackendArray: + if workspace is not None: + return workspace.residual(parameters, indices=indices) prediction = evaluate_model(model_parameters(parameters)) value = (prediction - images_array[indices]) * weights[indices] return value.reshape(parameters.shape[0], -1) @@ -760,6 +778,8 @@ def residual( def jacobian( parameters: BackendArray, *, indices: BackendArray ) -> BackendArray: + if workspace is not None: + return workspace.jacobian(parameters, indices=indices) assert isinstance(selected_model, GaussianDipoleModel) physical = model_parameters(parameters) value = selected_model._jacobian_positive_unchecked( @@ -773,7 +793,7 @@ def jacobian( jacobian_function = jacobian normal_equations = None - if gaussian_model and resolved.name == "cutile": + if gaussian_model and resolved.name == "cutile" and workspace is None: from ._cutile import gaussian_normal_equations def normal_equations( @@ -1018,6 +1038,25 @@ def _validate_execution_options( ) +def _validate_fusion( + fusion: bool, + *, + backend: BackendRequest, + gaussian_model: bool, + use_finite_difference: bool, +) -> None: + if not isinstance(fusion, bool): + raise TypeError("fusion must be a boolean") + if not fusion: + return + if backend != "cutile": + raise ValueError("fusion requires backend='cutile'") + if not gaussian_model: + raise ValueError("fusion supports only the Gaussian model") + if use_finite_difference: + raise ValueError("fusion does not support finite-difference fitting") + + def fit_dipoles( images: ArrayLike, *, @@ -1028,6 +1067,7 @@ def fit_dipoles( mode: FitMode = "difference", backend: BackendRequest = "auto", config: LMConfig | None = None, + fusion: bool = False, ) -> DipoleFitResult: """Fit batched Gaussian or sampled-stamp dipoles. @@ -1049,9 +1089,25 @@ def fit_dipoles( treats all three planes as independent; when difference equals positive minus negative, the reported degrees of freedom and uncertainties are not statistically calibrated for that dependence. + + ``fusion=True`` selects fused Gaussian residual and analytic-Jacobian + evaluation with ``backend="cutile"``. Normal equations retain CuPy's + contractions and the solver's convergence and diagnostics controls. """ _validate_execution_options(model, backend, config) + _validate_fusion( + fusion, + backend=backend, + gaussian_model=( + model == "gaussian" + if isinstance(model, str) + else isinstance(model, GaussianDipoleModel) + ), + use_finite_difference=( + config is not None and config.use_finite_difference + ), + ) resolved = resolve_backend(backend) backend_result = _fit_dipoles_backend( images, @@ -1062,6 +1118,7 @@ def fit_dipoles( mode=mode, resolved=resolved, config=config, + fusion=fusion, ) return _materialize_dipole_fit_result(backend_result) @@ -1075,11 +1132,14 @@ def fit_dipoles_device( variance: ArrayLike | None = None, mode: FitMode = "difference", config: LMConfig | None = None, - backend: Literal["cupy", "native"] = "cupy", + backend: Literal["cupy", "cutile", "native"] = "cupy", + fusion: bool = False, ) -> DeviceDipoleFitResult: """Fit dipoles while keeping every result array on the active GPU. - This experimental seam uses CuPy-owned buffers. The optional native + This experimental seam uses CuPy-owned buffers. The cuTile backend + supports Gaussian analytic fits; ``fusion=True`` selects fused residual + and analytic-Jacobian evaluation. The optional native backend runs Gaussian LM iterations in CUDA C++. NumPy inputs are copied to the active device. Existing CuPy inputs must already reside on that device and are checked before any input conversion. Input arrays are @@ -1098,8 +1158,20 @@ def fit_dipoles_device( """ _validate_execution_options(model, backend, config) - if backend not in {"cupy", "native"}: - raise ValueError("device fit backend must be cupy or native") + _validate_fusion( + fusion, + backend=backend, + gaussian_model=( + model == "gaussian" + if isinstance(model, str) + else isinstance(model, GaussianDipoleModel) + ), + use_finite_difference=( + config is not None and config.use_finite_difference + ), + ) + if backend not in {"cupy", "cutile", "native"}: + raise ValueError("device fit backend must be cupy, cutile or native") if backend == "native": from ._native import load_extension @@ -1134,6 +1206,7 @@ def fit_dipoles_device( mode=mode, resolved=resolve_backend(backend), config=config, + fusion=fusion, ) return DeviceDipoleFitResult( parameters=result.parameters, diff --git a/src/cuphoton/xfit/commands.py b/src/cuphoton/xfit/commands.py index 8d19d3c8..fda3e8ff 100644 --- a/src/cuphoton/xfit/commands.py +++ b/src/cuphoton/xfit/commands.py @@ -236,6 +236,7 @@ class FitDipolesCommand(ExecutorOptions, _ValidatedDatasetCommand): g_tol = None max_evaluations = None use_finite_difference = None + fusion = None chunk_size = None class ChunkSizeArg(PositiveIntegerInvariant): @@ -266,6 +267,12 @@ class ComputeDtypeArg(ComputeDtypeInvariant): _mandatory = False _default = "input" + class FusionArg(BoolInvariant): + _arg = "--fusion" + _help = "Fuse Gaussian evaluation with the cutile backend." + _mandatory = False + _default = False + class StampEvaluationArg(StampEvaluationInvariant): _arg = "--stamp-evaluation" _help = ( @@ -330,6 +337,17 @@ def _model_specification( ) def run(self) -> None: + from .api import _validate_fusion + + assert self.backend is not None + assert self.fusion is not None + self._call( + _validate_fusion, + self.fusion, + backend=self.backend, + gaussian_model=self.model == "gaussian", + use_finite_difference=bool(self.use_finite_difference), + ) executor_options = self.executor_options() if self.executor != "local": from cuphoton.core.executors import run_workload @@ -362,6 +380,7 @@ def run(self) -> None: "g_tol", "max_evaluations", "use_finite_difference", + "fusion", "fits_reader", ) } @@ -439,6 +458,7 @@ def run(self) -> None: mode=self.mode, backend=self.backend, config=config, + fusion=self.fusion, ) dtype = np.dtype(result.dtype) default_tolerance = float(np.sqrt(np.finfo(dtype).eps)) @@ -451,6 +471,7 @@ def run(self) -> None: "model": self.model, "mode": self.mode, "backend": self.backend, + "fusion": self.fusion, **( {"fits_reader": self.fits_reader} if dataset.input_sources diff --git a/src/cuphoton/xfit/executor.py b/src/cuphoton/xfit/executor.py index cbff8f8b..32b7436f 100644 --- a/src/cuphoton/xfit/executor.py +++ b/src/cuphoton/xfit/executor.py @@ -26,7 +26,12 @@ ) from cuphoton.core.fits_options import normalize_xdr_options -from .api import DipoleFitResult, _floating_dtype, fit_dipoles +from .api import ( + DipoleFitResult, + _floating_dtype, + _validate_fusion, + fit_dipoles, +) from .io import ( XFitDataset, _broadcast_image_auxiliary, @@ -62,6 +67,7 @@ "g_tol": None, "max_evaluations": None, "use_finite_difference": False, + "fusion": False, "fits_reader": "auto", "xdr_options": None, } @@ -101,6 +107,12 @@ def _plan_xfit_chunks( settings = json_mapping( {**_DEFAULTS, **fit_options}, field="xFit options" ) + _validate_fusion( + settings["fusion"], + backend=settings["backend"], + gaussian_model=settings["model"] == "gaussian", + use_finite_difference=settings["use_finite_difference"], + ) if settings["compute_dtype"] not in {"input", "float32", "float64"}: raise ValueError("unsupported xFit compute dtype") if settings["fits_reader"] not in {"auto", "astropy", "xdr"}: @@ -281,6 +293,7 @@ def run_item(self, item: WorkItem, item_dir: Path) -> Mapping[str, Any]: mode=self.settings["mode"], backend=self.settings["backend"], config=self.config, + fusion=self.settings["fusion"], ) solve_seconds = time.perf_counter() - solve_start _check_input_stat(self.options) @@ -443,6 +456,7 @@ def _effective_config( "model": settings["model"], "mode": settings["mode"], "backend": settings["backend"], + "fusion": settings["fusion"], **( {"fits_reader": settings["fits_reader"]} if dataset.input_sources diff --git a/src/cuphoton/xfit/solver.py b/src/cuphoton/xfit/solver.py index c73c4e81..e08ab527 100644 --- a/src/cuphoton/xfit/solver.py +++ b/src/cuphoton/xfit/solver.py @@ -8,6 +8,7 @@ from dataclasses import dataclass from enum import IntEnum +from functools import lru_cache from types import ModuleType from typing import Any, Protocol @@ -356,8 +357,8 @@ def _final_jacobian( def _factorize_jacobians( jacobians: BackendArray, ap: ModuleType -) -> tuple[BackendArray, BackendArray]: - """Return left singular vectors and values for covariance diagnostics.""" +) -> tuple[BackendArray, BackendArray, BackendArray]: + """Return unscaled rank factors and the matrix used to compute them.""" batch, parameters, observations = jacobians.shape module_name = type(jacobians).__module__.split(".", maxsplit=1)[0] @@ -371,11 +372,75 @@ def _factorize_jacobians( left_vectors, singular_values, _ = ap.linalg.svd( reduced, full_matrices=False ) - return left_vectors, singular_values + return left_vectors, singular_values, reduced left_vectors, singular_values, _ = ap.linalg.svd( jacobians, full_matrices=False ) - return left_vectors, singular_values + return left_vectors, singular_values, jacobians + + +def _covariance_from_factor( + factor: BackendArray, ap: ModuleType +) -> BackendArray: + """Equilibrate a full-rank parameter-major factor for covariance only.""" + + scale = ap.max(ap.abs(factor), axis=-1) + left_vectors, singular_values, _ = ap.linalg.svd( + factor / scale[..., None], full_matrices=False + ) + inverse_factor = (left_vectors / singular_values[..., None, :]) / scale[ + ..., None + ] + return inverse_factor @ inverse_factor.swapaxes(-2, -1) + + +@lru_cache(maxsize=2) +def _fp32_workspace_cast_kernel(ap: ModuleType, to_float32: bool) -> Any: + # CuPy appends -ftz=true even after a caller requests --ftz=false. + # Explicit PTX without .ftz preserves subnormals in both directions. + if to_float32: + return ap.ElementwiseKernel( + "float64 x", + "float32 y", + 'asm("cvt.rn.f32.f64 %0, %1;" : "=f"(y) : "d"(x));', + "xfit_workspace_float32", + ) + return ap.ElementwiseKernel( + "float32 x", + "float64 y", + 'asm("cvt.f64.f32 %0, %1;" : "=d"(y) : "f"(x));', + "xfit_workspace_float64", + ) + + +def _cast_fp32_workspace( + values: BackendArray, dtype: Any, ap: ModuleType +) -> BackendArray: + """Convert FP32 solver workspaces without flushing subnormal values.""" + + if ap is np: + return values.astype(dtype) + return _fp32_workspace_cast_kernel(ap, dtype == np.float32)(values) + + +def _normal_equations_from_jacobian( + jacobian: BackendArray, residual: BackendArray, ap: ModuleType +) -> tuple[BackendArray, BackendArray]: + """Accumulate FP32 products in FP64, then retain the input precision.""" + + widen = jacobian.dtype == np.float32 and residual.dtype == np.float32 + normal_jacobian = ( + _cast_fp32_workspace(jacobian, np.float64, ap) if widen else jacobian + ) + normal_residual = ( + _cast_fp32_workspace(residual, np.float64, ap) if widen else residual + ) + gradient = ap.einsum("knm,km->kn", normal_jacobian, normal_residual) + hessian = ap.einsum("knm,kpm->knp", normal_jacobian, normal_jacobian) + if widen: + gradient = _cast_fp32_workspace(gradient, np.float32, ap) + hessian = _cast_fp32_workspace(hessian, np.float32, ap) + return gradient, hessian def batched_levenberg_marquardt( @@ -390,6 +455,12 @@ def batched_levenberg_marquardt( interpreted as parameter-major, with shape ``(K, N, M)``. When it is absent, or finite differences are requested, the solver allocates the perturbation workspace in the same backend as ``x0``. + + Built-in FP32 normal equations accumulate in FP64 and are rounded back + before convergence checks. FP32 damped systems likewise use FP64 solve + intermediates and return FP32 steps. Model evaluations, solver state, + covariance and rank thresholds retain their input precision. A supplied + normal-equations callback controls its own accumulation precision. """ if not hasattr(x0, "ndim") or x0.ndim != 2: @@ -512,13 +583,13 @@ def finite_rows(value: BackendArray) -> BackendArray: fallback_jacobian = fallback_jacobian[valid_fallback] gradient = gradient.copy() hessian = hessian.copy() - gradient[recovered_positions] = ap.einsum( - "knm,km->kn", + ( + gradient[recovered_positions], + hessian[recovered_positions], + ) = _normal_equations_from_jacobian( fallback_jacobian, residual_active[recovered_positions], - ) - hessian[recovered_positions] = ap.einsum( - "knm,kpm->knp", fallback_jacobian, fallback_jacobian + ap, ) valid_normal_equations[recovered_positions] = True jacobians[recovered_indices] = fallback_jacobian @@ -566,8 +637,9 @@ def finite_rows(value: BackendArray) -> BackendArray: jacobian = jacobian[valid_jacobian] jacobians[active_indices] = jacobian jacobian_current[active_indices] = True - gradient = ap.einsum("knm,km->kn", jacobian, residual_active) - hessian = ap.einsum("knm,kpm->knp", jacobian, jacobian) + gradient, hessian = _normal_equations_from_jacobian( + jacobian, residual_active, ap + ) diagonal = ap.maximum(ap.diagonal(hessian, axis1=1, axis2=2), eps) scaled_gradient = ap.max(ap.abs(gradient) / ap.sqrt(diagonal), axis=1) residual_norm = ap.linalg.norm(residual_active, axis=1) @@ -592,10 +664,32 @@ def finite_rows(value: BackendArray) -> BackendArray: * diagonal[:, None, :] ) try: - step = ap.linalg.solve(system, -gradient[..., None])[..., 0] + if system.dtype == np.float32 and gradient.dtype == np.float32: + # NumPy already solves FP32 inputs with FP64 intermediates. + # Match that policy on the GPU before rounding the step. + step = _cast_fp32_workspace( + ap.linalg.solve( + _cast_fp32_workspace(system, np.float64, ap), + -_cast_fp32_workspace( + gradient[..., None], np.float64, ap + ), + )[..., 0], + np.float32, + ap, + ) + else: + step = ap.linalg.solve(system, -gradient[..., None])[..., 0] except Exception: status[active_indices] = int(LMStatus.SINGULAR) continue + # For (H + damping * D) step = -gradient, the quadratic model + # predicts -g.T step - step.T H step / 2. Use the equivalent + # O(n) expression with the damping that produced this step. + predicted_reduction = 0.5 * ap.einsum( + "kn,kn->k", + step, + damping[active_indices, None] * diagonal * step - gradient, + ) valid_step = finite_rows(step) status[active_indices[~valid_step]] = int(LMStatus.SINGULAR) active_indices = active_indices[valid_step] @@ -604,6 +698,7 @@ def finite_rows(value: BackendArray) -> BackendArray: x_active = x_active[valid_step] residual_active = residual_active[valid_step] step = step[valid_step] + predicted_reduction = predicted_reduction[valid_step] can_evaluate = evaluations[active_indices] < settings.max_evaluations status[active_indices[~can_evaluate]] = int(LMStatus.MAX_EVALUATIONS) @@ -613,6 +708,7 @@ def finite_rows(value: BackendArray) -> BackendArray: x_active = x_active[can_evaluate] residual_active = residual_active[can_evaluate] step = step[can_evaluate] + predicted_reduction = predicted_reduction[can_evaluate] trial_x = x_active + step trial_residual = _call_residual(problem, trial_x, active_indices) @@ -638,8 +734,16 @@ def finite_rows(value: BackendArray) -> BackendArray: damping[accepted_indices] * settings.damping_decrease, ) rejected_steps[accepted_indices] = 0 - f_converged = reduction[accepted] <= ( - settings.f_tol * ap.maximum(1.0, cost[accepted]) + # As in MINPACK, both actual and predicted progress must be + # small, without the model underpredicting by more than 2x. + prediction = predicted_reduction[accepted] + f_threshold = settings.f_tol * ap.maximum(1.0, cost[accepted]) + f_converged = ( + (reduction[accepted] <= f_threshold) + & ap.isfinite(prediction) + & (prediction > 0) + & (prediction <= f_threshold) + & (0.5 * reduction[accepted] <= prediction) ) x_converged = ap.linalg.norm(accepted_step, axis=1) <= ( settings.x_tol * (settings.x_tol + 1.0) @@ -819,7 +923,7 @@ def finite_rows(value: BackendArray) -> BackendArray: hessians[diagnostic_indices] = ap.einsum( "knm,kpm->knp", diagnostic_jacobians, diagnostic_jacobians ) - left_vectors, singular_values = _factorize_jacobians( + _, singular_values, covariance_factor = _factorize_jacobians( diagnostic_jacobians, ap ) tolerance = max(n, m) * eps * ap.max(singular_values, axis=1) @@ -828,14 +932,23 @@ def finite_rows(value: BackendArray) -> BackendArray: full_rank = rank_values == n full_rank_indices = diagnostic_indices[full_rank] if full_rank_indices.shape[0] > 0: - full_rank_vectors = left_vectors[full_rank] - full_rank_singular_values = singular_values[full_rank] - inverse_squared = 1.0 / ( - full_rank_singular_values * full_rank_singular_values - ) - covariance[full_rank_indices] = ( - full_rank_vectors * inverse_squared[:, None, :] - ) @ full_rank_vectors.transpose(0, 2, 1) + full_rank_factor = covariance_factor[full_rank] + try: + covariance[full_rank_indices] = _covariance_from_factor( + full_rank_factor, ap + ) + except Exception: + # A covariance-only failure must not replace successful + # unscaled rank decisions with a different factorization. + for index in range(full_rank_indices.shape[0]): + try: + covariance[full_rank_indices[index]] = ( + _covariance_from_factor( + full_rank_factor[index], ap + ) + ) + except Exception: + continue except Exception: # Preserve row-level failure isolation if a backend cannot perform # one of the batched linear-algebra operations. @@ -844,18 +957,15 @@ def finite_rows(value: BackendArray) -> BackendArray: jacobian = jacobians[row] gradients[row] = jacobian @ residuals[row] hessians[row] = jacobian @ jacobian.T - left_vectors, singular_values, _ = ap.linalg.svd( + _, singular_values, _ = ap.linalg.svd( jacobian, full_matrices=False ) tolerance = max(n, m) * eps * ap.max(singular_values) rank[row] = ap.sum(singular_values > tolerance) if _int(rank[row]) == n: - inverse_squared = 1.0 / ( - singular_values * singular_values + covariance[row] = _covariance_from_factor( + jacobian, ap ) - covariance[row] = ( - left_vectors * inverse_squared[None, :] - ) @ left_vectors.T except Exception: continue diff --git a/src/cuphoton/xscan/device_pipeline.py b/src/cuphoton/xscan/device_pipeline.py index 37754c1b..04b0247e 100644 --- a/src/cuphoton/xscan/device_pipeline.py +++ b/src/cuphoton/xscan/device_pipeline.py @@ -650,17 +650,28 @@ class DeviceXFitPipelineConfig: damping_decrease: float = 0.3 finite_difference_step: float | None = None use_finite_difference: bool = False - backend: Literal["cupy", "native"] = "cupy" + backend: Literal["cupy", "cutile", "native"] = "cupy" + fusion: bool = False def __post_init__(self) -> None: if self.model != "gaussian": raise ValueError("xfit model must be 'gaussian'") if self.mode != "difference": raise ValueError("xfit mode must be 'difference'") - if self.backend not in ("cupy", "native"): - raise ValueError("xfit backend must be 'cupy' or 'native'") + if self.backend not in ("cupy", "cutile", "native"): + raise ValueError( + "xfit backend must be 'cupy', 'cutile' or 'native'" + ) if not isinstance(self.use_finite_difference, bool): raise TypeError("xfit use_finite_difference must be boolean") + if not isinstance(self.fusion, bool): + raise TypeError("xfit fusion must be boolean") + if self.fusion and self.backend != "cutile": + raise ValueError("xfit fusion requires backend='cutile'") + if self.backend == "cutile" and self.use_finite_difference: + raise ValueError( + "xfit backend='cutile' does not support finite differences" + ) if self.backend == "native" and self.use_finite_difference: raise ValueError( "native xfit does not support finite differences" @@ -715,6 +726,8 @@ def to_payload(self) -> dict[str, Any]: payload = asdict(self) if self.backend == "cupy": del payload["backend"] + if not self.fusion: + del payload["fusion"] return payload def execution_payload(self) -> dict[str, Any]: @@ -723,6 +736,8 @@ def execution_payload(self) -> dict[str, Any]: payload: dict[str, Any] = {} if self.backend != "cupy": payload["backend"] = self.backend + if self.fusion: + payload["fusion"] = True return payload def solver_payload(self) -> dict[str, Any]: @@ -732,6 +747,7 @@ def solver_payload(self) -> dict[str, Any]: del payload["model"] del payload["mode"] payload.pop("backend", None) + payload.pop("fusion", None) return payload @classmethod @@ -740,9 +756,10 @@ def from_payload( ) -> DeviceXFitPipelineConfig: """Restore exact LM settings from JSON-compatible values.""" + expected = frozenset(value.name for value in fields(cls)) values = json_mapping(payload, field="device pipeline xfit config") values.setdefault("backend", "cupy") - expected = frozenset(value.name for value in fields(cls)) + values.setdefault("fusion", False) values = _require_exact_fields( values, expected=expected, @@ -1079,9 +1096,10 @@ def _device_pipeline_evidence_layout_contract( "mode": "difference", "dtype": "float64", } - if xfit["backend"] not in ("cupy", "native"): + if xfit["backend"] not in ("cupy", "cutile", "native"): raise ValueError( - "scientific evidence xfit.backend must be 'cupy' or 'native'" + "scientific evidence xfit.backend must be 'cupy', 'cutile' " + "or 'native'" ) for name, expected in expected_xfit.items(): if xfit[name] != expected: @@ -2158,6 +2176,7 @@ def _pack_scientific_evidence( candidate_count: int, kernel_shape: tuple[int, int], flux_conserve: bool, + xfit_backend: Literal["cupy", "cutile", "native"] = "cupy", ) -> _PackedScientificEvidence: """Pack all parity evidence on the producer stream as float64.""" @@ -2183,7 +2202,8 @@ def _pack_scientific_evidence( str(xfit_result.mode), str(xfit_result.dtype), ) - if xfit_contract[0] not in ("cupy", "native") or xfit_contract[1:] != ( + if xfit_contract != ( + xfit_backend, "levenberg-marquardt", "gaussian", "difference", @@ -3352,6 +3372,7 @@ def run_device_pipeline_item( candidate_count=len(item.candidates), kernel_shape=context.config.xpois.kernel_shape, flux_conserve=context.config.xpois.flux_conserve, + xfit_backend=context.config.xfit.backend, ), ) evidence_layout = _append_prediction_evidence_layout( diff --git a/src/cuphoton/xscan/pipeline_benchmark/runner.py b/src/cuphoton/xscan/pipeline_benchmark/runner.py index 543bc7b2..d0f2cc3f 100644 --- a/src/cuphoton/xscan/pipeline_benchmark/runner.py +++ b/src/cuphoton/xscan/pipeline_benchmark/runner.py @@ -383,6 +383,8 @@ def audit(args: SimpleNamespace) -> dict[str, Any]: if metadata["xpois"][f"xpois_{name}"] != value: raise ValueError(f"xPois {name} differs") evidence = result["scientific_evidence"] + if metadata["xfit"].get("backend") != evidence["xfit"]["backend"]: + raise ValueError("xFit backend differs") if ( metadata["xpois"]["basis_terms"] != evidence["xpois"]["basis_terms"] @@ -427,6 +429,8 @@ def provenance() -> dict[str, Any]: "xfit/api.py", "xfit/backend.py", "xfit/_native.py", + "xfit/_cutile.py", + "xfit/_cutile_fused.py", "xfit/models.py", "xfit/solver.py", "xscan/xfit_features.py", @@ -437,7 +441,7 @@ def provenance() -> dict[str, Any]: cupy_packages = importlib.metadata.packages_distributions().get( "cupy", [] ) - for name in ("cuphoton", "numpy", "torch", *cupy_packages): + for name in ("cuphoton", "numpy", "torch", "cuda-tile", *cupy_packages): try: packages[name] = importlib.metadata.version(name) except importlib.metadata.PackageNotFoundError: diff --git a/src/cuphoton/xscan/xfit_features.py b/src/cuphoton/xscan/xfit_features.py index aaa5d0bd..fd379d57 100644 --- a/src/cuphoton/xscan/xfit_features.py +++ b/src/cuphoton/xscan/xfit_features.py @@ -1292,10 +1292,12 @@ def _validate_device_feature_source( if result.solver != "levenberg-marquardt": raise ValueError("xFit device result solver is unsupported") if ( - result.backend not in {"cupy", "native"} + result.backend not in {"cupy", "cutile", "native"} or result.result_location != "device" ): raise ValueError("xFit result must be a CuPy device result") + if result.backend == "cutile" and result.model != "gaussian": + raise ValueError("cuTile xFit features require the Gaussian model") parameter_names = _validate_live_feature_source(result) if int(result.device_id) != active_device_id: raise ValueError( diff --git a/tests/core/test_cli_contract.py b/tests/core/test_cli_contract.py index ad4fdec1..ca2493aa 100644 --- a/tests/core/test_cli_contract.py +++ b/tests/core/test_cli_contract.py @@ -71,7 +71,7 @@ def test_public_command_surface_counts_are_exact() -> None: assert per_group == [ ("xdr", 1, 0, 17, 0), - ("xfit", 3, 3, 38, 1), + ("xfit", 3, 3, 39, 1), ("xpois", 7, 7, 162, 1), ("xscan", 44, 44, 229, 1), ("xrep", 6, 6, 123, 1), @@ -81,7 +81,7 @@ def test_public_command_surface_counts_are_exact() -> None: assert sum(item[1] for item in per_group) == 94 assert sum(item[1] + item[4] for item in per_group) == 99 assert sum(item[2] for item in per_group) == 91 - assert sum(item[3] for item in per_group) == 961 + assert sum(item[3] for item in per_group) == 962 def test_public_registry_order_and_component_derivations() -> None: diff --git a/tests/test_public_api.py b/tests/test_public_api.py index 6787184b..9d54d87e 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -123,6 +123,7 @@ def test_xfit_curated_exports_and_fit_signature() -> None: "mode", "backend", "config", + "fusion", ] assert list(inspect.signature(xfit.fit_dipoles_device).parameters) == [ "images", @@ -133,6 +134,7 @@ def test_xfit_curated_exports_and_fit_signature() -> None: "mode", "config", "backend", + "fusion", ] assert [field.name for field in fields(xfit.DipoleFitResult)] == [ "parameters", diff --git a/tests/xfit/test_commands.py b/tests/xfit/test_commands.py index f975d85c..e03dc38a 100644 --- a/tests/xfit/test_commands.py +++ b/tests/xfit/test_commands.py @@ -477,10 +477,12 @@ def test_artifacts_retain_archive_hash_captured_at_load_time( assert loaded_sha256 != file_sha256(input_path) +@pytest.mark.parametrize("fusion", [False, True]) def test_fit_dipoles_command_writes_complete_artifact_set( monkeypatch, tmp_path, capsys, + fusion, ) -> None: import cuphoton.xfit as xfit @@ -510,6 +512,7 @@ def fake_fit_dipoles(received_images, **kwargs): assert received_images.dtype == np.float64 assert kwargs["model"] == "gaussian" assert isinstance(kwargs["config"], FakeConfig) + assert kwargs["fusion"] is fusion return _fake_result(received_images) monkeypatch.setattr(xfit, "LMConfig", FakeConfig, raising=False) @@ -531,7 +534,8 @@ def fake_fit_dipoles(received_images, **kwargs): "--model", "gaussian", "--backend", - "numpy", + "cutile" if fusion else "numpy", + *(["--fusion"] if fusion else []), "--compute-dtype", "float64", ], @@ -557,6 +561,56 @@ def fake_fit_dipoles(received_images, **kwargs): "requested": "float64", "resolved": "float64", } + assert effective["fusion"] is fusion + assert effective["backend"] == ("cutile" if fusion else "numpy") + + +@pytest.mark.parametrize( + ("backend", "model", "extra_args", "message"), + [ + ("cupy", "gaussian", [], "requires backend='cutile'"), + ("cutile", "stamp", [], "only the Gaussian model"), + ( + "cutile", + "gaussian", + ["--use-finite-difference"], + "finite-difference fitting", + ), + ], +) +def test_fusion_command_rejects_invalid_settings_before_loading_input( + tmp_path, monkeypatch, capsys, backend, model, extra_args, message +) -> None: + from cuphoton.xfit.commands import FitDipolesCommand + + input_path = tmp_path / "input.npz" + _write_gaussian_input(input_path) + monkeypatch.setattr( + FitDipolesCommand, + "_load_dataset", + lambda *args, **kwargs: pytest.fail( + "invalid fusion settings read the input" + ), + ) + rc = run_component( + "xfit", + [ + "fit-dipoles", + "--input", + str(input_path), + "--output-dir", + str(tmp_path / "fit"), + "--backend", + backend, + "--model", + model, + "--fusion", + *extra_args, + ], + ) + assert rc == 1 + assert message in capsys.readouterr().err + assert not (tmp_path / "fit").exists() def test_fit_dipoles_command_runs_real_synthetic_gaussian_fit( @@ -754,6 +808,7 @@ def test_command_help_exposes_stable_safe_input_options(capsys) -> None: assert "vignetted" in captured.out assert "finite-volume" in captured.out assert "--use-finite-difference" in captured.out + assert "--fusion" in captured.out @pytest.mark.parametrize("runtime", ["dragon", "mpi"]) diff --git a/tests/xfit/test_cutile_fused.py b/tests/xfit/test_cutile_fused.py new file mode 100644 index 00000000..9803a041 --- /dev/null +++ b/tests/xfit/test_cutile_fused.py @@ -0,0 +1,649 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import re +import subprocess +import sys +from concurrent.futures import ThreadPoolExecutor +from dataclasses import fields +from threading import Barrier + +import numpy as np +import pytest + +import cuphoton.xfit.api as xfit_api +from cuphoton.xfit import GaussianDipoleModel, LMConfig, fit_dipoles + + +@pytest.fixture +def cutile_gpu(): + cp = pytest.importorskip("cupy") + pytest.importorskip("cuda.tile") + try: + count = cp.cuda.runtime.getDeviceCount() + except Exception as exc: + pytest.skip(f"CUDA runtime is unavailable: {exc}") + if count < 1: + pytest.skip("CUDA device is unavailable") + return cp + + +def _case(mode, dtype): + model = GaussianDipoleModel( + (11, 15), split_weights=(1.0, 0.25, 1.75), dtype=dtype + ) + truth = np.asarray( + [ + [5.0, 1.3, 1.8, 0.25, -2.0, 0.5, 2.2, -0.7], + [3.0, 1.6, 1.1, -0.3, -1.0, -1.0, 2.0, 1.0], + [4.2, 1.4, 2.0, 0.1, -1.7, 0.8, 1.9, -0.9], + ], + dtype=dtype, + ) + images = np.asarray(model.evaluate(truth, mode=mode)) + initial = truth + np.asarray( + [0.2, 0.1, -0.1, 0.03, 0.1, -0.1, -0.1, 0.1], dtype=dtype + ) + # The first row finishes before the other rows, exercising original + # candidate indices after the active batch has been compacted. + initial[0] = truth[0] + mask = np.ones(images.shape, dtype=bool) + mask.reshape(3, -1)[0, ::17] = False + mask.reshape(3, -1)[1, 3::19] = False + mask.reshape(3, -1)[2, 5::23] = False + variance = np.linspace(0.5, 2.0, images.size, dtype=dtype).reshape( + images.shape + ) + images[~mask] = np.nan + variance[~mask] = np.nan + return dict( + images=images, + model=model, + initial=initial, + mode=mode, + mask=mask, + variance=variance, + ) + + +def _assert_fit_equal(actual, expected, dtype): + tolerance = 3e-5 if dtype is np.float32 else 5e-12 + for name in ( + "status", + "converged", + "evaluations", + "valid_pixel_count", + "degrees_of_freedom", + "uncertainty_valid", + ): + np.testing.assert_array_equal( + getattr(actual, name), getattr(expected, name), err_msg=name + ) + assert actual.uncertainty_reason == expected.uncertainty_reason + for name in ( + "parameters", + "residual_norm", + "chi_square", + "valid_pixel_fraction", + "null_chi_square", + "delta_chi_square", + "fractional_null_improvement", + "reduced_chi_square", + "covariance", + "standard_errors", + "residuals", + ): + value = getattr(actual, name) + reference = getattr(expected, name) + np.testing.assert_array_equal( + np.isnan(value), np.isnan(reference), err_msg=name + ) + np.testing.assert_allclose( + value, + reference, + rtol=tolerance, + atol=tolerance, + equal_nan=True, + err_msg=name, + ) + + +def _fit_and_rank(cp, monkeypatch, **kwargs): + # Rank is a solver diagnostic; the public result reports its effect on + # covariance through uncertainty_reason instead of exposing the number. + solve = xfit_api.batched_levenberg_marquardt + ranks = [] + + def record_rank(*args, **options): + result = solve(*args, **options) + ranks.append(cp.asnumpy(result.rank).copy()) + return result + + with monkeypatch.context() as patch: + patch.setattr(xfit_api, "batched_levenberg_marquardt", record_rank) + result = fit_dipoles(**kwargs) + assert len(ranks) == 1 + return result, ranks[0] + + +@pytest.mark.parametrize("mode", ["difference", "split"]) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_fused_fit_preserves_weighted_compacted_results( + cutile_gpu, monkeypatch, mode, dtype +): + case = _case(mode, dtype) + fused, fused_rank = _fit_and_rank( + cutile_gpu, monkeypatch, **case, backend="cutile", fusion=True + ) + assert fused.evaluations[0] == 1 + assert np.all(fused.evaluations[1:] > 1) + assert fused.converged.all() + assert fused.uncertainty_valid.all() + for backend in ("cupy", "cutile"): + reference, rank = _fit_and_rank( + cutile_gpu, monkeypatch, **case, backend=backend + ) + _assert_fit_equal(fused, reference, dtype) + np.testing.assert_array_equal(fused_rank, rank) + + +@pytest.mark.parametrize("scenario", ["rank", "budget", "rejected-step"]) +def test_fused_fit_preserves_termination_and_failed_uncertainty( + cutile_gpu, monkeypatch, scenario +): + model = GaussianDipoleModel((11, 15), dtype=np.float64) + truth = np.asarray([[5.0, 1.3, 1.8, 0.25, -2.0, 0.5, 2.2, -0.7]]) + initial = truth.copy() + config = LMConfig() + if scenario == "rank": + truth[:, 0] = 0 + initial = truth.copy() + elif scenario == "budget": + initial[:, 0] += 0.2 + config = LMConfig(max_evaluations=1) + else: + initial = np.asarray([[1.0, 0.5, 0.7, 0.0, 0.0, 0.0, 0.5, 0.0]]) + config = LMConfig(max_evaluations=2, initial_damping=1e-8) + case = dict( + images=model.evaluate(truth), + model=model, + initial=initial, + config=config, + ) + fused, rank = _fit_and_rank( + cutile_gpu, monkeypatch, **case, backend="cutile", fusion=True + ) + unfused, unfused_rank = _fit_and_rank( + cutile_gpu, monkeypatch, **case, backend="cutile" + ) + _assert_fit_equal(fused, unfused, np.float64) + np.testing.assert_array_equal(rank, unfused_rank) + # Analytic final diagnostics do not consume residual evaluations, even + # when the normal-equation path needs to refresh its Jacobian. + reference, reference_rank = _fit_and_rank( + cutile_gpu, monkeypatch, **case, backend="cupy" + ) + _assert_fit_equal(fused, reference, np.float64) + np.testing.assert_array_equal(rank, reference_rank) + assert not fused.uncertainty_valid.any() + assert np.isnan(fused.covariance).all() + if scenario == "rank": + assert np.all(rank < 8) + assert fused.uncertainty_reason == ("rank-deficient Jacobian",) + elif scenario == "budget": + np.testing.assert_array_equal(fused.evaluations, [1]) + np.testing.assert_array_equal(rank, [-1]) + else: + np.testing.assert_array_equal(fused.evaluations, [2]) + np.testing.assert_array_equal(fused.status, ["max_evaluations"]) + # A rejected trial must preserve the accepted parameters/residuals. + np.testing.assert_allclose( + fused.parameters, initial, rtol=5e-12, atol=5e-12 + ) + + +@pytest.mark.parametrize("mode", ["difference", "split"]) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_fused_kernel_outputs_match_noisy_strided_inputs( + cutile_gpu, mode, dtype +): + from cuphoton.xfit._cutile_fused import GaussianWorkspace + + cp = cutile_gpu + rng = np.random.default_rng(1219) + shape = (8, 10) + physical = cp.asarray(_case(mode, dtype)["initial"]) + model = GaussianDipoleModel(shape, backend="cupy", dtype=dtype) + prediction = model.evaluate(physical, mode=mode).reshape(3, -1) + images = prediction + cp.asarray( + rng.normal(scale=0.2, size=prediction.shape), dtype=dtype + ) + weights = cp.asarray( + rng.uniform(0.25, 2.0, size=prediction.shape), dtype=dtype + ) + weights[:, ::7] = 0 + # Both inputs have non-unit observation strides; the active rows also + # arrive in a different order from their original image/weight storage. + images = cp.repeat(images, 2, axis=1)[:, ::2] + weights = cp.repeat(weights, 2, axis=1)[:, ::2] + assert not images.flags.c_contiguous + assert not weights.flags.c_contiguous + indices = cp.asarray([2, 0], dtype=cp.int64) + active = physical[indices].copy() + active[:, 0] *= 1.1 + parameters = active.copy() + parameters[:, 1:3] = cp.log(parameters[:, 1:3]) + expected_residuals = ( + model.evaluate(active, mode=mode).reshape(2, -1) - images[indices] + ) * weights[indices] + jacobian = model.jacobian(active, mode=mode).reshape(2, 8, -1) + jacobian[:, 1:3] *= active[:, 1:3, None] + jacobian *= weights[indices, None, :] + workspace = GaussianWorkspace( + images, weights, image_shape=shape, mode=mode + ) + residuals = workspace.residual(parameters, indices=indices) + actual_jacobian = workspace.jacobian(parameters, indices=indices) + tolerance = 3e-5 if dtype is np.float32 else 5e-12 + for actual, expected in ( + (residuals, expected_residuals), + (actual_jacobian, jacobian), + ): + np.testing.assert_allclose( + cp.asnumpy(actual), + cp.asnumpy(expected), + rtol=tolerance, + atol=tolerance, + ) + + +def test_fused_jacobian_preserves_cancellation_sensitive_hessian(cutile_gpu): + from cuphoton.xfit._cutile_fused import GaussianWorkspace + + cp = cutile_gpu + shape = (51, 51) + parameters = cp.asarray( + [[2000.0, 1.5, 2.25, 0.4, -2.0, 0.0, 2.0, 0.0]], + dtype=cp.float64, + ) + parameters[:, 1:3] = cp.log(parameters[:, 1:3]) + physical = parameters.copy() + physical[:, 1:3] = cp.exp(physical[:, 1:3]) + model = GaussianDipoleModel(shape, backend="cupy", dtype=np.float64) + expected_jacobian = model.jacobian(physical).reshape(1, 8, -1) + expected_jacobian[:, 1:3] *= physical[:, 1:3, None] + images = cp.zeros((1, shape[0] * shape[1]), dtype=cp.float64) + workspace = GaussianWorkspace( + images, cp.ones_like(images), image_shape=shape, mode="difference" + ) + jacobian = workspace.jacobian( + parameters, indices=cp.asarray([0], dtype=cp.int64) + ) + # Large symmetric stars make cancellation in Hessian cross terms + # sensitive to the separately rounded products in the analytic model. + expected_hessian = cp.einsum( + "knm,kpm->knp", expected_jacobian, expected_jacobian + ) + hessian = cp.einsum("knm,kpm->knp", jacobian, jacobian) + for actual, expected in ( + (jacobian, expected_jacobian), + (hessian, expected_hessian), + ): + np.testing.assert_allclose( + cp.asnumpy(actual), cp.asnumpy(expected), rtol=5e-12, atol=5e-12 + ) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_fused_invalid_sigma_trials_keep_accepted_residuals( + cutile_gpu, dtype +): + from cuphoton.xfit._cutile_fused import GaussianWorkspace + + cp = cutile_gpu + model = GaussianDipoleModel((7, 9), dtype=dtype) + physical = np.asarray([[5, 1.3, 1.8, 0.25, -2, 0.5, 2.2, -0.7]], dtype) + image = cp.asarray(model.evaluate(physical)).reshape(1, -1) + workspace = GaussianWorkspace( + image, cp.ones_like(image), image_shape=(7, 9), mode="difference" + ) + parameters = cp.asarray(physical) + parameters[:, 1:3] = cp.log(parameters[:, 1:3]) + indices = cp.asarray([0], dtype=cp.int64) + accepted = workspace.residual(parameters, indices=indices) + before = cp.asnumpy(accepted) + for log_sigma in (-1000.0, 1000.0): + trial = parameters.copy() + trial[:, 1:3] = log_sigma + invalid = workspace.residual(trial, indices=indices) + assert cp.isnan(invalid).all().item() + np.testing.assert_array_equal(cp.asnumpy(accepted), before) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("fusion", [False, True]) +def test_cutile_large_width_keeps_physical_derivative_order( + cutile_gpu, monkeypatch, dtype, fusion +): + from cuphoton.xfit._cutile import gaussian_normal_equations + from cuphoton.xfit._cutile_fused import GaussianWorkspace + + cp = cutile_gpu + width = 1e13 if dtype is np.float32 else 1e103 + initial = np.asarray([[1, width, 1, 0, -width, 0, width, 0]], dtype) + model = GaussianDipoleModel((5, 7), split_weights=(1, 1, 1), dtype=dtype) + images = np.asarray(model.evaluate(initial, mode="split")) * dtype(1.01) + mask = np.zeros_like(images, dtype=bool) + mask[..., 2, :] = True + parameters = cp.asarray(initial) + parameters[:, 1:3] = cp.log(parameters[:, 1:3]) + physical = parameters.copy() + physical[:, 1:3] = cp.exp(physical[:, 1:3]) + gpu_model = model.to_backend("cupy") + weights = cp.asarray(mask, dtype=dtype).reshape(1, -1) + data = cp.asarray(images).reshape(1, -1) + expected_residual = ( + gpu_model.evaluate(physical, mode="split").reshape(1, -1) - data + ) * weights + jacobian = gpu_model.jacobian(physical, mode="split").reshape(1, 8, -1) + jacobian[:, 1:3] *= physical[:, 1:3, None] + jacobian *= weights[:, None, :] + # The physical sigma**3 denominator overflows, while sigma**2 and + # every weighted Jacobian entry remain finite. Simplifying the chain + # rule to star*x*x/sigma**2 would introduce a spurious width direction. + assert cp.isfinite(jacobian).all().item() + assert (jacobian[:, 1] == 0).all().item() + workspace = GaussianWorkspace( + data, weights, image_shape=(5, 7), mode="split" + ) + indices = cp.asarray([0], dtype=cp.int64) + if fusion: + actual_jacobian = workspace.jacobian(parameters, indices=indices) + tolerance = 3e-5 if dtype is np.float32 else 5e-12 + np.testing.assert_allclose( + cp.asnumpy(actual_jacobian), + cp.asnumpy(jacobian), + rtol=tolerance, + atol=tolerance, + ) + assert cp.isfinite(actual_jacobian).all().item() + assert (actual_jacobian[:, 1] == 0).all().item() + else: + gradient, hessian = gaussian_normal_equations( + parameters, + expected_residual, + weights, + image_shape=(5, 7), + mode="split", + ) + assert cp.isfinite(gradient).all().item() + assert cp.isfinite(hessian).all().item() + assert (gradient[:, 1] == 0).all().item() + assert (hessian[:, 1, :] == 0).all().item() + case = dict( + images=images, + model=model, + initial=initial, + mask=mask, + mode="split", + config=LMConfig(max_evaluations=2, f_tol=0, x_tol=0, g_tol=0), + ) + reference, reference_rank = _fit_and_rank( + cp, monkeypatch, **case, backend="cupy" + ) + fused, rank = _fit_and_rank( + cp, monkeypatch, **case, backend="cutile", fusion=fusion + ) + np.testing.assert_array_equal(reference.status, ["max_evaluations"]) + np.testing.assert_array_equal(reference.evaluations, [2]) + _assert_fit_equal(fused, reference, dtype) + np.testing.assert_array_equal(rank, reference_rank) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("masked", [False, True]) +@pytest.mark.parametrize("fusion", [False, True]) +def test_cutile_tiny_width_preserves_invalid_derivative( + cutile_gpu, monkeypatch, dtype, masked, fusion +): + from cuphoton.xfit._cutile import gaussian_normal_equations + from cuphoton.xfit._cutile_fused import GaussianWorkspace + + cp = cutile_gpu + width = 1e-16 if dtype is np.float32 else 1e-110 + initial = np.asarray([[1, width, 1, 0, -1, 0, 1, 0]], dtype) + images = np.full((1, 2, 2), np.nan if masked else 0, dtype=dtype) + mask = np.zeros_like(images, dtype=bool) if masked else None + parameters = cp.asarray(initial) + parameters[:, 1:3] = cp.log(parameters[:, 1:3]) + # Every model pixel is zero, but the physical width derivative is 0/0 + # because sigma**3 underflows. Zero mask weights preserve that NaN. + weights = cp.full((1, 4), 0 if masked else 1, dtype=dtype) + workspace = GaussianWorkspace( + cp.zeros((1, 4), dtype=dtype), + weights, + image_shape=(2, 2), + mode="difference", + ) + indices = cp.asarray([0], dtype=cp.int64) + residual = workspace.residual(parameters, indices=indices) + assert (residual == 0).all().item() + if fusion: + jacobian = workspace.jacobian(parameters, indices=indices) + physical = parameters.copy() + physical[:, 1:3] = cp.exp(physical[:, 1:3]) + model = GaussianDipoleModel((2, 2), backend="cupy", dtype=dtype) + expected_jacobian = model.jacobian(physical).reshape(1, 8, 4) + expected_jacobian[:, 1:3] *= physical[:, 1:3, None] + expected_jacobian *= weights[:, None, :] + assert not cp.isfinite(jacobian).all().item() + np.testing.assert_array_equal( + cp.asnumpy(cp.isnan(jacobian)), + cp.asnumpy(cp.isnan(expected_jacobian)), + ) + tolerance = 3e-5 if dtype is np.float32 else 5e-12 + np.testing.assert_allclose( + cp.asnumpy(jacobian), + cp.asnumpy(expected_jacobian), + rtol=tolerance, + atol=tolerance, + equal_nan=True, + ) + else: + gradient, hessian = gaussian_normal_equations( + parameters, + residual, + weights, + image_shape=(2, 2), + mode="difference", + ) + assert not cp.isfinite(gradient).all().item() + assert not cp.isfinite(hessian).all().item() + case = dict( + images=images, + model="gaussian", + initial=initial, + mask=mask, + config=LMConfig(max_evaluations=2), + ) + reference, reference_rank = _fit_and_rank( + cp, monkeypatch, **case, backend="cupy" + ) + fused, rank = _fit_and_rank( + cp, monkeypatch, **case, backend="cutile", fusion=fusion + ) + np.testing.assert_array_equal(reference.status, ["invalid_residual"]) + np.testing.assert_array_equal(reference.evaluations, [1]) + _assert_fit_equal(fused, reference, dtype) + np.testing.assert_array_equal(rank, reference_rank) + + +@pytest.mark.parametrize( + ("dtype", "amplitude"), [(np.float32, 1e25), (np.float64, 1e200)] +) +@pytest.mark.parametrize("nonzero_residual", [False, True]) +@pytest.mark.parametrize("fusion", [False, True]) +def test_cutile_preserves_finite_jacobian_overflow_results( + cutile_gpu, monkeypatch, dtype, amplitude, nonzero_residual, fusion +): + cp = cutile_gpu + model = GaussianDipoleModel((5, 7), backend="cupy", dtype=dtype) + # Coincident stars keep the model finite and exactly zero. Only the + # middle row's finite position derivatives overflow their products. + initial = np.asarray( + [ + [5, 1, 1, 0, 0, 0, 0, 0], + [amplitude, 1, 1, 0, 0, 0, 0, 0], + [3, 1, 1, 0, 0, 0, 0, 0], + ], + dtype=dtype, + ) + images = np.zeros((3, 5, 7), dtype=dtype) + if nonzero_residual: + images[1, 1, 2] = amplitude + jacobian = model.jacobian(cp.asarray(initial)).reshape(3, 8, -1) + assert cp.isfinite(jacobian).all().item() + hessian = cp.einsum("knm,kpm->knp", jacobian, jacobian) + np.testing.assert_array_equal( + cp.asnumpy(cp.isfinite(hessian).all(axis=(1, 2))), + [True, False, True], + ) + case = dict(images=images, model=model, initial=initial) + reference, reference_rank = _fit_and_rank( + cp, monkeypatch, **case, backend="cupy" + ) + actual, rank = _fit_and_rank( + cp, monkeypatch, **case, backend="cutile", fusion=fusion + ) + assert reference.status.tolist() == [ + "converged_g_tol", + "singular" if nonzero_residual else "converged_g_tol", + "converged_g_tol", + ] + np.testing.assert_array_equal(reference.evaluations, [1, 1, 1]) + np.testing.assert_array_equal(actual.parameters, reference.parameters) + _assert_fit_equal(actual, reference, dtype) + np.testing.assert_array_equal(rank, reference_rank) + + +def test_fused_concurrent_workers_keep_independent_results(cutile_gpu): + cp = cutile_gpu + cases = [_case("split", np.float64) for _ in range(2)] + cases[1]["images"] = cases[1]["images"] * 1.25 + cases[1]["initial"] = cases[1]["initial"].copy() + cases[1]["initial"][:, 0] *= 1.25 + # Warm all specializations on the default stream before testing worker + # ownership; compilation concurrency is a separate concern. + with cp.cuda.Stream.null: + expected = [fit_dipoles(**case, backend="cupy") for case in cases] + for case in cases: + fit_dipoles(**case, backend="cutile", fusion=True) + cp.cuda.Stream.null.synchronize() + barrier = Barrier(2) + + def worker(index): + pool = cp.cuda.MemoryPool() + stream = cp.cuda.Stream(non_blocking=True) + retained = [] + with cp.cuda.Device(0), stream, cp.cuda.using_allocator(pool.malloc): + barrier.wait(timeout=30) + for _ in range(3): + result = fit_dipoles( + **cases[index], backend="cutile", fusion=True + ) + snapshots = { + field.name: value.copy() + for field in fields(result) + if isinstance( + value := getattr(result, field.name), np.ndarray + ) + } + retained.append((result, snapshots)) + stream.synchronize() + pool.free_all_blocks() + for result, snapshots in retained: + _assert_fit_equal(result, expected[index], np.float64) + for name, snapshot in snapshots.items(): + np.testing.assert_array_equal(getattr(result, name), snapshot) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(worker, index) for index in range(2)] + for future in futures: + future.result(timeout=180) + + +def test_fused_cli_matches_unfused_artifacts(cutile_gpu, tmp_path): + import pyarrow.parquet as pq + import yaml + + case = _case("difference", np.float64) + input_path = tmp_path / "synthetic.npz" + np.savez( + input_path, + candidate_id=np.asarray(["a", "b", "c"]), + **{ + name: case[name] + for name in ("images", "initial", "mask", "variance") + }, + ) + outputs = [] + for fusion in (False, True): + output = tmp_path / ("fused" if fusion else "unfused") + command = [ + sys.executable, + "-m", + "cuphoton", + "xfit", + "fit-dipoles", + "--input", + str(input_path), + "--output-dir", + str(output), + "--model", + "gaussian", + "--backend", + "cutile", + ] + if fusion: + command.append("--fusion") + completed = subprocess.run( + command, text=True, capture_output=True, timeout=240, check=False + ) + assert completed.returncode == 0, completed.stdout + completed.stderr + assert not re.search( + r"CUDA_ERROR_|Exception ignored (?:in:|while)", + completed.stdout + completed.stderr, + ), completed.stdout + completed.stderr + effective = yaml.safe_load( + (output / "effective-config.yaml").read_text() + ) + assert effective["fusion"] is fusion + outputs.append(output) + reference, fused = [pq.read_table(p / "fits.parquet") for p in outputs] + for name in ("candidate_id", "status", "evaluations", "converged"): + assert reference[name].to_pylist() == fused[name].to_pylist() + for name in case["model"].parameter_names: + np.testing.assert_allclose( + fused[name].to_numpy(), + reference[name].to_numpy(), + rtol=5e-12, + atol=5e-12, + ) + with ( + np.load(outputs[0] / "fit-arrays.npz", allow_pickle=False) as before, + np.load(outputs[1] / "fit-arrays.npz", allow_pickle=False) as after, + ): + for name in ("covariance", "residuals"): + np.testing.assert_array_equal( + np.isnan(after[name]), np.isnan(before[name]) + ) + np.testing.assert_allclose( + after[name], + before[name], + rtol=5e-12, + atol=5e-12, + equal_nan=True, + ) diff --git a/tests/xfit/test_executor.py b/tests/xfit/test_executor.py index 6f058f81..71967854 100644 --- a/tests/xfit/test_executor.py +++ b/tests/xfit/test_executor.py @@ -13,6 +13,7 @@ import numpy as np import pyarrow.parquet as pq import pytest +import yaml from cuphoton.core.artifacts import file_sha256 from cuphoton.xfit import ( @@ -103,8 +104,9 @@ def test_chunk_size_is_part_of_configuration_identity(tmp_path): assert first["configuration_sha256"] != second["configuration_sha256"] +@pytest.mark.parametrize("fusion", [False, True]) def test_finalizer_retains_planning_input_and_still_checks_content( - tmp_path, monkeypatch + tmp_path, monkeypatch, fusion ): path = _input(tmp_path) loads = [] @@ -115,8 +117,11 @@ def load(*args, **kwargs): return original_load(*args, **kwargs) def cpu_fit(*args, **kwargs): - result = fit_dipoles(*args, **{**kwargs, "backend": "numpy"}) - return replace(result, backend="cupy", device="cuda:0") + assert kwargs["fusion"] is fusion + result = fit_dipoles( + *args, **{**kwargs, "backend": "numpy", "fusion": False} + ) + return replace(result, backend=kwargs["backend"], device="cuda:0") monkeypatch.setattr(executor, "load_xfit_dataset", load) monkeypatch.setattr(executor, "fit_dipoles", cpu_fit) @@ -124,7 +129,11 @@ def cpu_fit(*args, **kwargs): spec = executor.prepare_xfit_workload( input_path=path, chunk_size=2, - fit_options={"backend": "cupy", "max_evaluations": 2}, + fit_options={ + "backend": "cutile" if fusion else "cupy", + "fusion": fusion, + "max_evaluations": 2, + }, ) worker = spec.worker_factory(spec.options_payload) for index in range(2): @@ -134,6 +143,10 @@ def cpu_fit(*args, **kwargs): worker.run_item(item, round_dir / "items" / item.item_id) records.append({"item_id": item.item_id, "status": "success"}) assert spec.finalize_round(round_dir, records)["candidate_count"] == 3 + effective = yaml.safe_load( + (round_dir / "scientific" / "effective-config.yaml").read_text() + ) + assert effective["fusion"] is fusion assert len(loads) == 2 # One planning load and one worker load. before = path.stat() changed = bytearray(path.read_bytes()) @@ -146,18 +159,46 @@ def cpu_fit(*args, **kwargs): worker.close() -def test_preflight_accepts_analytic_gaussian_cutile(tmp_path): +@pytest.mark.parametrize("fusion", [False, True]) +def test_preflight_accepts_analytic_gaussian_cutile(tmp_path, fusion): spec = executor.prepare_xfit_workload( - input_path=_input(tmp_path), fit_options={"backend": "cutile"} + input_path=_input(tmp_path), + fit_options={"backend": "cutile", "fusion": fusion}, ) assert spec.backend == "cutile" + assert spec.options_payload["fit_options"]["fusion"] is fusion + + +@pytest.mark.parametrize( + ("backend", "fusion", "message"), + [ + ("cupy", True, "requires backend='cutile'"), + ("cutile", "yes", "must be a boolean"), + ], +) +def test_preflight_rejects_invalid_fusion_before_input_reads( + tmp_path, monkeypatch, backend, fusion, message +): + monkeypatch.setattr( + executor, + "load_xfit_dataset", + lambda *args, **kwargs: pytest.fail( + "invalid fusion settings read the input" + ), + ) + with pytest.raises((ValueError, TypeError), match=message): + executor.prepare_xfit_workload( + input_path=tmp_path / "missing.npz", + fit_options={"backend": backend, "fusion": fusion}, + ) @pytest.mark.parametrize( "runtime,rank", [("dragon", 0), ("mpi", 0), ("mpi", 1)] ) +@pytest.mark.parametrize("fusion", [False, True]) def test_cli_dispatches_collective_preflight_and_preserves_options( - tmp_path, monkeypatch, capsys, runtime, rank + tmp_path, monkeypatch, capsys, runtime, rank, fusion ): from types import SimpleNamespace @@ -197,7 +238,8 @@ def run(**kwargs): "--model", "gaussian", "--backend", - "cupy", + "cutile" if fusion else "cupy", + *(["--fusion"] if fusion else []), "--executor", runtime, "--chunk-size", @@ -222,6 +264,7 @@ def run(**kwargs): assert kwargs["benchmark"].measure_rounds == 2 assert len(spec.items) == 2 assert spec.options_payload["fit_options"]["max_evaluations"] == 10 + assert spec.options_payload["fit_options"]["fusion"] is fusion assert json.loads(capsys.readouterr().out)["status"] == "success" diff --git a/tests/xfit/test_models.py b/tests/xfit/test_models.py index 13adbc3d..3fb8569a 100644 --- a/tests/xfit/test_models.py +++ b/tests/xfit/test_models.py @@ -645,16 +645,24 @@ def test_vignetted_bilinear_preserves_visible_requested_flux() -> None: assert vignetted.to_backend("numpy").evaluation == "bilinear-vignetted" -def test_rank_deficient_fit_reports_invalid_public_uncertainty() -> None: - model = GaussianDipoleModel((7, 9), dtype=np.float64) - initial = np.asarray([[0.0, 1.0, 1.0, 0.0, -1.0, 0.0, 1.0, 0.0]]) +@pytest.mark.parametrize("backend", ["numpy", "cupy"]) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_rank_deficient_fit_reports_invalid_public_uncertainty( + backend, dtype +) -> None: + if backend == "cupy": + _require_cupy_device() + model = GaussianDipoleModel((7, 9), dtype=dtype) + initial = np.asarray( + [[0.0, 1.0, 1.0, 0.0, -1.0, 0.0, 1.0, 0.0]], dtype=dtype + ) images = model.evaluate(initial) result = fit_dipoles( images, model=model, initial=initial, - backend="numpy", + backend=backend, ) assert result.converged.all() @@ -662,6 +670,7 @@ def test_rank_deficient_fit_reports_invalid_public_uncertainty() -> None: assert result.uncertainty_reason == ("rank-deficient Jacobian",) assert np.isnan(result.covariance).all() assert np.isnan(result.standard_errors).all() + assert result.covariance.dtype == dtype def test_stamp_multi_mode_basis_requires_explicit_weights() -> None: @@ -920,14 +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"), - ("backend", "numpy"), - ("result_location", "host"), - ): - with pytest.raises(ValueError): - replace(device, **{name: value}) + for name in ("schema", "solver", "result_location"): + # Python 3.14 changed replace()'s init=False error to TypeError. + with pytest.raises( + (ValueError, TypeError), match=rf"field {name}.*init=False" + ): + replace(device, **{name: "other"}) + with pytest.raises(ValueError, match="backend"): + replace(device, backend="numpy") with pytest.raises(ValueError, match="CUDA device"): replace(device, device_id=device.device_id + 1) with pytest.raises(TypeError, match=device.dtype): @@ -1097,6 +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: @@ -1222,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] @@ -1543,3 +1620,292 @@ def test_cutile_preserves_rank_and_budget_failure_semantics() -> None: assert np.array_equal(cutile_budget.status, cupy_budget.status) assert np.array_equal(cutile_budget.evaluations, cupy_budget.evaluations) assert cutile_budget.uncertainty_reason == cupy_budget.uncertainty_reason + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("fusion", [False, True]) +def test_cutile_device_results_survive_worker_reuse(dtype, fusion) -> None: + from concurrent.futures import ThreadPoolExecutor + from dataclasses import fields + from threading import Barrier + + cp = _require_cutile() + pytest.importorskip("torch") + + from cuphoton.xscan.xfit_features import ( + transform_xfit_result_features_device, + ) + + device_id = int(cp.cuda.Device().id) + model = GaussianDipoleModel((11, 15), dtype=dtype) + cases = [] + for scale in (1.0, 1.5): + truth = _gaussian_truth(dtype) + truth[:, 0] *= scale + images = model.evaluate(truth, mode="difference") + initial = truth + np.asarray( + [0.2, 0.1, -0.1, 0.03, 0.1, -0.1, -0.1, 0.1], dtype=dtype + ) + initial[0] = truth[0] + cases.append( + dict( + images=images, + model=model, + initial=initial, + variance=np.full_like(images, 0.5), + mode="difference", + ) + ) + + def snapshot(result): + return { + field.name: cp.asnumpy(value) + for field in fields(result) + if isinstance(value := getattr(result, field.name), cp.ndarray) + } + + def features(result): + transformed = transform_xfit_result_features_device( + result, image_shape=(11, 15), variance_present=True + ) + assert transformed.backend == "cupy" + return transformed + + expected = [] + with cp.cuda.Stream.null: + for case in cases: + reference = fit_dipoles_device(**case) + expected.append( + (snapshot(reference), cp.asnumpy(features(reference).values)) + ) + fit_dipoles_device(**case, backend="cutile", fusion=fusion) + cp.cuda.Stream.null.synchronize() + barrier = Barrier(2) + tolerance = 3e-5 if dtype is np.float32 else 5e-12 + + def worker(index): + with cp.cuda.Device(device_id): + pool = cp.cuda.MemoryPool() + stream = cp.cuda.Stream(non_blocking=True) + retained = [] + with stream, cp.cuda.using_allocator(pool.malloc): + barrier.wait(timeout=30) + for iteration in range(2): + case_index = (index + iteration) % len(cases) + result = fit_dipoles_device( + **cases[case_index], backend="cutile", fusion=fusion + ) + assert result.backend == "cutile" + assert result.device_id == device_id + transformed = features(result) + arrays = snapshot(result) + reference, reference_features = expected[case_index] + assert arrays.keys() == reference.keys() + for name, actual in arrays.items(): + target = reference[name] + if actual.dtype.kind in "fc": + np.testing.assert_array_equal( + np.isnan(actual), np.isnan(target) + ) + np.testing.assert_allclose( + actual, target, rtol=tolerance, atol=tolerance + ) + else: + np.testing.assert_array_equal(actual, target) + feature_values = cp.asnumpy(transformed.values) + # Canonical xScan features have float32 storage even + # when the fit was computed in float64. + np.testing.assert_allclose( + feature_values, + reference_features, + rtol=3e-5, + atol=3e-5, + ) + retained.append( + (result, arrays, transformed, feature_values) + ) + stream.synchronize() + pool.free_all_blocks() + for result, arrays, transformed, feature_values in retained: + for name, actual in snapshot(result).items(): + np.testing.assert_array_equal(actual, arrays[name]) + np.testing.assert_array_equal( + cp.asnumpy(transformed.values), feature_values + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(worker, index) for index in range(2)] + for future in futures: + future.result(timeout=180) + + +@pytest.mark.parametrize( + "variation", ["original", "nearby-noise", "masked-nonfinite"] +) +def test_fp32_public_covariance_matches_with_variance( + variation, +) -> None: + _require_cupy_device() + truth = np.asarray( + [ + [4.0, 1.2, 1.6, 0.13, -1.0, 0.2, 1.0, -0.2], + [3.1, 1.4, 1.1, -0.17, -0.8, -0.3, 0.9, 0.35], + ], + dtype=np.float32, + ) + model = GaussianDipoleModel((7, 9), dtype=np.float32) + images = model.evaluate(truth) + images += ( + np.random.default_rng(20260929) + .normal(0, 0.004, images.shape) + .astype(np.float32) + ) + if variation == "nearby-noise": + images += ( + np.random.default_rng(20261013) + .normal(0, 1e-4, images.shape) + .astype(np.float32) + ) + weights = np.linspace(0.7, 1.6, images.size, dtype=np.float32).reshape( + images.shape + ) + weights.reshape(2, -1)[:, ::11] = 0 + mask = weights != 0 + images[~mask] = np.nan + variance = 1 / np.where(mask, weights, 1) ** 2 + if variation == "masked-nonfinite": + variance[~mask] = np.inf + initial = truth + np.asarray( + [0.1, 0.04, -0.04, 0.015, 0.05, -0.04, -0.05, 0.04], + dtype=np.float32, + ) + originals = [value.copy() for value in (images, initial, variance, mask)] + results = [] + for backend in ("numpy", "cupy"): + results.append( + fit_dipoles( + images, + model=model, + initial=initial, + mask=mask, + variance=variance, + config=LMConfig(max_evaluations=80), + backend=backend, + ) + ) + with pytest.raises(ValueError, match="finite on included pixels"): + fit_dipoles(images, model=model, initial=initial, backend=backend) + cpu, gpu = results + assert cpu.uncertainty_valid.all() and gpu.uncertainty_valid.all() + assert gpu.parameters.dtype == gpu.covariance.dtype == np.float32 + for field, expected in vars(cpu).items(): + if field in {"backend", "device"}: + continue + actual = getattr(gpu, field) + if isinstance(expected, np.ndarray): + assert actual.dtype == expected.dtype + if isinstance(expected, np.ndarray) and expected.dtype.kind == "f": + np.testing.assert_allclose( + actual, expected, rtol=3e-5, atol=3e-5, equal_nan=True + ) + else: + np.testing.assert_array_equal(actual, expected) + for actual, expected in zip( + (images, initial, variance, mask), originals, strict=True + ): + assert actual.tobytes() == expected.tobytes() + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_fp32_workspace_cast_preserves_subnormals_on_device(dtype) -> None: + cp = _require_cupy_device() + values = np.asarray( + [ + 0, + 0x80000000, + 1, + 0x80000001, + 16384, + 0x80004000, + 0x7FFFFF, + 0x807FFFFF, + 0x800000, + 0x80800000, + 0x3F800000, + 0xBF800000, + 0x7F7FFFFF, + 0xFF7FFFFF, + 0x7F800000, + 0xFF800000, + 0x7FC00000, + ], + dtype=np.uint32, + ).view(np.float32) + if dtype == np.float32: + tiny = float(np.finfo(np.float32).smallest_subnormal) + values = np.concatenate( + ( + values.astype(np.float64), + np.asarray( + [ + tiny * 0.5, + tiny * 1.5, + tiny * 2.5, + -tiny * 0.5, + -tiny * 1.5, + -tiny * 2.5, + ] + ), + ) + ) + expected = values.astype(dtype) + stream = cp.cuda.Stream(non_blocking=True) + with stream: + device = cp.asarray(np.repeat(values, 2))[::2] + actual_device = xfit_solver._cast_fp32_workspace(device, dtype, cp) + stream.synchronize() + actual = cp.asnumpy(actual_device) + assert not device.flags.c_contiguous + assert actual.shape == expected.shape + assert actual.dtype == expected.dtype + np.testing.assert_array_equal(np.isnan(actual), np.isnan(expected)) + np.testing.assert_array_equal(np.isposinf(actual), np.isposinf(expected)) + np.testing.assert_array_equal(np.isneginf(actual), np.isneginf(expected)) + finite = np.isfinite(expected) + bits = np.uint32 if dtype == np.float32 else np.uint64 + np.testing.assert_array_equal( + actual[finite].view(bits), expected[finite].view(bits) + ) + assert cp.asnumpy(device).tobytes() == values.tobytes() + + +@pytest.mark.parametrize( + ("jacobian_scale", "residual_scale"), + [(2.0**-65, 2.0**-70), (2.0**-135, 2.0**100)], +) +def test_fp32_normal_equations_preserve_subnormal_contributions( + jacobian_scale, residual_scale +) -> None: + cp = _require_cupy_device() + jacobian = np.full((2, 2, 16), jacobian_scale, dtype=np.float32) + residual = np.full((2, 16), residual_scale, dtype=np.float32) + expected_gradient = np.full( + (2, 2), 16 * jacobian_scale * residual_scale, dtype=np.float32 + ) + expected_hessian = np.full( + (2, 2, 2), 16 * jacobian_scale**2, dtype=np.float32 + ) + for ap in (np, cp): + gradient, hessian = xfit_solver._normal_equations_from_jacobian( + ap.asarray(jacobian), ap.asarray(residual), ap + ) + if ap is cp: + gradient, hessian = cp.asnumpy(gradient), cp.asnumpy(hessian) + for actual, expected in ( + (gradient, expected_gradient), + (hessian, expected_hessian), + ): + assert actual.dtype == np.float32 + np.testing.assert_array_equal( + actual.view(np.uint32), expected.view(np.uint32) + ) diff --git a/tests/xfit/test_solver.py b/tests/xfit/test_solver.py index 079bf479..2b6abd25 100644 --- a/tests/xfit/test_solver.py +++ b/tests/xfit/test_solver.py @@ -71,6 +71,60 @@ def jacobian(x, *, indices, out=None): assert np.isfinite(result.covariance).all() +def test_f_tol_continues_after_poorly_predicted_nonlinear_progress() -> None: + coefficient = 1.001 - 1.0e-8 + + def residual(x, *, indices): + del indices + return 1.0 - x + coefficient * x * x + + def jacobian(x, *, indices): + del indices + return (-1.0 + 2.0 * coefficient * x)[:, :, None] + + # The first damped step is almost one: actual cost improvement is about + # 1e-8, although the linear model predicts about 0.5. Small actual + # progress alone would stop far from this positive parabola's minimum. + problem = BatchedLeastSquaresProblem(residual, jacobian) + initial = np.zeros((1, 1)) + limited = batched_levenberg_marquardt( + problem, initial, config=LMConfig(max_evaluations=2) + ) + assert limited.status.tolist() == [LMStatus.MAX_EVALUATIONS] + assert limited.evaluations.tolist() == [2] + np.testing.assert_allclose(limited.parameters, 1.0 / 1.001) + + result = batched_levenberg_marquardt(problem, initial) + + assert result.converged.all() + assert result.evaluations[0] > 2 + np.testing.assert_allclose( + result.parameters, 1.0 / (2.0 * coefficient), atol=5.0e-5, rtol=0 + ) + np.testing.assert_allclose( + result.residuals, 1.0 - 1.0 / (4.0 * coefficient), atol=1.0e-8 + ) + + +def test_f_tol_still_accepts_small_well_predicted_progress() -> None: + def residual(x, *, indices): + del indices + return np.concatenate((x, np.ones_like(x)), axis=1) + + def jacobian(x, *, indices): + del indices + return np.broadcast_to(np.asarray([[[1.0, 0.0]]]), (len(x), 1, 2)) + + result = batched_levenberg_marquardt( + BatchedLeastSquaresProblem(residual, jacobian), + np.asarray([[1.0e-5]]), + ) + + assert result.status.tolist() == [LMStatus.CONVERGED_F_TOL] + assert result.evaluations.tolist() == [2] + np.testing.assert_allclose(result.parameters, 1.0e-5 / 1001.0) + + def test_active_batch_is_compacted_after_independent_termination() -> None: target = np.asarray([[0.0], [2.0], [-3.0]]) seen_indices: list[tuple[int, ...]] = [] @@ -516,7 +570,48 @@ def jacobian(x, *, indices): ) -def test_final_diagnostics_factorize_batch_once(monkeypatch) -> None: +def test_ill_scaled_covariance_preserves_unscaled_rank_policy() -> None: + orthogonal = ( + np.asarray( + [[1, 1, 1, 1], [1, -1, 1, -1], [1, 1, -1, -1], [1, -1, -1, 1]], + dtype=np.float64, + ) + / 2 + ) + singular_values = np.asarray([1, 0.5, 0.25, 2.0**-8]) + scales = np.asarray([2.0**-24, 1, 32, 2.0**-8]) + factor = scales[:, None] * orthogonal * singular_values[None, :] + matrix = np.tile(factor, (1, 3)) + matrices = np.stack((matrix, matrix.copy())) + # The second row remains numerically rank deficient under the existing + # unscaled threshold, even though equilibration can recover its rank. + matrices[1, 0] *= 2.0**-36 + expected = ( + ((orthogonal / singular_values**2) @ orthogonal.T) + / scales[:, None] + / scales[None, :] + / 3 + ) + + def residual(x, *, indices): + return np.einsum("kn,knm->km", x, matrices[indices]) + + def jacobian(x, *, indices): + del x + return matrices[indices] + + result = batched_levenberg_marquardt( + BatchedLeastSquaresProblem(residual, jacobian), np.zeros((2, 4)) + ) + + assert result.rank.tolist() == [4, 3] + np.testing.assert_allclose( + result.covariance[0], expected, rtol=2.0e-8, atol=2.0e-8 + ) + assert np.isnan(result.covariance[1]).all() + + +def test_final_diagnostics_factorize_each_stage_as_batch(monkeypatch) -> None: target = np.asarray( [[1.0, -2.0], [-3.0, 4.0], [5.0, -6.0]], dtype=np.float64 ) @@ -547,7 +642,7 @@ def recording_svd(value, *args, **kwargs): target.copy(), ) - assert seen_shapes == [(3, 2, 3)] + assert seen_shapes == [(3, 2, 3), (3, 2, 3)] assert result.converged.all() assert result.rank.tolist() == [2, 2, 2] assert np.isfinite(result.covariance).all() @@ -583,8 +678,10 @@ def jacobian(x, *, indices): assert np.isnan(result.gradient[1]).all() +@pytest.mark.parametrize("failing_stage", ["rank", "covariance"]) def test_final_diagnostics_fall_back_to_rows_when_batched_svd_fails( monkeypatch, + failing_stage, ) -> None: target = np.asarray( [[1.0, -2.0], [-3.0, 4.0], [5.0, -6.0]], dtype=np.float64 @@ -611,23 +708,40 @@ def jacobian(x, *, indices): def row_only_svd(value, *args, **kwargs): seen_shapes.append(value.shape) - if value.ndim == 3: + if value.ndim == 3 and ( + failing_stage == "rank" or len(seen_shapes) == 2 + ): raise RuntimeError("batched svd is unavailable") + if failing_stage == "covariance" and len(seen_shapes) == 4: + raise RuntimeError("row 1 covariance factorization failed") return original_svd(value, *args, **kwargs) monkeypatch.setattr(np.linalg, "svd", row_only_svd) result = batched_levenberg_marquardt(problem, target.copy()) - assert seen_shapes == [(3, 2, 3), (2, 3), (2, 3), (2, 3)] + if failing_stage == "rank": + assert seen_shapes == [(3, 2, 3), *[(2, 3)] * 6] + else: + # The successful batched rank factorization is not repeated per row. + assert seen_shapes == [(3, 2, 3), (3, 2, 3), *[(2, 3)] * 3] assert result.rank.tolist() == [2, 2, 2] assert np.array_equal(result.status, expected.status) - for field in ("gradient", "gn_hessian", "covariance"): + for field in ("gradient", "gn_hessian"): assert np.allclose( getattr(result, field), getattr(expected, field), rtol=1.0e-12, atol=1.0e-12, ) + if failing_stage == "covariance": + expected.covariance[1] = np.nan + np.testing.assert_allclose( + result.covariance, + expected.covariance, + rtol=1.0e-12, + atol=1.0e-12, + equal_nan=True, + ) @pytest.mark.parametrize("dtype", [np.float32, np.float64]) diff --git a/tests/xscan/test_device_pipeline.py b/tests/xscan/test_device_pipeline.py index a4f41f7a..40314b00 100644 --- a/tests/xscan/test_device_pipeline.py +++ b/tests/xscan/test_device_pipeline.py @@ -795,6 +795,20 @@ def test_scientific_evidence_config_binds_xpois_metadata_and_layout( ) +@pytest.mark.parametrize("backend", ["cupy", "cutile", "native"]) +def test_scientific_evidence_binds_selected_xfit_backend( + tmp_path: Path, backend: str +) -> None: + config = _evidence_config(tmp_path) + config = replace(config, xfit=replace(config.xfit, backend=backend)) + evidence = _scientific_evidence_payload() + evidence["xfit"]["backend"] = backend + decode_device_pipeline_evidence(evidence, config=config) + evidence["xfit"]["backend"] = "cutile" if backend == "cupy" else "cupy" + with pytest.raises(ValueError, match="xFit backend does not match"): + decode_device_pipeline_evidence(evidence, config=config) + + def test_scientific_evidence_config_binds_flux_basis_order_and_reference( tmp_path: Path, ) -> None: @@ -1366,7 +1380,9 @@ def fit(images: Any, **kwargs: Any) -> Any: assert kwargs == { "model": "gaussian", "mode": "difference", + "backend": "cupy", "config": context.solver_config, + "fusion": False, } stage_calls.append("xfit") return _fake_xfit_result(cp, 2) @@ -1748,6 +1764,50 @@ def test_changed_input_fails_before_device_ingress( _read_item_inputs(item) +def test_xfit_backend_defaults_preserve_legacy_payload_and_hash( + tmp_path: Path, +) -> None: + positional = DeviceXFitPipelineConfig("gaussian", "difference", 1e-6) + assert positional.f_tol == 1e-6 + assert positional.backend == "cupy" + assert positional.fusion is False + config = _evidence_config(tmp_path) + legacy = config.to_payload() + assert "backend" not in legacy["xfit"] + assert "fusion" not in legacy["xfit"] + explicit = config.to_payload() + explicit["xfit"].update(backend="cupy", fusion=False) + restored = DevicePipelineConfig.from_payload(explicit) + assert restored.to_payload() == legacy + assert restored.configuration_sha256 == config.configuration_sha256 + + fused = replace( + config, xfit=replace(config.xfit, backend="cutile", fusion=True) + ) + assert DevicePipelineConfig.from_payload(fused.to_payload()) == fused + assert fused.configuration_sha256 != config.configuration_sha256 + assert fused.xfit.solver_payload() == config.xfit.solver_payload() + + +@pytest.mark.parametrize( + ("options", "message"), + [ + ({"backend": "numpy"}, "xfit backend"), + ({"fusion": 1}, "fusion must be boolean"), + ({"fusion": True}, "fusion requires backend='cutile'"), + ( + {"backend": "cutile", "use_finite_difference": True}, + "does not support finite differences", + ), + ], +) +def test_xfit_config_rejects_unsupported_backend_options( + options: dict[str, Any], message: str +) -> None: + with pytest.raises((TypeError, ValueError), match=message): + DeviceXFitPipelineConfig(**options) + + def test_xfit_config_rejects_values_the_solver_cannot_use() -> None: with pytest.raises(ValueError, match="model"): DeviceXFitPipelineConfig(model="stamp") # type: ignore[arg-type] diff --git a/tests/xscan/test_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) 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