From 13b8a5882d5d8bb1170dcef1741b16368c0e90a3 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 30 Sep 2026 05:51:21 -0700 Subject: [PATCH 1/4] Add optional Numba-CUDA-MLIR Gaussian xFit backend Fuse Gaussian model and analytic Jacobian evaluation while preserving the reference weighted normal contractions and solver behavior. Keep CuPy as the default and expose the optional backend through xFit and the device pipeline, with real CLI and scientific coverage. Signed-off-by: Trent Nelson --- .github/workflows/ci.yml | 24 +- README.md | 1 + THIRD_PARTY_NOTICES.md | 2 + docs/components/xfit.md | 40 ++ examples/xfit/benchmark_gaussian.py | 2 +- pyproject.toml | 6 +- src/cuphoton/core/bulk.py | 2 +- src/cuphoton/core/dragon.py | 4 +- src/cuphoton/core/execution.py | 12 +- src/cuphoton/core/runtime.py | 14 +- src/cuphoton/xfit/_numba_mlir.py | 390 +++++++++++ src/cuphoton/xfit/_types.py | 8 +- src/cuphoton/xfit/api.py | 82 ++- src/cuphoton/xfit/backend.py | 13 +- src/cuphoton/xfit/commands.py | 3 +- src/cuphoton/xfit/executor.py | 10 +- src/cuphoton/xscan/device_pipeline.py | 30 +- .../xscan/pipeline_benchmark/runner.py | 9 +- src/cuphoton/xscan/xfit_features.py | 6 +- tests/core/test_bulk.py | 5 +- tests/core/test_dragon.py | 14 +- tests/core/test_runtime.py | 43 ++ tests/xfit/test_executor.py | 25 +- tests/xfit/test_models.py | 1 + tests/xfit/test_numba_mlir.py | 617 ++++++++++++++++++ tests/xfit/test_numba_mlir_commands.py | 210 ++++++ tests/xscan/test_numba_mlir_pipeline.py | 392 +++++++++++ tests/xscan/test_pipeline_benchmark.py | 6 +- uv.lock | 64 +- 29 files changed, 1948 insertions(+), 87 deletions(-) create mode 100644 src/cuphoton/xfit/_numba_mlir.py create mode 100644 tests/xfit/test_numba_mlir.py create mode 100644 tests/xfit/test_numba_mlir_commands.py create mode 100644 tests/xscan/test_numba_mlir_pipeline.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5c7966cc..e6c12549 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -109,7 +109,9 @@ jobs: version: "0.12.1" enable-cache: false - run: uv python install 3.12 - - run: uv sync --locked --python 3.12 --extra dev --extra gpu --extra viz + - run: >- + uv sync --locked --python 3.12 --extra dev --extra gpu --extra viz + --extra numba-cuda-mlir - name: Build native xFit with CUDA 13 compiler wheels env: CUPHOTON_XFIT_BUILD_EXT: "1" @@ -192,6 +194,25 @@ jobs: assert not report.findall(".//error") print(f"All {len(cases)} native xFit tests passed without skips") PY + - name: Run Numba-CUDA-MLIR scientific and CLI tests + run: >- + uv run --no-sync pytest -q -rs + --junitxml="${RUNNER_TEMP}/xfit-numba-mlir.xml" + tests/xfit/test_numba_mlir.py tests/xfit/test_numba_mlir_commands.py + tests/xscan/test_numba_mlir_pipeline.py::test_numba_pipeline_cli_runs_real_stages_and_matches_cupy + - name: Require Numba-CUDA-MLIR tests to execute + run: | + uv run --no-sync python - <<'PY' + import os + from pathlib import Path + import xml.etree.ElementTree as ET + + report = ET.parse(Path(os.environ["RUNNER_TEMP"]) / "xfit-numba-mlir.xml") + assert report.findall(".//testcase"), "No Numba-CUDA-MLIR tests ran" + assert not report.findall(".//skipped"), "GPU tests must not skip" + assert not report.findall(".//failure") + assert not report.findall(".//error") + PY - name: Run synthetic GPU quickstarts run: >- uv run --no-sync python examples/run_quickstarts.py @@ -204,6 +225,7 @@ jobs: path: | ${{ runner.temp }}/xfit-gpu.xml ${{ runner.temp }}/xfit-native-gpu.xml + ${{ runner.temp }}/xfit-numba-mlir.xml ${{ runner.temp }}/cuphoton-quickstart/**/summary.json if-no-files-found: error retention-days: 14 diff --git a/README.md b/README.md index 095db8f0..c78ab280 100644 --- a/README.md +++ b/README.md @@ -173,6 +173,7 @@ upstream release has no Python 3.14 wheel. | `torch` | PyTorch workflows that can be forced to CPU execution | | `gpu` | The `io` and `photometry` extras plus CUDA 13 PyTorch and Numba-CUDA | | `cutile` | Experimental `cuda.tile` backend and CuPy | +| `numba-cuda-mlir` | Gaussian xFit kernels compiled with Numba-CUDA-MLIR and CuPy | | `mpi` | mpi4py bindings for an existing MPI runtime | | `dragon` | DragonHPC runtime on Python 3.12 or 3.13 | | `viz` | Bokeh reviews and Pillow image outputs | diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index eb923b66..63cb976d 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -60,6 +60,8 @@ more than one license. | `gpu` | `torch>=2.13,<3` | `2.13.0` | `Apache-2.0 AND Apache-2.0 WITH LLVM-exception AND BSD-2-Clause AND BSD-3-Clause AND BSL-1.0 AND MIT` | [PyTorch](https://github.com/pytorch/pytorch) | `uv / PyPI` | | `cutile` | `cuda-tile>=1.6,<2` | `1.6.0` | `Apache-2.0` | [CUDA Tile](https://github.com/NVIDIA/cutile-python) | `uv / PyPI` | | `cutile` | `cupy-cuda13x[ctk]>=14,<15` | `14.1.1` | `MIT`; the `ctk` extra installs separately licensed NVIDIA CUDA component wheels | [CuPy](https://github.com/cupy/cupy) | `uv / PyPI` | +| `numba-cuda-mlir` | `numba-cuda-mlir[cu13]>=0.5.4,<0.6` | `0.5.4` | `Apache-2.0` with bundled third-party notices; the `cu13` extra installs separately licensed NVIDIA CUDA component wheels | [Numba-CUDA-MLIR](https://github.com/NVIDIA/numba-cuda-mlir) | `uv / PyPI` | +| `numba-cuda-mlir` | `cupy-cuda13x[ctk]>=14,<15` | `14.1.1` | `MIT`; the `ctk` extra installs separately licensed NVIDIA CUDA component wheels | [CuPy](https://github.com/cupy/cupy) | `uv / PyPI` | | `dev` | `setuptools>=83.0.0` | `83.0.0` | `MIT` | [setuptools](https://github.com/pypa/setuptools) | `uv / PyPI` | | `dev` | `pre-commit>=4.0` | `4.6.0` | `MIT` | [pre-commit](https://github.com/pre-commit/pre-commit) | `uv / PyPI` | | `dev` | `virtualenv>=21.7.12` | `21.7.12` | `MIT` | [virtualenv](https://github.com/pypa/virtualenv) | `uv / PyPI; pre-commit environment creation` | diff --git a/docs/components/xfit.md b/docs/components/xfit.md index 181a2589..848aac94 100644 --- a/docs/components/xfit.md +++ b/docs/components/xfit.md @@ -209,6 +209,46 @@ from end-to-end fit and process-communication timings. The default per-round wait is 180 seconds. `--timeout-seconds` allows longer cold compilation; worker cleanup may take additional time after a timeout. +### Optional Numba-CUDA-MLIR backend + +`backend="numba-cuda-mlir"` uses Python-authored CUDA kernels that fuse +Gaussian model evaluation and weighting in separate kernels for residuals +and the analytic Jacobian. CuPy contracts the weighted Jacobian to form the +gradient and Gauss–Newton Hessian. Install the optional +`cuphoton[numba-cuda-mlir]` extra on Linux with CUDA 13. It supports float32 +and float64 computation in difference and split modes, including masks and +variance weights. The existing LM solver and final analytic-Jacobian +diagnostics are retained. Sampled-stamp models and finite differences are +rejected. Automatic backend selection continues to choose NumPy or CuPy. + +Select it in the CLI with: + +```bash +cuphoton xfit fit-dipoles --input dipoles.npz --output-dir fit-run \ + --model gaussian --backend numba-cuda-mlir --compute-dtype float64 +``` + +`fit_dipoles_device(..., backend="numba-cuda-mlir")` returns CuPy-owned +arrays and reports the selected execution backend. Inputs and consumers use +the current CuPy stream; establish an event dependency before passing data +produced on another stream. Each worker thread reuses compiled kernels for +repeated fits of the same dtype and stamp dimensions. Allow for compilation +when warming a new worker or stamp shape. Each fit call retains a buffer +for the weighted Jacobian with eight values per observation and candidate, +in the selected compute dtype. Split mode has three observations per spatial +pixel. Account for this device-memory cost when choosing batch sizes. + +For a persistent device pipeline or pipeline benchmark, set +`"backend": "numba-cuda-mlir"` in the configuration's `xfit` object. + +Compare warmed CuPy and Numba-CUDA-MLIR fits on synthetic inputs with: + +```bash +uv run --locked --extra numba-cuda-mlir \ + python examples/xfit/benchmark_gaussian.py \ + --backend cupy --backend numba-cuda-mlir --dtype float64 +``` + Split mode uses diagonal per-plane weights. When the difference plane is derived from the positive and negative planes, those residuals are correlated; statistical calibration of the reported covariance requires a caller-supplied diff --git a/examples/xfit/benchmark_gaussian.py b/examples/xfit/benchmark_gaussian.py index 3a0b2189..e52cbeb9 100644 --- a/examples/xfit/benchmark_gaussian.py +++ b/examples/xfit/benchmark_gaussian.py @@ -184,7 +184,7 @@ def main() -> None: parser.add_argument( "--backend", action="append", - choices=("numpy", "cupy", "cutile", "native"), + choices=("numpy", "cupy", "cutile", "native", "numba-cuda-mlir"), ) parser.add_argument("--compare-native", action="store_true") parser.add_argument("--workers", type=int, default=1) diff --git a/pyproject.toml b/pyproject.toml index f9f35cef..52331700 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,6 +67,10 @@ cutile = [ "cuda-tile>=1.6,<2; platform_system == 'Linux'", "cupy-cuda13x[ctk]>=14,<15; platform_system == 'Linux'", ] +numba-cuda-mlir = [ + "numba-cuda-mlir[cu13]>=0.5.4,<0.6; platform_system == 'Linux'", + "cupy-cuda13x[ctk]>=14,<15; platform_system == 'Linux'", +] mpi = [ "mpi4py>=4.1.2,<5", ] @@ -163,7 +167,7 @@ show_error_codes = true module = [ "astropy.*", "bokeh.*", "cuda.*", "cupy.*", "cupyx.*", "dragon.*", "h5py.*", "kvikio.*", "mpi4py.*", "numba.*", "nvidia.*", "nvtx.*", - "pandas.*", "photutils.*", "PIL.*", "pyarrow.*", "pybind11.*", + "numba_cuda_mlir.*", "pandas.*", "photutils.*", "PIL.*", "pyarrow.*", "pybind11.*", "scipy.*", "setuptools.*", "torch.*", ] follow_imports = "skip" diff --git a/src/cuphoton/core/bulk.py b/src/cuphoton/core/bulk.py index da3c4d03..eefac187 100644 --- a/src/cuphoton/core/bulk.py +++ b/src/cuphoton/core/bulk.py @@ -344,7 +344,7 @@ def _require_unique_item_ids(items: Sequence[WorkItem]) -> None: def collect_gpu_identity(backend: str) -> dict[str, Any]: """Collect stable device identity after caller-established placement.""" - if backend in {"auto", "cupy", "cutile", "native"}: + if backend in {"auto", "cupy", "cutile", "native", "numba-cuda-mlir"}: try: return _collect_cupy_identity() except (ImportError, OSError, RuntimeError): diff --git a/src/cuphoton/core/dragon.py b/src/cuphoton/core/dragon.py index c825e92a..0d709986 100644 --- a/src/cuphoton/core/dragon.py +++ b/src/cuphoton/core/dragon.py @@ -330,7 +330,9 @@ def _valid_shard_provenance( return False gpu_backend = gpu.get("backend") expected_identity_backend = ( - "cupy" if backend in {"cupy", "cutile", "native"} else backend + "cupy" + if backend in {"cupy", "cutile", "native", "numba-cuda-mlir"} + else backend ) return ( isinstance(gpu_backend, str) diff --git a/src/cuphoton/core/execution.py b/src/cuphoton/core/execution.py index 1d75c911..a7a1fe28 100644 --- a/src/cuphoton/core/execution.py +++ b/src/cuphoton/core/execution.py @@ -43,9 +43,10 @@ class Worker(Protocol): """Worker-local state constructed after placement and CUDA binding. - ``gpu_identity`` must contain ``backend`` (``cupy`` for ``cutile`` - or ``native``), ``device_index=0``, ``identity_error=None``, and at least - one non-empty physical ``uuid`` or ``pci_bus_id`` string. + ``gpu_identity`` must contain ``backend`` (``cupy`` for ``cutile``, + ``native``, or ``numba-cuda-mlir``), ``device_index=0``, + ``identity_error=None``, and at least one non-empty physical ``uuid`` + or ``pci_bus_id`` string. """ @property @@ -118,6 +119,7 @@ def __post_init__(self) -> None: "cupy", "cutile", "native", + "numba-cuda-mlir", "numba-cuda", "torch", }: @@ -417,7 +419,9 @@ def audit_worker_provenance( invalid.append("cuda_visible_devices") ids: frozenset[tuple[str, str]] = frozenset() expected_backend = ( - "cupy" if backend in {"cutile", "native"} else backend + "cupy" + if backend in {"cutile", "native", "numba-cuda-mlir"} + else backend ) if not isinstance(gpu, Mapping): invalid.append("gpu") diff --git a/src/cuphoton/core/runtime.py b/src/cuphoton/core/runtime.py index 155330bb..96367f4e 100644 --- a/src/cuphoton/core/runtime.py +++ b/src/cuphoton/core/runtime.py @@ -36,8 +36,12 @@ def runtime_metadata( if backend == "torch": _add_torch_metadata(metadata, resolved_device) - elif backend in {"cupy", "cutile", "native"}: + elif backend in {"cupy", "cutile", "native", "numba-cuda-mlir"}: _add_cupy_metadata(metadata) + if backend == "numba-cuda-mlir": + metadata["numba_cuda_mlir_version"] = _package_version( + "numba-cuda-mlir" + ) elif backend == "numba-cuda": _add_numba_metadata(metadata) return metadata @@ -59,7 +63,13 @@ def _device_type(backend: str, device: str | None) -> str: return "cpu" if backend == "cpu": return "cpu" - if backend in {"cupy", "cutile", "native", "numba-cuda"}: + if backend in { + "cupy", + "cutile", + "native", + "numba-cuda", + "numba-cuda-mlir", + }: return "cuda" return device or "unknown" diff --git a/src/cuphoton/xfit/_numba_mlir.py b/src/cuphoton/xfit/_numba_mlir.py new file mode 100644 index 00000000..bc7fbbd3 --- /dev/null +++ b/src/cuphoton/xfit/_numba_mlir.py @@ -0,0 +1,390 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Optional Numba-CUDA-MLIR Gaussian kernels using CuPy-owned storage.""" + +from __future__ import annotations + +import math +import threading +from collections import OrderedDict +from typing import Any + +import numpy as np + +from ._types import FitMode + +_THREADS = 128 +_KERNEL_CACHE_SIZE = 16 +_worker_dispatchers = threading.local() + + +def _kernels( + cuda: Any, dtype: Any, height: int, width: int +) -> tuple[Any, Any]: + """Construct independent dispatchers with their own launch state.""" + + pixels = height * width + libdevice = cuda.libdevice + if dtype == np.float32: + add = libdevice.fadd_rn + multiply = libdevice.fmul_rn + divide = libdevice.fdiv_rn + power = libdevice.powf + else: + add = libdevice.dadd_rn + multiply = libdevice.dmul_rn + divide = libdevice.ddiv_rn + power = libdevice.pow + + # These kernels need no experimental AST features. An explicit option + # avoids optional-module auto-detection during concurrent startup. + @cuda.jit( + device=True, + inline=True, + fastmath=False, + experimental_ast_transforms=False, + ) + def components(parameters, row, observation): + amplitude = parameters[row, 0] + sigma_x = math.exp(parameters[row, 1]) + sigma_y = math.exp(parameters[row, 2]) + if not math.isfinite(sigma_x) or sigma_x <= 0: + sigma_x = dtype(math.nan) + if not math.isfinite(sigma_y) or sigma_y <= 0: + sigma_y = dtype(math.nan) + cosine = math.cos(parameters[row, 3]) + sine = math.sin(parameters[row, 3]) + inv_x2 = divide(dtype(1), multiply(sigma_x, sigma_x)) + inv_y2 = divide(dtype(1), multiply(sigma_y, sigma_y)) + pixel = observation % pixels + x = add(dtype(pixel % width), -dtype((width - 1) / 2)) + y = add(dtype(pixel // width), -dtype((height - 1) / 2)) + xp = add(x, -parameters[row, 4]) + yp = add(y, -parameters[row, 5]) + xn = add(x, -parameters[row, 6]) + yn = add(y, -parameters[row, 7]) + px = add(multiply(cosine, xp), multiply(sine, yp)) + py = add(multiply(cosine, yp), -multiply(sine, xp)) + nx = add(multiply(cosine, xn), multiply(sine, yn)) + ny = add(multiply(cosine, yn), -multiply(sine, xn)) + pu = math.exp( + multiply( + dtype(-0.5), + add( + multiply(multiply(px, px), inv_x2), + multiply(multiply(py, py), inv_y2), + ), + ) + ) + nu = math.exp( + multiply( + dtype(-0.5), + add( + multiply(multiply(nx, nx), inv_x2), + multiply(multiply(ny, ny), inv_y2), + ), + ) + ) + positive = multiply(amplitude, pu) + negative = multiply(amplitude, nu) + sigma_x_cubed = dtype(power(sigma_x, dtype(3))) + sigma_y_cubed = dtype(power(sigma_y, dtype(3))) + p1 = divide(multiply(multiply(positive, px), px), sigma_x_cubed) + p2 = divide(multiply(multiply(positive, py), py), sigma_y_cubed) + p3 = multiply( + multiply(multiply(-positive, px), py), add(inv_x2, -inv_y2) + ) + p4 = multiply( + positive, + add( + multiply(multiply(cosine, px), inv_x2), + -multiply(multiply(sine, py), inv_y2), + ), + ) + p5 = multiply( + positive, + add( + multiply(multiply(sine, px), inv_x2), + multiply(multiply(cosine, py), inv_y2), + ), + ) + n1 = divide(multiply(multiply(negative, nx), nx), sigma_x_cubed) + n2 = divide(multiply(multiply(negative, ny), ny), sigma_y_cubed) + n3 = multiply( + multiply(multiply(-negative, nx), ny), add(inv_x2, -inv_y2) + ) + n4 = multiply( + negative, + add( + multiply(multiply(cosine, nx), inv_x2), + -multiply(multiply(sine, ny), inv_y2), + ), + ) + n5 = multiply( + negative, + add( + multiply(multiply(sine, nx), inv_x2), + multiply(multiply(cosine, ny), inv_y2), + ), + ) + plane = observation // pixels + zero = dtype(0) + if plane == 1: + return ( + positive, + pu, + multiply(p1, sigma_x), + multiply(p2, sigma_y), + p3, + p4, + p5, + zero, + zero, + ) + if plane == 2: + return ( + negative, + nu, + multiply(n1, sigma_x), + multiply(n2, sigma_y), + n3, + zero, + zero, + n4, + n5, + ) + return ( + add(positive, -negative), + add(pu, -nu), + multiply(add(p1, -n1), sigma_x), + multiply(add(p2, -n2), sigma_y), + add(p3, -n3), + p4, + p5, + -n4, + -n5, + ) + + @cuda.jit(fastmath=False, experimental_ast_transforms=False) + def residual_kernel(parameters, indices, images, weights, residuals): + row = cuda.blockIdx.x + original = indices[row] + for observation in range(cuda.threadIdx.x, images.shape[1], _THREADS): + prediction = components(parameters, row, observation)[0] + residuals[row, observation] = ( + prediction - images[original, observation] + ) * weights[original, observation] + + @cuda.jit(fastmath=False, experimental_ast_transforms=False) + def jacobian_kernel(parameters, indices, weights, jacobian): + row = cuda.blockIdx.x + original = indices[row] + for observation in range( + cuda.threadIdx.x, weights.shape[1], _THREADS + ): + sample = components(parameters, row, observation) + weight = weights[original, observation] + for parameter in range(8): + jacobian[row, parameter, observation] = ( + sample[parameter + 1] * weight + ) + + return residual_kernel, jacobian_kernel + + +def _cached_kernels( + cuda: Any, dtype: Any, device: int, height: int, width: int +) -> tuple[Any, Any]: + """Reuse a worker's dispatchers for stable dtype and image dimensions. + + Each thread retains at most sixteen device/shape/dtype specializations. + Fit buffers and their stream lifetimes remain workspace-local. + """ + + cache = getattr(_worker_dispatchers, "cache", None) + if cache is None: + cache = OrderedDict() + _worker_dispatchers.cache = cache + key = (device, dtype.str, height, width) + dispatchers = cache.get(key) + if dispatchers is None: + dispatchers = _kernels(cuda, dtype.type, height, width) + cache[key] = dispatchers + if len(cache) > _KERNEL_CACHE_SIZE: + cache.popitem(last=False) + cache.move_to_end(key) + return dispatchers + + +class GaussianWorkspace: + """Own one fit's buffers and owners, using worker-local dispatchers. + + Images and weights must already be sanitized and ready on CuPy's current + stream. Every callback and its consumers must use that same stream. A + workspace belongs to one fit and must not be called concurrently. Use it + as a context manager: closing waits before releasing queued work's owners. + The weighted Jacobian is scratch storage reused on that same stream. + """ + + def __init__( + self, + images: Any, + weights: Any, + *, + image_shape: tuple[int, int], + mode: FitMode, + ) -> None: + try: + import cupy as cp + from numba_cuda_mlir import cuda + except ImportError as exc: + raise ImportError( + "Numba-CUDA-MLIR Gaussian kernels require CuPy and " + "numba-cuda-mlir" + ) from exc + if mode not in {"difference", "split"}: + raise ValueError("mode must be 'difference' or 'split'") + height, width = image_shape + if height < 1 or width < 1: + raise ValueError("image dimensions must be positive") + shape = ( + (height, width) if mode == "difference" else (3, height, width) + ) + if images.ndim != len(shape) + 1 or images.shape[1:] != shape: + raise ValueError("images do not match image_shape and mode") + if weights.shape != images.shape or weights.dtype != images.dtype: + raise ValueError("weights must have the images' shape and dtype") + if images.dtype not in (np.dtype("float32"), np.dtype("float64")): + raise TypeError("Numba-CUDA-MLIR supports float32 and float64") + self._cp = cp + self._cuda = cuda + self._stream = cp.cuda.get_current_stream() + self._device = cp.cuda.Device().id + for value in (images, weights): + if ( + not isinstance(value, cp.ndarray) + or value.device.id != self._device + ): + raise ValueError( + "images and weights must be on the active GPU" + ) + if not value.flags.c_contiguous: + raise ValueError("images and weights must be contiguous") + self._external_stream = cuda.external_stream(self._stream.ptr) + self._images = images.reshape(images.shape[0], -1) + self._weights = weights.reshape(weights.shape[0], -1) + self._dtype = images.dtype + self._jacobian = cp.empty( + (images.shape[0], 8, self._images.shape[1]), dtype=self._dtype + ) + self._stable_wrappers = { + id(value): cuda.as_cuda_array(value, sync=False) + for value in (self._images, self._weights, self._jacobian) + } + self._residual_kernel, self._jacobian_kernel = _cached_kernels( + cuda, images.dtype, self._device, height, width + ) + self._pending: list[tuple[Any, ...]] = [] + self._closed = False + + def _check(self, parameters: Any, indices: Any) -> None: + if self._closed: + raise RuntimeError("Gaussian workspace is closed") + if ( + self._cp.cuda.Device().id != self._device + or self._cp.cuda.get_current_stream().ptr != self._stream.ptr + ): + raise RuntimeError( + "Gaussian workspace requires its original CUDA stream" + ) + if ( + parameters.shape != (indices.size, 8) + or parameters.dtype != self._dtype + or indices.ndim != 1 + or indices.dtype != np.dtype("int64") + or not 0 < indices.size <= self._images.shape[0] + ): + raise ValueError( + "incompatible Gaussian parameters or row indices" + ) + for value in (parameters, indices): + if ( + value.device.id != self._device + or not value.flags.c_contiguous + ): + raise ValueError( + "callback arrays must be contiguous on the active GPU" + ) + if self._stream.done: + self._pending.clear() + + def _launch( + self, kernel: Any, rows: int, arrays: tuple[Any, ...] + ) -> None: + wrappers = tuple( + self._stable_wrappers[id(value)] + if id(value) in self._stable_wrappers + else self._cuda.as_cuda_array(value, sync=False) + for value in arrays + ) + # Retain owners before enqueue, including when a launch raises. + self._pending.append((arrays, wrappers)) + kernel[rows, _THREADS, self._external_stream](*wrappers) + + def residual(self, parameters: Any, *, indices: Any) -> Any: + """Return fresh weighted residuals in compacted candidate order.""" + + self._check(parameters, indices) + result = self._cp.empty( + (parameters.shape[0], self._images.shape[1]), dtype=self._dtype + ) + self._launch( + self._residual_kernel, + parameters.shape[0], + (parameters, indices, self._images, self._weights, result), + ) + return result + + def normal_equations( + self, parameters: Any, residuals: Any, *, indices: Any + ) -> tuple[Any, Any]: + """Return weighted gradient and Hessian in log-width coordinates.""" + + self._check(parameters, indices) + if ( + residuals.shape != (parameters.shape[0], self._images.shape[1]) + or residuals.dtype != self._dtype + or residuals.device.id != self._device + or not residuals.flags.c_contiguous + ): + raise ValueError("incompatible Gaussian residual array") + self._launch( + self._jacobian_kernel, + parameters.shape[0], + (parameters, indices, self._weights, self._jacobian), + ) + jacobian = self._jacobian[: parameters.shape[0]] + self._pending.append(((residuals, jacobian), ())) + gradient = self._cp.einsum("knm,km->kn", jacobian, residuals) + self._pending.append(((gradient,), ())) + hessian = self._cp.einsum("knm,kpm->knp", jacobian, jacobian) + self._pending.append(((hessian,), ())) + return gradient, hessian + + def close(self) -> None: + """Complete queued work before releasing any retained input owners.""" + + if not self._closed: + self._stream.synchronize() + self._pending.clear() + self._closed = True + + def __enter__(self) -> GaussianWorkspace: + if self._closed: + raise RuntimeError("Gaussian workspace is closed") + return self + + def __exit__(self, *exc: Any) -> None: + self.close() diff --git a/src/cuphoton/xfit/_types.py b/src/cuphoton/xfit/_types.py index 07ca7529..bef5b8db 100644 --- a/src/cuphoton/xfit/_types.py +++ b/src/cuphoton/xfit/_types.py @@ -14,9 +14,11 @@ # Keep aliases eagerly expanded for public get_type_hints/get_args users. FitMode: TypeAlias = Literal["difference", "split"] BackendRequest: TypeAlias = Literal[ - "auto", "numpy", "cupy", "cutile", "native" + "auto", "numpy", "cupy", "cutile", "native", "numba-cuda-mlir" +] +ResolvedBackend: TypeAlias = Literal[ + "numpy", "cupy", "cutile", "native", "numba-cuda-mlir" ] -ResolvedBackend: TypeAlias = Literal["numpy", "cupy", "cutile", "native"] ModelName: TypeAlias = Literal["gaussian", "stamp"] StampEvaluation: TypeAlias = Literal[ "bilinear", "bilinear-vignetted", "finite-volume" @@ -26,7 +28,7 @@ FIT_MODES: frozenset[FitMode] = frozenset(("difference", "split")) BACKEND_REQUESTS: frozenset[BackendRequest] = frozenset( - ("auto", "numpy", "cupy", "cutile", "native") + ("auto", "numpy", "cupy", "cutile", "native", "numba-cuda-mlir") ) MODEL_NAMES: frozenset[ModelName] = frozenset(("gaussian", "stamp")) STAMP_EVALUATIONS: frozenset[StampEvaluation] = frozenset( diff --git a/src/cuphoton/xfit/api.py b/src/cuphoton/xfit/api.py index 8e6fa971..f33c556e 100644 --- a/src/cuphoton/xfit/api.py +++ b/src/cuphoton/xfit/api.py @@ -131,14 +131,14 @@ class DeviceDipoleFitResult: init=False, default="levenberg-marquardt", ) - backend: Literal["cupy", "native"] = "cupy" + backend: Literal["cupy", "native", "numba-cuda-mlir"] = "cupy" result_location: Literal["device"] = field( init=False, default="device", ) def __post_init__(self) -> None: - if self.backend not in {"cupy", "native"}: + if self.backend not in {"cupy", "native", "numba-cuda-mlir"}: raise ValueError("device fit backend must be cupy or native") if isinstance(self.device_id, bool) or not isinstance( self.device_id, (int, np.integer) @@ -791,29 +791,41 @@ def normal_equations( mode=mode, ) - problem = BatchedLeastSquaresProblem( - residual, - jacobian_function, - normal_equations=normal_equations, - ) - if resolved.name == "native": - from ._native import fit_gaussian + execution: Any = nullcontext() + if resolved.name == "numba-cuda-mlir": + from ._numba_mlir import GaussianWorkspace - low_level = fit_gaussian( - problem, - solver_initial, - images_array, - weights, - image_shape=(height, width), - mode=mode, - config=config, + execution = GaussianWorkspace( + images_array, weights, image_shape=(height, width), mode=mode ) - else: - low_level = batched_levenberg_marquardt( - problem, - solver_initial, - config=config, + with execution as workspace: + problem = BatchedLeastSquaresProblem( + residual if workspace is None else workspace.residual, + jacobian_function, + normal_equations=( + normal_equations + if workspace is None + else workspace.normal_equations + ), ) + if resolved.name == "native": + from ._native import fit_gaussian + + low_level = fit_gaussian( + problem, + solver_initial, + images_array, + weights, + image_shape=(height, width), + mode=mode, + config=config, + ) + else: + low_level = batched_levenberg_marquardt( + problem, + solver_initial, + config=config, + ) parameters_backend = model_parameters(low_level.parameters) if gaussian_model: # An ellipse orientation is periodic modulo pi. Keep portable fit @@ -1005,7 +1017,7 @@ def _validate_execution_options( raise ValueError( "native xFit does not support custom Gaussian models" ) - if backend in {"cutile", "native"}: + if backend in {"cutile", "native", "numba-cuda-mlir"}: if isinstance(model, StampDipoleModel) or model == "stamp": raise ValueError( f"backend={backend!r} currently supports only " @@ -1075,15 +1087,17 @@ def fit_dipoles_device( variance: ArrayLike | None = None, mode: FitMode = "difference", config: LMConfig | None = None, - backend: Literal["cupy", "native"] = "cupy", + backend: Literal["cupy", "native", "numba-cuda-mlir"] = "cupy", ) -> DeviceDipoleFitResult: """Fit dipoles while keeping every result array on the active GPU. - This experimental seam uses CuPy-owned buffers. 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 - borrowed and are not mutated. + Results are CuPy arrays. The optional native backend runs Gaussian LM + iterations in CUDA C++. The optional ``numba-cuda-mlir`` backend fuses + Gaussian model evaluation and weighting in separate kernels for residuals + and the analytic Jacobian, then uses CuPy contractions for the normal + equations. 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 borrowed and are not mutated. The solver may copy per-fit status, evaluation counts and boolean diagnostic masks to the host for control flow. Parameters, residuals, @@ -1098,13 +1112,17 @@ 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") + if backend not in {"cupy", "native", "numba-cuda-mlir"}: + raise ValueError("unsupported device xFit backend") if backend == "native": from ._native import load_extension load_extension() - cp = _load_cupy() + cp = ( + resolve_backend(backend).module + if backend == "numba-cuda-mlir" + else _load_cupy() + ) active_device_id = _active_cupy_device_id(cp) for name, value in ( ("images", images), diff --git a/src/cuphoton/xfit/backend.py b/src/cuphoton/xfit/backend.py index 8b27c50d..77b47b3f 100644 --- a/src/cuphoton/xfit/backend.py +++ b/src/cuphoton/xfit/backend.py @@ -30,7 +30,7 @@ def _cupy_backend( try: import cupy as cp except ImportError as exc: - extra = "cutile" if name == "cutile" else "gpu" + extra = name if name in {"cutile", "numba-cuda-mlir"} else "gpu" raise RuntimeError( f"CuPy is not installed; install cuphoton with the {extra} extra" ) from exc @@ -72,10 +72,19 @@ def resolve_backend(name: BackendRequest = "auto") -> Backend: load_extension() backend = _cupy_backend(require_device=True, name="native") return backend + if name == "numba-cuda-mlir": + try: + from numba_cuda_mlir import cuda # noqa: F401 + except ImportError as exc: + raise RuntimeError( + "Numba-CUDA-MLIR is not installed; install cuphoton with " + "the numba-cuda-mlir extra" + ) from exc + return _cupy_backend(require_device=True, name=name) if name != "auto": raise ValueError( f"unsupported xFit backend {name!r}; expected auto, " - "numpy, cupy, cutile, or native" + "numpy, cupy, cutile, native, or numba-cuda-mlir" ) try: return _cupy_backend(require_device=True) diff --git a/src/cuphoton/xfit/commands.py b/src/cuphoton/xfit/commands.py index 8d19d3c8..83cb6c00 100644 --- a/src/cuphoton/xfit/commands.py +++ b/src/cuphoton/xfit/commands.py @@ -252,7 +252,8 @@ class OutputDirArg(PathSpecInvariant): class BackendArg(BackendInvariant): _arg = "--backend" _help = ( - "Fit backend: auto, numpy, cupy, cutile, or native. " + "Fit backend: auto, numpy, cupy, cutile, native, " + "or numba-cuda-mlir. " "[default: %default]" ) _mandatory = False diff --git a/src/cuphoton/xfit/executor.py b/src/cuphoton/xfit/executor.py index cbff8f8b..a6f3c6d7 100644 --- a/src/cuphoton/xfit/executor.py +++ b/src/cuphoton/xfit/executor.py @@ -613,15 +613,17 @@ def prepare_xfit_workload( settings = {**_DEFAULTS, **fit_options} backend = settings["backend"] - if backend not in {"cupy", "cutile", "native"}: + if backend not in {"cupy", "cutile", "native", "numba-cuda-mlir"}: raise ValueError( - "distributed xFit requires backend cupy, cutile, or native" + "distributed xFit requires backend cupy, cutile, native, or " + "numba-cuda-mlir" ) - if backend in {"cutile", "native"} and settings["model"] == "stamp": + specialized = backend in {"cutile", "native", "numba-cuda-mlir"} + if specialized and settings["model"] == "stamp": raise ValueError( f"backend={backend!r} currently supports only the Gaussian model" ) - if backend in {"cutile", "native"} and settings["use_finite_difference"]: + if specialized and settings["use_finite_difference"]: raise ValueError( f"backend={backend!r} does not support finite-difference fitting" ) diff --git a/src/cuphoton/xscan/device_pipeline.py b/src/cuphoton/xscan/device_pipeline.py index 37754c1b..d9ecd8dc 100644 --- a/src/cuphoton/xscan/device_pipeline.py +++ b/src/cuphoton/xscan/device_pipeline.py @@ -637,7 +637,7 @@ def from_payload( @dataclass(frozen=True, slots=True) class DeviceXFitPipelineConfig: - """Levenberg-Marquardt controls for Gaussian difference fitting.""" + """Backend and LM controls for Gaussian difference fitting.""" model: Literal["gaussian"] = "gaussian" mode: Literal["difference"] = "difference" @@ -650,20 +650,25 @@ 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", "native", "numba-cuda-mlir"] = "cupy" 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", "native", "numba-cuda-mlir"): + raise ValueError( + "xfit backend must be 'cupy', 'native', or 'numba-cuda-mlir'" + ) if not isinstance(self.use_finite_difference, bool): raise TypeError("xfit use_finite_difference must be boolean") - if self.backend == "native" and self.use_finite_difference: + if ( + self.backend in {"native", "numba-cuda-mlir"} + and self.use_finite_difference + ): raise ValueError( - "native xfit does not support finite differences" + f"{self.backend} xfit does not support finite differences" ) for name in ("f_tol", "x_tol", "g_tol", "finite_difference_step"): value = getattr(self, name) @@ -738,7 +743,7 @@ def solver_payload(self) -> dict[str, Any]: def from_payload( cls, payload: Mapping[str, Any] ) -> DeviceXFitPipelineConfig: - """Restore exact LM settings from JSON-compatible values.""" + """Restore settings, preserving CuPy for older configuration files.""" values = json_mapping(payload, field="device pipeline xfit config") values.setdefault("backend", "cupy") @@ -1079,10 +1084,8 @@ def _device_pipeline_evidence_layout_contract( "mode": "difference", "dtype": "float64", } - if xfit["backend"] not in ("cupy", "native"): - raise ValueError( - "scientific evidence xfit.backend must be 'cupy' or 'native'" - ) + if xfit["backend"] not in ("cupy", "native", "numba-cuda-mlir"): + raise ValueError("scientific evidence xfit.backend is unsupported") for name, expected in expected_xfit.items(): if xfit[name] != expected: raise ValueError( @@ -2158,6 +2161,7 @@ def _pack_scientific_evidence( candidate_count: int, kernel_shape: tuple[int, int], flux_conserve: bool, + xfit_backend: Literal["cupy", "native", "numba-cuda-mlir"] = "cupy", ) -> _PackedScientificEvidence: """Pack all parity evidence on the producer stream as float64.""" @@ -2183,7 +2187,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 +3357,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..c654d95f 100644 --- a/src/cuphoton/xscan/pipeline_benchmark/runner.py +++ b/src/cuphoton/xscan/pipeline_benchmark/runner.py @@ -427,6 +427,7 @@ def provenance() -> dict[str, Any]: "xfit/api.py", "xfit/backend.py", "xfit/_native.py", + "xfit/_numba_mlir.py", "xfit/models.py", "xfit/solver.py", "xscan/xfit_features.py", @@ -437,7 +438,13 @@ 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", + "numba-cuda-mlir", + *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..73f1f7da 100644 --- a/src/cuphoton/xscan/xfit_features.py +++ b/src/cuphoton/xscan/xfit_features.py @@ -1292,10 +1292,14 @@ 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", "native", "numba-cuda-mlir"} or result.result_location != "device" ): raise ValueError("xFit result must be a CuPy device result") + if result.backend == "numba-cuda-mlir" and result.model != "gaussian": + raise ValueError( + "Numba-CUDA-MLIR xFit results require a 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_bulk.py b/tests/core/test_bulk.py index 7e727067..586d7f1d 100644 --- a/tests/core/test_bulk.py +++ b/tests/core/test_bulk.py @@ -353,8 +353,10 @@ def PCI_DOMAIN_ID(self): ) +@pytest.mark.parametrize("backend", ["cupy", "cutile", "numba-cuda-mlir"]) def test_cupy_gpu_identity_uses_normalized_pci_address( monkeypatch: pytest.MonkeyPatch, + backend: str, ) -> None: runtime = SimpleNamespace( getDevice=lambda: 0, @@ -370,8 +372,9 @@ def test_cupy_gpu_identity_uses_normalized_pci_address( SimpleNamespace(cuda=SimpleNamespace(runtime=runtime)), ) - identity = collect_gpu_identity("cupy") + identity = collect_gpu_identity(backend) + assert identity["backend"] == "cupy" assert identity["pci_bus_id"] == "00000000:C1:00.0" assert identity["identity_warnings"] == [] diff --git a/tests/core/test_dragon.py b/tests/core/test_dragon.py index 9aa812f9..499d8bbe 100644 --- a/tests/core/test_dragon.py +++ b/tests/core/test_dragon.py @@ -34,7 +34,7 @@ def __init__(self, options): _STATE.factories.append(self.placement.worker_id) _STATE.instances.append(self) self.gpu_identity = { - "backend": options["backend"], + "backend": "cupy", "device_index": 0, "name": "fake-gpu", "uuid": ( @@ -256,6 +256,7 @@ def _run( workers_per_gpu=1, items=None, mps_pipe_directory=None, + backend="cupy", ): spec = WorkloadSpec( items=items @@ -263,11 +264,11 @@ def _run( WorkItem("one", {"value": "x" * 300_000 if large else 1}, 12), WorkItem("two", {"value": 2}, 6), ), - options_payload={"backend": "cupy"}, + options_payload={"backend": backend}, manifest_payload={"schema": "example"}, input_identity_payload={"schema": "example-identity"}, manifest_sha256="0" * 64, - backend="cupy", + backend=backend, worker_factory=_factory, success_record_validator=validator, finalize_round=finalizer, @@ -591,9 +592,12 @@ def test_native_workers_keep_spawn_children_on_standard_multiprocessing( assert completed.returncode == 0, completed.stderr -def test_ordinary_workload_retains_root_artifacts(monkeypatch, tmp_path): +@pytest.mark.parametrize("backend", ["cupy", "cutile", "numba-cuda-mlir"]) +def test_ordinary_workload_retains_root_artifacts( + monkeypatch, tmp_path, backend +): state = _install_runtime(monkeypatch) - result = _run(tmp_path) + result = _run(tmp_path, backend=backend) assert result.status == "success", result.summary assert "benchmark" not in result.summary assert (result.run_dir / "records/one.json").is_file() diff --git a/tests/core/test_runtime.py b/tests/core/test_runtime.py index 57d12336..f9e531cb 100644 --- a/tests/core/test_runtime.py +++ b/tests/core/test_runtime.py @@ -4,7 +4,11 @@ from __future__ import annotations +import sys +from types import SimpleNamespace + from cuphoton import __version__ +from cuphoton.core import runtime from cuphoton.core.runtime import runtime_metadata @@ -17,3 +21,42 @@ def test_cpu_runtime_metadata_is_self_describing() -> None: assert metadata["dtype"] == "float64" assert metadata["python_version"] assert metadata["numpy_version"] + + +def test_numba_mlir_runtime_keeps_cupy_device_and_compiler_version( + monkeypatch, +) -> None: + cuda_runtime = SimpleNamespace( + getDevice=lambda: 0, + getDeviceProperties=lambda index: { + "name": b"test GPU", + "major": 12, + "minor": 0, + }, + driverGetVersion=lambda: 13000, + runtimeGetVersion=lambda: 13000, + ) + monkeypatch.setitem( + sys.modules, + "cupy", + SimpleNamespace( + __version__="14.0.0", cuda=SimpleNamespace(runtime=cuda_runtime) + ), + ) + monkeypatch.setattr( + runtime, + "_package_version", + lambda distribution: { + "numpy": "2.3.0", + "numba-cuda-mlir": "0.5.4", + }.get(distribution), + ) + + metadata = runtime_metadata(backend="numba-cuda-mlir", dtype="float64") + + assert metadata["backend"] == "numba-cuda-mlir" + assert metadata["device"] == "cuda" + assert metadata["device_name"] == "test GPU" + assert metadata["cupy_version"] == "14.0.0" + assert metadata["numba_cuda_mlir_version"] == "0.5.4" + assert metadata["compute_capability"] == "12.0" diff --git a/tests/xfit/test_executor.py b/tests/xfit/test_executor.py index 6f058f81..53e5cd8c 100644 --- a/tests/xfit/test_executor.py +++ b/tests/xfit/test_executor.py @@ -65,8 +65,9 @@ def _input(tmp_path, *, mode="difference", auxiliary="candidate"): ("gaussian", True, "does not support finite-difference fitting"), ], ) -def test_preflight_rejects_unsupported_cutile_settings( - tmp_path, monkeypatch, model, finite_difference, message +@pytest.mark.parametrize("backend", ["cutile", "numba-cuda-mlir"]) +def test_preflight_rejects_unsupported_specialized_settings( + tmp_path, monkeypatch, model, finite_difference, message, backend ): def unexpected_load(*args, **kwargs): pytest.fail("invalid settings must fail before reading input") @@ -76,7 +77,7 @@ def unexpected_load(*args, **kwargs): executor.prepare_xfit_workload( input_path=tmp_path / "missing.npz", fit_options={ - "backend": "cutile", + "backend": backend, "model": model, "use_finite_difference": finite_difference, }, @@ -85,9 +86,7 @@ def unexpected_load(*args, **kwargs): @pytest.mark.parametrize("backend", ["auto", "numpy"]) def test_preflight_rejects_cpu_backend_before_input_reads(tmp_path, backend): - with pytest.raises( - ValueError, match="requires backend cupy, cutile, or native" - ): + with pytest.raises(ValueError, match="requires backend cupy"): executor.prepare_xfit_workload( input_path=tmp_path / "missing.npz", fit_options={"backend": backend}, @@ -146,18 +145,20 @@ def cpu_fit(*args, **kwargs): worker.close() -def test_preflight_accepts_analytic_gaussian_cutile(tmp_path): +@pytest.mark.parametrize("backend", ["cutile", "numba-cuda-mlir"]) +def test_preflight_accepts_analytic_gaussian_specialized(tmp_path, backend): spec = executor.prepare_xfit_workload( - input_path=_input(tmp_path), fit_options={"backend": "cutile"} + input_path=_input(tmp_path), fit_options={"backend": backend} ) - assert spec.backend == "cutile" + assert spec.backend == backend @pytest.mark.parametrize( "runtime,rank", [("dragon", 0), ("mpi", 0), ("mpi", 1)] ) +@pytest.mark.parametrize("backend", ["cupy", "numba-cuda-mlir"]) def test_cli_dispatches_collective_preflight_and_preserves_options( - tmp_path, monkeypatch, capsys, runtime, rank + tmp_path, monkeypatch, capsys, runtime, rank, backend ): from types import SimpleNamespace @@ -197,7 +198,7 @@ def run(**kwargs): "--model", "gaussian", "--backend", - "cupy", + backend, "--executor", runtime, "--chunk-size", @@ -221,6 +222,8 @@ def run(**kwargs): assert kwargs["benchmark"].warmup_rounds == 1 assert kwargs["benchmark"].measure_rounds == 2 assert len(spec.items) == 2 + assert spec.backend == backend + assert spec.options_payload["fit_options"]["backend"] == backend assert spec.options_payload["fit_options"]["max_evaluations"] == 10 assert json.loads(capsys.readouterr().out)["status"] == "success" diff --git a/tests/xfit/test_models.py b/tests/xfit/test_models.py index 13adbc3d..a8eeb2a7 100644 --- a/tests/xfit/test_models.py +++ b/tests/xfit/test_models.py @@ -211,6 +211,7 @@ def test_public_fit_option_hints_are_literal_and_cupy_optional() -> None: "cupy", "cutile", "native", + "numba-cuda-mlir", } diff --git a/tests/xfit/test_numba_mlir.py b/tests/xfit/test_numba_mlir.py new file mode 100644 index 00000000..fdc09f66 --- /dev/null +++ b/tests/xfit/test_numba_mlir.py @@ -0,0 +1,617 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Scientific and stream-lifetime coverage for the optional MLIR backend.""" + +from __future__ import annotations + +import gc +import sys +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier +from types import ModuleType + +import numpy as np +import pytest + +from cuphoton.xfit import ( + GaussianDipoleModel, + LMConfig, + LMStatus, + fit_dipoles, + fit_dipoles_device, +) + +_BACKEND = "numba-cuda-mlir" + + +@pytest.fixture(scope="module") +def gpu(): + cp = pytest.importorskip("cupy") + pytest.importorskip("numba_cuda_mlir") + try: + count = cp.cuda.runtime.getDeviceCount() + except cp.cuda.runtime.CUDARuntimeError as exc: + if exc.status != 100: # cudaErrorNoDevice + raise + pytest.skip("CUDA device is unavailable") + if count < 1: + pytest.skip("CUDA device is unavailable") + return cp + + +def test_dispatcher_creation_ignores_partial_experimental_import(monkeypatch): + cuda = pytest.importorskip("numba_cuda_mlir.cuda") + from cuphoton.xfit._numba_mlir import _kernels + + # Load compiler prerequisites through the real public JIT, without + # compiling a signature or initializing a device. + _kernels(cuda, np.float64, 7, 9) + name = "numba_cuda_mlir.cuda.experimental" + # Another thread can publish this module before defining its helpers. + monkeypatch.setitem(sys.modules, name, ModuleType(name)) + for kernel in _kernels(cuda, np.float64, 7, 9): + assert not kernel.signatures + + +def _truth(dtype=np.float64): + return np.asarray( + [ + [5.0, 1.3, 1.8, 0.25, -2.0, 0.5, 2.2, -0.7], + [3.0, 1.6, 1.1, -0.3, -1.0, -1.0, 2.0, 1.0], + [4.2, 1.4, 2.0, 0.1, -1.7, 0.8, 1.9, -0.9], + ], + dtype=dtype, + ) + + +def _tolerance(dtype): + return 3.0e-5 if np.dtype(dtype) == np.dtype(np.float32) else 5.0e-12 + + +def _assert_fit_equal(actual, expected, dtype): + 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 + tolerance = _tolerance(dtype) + for name in ( + "parameters", + "residuals", + "residual_norm", + "chi_square", + "valid_pixel_fraction", + "null_chi_square", + "delta_chi_square", + "fractional_null_improvement", + "reduced_chi_square", + "covariance", + "standard_errors", + ): + np.testing.assert_allclose( + getattr(actual, name), + getattr(expected, name), + rtol=tolerance, + atol=tolerance, + equal_nan=True, + err_msg=name, + ) + + +@pytest.mark.parametrize("mode", ["difference", "split"]) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_gaussian_fit_preserves_scientific_outputs(gpu, mode, dtype): + truth = _truth(dtype)[:2] + model = GaussianDipoleModel((11, 15), dtype=dtype) + images = model.evaluate(truth, mode=mode) + initial = truth + np.asarray( + [0.2, 0.1, -0.1, 0.03, 0.1, -0.1, -0.1, 0.1], dtype=dtype + ) + # Explicit variance keeps covariance appreciable for noiseless images. + options = dict( + model=model, + initial=initial, + variance=np.full(images.shape, 0.5, dtype=dtype), + mode=mode, + ) + expected = fit_dipoles(images, backend="cupy", **options) + actual = fit_dipoles(images, backend=_BACKEND, **options) + + assert actual.backend == _BACKEND + assert actual.converged.all() + assert actual.uncertainty_valid.all() + assert np.abs(expected.covariance).max() > 1.0e-3 + _assert_fit_equal(actual, expected, dtype) + + +@pytest.mark.parametrize("mode", ["difference", "split"]) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_fused_boundary_preserves_weights_and_compacted_indices( + gpu, mode, dtype +): + from cuphoton.xfit._numba_mlir import GaussianWorkspace + + cp = gpu + rng = np.random.default_rng(1219) + physical = cp.asarray(_truth(dtype)) + parameters = physical.copy() + parameters[:, 1:3] = cp.log(physical[:, 1:3]) + model = GaussianDipoleModel((8, 11), backend="cupy", dtype=dtype) + prediction = model.evaluate(physical, mode=mode) + 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.reshape(3, -1)[:, ::7] = 0 + residuals = ((prediction - images) * weights).reshape(3, -1) + jacobian = model.jacobian(physical, mode=mode) + chain_shape = (3, 2) + (1,) * (jacobian.ndim - 2) + jacobian[:, 1:3] *= physical[:, 1:3].reshape(chain_shape) + jacobian = (jacobian * weights[:, None, ...]).reshape(3, 8, -1) + tolerance = _tolerance(dtype) + + # FP64 inputs and sums are deliberately not exactly representable in + # FP32: reductions must not narrow their intermediate values. + with GaussianWorkspace( + images, weights, image_shape=(8, 11), mode=mode + ) as workspace: + for rows in ((0, 1, 2), (2, 0), (1,)): + indices = cp.asarray(rows, dtype=cp.int64) + active = parameters[indices] + actual_residual = workspace.residual(active, indices=indices) + gradient, hessian = workspace.normal_equations( + active, actual_residual, indices=indices + ) + expected_gradient = cp.einsum( + "knm,km->kn", jacobian[indices], residuals[indices] + ) + expected_hessian = cp.einsum( + "knm,kpm->knp", jacobian[indices], jacobian[indices] + ) + for actual, expected in ( + (actual_residual, residuals[indices]), + (gradient, expected_gradient), + (hessian, expected_hessian), + ): + np.testing.assert_allclose( + cp.asnumpy(actual), + cp.asnumpy(expected), + rtol=tolerance, + atol=tolerance, + ) + + +def test_fused_boundary_preserves_cancelling_hessian_terms(gpu): + from cuphoton.xfit._numba_mlir import GaussianWorkspace + + cp = gpu + shape = (51, 51) + parameters = cp.asarray( + [[2**20, 8.0, 4.0, 0.0, -0.25, 0.0, 0.25, 0.0]], + dtype=cp.float64, + ) + parameters[:, 1:3] = cp.log(parameters[:, 1:3]) + images = cp.zeros((1, *shape), dtype=cp.float64) + coordinate = cp.arange(51, dtype=cp.float64) - 25 + radius_squared = coordinate[:, None] ** 2 + coordinate[None, :] ** 2 + weights = (1 + radius_squared[None, ...] / 4096).copy() + indices = cp.asarray([0], dtype=cp.int64) + + physical = parameters.copy() + physical[:, 1:3] = cp.exp(parameters[:, 1:3]) + model = GaussianDipoleModel(shape, backend="cupy", dtype=np.float64) + jacobian = model.jacobian(physical, mode="difference") + jacobian[:, 1:3] *= physical[:, 1:3, None, None] + jacobian = (jacobian * weights[:, None, ...]).reshape(1, 8, -1) + residual = cp.ones((1, 51 * 51), dtype=cp.float64) + expected_gradient = cp.einsum("knm,km->kn", jacobian, residual) + expected_hessian = cp.einsum("knm,kpm->knp", jacobian, jacobian) + + # Symmetry cancels products of bright-lobe position derivatives. + absolute_sum = float( + cp.sum(cp.abs(jacobian[0, 4] * jacobian[0, 5])).item() + ) + assert absolute_sum > 1e8 + assert abs(float(expected_hessian[0, 4, 5].item())) < 1e-12 * absolute_sum + + with GaussianWorkspace( + images, weights, image_shape=shape, mode="difference" + ) as workspace: + gradient, hessian = workspace.normal_equations( + parameters, residual, indices=indices + ) + for actual, expected in ( + (gradient, expected_gradient), + (hessian, expected_hessian), + ): + assert actual.shape == expected.shape + assert actual.dtype == expected.dtype + np.testing.assert_allclose( + cp.asnumpy(actual), + cp.asnumpy(expected), + rtol=_tolerance(np.float64), + atol=_tolerance(np.float64), + equal_nan=True, + ) + + +def test_fit_masks_variance_and_independent_termination(gpu): + model = GaussianDipoleModel( + (11, 15), split_weights=(1.0, 0.25, 1.75), dtype=np.float64 + ) + truth = _truth() + images = model.evaluate(truth, mode="split") + initial = truth + np.asarray([0.2, 0.1, -0.1, 0.03, 0.1, -0.1, -0.1, 0.1]) + initial[0] = truth[0] + mask = np.ones(images.shape, dtype=bool) + mask[0, 0, 0, :3] = False + mask[1, 1, 2:4, 5:8] = False + mask[2, 2, 7:, 10:] = False + y, x = np.mgrid[:11, :15] + variance = np.empty(images.shape) + for candidate in range(3): + for plane in range(3): + variance[candidate, plane] = ( + 0.75 + 0.03 * x + 0.05 * y + 0.2 * candidate + 0.1 * plane + ) + images[~mask] = np.nan + variance[~mask] = np.nan + options = dict( + model=model, + initial=initial, + mask=mask, + variance=variance, + mode="split", + ) + expected = fit_dipoles(images, backend="cupy", **options) + actual = fit_dipoles(images, backend=_BACKEND, **options) + + assert actual.evaluations[0] == 1 + assert np.all(actual.evaluations[1:] > 1) + _assert_fit_equal(actual, expected, np.float64) + + +@pytest.mark.parametrize("failure", ["rank", "budget"]) +def test_fit_preserves_failure_diagnostics(gpu, failure): + model = GaussianDipoleModel((7, 9)) + truth = _truth()[:1] + if failure == "rank": + truth[:, 0] = 0 + initial = truth.copy() + config = LMConfig() + else: + initial = truth + np.asarray( + [0.3, 0.1, -0.1, 0.05, 0.1, -0.1, -0.1, 0.1] + ) + config = LMConfig(max_evaluations=1) + images = model.evaluate(truth) + options = dict(model=model, initial=initial, config=config) + expected = fit_dipoles(images, backend="cupy", **options) + actual = fit_dipoles(images, backend=_BACKEND, **options) + + assert not actual.uncertainty_valid.any() + assert np.isnan(actual.covariance).all() + if failure == "budget": + assert actual.status.tolist() == ["max_evaluations"] + assert actual.evaluations.tolist() == [1] + else: + assert actual.uncertainty_reason == ("rank-deficient Jacobian",) + _assert_fit_equal(actual, expected, np.float64) + + +@pytest.mark.parametrize( + "log_width", [1000.0, -1000.0, np.inf, -np.inf, np.nan] +) +def test_nonfinite_trial_widths_produce_invalid_residuals(gpu, log_width): + from cuphoton.xfit._numba_mlir import GaussianWorkspace + + cp = gpu + model = GaussianDipoleModel((7, 9), backend="cupy") + physical = cp.asarray(_truth()[:1]) + parameters = physical.copy() + parameters[:, 1:3] = log_width + physical[:, 1:3] = cp.exp(parameters[:, 1:3]) + images = cp.zeros((1, 7, 9), dtype=cp.float64) + weights = cp.ones_like(images) + weights[:, 0, 0] = 0 + expected = ( + model._evaluate_positive_unchecked(physical, mode="difference") + * weights + ).reshape(1, -1) + with GaussianWorkspace( + images, weights, image_shape=(7, 9), mode="difference" + ) as workspace: + indices = cp.asarray([0], dtype=cp.int64) + actual = workspace.residual(parameters, indices=indices) + gradient, hessian = workspace.normal_equations( + parameters, actual, indices=indices + ) + np.testing.assert_array_equal( + cp.asnumpy(actual), cp.asnumpy(expected) + ) + assert not cp.isfinite(actual).any().item() + assert not cp.isfinite(gradient).any().item() + assert not cp.isfinite(hessian).any().item() + + +@pytest.mark.parametrize("mode", ["difference", "split"]) +@pytest.mark.parametrize( + ("dtype", "sigma"), [(np.float32, 1e-16), (np.float64, 1e-110)] +) +def test_finite_tiny_widths_preserve_sigma_cube_semantics( + gpu, mode, dtype, sigma +): + from cuphoton.xfit._numba_mlir import GaussianWorkspace + + cp = gpu + initial = np.asarray( + [[1, sigma, sigma, 0, -0.25, 0.25, 0.25, -0.25]], dtype=dtype + ) + parameters = cp.asarray(initial) + parameters[:, 1:3] = cp.log(parameters[:, 1:3]) + physical = parameters.copy() + physical[:, 1:3] = cp.exp(parameters[:, 1:3]) + model = GaussianDipoleModel((2, 2), backend="cupy", dtype=dtype) + prediction = model.evaluate(physical, mode=mode) + images = cp.zeros_like(prediction) + mask = cp.ones(images.shape, dtype=cp.bool_) + mask.reshape(1, -1)[:, 0] = False + weights = mask.astype(dtype) + if mode == "split": + weights *= cp.asarray(model.split_weights, dtype=dtype).reshape( + 1, 3, 1, 1 + ) + residuals = (prediction * weights).reshape(1, -1) + jacobian = model.jacobian(physical, mode=mode) + chain_shape = (1, 2) + (1,) * (jacobian.ndim - 2) + jacobian[:, 1:3] *= physical[:, 1:3].reshape(chain_shape) + jacobian = (jacobian * weights[:, None, ...]).reshape(1, 8, -1) + gradient = cp.einsum("knm,km->kn", jacobian, residuals) + hessian = cp.einsum("knm,kpm->knp", jacobian, jacobian) + + # Finite positive widths yield finite residuals, but sigma**3 rounds to + # zero. Preserve the reference's 0/0 derivatives, including masked pixels; + # cancelling sigma algebraically would change the solver's failure status. + assert cp.isfinite(physical).all().item() + assert cp.all(physical[:, 1:3] ** 3 == 0).item() + assert cp.isfinite(residuals).all().item() + assert not cp.isfinite(jacobian).all().item() + with GaussianWorkspace( + images, weights, image_shape=(2, 2), mode=mode + ) as workspace: + indices = cp.asarray([0], dtype=cp.int64) + actual_residuals = workspace.residual(parameters, indices=indices) + actual_gradient, actual_hessian = workspace.normal_equations( + parameters, actual_residuals, indices=indices + ) + for actual, expected in ( + (actual_residuals, residuals), + (actual_gradient, gradient), + (actual_hessian, hessian), + ): + np.testing.assert_array_equal( + cp.asnumpy(actual), cp.asnumpy(expected) + ) + + options = dict(model=model, initial=initial, mask=mask, mode=mode) + expected_fit = fit_dipoles(images, backend="cupy", **options) + actual_fit = fit_dipoles(images, backend=_BACKEND, **options) + assert expected_fit.status.tolist() == ["invalid_residual"] + _assert_fit_equal(actual_fit, expected_fit, dtype) + + +@pytest.mark.parametrize( + ("dtype", "amplitude"), [(np.float32, 1e25), (np.float64, 1e200)] +) +@pytest.mark.parametrize("nonzero_residual", [False, True]) +def test_finite_jacobian_overflow_preserves_solver_outcomes( + gpu, dtype, amplitude, nonzero_residual +): + cp = gpu + model = GaussianDipoleModel((5, 7), backend="cupy", dtype=dtype) + # Coincident stars have zero difference but finite position derivatives. + # Only the middle row overflows the normal equations. + 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], + ) + + options = dict(model=model, initial=initial) + expected = fit_dipoles(images, backend="cupy", **options) + actual = fit_dipoles(images, backend=_BACKEND, **options) + assert expected.status.tolist() == [ + "converged_g_tol", + "singular" if nonzero_residual else "converged_g_tol", + "converged_g_tol", + ] + np.testing.assert_array_equal(expected.evaluations, [1, 1, 1]) + np.testing.assert_array_equal(actual.parameters, expected.parameters) + _assert_fit_equal(actual, expected, dtype) + + +def test_device_fit_retains_results_after_repeated_concurrent_calls(gpu): + cp = gpu + workers = 4 + barrier = Barrier(workers) + model = GaussianDipoleModel((11, 15)) + truth = _truth()[:2] + images = model.evaluate(truth) + initial = truth + np.asarray([0.2, 0.1, -0.1, 0.03, 0.1, -0.1, -0.1, 0.1]) + cases = {} + for worker in range(workers): + for iteration in range(3): + scale = 1.0 + 0.1 * worker + 0.05 * iteration + scaled_initial = initial.copy() + scaled_initial[:, 0] *= scale + options = dict( + initial=scaled_initial, + variance=np.ones_like(images) * scale**2, + ) + expected = fit_dipoles( + images * scale, model="gaussian", backend=_BACKEND, **options + ) + cases[worker, iteration] = (images * scale, options, expected) + device_id = cp.cuda.runtime.getDevice() + + def fit_on_owned_stream(worker): + with cp.cuda.Device(device_id), cp.cuda.Stream(non_blocking=True): + barrier.wait(timeout=120) + results = [] + for iteration in range(3): + # Different amplitudes expose accidental cross-worker or + # cross-call reuse of model observations and scratch arrays. + scaled_images, options, _ = cases[worker, iteration] + result = fit_dipoles_device( + cp.asarray(scaled_images), + model="gaussian", + initial=cp.asarray(options["initial"]), + variance=cp.asarray(options["variance"]), + backend=_BACKEND, + ) + results.append(result) + gc.collect() + portable = [] + for result in results: + assert result.backend == _BACKEND + assert isinstance(result.parameters, cp.ndarray) + portable.append( + ( + cp.asnumpy(result.parameters), + cp.asnumpy(result.status_codes), + cp.asnumpy(result.evaluations), + cp.asnumpy(result.covariance), + ) + ) + return portable + + with ThreadPoolExecutor(max_workers=workers) as executor: + results = list(executor.map(fit_on_owned_stream, range(workers))) + for worker, worker_results in enumerate(results): + for iteration, result in enumerate(worker_results): + parameters, status_codes, evaluations, covariance = result + _, _, expected = cases[worker, iteration] + # Scheduling alone must not change any numerical result. + np.testing.assert_array_equal(parameters, expected.parameters) + np.testing.assert_array_equal(covariance, expected.covariance) + np.testing.assert_array_equal(evaluations, expected.evaluations) + status = [ + LMStatus(int(code)).name.lower() for code in status_codes + ] + assert status == expected.status.tolist() + + +def test_workspace_orders_delayed_producer_and_retains_input_owners(gpu): + from cuphoton.xfit._numba_mlir import GaussianWorkspace + + cp = gpu + shape = (8, 11) + model = GaussianDipoleModel(shape) + truth = _truth()[:2] + expected = model.evaluate(truth) + jacobian = model.jacobian(truth) + jacobian[:, 1:3] *= truth[:, 1:3, None, None] + jacobian = jacobian.reshape(2, 8, -1) + expected_gradient = np.einsum( + "knm,km->kn", jacobian, expected.reshape(2, -1) + ) + expected_hessian = np.einsum("knm,kpm->knp", jacobian, jacobian) + log_parameters = truth.copy() + log_parameters[:, 1:3] = np.log(log_parameters[:, 1:3]) + delay = cp.RawKernel( + r""" + extern "C" __global__ void delay(unsigned int repetitions) { + for (unsigned int i = 0; i < repetitions; ++i) { + __nanosleep(1000000); + } + } + """, + "delay", + ) + producer = cp.cuda.Stream(non_blocking=True) + consumer = cp.cuda.Stream(non_blocking=True) + with consumer: + delay((1,), (1,), (np.uint32(0),)) + with GaussianWorkspace( + cp.zeros_like(cp.asarray(expected)), + cp.ones_like(cp.asarray(expected)), + image_shape=shape, + mode="difference", + ) as warm: + parameters = cp.asarray(log_parameters) + indices = cp.arange(2, dtype=cp.int64) + residuals = warm.residual(parameters, indices=indices) + warm.normal_equations(parameters, residuals, indices=indices) + + with producer: + source_parameters = cp.asarray(log_parameters) + source_indices = cp.arange(2, dtype=cp.int64) + producer.synchronize() + for _ in range(3): + with producer: + images = cp.empty(expected.shape, dtype=cp.float64) + weights = cp.empty_like(images) + parameters = cp.empty(log_parameters.shape, dtype=cp.float64) + indices = cp.empty(2, dtype=cp.int64) + delay((1,), (1,), (np.uint32(100),)) + images.fill(0) + weights.fill(1) + cp.copyto(parameters, source_parameters) + cp.copyto(indices, source_indices) + ready = cp.cuda.Event() + ready.record(producer) + with consumer: + consumer.wait_event(ready) + with GaussianWorkspace( + images, weights, image_shape=shape, mode="difference" + ) as workspace: + residuals = workspace.residual(parameters, indices=indices) + gradient, hessian = workspace.normal_equations( + parameters, residuals, indices=indices + ) + del images, weights, parameters, indices + gc.collect() + # Closing waits for pending work before relinquishing owners; + # the returned arrays remain usable after the workspace exits. + np.testing.assert_allclose( + cp.asnumpy(residuals).reshape(expected.shape), + expected, + rtol=5.0e-12, + atol=5.0e-12, + ) + np.testing.assert_allclose( + cp.asnumpy(gradient), + expected_gradient, + rtol=5.0e-12, + atol=5.0e-12, + ) + np.testing.assert_allclose( + cp.asnumpy(hessian), + expected_hessian, + rtol=5.0e-12, + atol=5.0e-12, + ) diff --git a/tests/xfit/test_numba_mlir_commands.py b/tests/xfit/test_numba_mlir_commands.py new file mode 100644 index 00000000..29da6cd6 --- /dev/null +++ b/tests/xfit/test_numba_mlir_commands.py @@ -0,0 +1,210 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import sys + +import numpy as np +import pyarrow.parquet as pq +import pytest +import yaml + +from cuphoton.core.cli import run_component +from cuphoton.xfit import GaussianDipoleModel, StampDipoleModel + + +def _write_input(path, *, mode="difference", model="gaussian"): + if model == "gaussian": + selected = 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], + [3.0, 1.6, 1.1, -0.3, -1.0, -1.0, 2.0, 1.0], + ] + ) + initial = truth + np.asarray( + [0.2, 0.1, -0.1, 0.03, 0.1, -0.1, -0.1, 0.1] + ) + extra = {} + else: + y, x = np.mgrid[-3:4, -2:3] + basis = np.exp(-0.5 * (x * x + (y / 1.2) ** 2)) + selected = StampDipoleModel(basis, image_shape=(11, 15)) + truth = np.asarray([[-2.1, 0.6, 2.2, -0.4, 5.0]]) + initial = truth + extra = {"stamp_basis": basis} + images = selected.evaluate(truth, mode=mode) + mask = np.ones(images.shape, dtype=bool) + mask[..., 0, 0] = False + np.savez_compressed( + path, + candidate_id=np.asarray( + [f"candidate-{i}" for i in range(len(truth))] + ), + images=images, + initial=initial, + mask=mask, + variance=np.full_like(images, 0.5), + **extra, + ) + + +def _arguments( + input_path, output_dir, *, backend, model="gaussian", mode="difference" +): + return [ + "fit-dipoles", + "--input", + str(input_path), + "--output-dir", + str(output_dir), + "--model", + model, + "--mode", + mode, + "--backend", + backend, + ] + + +@pytest.mark.parametrize("mode", ["difference", "split"]) +def test_numba_mlir_cli_fits_and_writes_matching_artifacts( + tmp_path, capsys, mode +): + cp = pytest.importorskip("cupy") + pytest.importorskip("numba_cuda_mlir.cuda") + try: + count = cp.cuda.runtime.getDeviceCount() + except cp.cuda.runtime.CUDARuntimeError as exc: + if exc.status != 100: # cudaErrorNoDevice + raise + pytest.skip("CUDA device is unavailable") + if count < 1: + pytest.skip("CUDA device is unavailable") + input_path = tmp_path / "synthetic.npz" + _write_input(input_path, mode=mode) + tables = {} + for backend in ("cupy", "numba-cuda-mlir"): + output = tmp_path / backend + rc = run_component( + "xfit", _arguments(input_path, output, backend=backend, mode=mode) + ) + captured = capsys.readouterr() + assert rc == 0, captured.err + assert captured.err == "" + summary = json.loads(captured.out) + assert summary["backend"] == backend + assert summary["requested_backend"] == backend + assert summary["metrics"]["converged_count"] == 2 + assert json.loads((output / "summary.json").read_text()) == summary + config = yaml.safe_load( + (output / "effective-config.yaml").read_text() + ) + assert config["backend"] == backend + assert config["mode"] == mode + tables[backend] = pq.read_table(output / "fits.parquet").to_pydict() + assert tables[backend]["backend"] == [backend, backend] + + reference = tables["cupy"] + result = tables["numba-cuda-mlir"] + for field in ( + "candidate_id", + "status", + "converged", + "evaluations", + "valid_pixel_count", + "degrees_of_freedom", + "uncertainty_valid", + "uncertainty_reason", + ): + assert result[field] == reference[field] + assert reference["uncertainty_valid"] == [True, True] + assert min(reference["evaluations"]) > 1 + for field in ( + *GaussianDipoleModel.parameter_names, + "chi_square", + "residual_norm", + ): + np.testing.assert_allclose( + result[field], reference[field], rtol=5e-12, atol=5e-12 + ) + with ( + np.load( + tmp_path / "cupy" / "fit-arrays.npz", allow_pickle=False + ) as baseline, + np.load( + tmp_path / "numba-cuda-mlir" / "fit-arrays.npz", + allow_pickle=False, + ) as fused, + ): + np.testing.assert_array_equal( + fused["candidate_id"], baseline["candidate_id"] + ) + for field in ("covariance", "residuals"): + np.testing.assert_allclose( + fused[field], baseline[field], rtol=5e-12, atol=5e-12 + ) + assert np.abs(baseline["covariance"]).max() > 1e-3 + + +def test_numba_mlir_cli_reports_missing_dependency_before_cuda( + tmp_path, capsys, monkeypatch +): + import cuphoton.xfit.backend as backend_module + + input_path = tmp_path / "synthetic.npz" + output = tmp_path / "fit" + _write_input(input_path) + monkeypatch.setitem(sys.modules, "numba_cuda_mlir", None) + monkeypatch.setattr( + backend_module, + "_cupy_backend", + lambda **kwargs: pytest.fail("missing MLIR dependency reached CUDA"), + ) + rc = run_component( + "xfit", _arguments(input_path, output, backend="numba-cuda-mlir") + ) + captured = capsys.readouterr() + assert rc == 1 + assert "Numba-CUDA-MLIR is not installed" in captured.err + assert "numba-cuda-mlir extra" in captured.err + assert not output.exists() + + +@pytest.mark.parametrize( + ("model", "options", "message"), + [ + ("stamp", [], "supports only the Gaussian model"), + ( + "gaussian", + ["--use-finite-difference"], + "does not support finite-difference", + ), + ("gaussian", ["--max-evaluations", "0"], "must be at least 1"), + ], +) +def test_numba_mlir_cli_rejects_invalid_requests_before_cuda( + tmp_path, capsys, monkeypatch, model, options, message +): + import cuphoton.xfit.api as api + + input_path = tmp_path / "synthetic.npz" + output = tmp_path / "fit" + _write_input(input_path, model=model) + monkeypatch.setattr( + api, + "resolve_backend", + lambda *args, **kwargs: pytest.fail("invalid request reached CUDA"), + ) + rc = run_component( + "xfit", + _arguments(input_path, output, backend="numba-cuda-mlir", model=model) + + options, + ) + captured = capsys.readouterr() + assert rc != 0 + assert message in captured.err + assert not output.exists() diff --git a/tests/xscan/test_numba_mlir_pipeline.py b/tests/xscan/test_numba_mlir_pipeline.py new file mode 100644 index 00000000..54ae5d9f --- /dev/null +++ b/tests/xscan/test_numba_mlir_pipeline.py @@ -0,0 +1,392 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from contextlib import nullcontext +from types import SimpleNamespace + +import numpy as np +import pytest + +import cuphoton.xfit as xfit +import cuphoton.xscan.device_pipeline as pipeline +import cuphoton.xscan.xfit_features as features +from cuphoton.xscan.pipeline_benchmark import stages + + +def _config(tmp_path, *, backend="cupy"): + return pipeline.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=(3, 3), + decision_threshold=0.5, + xpois=pipeline.DeviceXPOISPipelineConfig( + kernel_shape=(3, 3), basis_sigmas=(0.8,), basis_degrees=(0,) + ), + xfit=pipeline.DeviceXFitPipelineConfig(backend=backend), + ) + + +def test_pipeline_backend_payload_preserves_legacy_default(tmp_path): + default = _config(tmp_path) + legacy = default.to_payload() + assert "backend" not in legacy["xfit"] + restored = pipeline.DevicePipelineConfig.from_payload(legacy) + assert restored.xfit.backend == "cupy" + assert restored.configuration_sha256 == default.configuration_sha256 + legacy["xfit"]["backend"] = "cupy" + assert ( + pipeline.DevicePipelineConfig.from_payload( + legacy + ).configuration_sha256 + == default.configuration_sha256 + ) + selected = _config(tmp_path, backend="numba-cuda-mlir") + payload = selected.to_payload() + assert payload["xfit"]["backend"] == "numba-cuda-mlir" + assert pipeline.DevicePipelineConfig.from_payload(payload) == selected + assert selected.configuration_sha256 != default.configuration_sha256 + assert "backend" not in selected.xfit.solver_payload() + assert xfit.LMConfig(**selected.xfit.solver_payload()) == xfit.LMConfig() + with pytest.raises(ValueError, match="xfit backend"): + pipeline.DeviceXFitPipelineConfig(backend="numba-cuda") + with pytest.raises(ValueError, match="finite-difference"): + pipeline.DeviceXFitPipelineConfig( + backend="numba-cuda-mlir", use_finite_difference=True + ) + + +@pytest.mark.parametrize("backend", ["cupy", "numba-cuda-mlir"]) +def test_pipeline_and_separated_stage_forward_backend( + tmp_path, monkeypatch, backend +): + config = _config(tmp_path, backend=backend) + descriptor = pipeline.NpyArrayDescriptor( + path=str(tmp_path / "image.npy"), + sha256="c" * 64, + shape=(5, 5), + dtype="float64", + ) + item = pipeline.DevicePipelineItem( + item_id="pair", + reference=descriptor, + target=descriptor, + candidates=(pipeline.DevicePipelineCandidate("candidate", 2, 2, 0),), + ) + difference = np.zeros((1, 3, 3)) + calls = [] + + class ReachedFit(Exception): + pass + + def fit(images, **kwargs): + assert images is difference + assert kwargs.get("backend", "cupy") == backend + assert kwargs["model"] == "gaussian" + assert kwargs["mode"] == "difference" + assert kwargs["config"] == xfit.LMConfig() + calls.append(kwargs) + raise ReachedFit + + class Stream(nullcontext): + def synchronize(self): + pass + + cp = SimpleNamespace( + asarray=np.asarray, + cuda=SimpleNamespace(Device=lambda: Stream()), + ) + context = pipeline.DeviceWorkerContext( + config=config, + cp=cp, + torch=None, + model=None, + performance=None, + feature_variance_present=False, + device="cuda:0", + producer_stream=Stream(), + consumer_stream=Stream(), + solver_config=xfit.LMConfig(), + load_seconds=0, + ) + monkeypatch.setattr( + pipeline.DeviceWorkerContext, + "_validate_active_device", + lambda self: None, + ) + monkeypatch.setattr( + pipeline, "_validate_item_contract", lambda *a, **k: None + ) + monkeypatch.setattr( + pipeline, + "_read_item_inputs", + lambda item: { + "reference": np.zeros((5, 5)), + "target": np.zeros((5, 5)), + }, + ) + monkeypatch.setattr( + pipeline, "_read_fits_device_inputs", lambda *a, **k: ({}, ()) + ) + monkeypatch.setattr( + pipeline, "_extract_stamps", lambda **k: (None, difference) + ) + monkeypatch.setattr( + pipeline, + "_load_pipeline_functions", + lambda: SimpleNamespace( + gaussian_basis_component=lambda **kwargs: None, + solve_constant_kernel_device=lambda *a, **k: SimpleNamespace( + residual=None + ), + fit_dipoles_device=fit, + ), + ) + with pytest.raises(ReachedFit): + pipeline.run_device_pipeline_item(item, context) + + monkeypatch.setattr(xfit, "fit_dipoles_device", fit) + monkeypatch.setattr(stages, "_read_array", lambda *args: difference) + priors = {"xpois": {"items": [{"artifacts": {"difference": None}}]}} + with pytest.raises(ReachedFit): + stages._run_item( + "xfit", + item, + config, + {"cp": cp, "torch": None}, + tmp_path, + priors, + 0, + ) + assert len(calls) == 2 + + +@pytest.mark.parametrize("requested", ["cupy", "numba-cuda-mlir"]) +def test_pipeline_evidence_rejects_backend_different_from_config( + tmp_path, requested +): + actual = "cupy" if requested == "numba-cuda-mlir" else "numba-cuda-mlir" + result = SimpleNamespace( + parameter_names=xfit.GaussianDipoleModel.parameter_names, + backend=actual, + solver="levenberg-marquardt", + model="gaussian", + mode="difference", + dtype="float64", + ) + with pytest.raises(ValueError, match="xFit evidence contract"): + pipeline._pack_scientific_evidence( + cp=None, + xpois_result=None, + xfit_result=result, + features=SimpleNamespace(feature_names=features.FEATURE_NAMES), + candidate_count=1, + kernel_shape=(3, 3), + flux_conserve=False, + xfit_backend=requested, + ) + with pytest.raises( + ValueError, match="xfit.backend does not match config" + ): + pipeline._validate_device_pipeline_evidence_config( + {"xfit": {"backend": actual}}, + {}, + config=_config(tmp_path, backend=requested), + ) + + +def test_numba_pipeline_feature_source_preserves_array_checks(): + class DeviceArray(np.ndarray): + @property + def device(self): + return SimpleNamespace(id=0) + + def array(shape, dtype="float64"): + return np.zeros(shape, dtype=dtype).view(DeviceArray) + + result = SimpleNamespace( + schema="cuphoton.xfit.device-fit-result/v1", + backend="numba-cuda-mlir", + solver="levenberg-marquardt", + result_location="device", + model="gaussian", + mode="difference", + dtype="float64", + device_id=0, + parameter_names=xfit.GaussianDipoleModel.parameter_names, + parameters=array((1, 8)), + standard_errors=array((1, 8)), + covariance=array((1, 8, 8)), + converged=array((1,), bool), + uncertainty_valid=array((1,), bool), + degrees_of_freedom=array((1,), "int64"), + valid_pixel_fraction=array((1,)), + fractional_null_improvement=array((1,)), + delta_chi_square=array((1,)), + reduced_chi_square=array((1,)), + ) + cp = SimpleNamespace(ndarray=DeviceArray) + names = features._validate_device_feature_source( + result, cp=cp, active_device_id=0 + ) + assert names == xfit.GaussianDipoleModel.parameter_names + with pytest.raises(ValueError, match="active device is 1"): + features._validate_device_feature_source( + result, cp=cp, active_device_id=1 + ) + result.parameters = np.zeros((1, 8)) + with pytest.raises(TypeError, match="parameters must be a CuPy array"): + features._validate_device_feature_source( + result, cp=cp, active_device_id=0 + ) + + +def test_numba_pipeline_cli_runs_real_stages_and_matches_cupy( + tmp_path, capsys +): + import json + from dataclasses import replace + + from cuphoton.core.cli import run_component + from cuphoton.xscan.pipeline_benchmark.fixture import prepare_fixture + from cuphoton.xscan.pipeline_benchmark.runner import read_inputs + + cp = pytest.importorskip("cupy") + torch = pytest.importorskip("torch") + pytest.importorskip("numba_cuda_mlir.cuda") + try: + count = cp.cuda.runtime.getDeviceCount() + except cp.cuda.runtime.CUDARuntimeError as exc: + if exc.status != 100: + raise + pytest.skip("CUDA device is unavailable") + if count < 1 or not torch.cuda.is_available(): + pytest.skip("CUDA device is unavailable") + + config_path, items_path = prepare_fixture( + tmp_path / "inputs", + images=1, + image_size=48, + candidates=2, + stamp_size=9, + seed=2026, + device="cuda:0", + ) + config, items = read_inputs(config_path, items_path) + assert "backend" not in config.to_payload()["xfit"] + exact_fields = { + "xfit.status_codes", + "xfit.converged", + "xfit.evaluations", + "xfit.valid_pixel_count", + "xfit.valid_pixel_fraction", + "xfit.null_chi_square", + "xfit.degrees_of_freedom", + "xfit.uncertainty_valid", + "xfit.uncertainty_reason_codes", + } + reference = reference_payload = None + for backend in ("cupy", "numba-cuda-mlir"): + selected = replace(config, xfit=replace(config.xfit, backend=backend)) + selected_path = tmp_path / f"{backend}.json" + selected_path.write_text(json.dumps(selected.to_payload())) + output = tmp_path / backend + rc = run_component( + "xscan", + [ + "benchmark-pipeline", + "--stage", + "pipeline", + "--config", + str(selected_path), + "--items", + str(items_path), + "--output", + str(output), + "--warmup", + "1", + "--repeat", + "1", + ], + ) + captured = capsys.readouterr() + assert rc == 0, captured.err + summary = json.loads((output / "summary.json").read_text()) + assert [row["completed_items"] for row in summary["rounds"]] == [1, 1] + previous = None + for index in range(2): + result = json.loads( + (output / f"round-{index:03d}" / "item-0000.json").read_text() + ) + assert ( + result["configuration_sha256"] + == selected.configuration_sha256 + ) + evidence = result["scientific_evidence"] + assert evidence["xfit"]["backend"] == backend + arrays = pipeline.decode_device_pipeline_evidence( + evidence, config=selected + ) + assert set(arrays) == set(stages.SCIENCE_KEYS) + assert len(result["predictions"]) == len(items[0].candidates) == 2 + assert result["transfers"]["pipeline_full_array_d2h_bytes"] == 0 + assert result["transfers"]["pipeline_compact_h2d_bytes"] == 0 + assert result["transfers"]["terminal_d2h_calls"] == 1 + if previous is not None: + for name, values in arrays.items(): + np.testing.assert_array_equal(values, previous[name]) + previous = arrays + if reference is None: + reference, reference_payload = arrays, result + for key in ( + "item_id", + "input_sha256", + "checkpoint_sha256", + "feature_schema_sha256", + "xpois", + "transfers", + ): + assert result[key] == reference_payload[key] + for actual, expected in zip( + result["predictions"], + reference_payload["predictions"], + strict=True, + ): + assert actual.keys() == expected.keys() + for key in actual: + if key in {"logit", "probability"}: + np.testing.assert_allclose( + actual[key], expected[key], rtol=3e-5, atol=3e-5 + ) + else: + assert actual[key] == expected[key] + for name, actual in arrays.items(): + expected = reference[name] + assert ( + actual.shape == expected.shape + and actual.dtype == expected.dtype + ) + if name.startswith("xpois.") or name in exact_fields: + np.testing.assert_array_equal(actual, expected) + continue + for mask in (np.isnan, np.isposinf, np.isneginf): + np.testing.assert_array_equal( + mask(actual), mask(expected) + ) + tolerance = ( + 3e-5 + if name == "xfit.features" or name.startswith("xscan.") + else 5e-12 + ) + np.testing.assert_allclose( + actual, + expected, + rtol=tolerance, + atol=tolerance, + equal_nan=True, + ) diff --git a/tests/xscan/test_pipeline_benchmark.py b/tests/xscan/test_pipeline_benchmark.py index fdd0d949..b11193c3 100644 --- a/tests/xscan/test_pipeline_benchmark.py +++ b/tests/xscan/test_pipeline_benchmark.py @@ -436,7 +436,7 @@ def child(command, log, *, timeout): ) -def test_provenance_discovers_the_installed_cupy_distribution( +def test_provenance_discovers_installed_cuda_distributions( monkeypatch, ) -> None: monkeypatch.setattr( @@ -447,4 +447,6 @@ def test_provenance_discovers_the_installed_cupy_distribution( monkeypatch.setattr( benchmark.importlib.metadata, "version", lambda name: "test-" + name ) - assert benchmark.provenance()["packages"]["cupy"] == "test-cupy" + packages = benchmark.provenance()["packages"] + assert packages["cupy"] == "test-cupy" + assert packages["numba-cuda-mlir"] == "test-numba-cuda-mlir" diff --git a/uv.lock b/uv.lock index 1c6b4a02..f937882a 100644 --- a/uv.lock +++ b/uv.lock @@ -355,6 +355,12 @@ cusparse = [ { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] +nvcc = [ + { name = "nvidia-cuda-crt", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cuda-nvcc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cuda-runtime", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvvm", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] nvjitlink = [ { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] @@ -422,6 +428,10 @@ io = [ mpi = [ { name = "mpi4py" }, ] +numba-cuda-mlir = [ + { name = "cupy-cuda13x", extra = ["ctk"] }, + { name = "numba-cuda-mlir", extra = ["cu13"] }, +] photometry = [ { name = "photutils" }, ] @@ -443,6 +453,7 @@ requires-dist = [ { name = "cuphoton", extras = ["io", "photometry"], marker = "extra == 'gpu'" }, { name = "cupy-cuda13x", extras = ["ctk"], marker = "sys_platform == 'linux' and extra == 'cutile'", specifier = ">=14,<15" }, { name = "cupy-cuda13x", extras = ["ctk"], marker = "sys_platform == 'linux' and extra == 'io'", specifier = ">=14,<15" }, + { name = "cupy-cuda13x", extras = ["ctk"], marker = "sys_platform == 'linux' and extra == 'numba-cuda-mlir'", specifier = ">=14,<15" }, { name = "dragonhpc", marker = "extra == 'dragon'", specifier = ">=0.14.2,<0.15" }, { name = "h5py", specifier = ">=3.10" }, { name = "kvikio-cu13", marker = "sys_platform == 'linux' and extra == 'io'", specifier = "==26.6.*" }, @@ -451,6 +462,7 @@ requires-dist = [ { name = "mypy", marker = "extra == 'dev'", specifier = "==1.19.1" }, { name = "numba", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = ">=0.61,<0.66" }, { name = "numba-cuda", extras = ["cu13"], marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = ">=0.30,<0.31" }, + { name = "numba-cuda-mlir", extras = ["cu13"], marker = "sys_platform == 'linux' and extra == 'numba-cuda-mlir'", specifier = ">=0.5.4,<0.6" }, { name = "numexpr", specifier = ">=2.10" }, { name = "numpy", specifier = ">=2.0,<2.6" }, { name = "nvidia-libnvcomp-cu13", marker = "sys_platform == 'linux' and extra == 'io'", specifier = "==5.2.*" }, @@ -471,7 +483,7 @@ requires-dist = [ { name = "types-pyyaml", marker = "extra == 'dev'", specifier = ">=6.0.12" }, { name = "virtualenv", marker = "extra == 'dev'", specifier = ">=21.7.12" }, ] -provides-extras = ["photometry", "torch", "viz", "io", "gpu", "cutile", "mpi", "dragon", "dev"] +provides-extras = ["photometry", "torch", "viz", "io", "gpu", "cutile", "numba-cuda-mlir", "mpi", "dragon", "dev"] [[package]] name = "cupy-cuda13x" @@ -876,6 +888,33 @@ cu13 = [ { name = "nvidia-nvjitlink" }, ] +[[package]] +name = "numba-cuda-mlir" +version = "0.5.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-bindings" }, + { name = "cuda-core" }, + { name = "numpy" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/a6/65/591bc6509f0715d3b8a817e8da53041c212bcc8505c503c50c3833ff2412/numba_cuda_mlir-0.5.4-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:39dac9cabea7fd95aa9f322c6b160af814c03ad3427bb3646fbe2ce00ac2eeed", size = 67214375, upload-time = "2026-09-29T18:56:38.13Z" }, + { url = "https://files.pythonhosted.org/packages/39/f3/194607f9a30bc28f18f5320ed3c8c10637c3378c2f27974b42005552f545/numba_cuda_mlir-0.5.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:327206fdfb112543cbae36c5f8961e5776a8189df19a03df0c6abb342a02c3f4", size = 71151884, upload-time = "2026-09-29T18:56:46.85Z" }, + { url = "https://files.pythonhosted.org/packages/95/09/fb8327306e09cf3e4dd076e6bf3ecc5382336d9259bbf026b4e62e41df5c/numba_cuda_mlir-0.5.4-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4d8b9fda2800adf48dfd0cfb49b6cd71e9abab5b8d67856306233500bba71aaa", size = 67214043, upload-time = "2026-09-29T18:57:01.832Z" }, + { url = "https://files.pythonhosted.org/packages/7a/dc/e1640872b2e0b9876ca2b004e60895794d0ec30fe73a4dbd184e94d3ce14/numba_cuda_mlir-0.5.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:958adab0338aaa7e6c7e39d028812f7506e185f9c763756b95bbdc07ee3ba21a", size = 71151839, upload-time = "2026-09-29T18:57:09.87Z" }, + { url = "https://files.pythonhosted.org/packages/7d/9a/56d536da43d9d92986a707bdf2ac8a29f4942b2aa984f5f4daddcb4d497d/numba_cuda_mlir-0.5.4-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5bd42215cf146d61eefbd0e848d0e2d4f3f48009fe21e5ae58174f984cca7f8e", size = 67217352, upload-time = "2026-09-29T18:57:23.633Z" }, + { url = "https://files.pythonhosted.org/packages/08/07/6def3acd0a934b91fa0053916da1f05d391e73eeea8750ec60b888d43b57/numba_cuda_mlir-0.5.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:669a14c69524e62169ca00c5a46c327c04141cb55676cde2ed01ef3774944053", size = 71151154, upload-time = "2026-09-29T18:57:31.576Z" }, + { url = "https://files.pythonhosted.org/packages/9d/0e/9e6732aea62b93deea1f57bf4319d9fcb94e36cce36ec9600ee41e677db2/numba_cuda_mlir-0.5.4-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b1ba33820a4a3948c91a2eacecaf217e39ac9bdc528b680f473587c893028ffc", size = 67237258, upload-time = "2026-09-29T18:57:47.372Z" }, + { url = "https://files.pythonhosted.org/packages/91/44/f74be3d3d914a1f0d2024a3929ee8bec66d149c4f174a7dcb4cf4e96bd9a/numba_cuda_mlir-0.5.4-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:938f671ed6a0fba6104965b5898493240f89cf4a8aaae49baa69bfdfdf1e4c9b", size = 71156253, upload-time = "2026-09-29T18:57:56.7Z" }, +] + +[package.optional-dependencies] +cu13 = [ + { name = "cuda-bindings" }, + { name = "cuda-toolkit", extra = ["cudart", "nvcc", "nvrtc"] }, + { name = "nvidia-nvjitlink" }, +] + [[package]] name = "numexpr" version = "2.14.1" @@ -956,6 +995,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ab/fb/0384bb2129bed6b1b39f8e44471c615ab5ab29b7e55817538a1d390d8f84/nvidia_cuda_cccl-13.0.85-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e0da7ad981f3a8aff08241b5bfc1af868742a63e2762f53a5171c492ef242649", size = 3120526, upload-time = "2025-09-04T08:25:25.144Z" }, ] +[[package]] +name = "nvidia-cuda-crt" +version = "13.0.88" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/b9/2cb230193e1570221eee8c9739f965bb4874bad44ee4c3373a5860d24c90/nvidia_cuda_crt-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ee2ea2a97073e02ee62bb27841f437332be2c248e3eac013df07997ada39c003", size = 134086, upload-time = "2025-09-04T08:25:33.024Z" }, + { url = "https://files.pythonhosted.org/packages/05/69/a1ec4d9f0747d85964206feb47bf359f48a2d6af74c0add8abba6efe3dda/nvidia_cuda_crt-13.0.88-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2c8043c7c9e02492716426e9919fc78d2c5b3b2a7a768a88e952676b08aa55a4", size = 134086, upload-time = "2025-09-04T08:25:40.741Z" }, +] + [[package]] name = "nvidia-cuda-cupti" version = "13.0.85" @@ -965,6 +1013,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:4eb01c08e859bf924d222250d2e8f8b8ff6d3db4721288cf35d14252a4d933c8", size = 10715597, upload-time = "2025-09-04T08:26:51.312Z" }, ] +[[package]] +name = "nvidia-cuda-nvcc" +version = "13.0.88" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cuda-crt" }, + { name = "nvidia-cuda-runtime" }, + { name = "nvidia-nvvm" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/9f/06/996d5cdc5ea45fb4a6111a1be4f0caf6556c0cb1bf9684a7252d8771797a/nvidia_cuda_nvcc-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c7ff28f86a24effdc6c034fa15230c549a273e4771b10a7fec14996f8cf3307f", size = 32455430, upload-time = "2025-09-04T08:27:27.39Z" }, + { url = "https://files.pythonhosted.org/packages/71/8b/a546c12881fffeba927d810598987df25d74b8b241788c7db8dfc93b0173/nvidia_cuda_nvcc-13.0.88-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:56fe502eb77625a12f25172caa3cdddb4e4c8ba2c8c17dba44b164761b380f03", size = 37384532, upload-time = "2025-09-04T08:27:37.916Z" }, +] + [[package]] name = "nvidia-cuda-nvrtc" version = "13.0.88" From 01bff98c92cc6d6f0c0a1cce2f8fc9d046b1f577 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 30 Sep 2026 06:43:15 -0700 Subject: [PATCH 2/4] Update public contracts for the optional xFit backend Include the Numba-CUDA-MLIR installation extra and device-fit backend argument in the existing public contract assertions. Signed-off-by: Trent Nelson --- tests/test_package_layout.py | 1 + tests/test_public_api.py | 1 + 2 files changed, 2 insertions(+) diff --git a/tests/test_package_layout.py b/tests/test_package_layout.py index 02c05011..90514d2c 100644 --- a/tests/test_package_layout.py +++ b/tests/test_package_layout.py @@ -70,6 +70,7 @@ def test_distribution_metadata_declares_supported_profiles() -> None: "gpu", "io", "mpi", + "numba-cuda-mlir", "photometry", "torch", "viz", diff --git a/tests/test_public_api.py b/tests/test_public_api.py index 6787184b..d376d2ea 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -131,6 +131,7 @@ def test_xfit_curated_exports_and_fit_signature() -> None: "mask", "variance", "mode", + "backend", "config", "backend", ] From 6f607d0859bfa3185dcd910ad13ee4a9f13fb71b Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Fri, 9 Oct 2026 12:19:06 -0700 Subject: [PATCH 3/4] Reconcile optional xFit backend contracts Keep the device signature and shared pipeline diagnostics consistent when native and Numba-CUDA-MLIR execution are both available. Signed-off-by: Trent Nelson --- src/cuphoton/xfit/api.py | 2 +- tests/test_public_api.py | 1 - tests/xscan/test_numba_mlir_pipeline.py | 4 ++-- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/cuphoton/xfit/api.py b/src/cuphoton/xfit/api.py index f33c556e..fec21b7e 100644 --- a/src/cuphoton/xfit/api.py +++ b/src/cuphoton/xfit/api.py @@ -139,7 +139,7 @@ class DeviceDipoleFitResult: def __post_init__(self) -> None: if self.backend not in {"cupy", "native", "numba-cuda-mlir"}: - raise ValueError("device fit backend must be cupy or native") + raise ValueError("unsupported device xFit backend") if isinstance(self.device_id, bool) or not isinstance( self.device_id, (int, np.integer) ): diff --git a/tests/test_public_api.py b/tests/test_public_api.py index d376d2ea..6787184b 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -131,7 +131,6 @@ def test_xfit_curated_exports_and_fit_signature() -> None: "mask", "variance", "mode", - "backend", "config", "backend", ] diff --git a/tests/xscan/test_numba_mlir_pipeline.py b/tests/xscan/test_numba_mlir_pipeline.py index 54ae5d9f..0725f658 100644 --- a/tests/xscan/test_numba_mlir_pipeline.py +++ b/tests/xscan/test_numba_mlir_pipeline.py @@ -55,7 +55,7 @@ def test_pipeline_backend_payload_preserves_legacy_default(tmp_path): assert xfit.LMConfig(**selected.xfit.solver_payload()) == xfit.LMConfig() with pytest.raises(ValueError, match="xfit backend"): pipeline.DeviceXFitPipelineConfig(backend="numba-cuda") - with pytest.raises(ValueError, match="finite-difference"): + with pytest.raises(ValueError, match="finite differences"): pipeline.DeviceXFitPipelineConfig( backend="numba-cuda-mlir", use_finite_difference=True ) @@ -191,7 +191,7 @@ def test_pipeline_evidence_rejects_backend_different_from_config( xfit_backend=requested, ) with pytest.raises( - ValueError, match="xfit.backend does not match config" + ValueError, match="xFit backend does not match config" ): pipeline._validate_device_pipeline_evidence_config( {"xfit": {"backend": actual}}, From d940e5399de303e314dbef67bc6322b17c8af45d Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Fri, 9 Oct 2026 12:20:43 -0700 Subject: [PATCH 4/4] Type the optional Numba CUDA runtime boundary Allow dynamic CUDA array and stream exports in GPU-enabled type checks. Signed-off-by: Trent Nelson --- src/cuphoton/xfit/_numba_mlir.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/cuphoton/xfit/_numba_mlir.py b/src/cuphoton/xfit/_numba_mlir.py index bc7fbbd3..f97ae609 100644 --- a/src/cuphoton/xfit/_numba_mlir.py +++ b/src/cuphoton/xfit/_numba_mlir.py @@ -238,12 +238,13 @@ def __init__( ) -> None: try: import cupy as cp - from numba_cuda_mlir import cuda + from numba_cuda_mlir import cuda as cuda_module except ImportError as exc: raise ImportError( "Numba-CUDA-MLIR Gaussian kernels require CuPy and " "numba-cuda-mlir" ) from exc + cuda: Any = cuda_module if mode not in {"difference", "split"}: raise ValueError("mode must be 'difference' or 'split'") height, width = image_shape