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..f97ae609 --- /dev/null +++ b/src/cuphoton/xfit/_numba_mlir.py @@ -0,0 +1,391 @@ +# 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 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 + 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..fec21b7e 100644 --- a/src/cuphoton/xfit/api.py +++ b/src/cuphoton/xfit/api.py @@ -131,15 +131,15 @@ 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"}: - raise ValueError("device fit backend must be cupy or native") + if self.backend not in {"cupy", "native", "numba-cuda-mlir"}: + raise ValueError("unsupported device xFit backend") 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/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/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..0725f658 --- /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 differences"): + 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"