Add an optional Numba-CUDA-MLIR Gaussian xFit backend

.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

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 |

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` |

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

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)

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"

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):

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)

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")

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"