From 0e18a23e89a9d8f1ec75aaa15e951fe7aa9d9e36 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sun, 20 Sep 2026 07:06:05 -0700 Subject: [PATCH 1/2] Add CUDA transfer and matmul kernels for the normalized views `view.to_gpu()` moves a normalized view onto an NVIDIA GPU as a `CudaNormalizedView`, supporting the same `@`/`__rmatmul__` the host view does. Every floating-point array on the device is float32. The transfer keeps the value-compressed layout: it uploads the same major_ptr/values/value_ptr/indices the host holds plus the O(n_rows + n_cols) statistics, never an expanded one-float-per-nonzero copy. That matters more on a device than on a host, where the whole reason to hold a count matrix in VCS form is that the expanded form does not fit. Keeping the layout is what forces custom kernels -- cuSPARSE reads CSR and CSC only, so nothing in CuPy can walk the two-level structure. There are four, the same two binary choices the host kernels make (which axis the statistics hang off, and whether the output's disjoint axis is the major one), generated from shared fragments rather than written out four times. Each is one warp per major slice, grid-strided, with the warp's 32 lanes split between the dense width and the slice's nonzeros so a PARAFAC2-shaped width of 8 does not idle 24 of them. Aligned directions accumulate in registers and reduce with __shfl_down_sync; misaligned ones atomicAdd, which is fast here but makes them nondeterministic in the last bits. Benchmarked against the alternative they exist to beat -- expanding to a CuPy CSR and calling cuSPARSE -- at 60k x 2k, 6M nonzeros, width 16. The kernels win every direction while holding 0.51x-0.62x the device memory, since this is bandwidth-bound: 0.56x-0.86x for `VCSR @ B` and `VCSC @ B`, and 0.17x-0.32x for the `dense @ sparse` directions, which cuSPARSE handles poorly. Ceilings are recorded on the slower of the two development GPUs. Two traps found while making that baseline honest, both of which would otherwise have flattered the kernels: - A VCS slice stores its indices grouped by the value they share, never ascending, so the expanded matrix is not canonical. cuSPARSE measured ~6x slower on it, so `to_cupy_sparse` sorts by default. - `B @ csc` measured ~5x `B @ csr`, so `to_cupy_sparse` grew a `format` parameter and the benchmark pins the baseline to CSR even for a VCSC-backed view. CI gains a `cuda` job on a self-hosted GPU runner, guarded so a forked PR cannot reach it and asserting a device is visible so it cannot silently pass by skipping. The three CPU jobs now pass `--no-extra cuda` to avoid the ~1 GB CuPy wheel. Co-Authored-By: Claude Opus 5 --- .github/workflows/ci.yml | 54 +++- README.md | 38 ++- benchmarks/README.md | 38 ++- benchmarks/baselines.json | 36 ++- benchmarks/cases.py | 127 +++++++- benchmarks/harness.py | 23 ++ benchmarks/run.py | 10 +- docs/api.md | 7 + docs/usage.md | 61 ++++ pyproject.toml | 10 + src/vsparse/__init__.py | 3 + src/vsparse/_cuda.py | 592 ++++++++++++++++++++++++++++++++++++ src/vsparse/_norm_common.py | 16 + tests/test_cuda.py | 189 ++++++++++++ uv.lock | 185 ++++++++++- 15 files changed, 1378 insertions(+), 11 deletions(-) create mode 100644 src/vsparse/_cuda.py create mode 100644 tests/test_cuda.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 34f61ba..10fcebb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -25,7 +25,9 @@ jobs: python-version: "3.12" - name: Install dependencies - run: uv sync --all-extras --dev + # --no-extra cuda: the CuPy wheel is ~1 GB and useless without a GPU. + # The `cuda` job below installs it on the self-hosted runner instead. + run: uv sync --all-extras --no-extra cuda --dev - name: Lint with Ruff run: uv run ruff check . @@ -57,7 +59,9 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install dependencies - run: uv sync --all-extras --dev + # --no-extra cuda: the CuPy wheel is ~1 GB and useless without a GPU. + # The `cuda` job below installs it on the self-hosted runner instead. + run: uv sync --all-extras --no-extra cuda --dev - name: Run pytest run: uv run pytest --cov @@ -76,7 +80,9 @@ jobs: python-version: "3.12" - name: Install dependencies - run: uv sync --all-extras --dev + # --no-extra cuda: the CuPy wheel is ~1 GB and useless without a GPU. + # The `cuda` job below installs it on the self-hosted runner instead. + run: uv sync --all-extras --no-extra cuda --dev # Fails if a gated metric exceeds its ceiling in benchmarks/baselines.json. # Memory and layout numbers are deterministic and gated tightly; timing is @@ -84,3 +90,45 @@ jobs: # because a shared runner's absolute speed means nothing. See benchmarks/README.md. - name: Run benchmark gate run: uv run python -m benchmarks.run --set fast + + cuda: + name: CUDA Tests and Benchmarks + # A self-hosted runner with an NVIDIA GPU and a driver new enough for the + # `cupy-cuda13x` wheel. Labelled rather than named so more than one machine + # can serve the job. Forked PRs never see this runner, and `pull_request` + # from a fork cannot reach a self-hosted runner anyway -- but the guard is + # explicit so an untrusted PR can never run code on our own hardware. + if: github.event_name == 'push' || github.event.pull_request.head.repo.full_name == github.repository + runs-on: [self-hosted, linux, x64, gpu, cuda] + steps: + - name: Checkout code + uses: actions/checkout@v7 + + - name: Install uv + uses: astral-sh/setup-uv@v10.0.1 + with: + enable-cache: true + python-version: "3.12" + + - name: Show the device the job landed on + run: nvidia-smi --query-gpu=name,driver_version,compute_cap,memory.total --format=csv + + - name: Install dependencies + run: uv sync --all-extras --dev + + # tests/test_cuda.py skips itself if CuPy cannot see a device, which on + # this runner would mean the job silently passed without testing + # anything. Assert the device is really there first. + - name: Assert a CUDA device is visible + run: uv run python -c "import vsparse, sys; sys.exit(0 if vsparse.cuda_is_available() else 'no CUDA device visible to CuPy')" + + - name: Run the CUDA tests + run: uv run pytest -m cuda + + # Same gate as the CPU benchmark job, over benchmarks/baselines.json. + # `cuda_device_bytes_ratio_vs_csr` is deterministic and gated tightly; + # the timing ratio is against cuSPARSE on the same device in the same + # process, which cancels most of the difference between GPU models, and + # is gated loosely for what it does not. See benchmarks/README.md. + - name: Run the CUDA benchmark gate + run: uv run python -m benchmarks.run --set cuda diff --git a/README.md b/README.md index 1ef2860..d0dc2ae 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,14 @@ pip install vsparse uv add vsparse ``` +For the CUDA support described below (needs an NVIDIA GPU; pulls in CuPy): + +```sh +pip install "vsparse[cuda]" +# or with uv +uv add "vsparse[cuda]" +``` + From source: ```sh @@ -140,14 +148,42 @@ adata_norm = vsparse.load_and_normalize( ) ``` +### CUDA (`to_gpu`) + +A normalized view moves onto an NVIDIA GPU with `.to_gpu()`, keeping its value-compressed layout there rather than expanding to one float per nonzero — so a matrix that only fits in device memory compressed still fits. Everything floating-point on the device is float32. + +```python +gpu = vsparse.VCSRArray.from_scipy(counts).normalized("parafac2").to_gpu() +out = gpu @ B # float32 CuPy array; B @ gpu likewise +``` + +cuSPARSE reads only CSR/CSC, so walking the VCS layout on the device needs custom kernels; `vsparse._cuda` supplies four, one per (format, direction). Against the alternative — expanding to a sorted CuPy CSR and calling cuSPARSE — they measured (RTX 5080, 60k x 2k, 6M nonzeros, width 16): + +| direction | kernel / cuSPARSE | device bytes vs CSR | +| --- | --- | --- | +| `VCSR @ B` | 0.71x–0.86x | 0.62x | +| `VCSC @ B` | 0.83x | 0.51x | +| `B @ VCSC` | 0.32x | 0.51x | +| `B @ VCSR` | 0.23x | 0.62x | + +`CudaNormalizedView.to_cupy_sparse()` builds that CSR baseline if you want it. + See `docs/` for full usage guides and API documentation. ## Development ```sh -uv sync --all-extras --dev +uv sync --all-extras --dev # add --no-extra cuda to skip the ~1 GB CuPy wheel uv run pytest uv run ruff check . uv run ty check uv run sphinx-build -b html docs docs/_build/html ``` + +The CUDA tests and benchmarks need an NVIDIA GPU; without one `tests/test_cuda.py` +skips itself. On a GPU machine: + +```sh +uv run pytest -m cuda +uv run python -m benchmarks.run --set cuda +``` diff --git a/benchmarks/README.md b/benchmarks/README.md index 0979092..e1bf960 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -6,6 +6,7 @@ hand. ```sh uv run python -m benchmarks.run --set fast # run + compare (what CI does) uv run python -m benchmarks.run --set slow # the larger cases +uv run python -m benchmarks.run --set cuda # the GPU cases (needs a device) uv run python -m benchmarks.run --case matvec_vs_scipy uv run python -m benchmarks.run --set fast --record # rewrite baselines.json ``` @@ -40,11 +41,44 @@ recorded for context. Memory is no longer measured here; see below. +## The CUDA set + +`--set cuda` measures `vsparse._cuda`'s kernels against the alternative they +exist to beat: expanding the value-compressed layout to one float per nonzero +and letting cuSPARSE do the product. It runs on the self-hosted GPU runner, not +the shared ones, and needs `uv sync --all-extras` (the CPU jobs pass +`--no-extra cuda`). + +The baseline is deliberately the *best* materialized option rather than the +cheapest to produce, so the comparison is not a strawman: CSR in every case, +including for a VCSC-backed view whose natural materialization is CSC, and +index-sorted. Both matter more than expected -- `B @ csc` measured ~5x +`B @ csr`, and an unsorted CSR ~6x a sorted one -- so `to_cupy_sparse` sorts by +default and the cases pass `format="csr"`. + +Three metrics, of which two are gated: + +- `cuda_time_ratio_kernel_over_csr` -- steady-state throughput, ours over + cuSPARSE's on the already-built CSR. Gated loosely, for the same reason the + scipy ratios are: it is a ratio measured on the same device in the same + process, which cancels most of the difference between GPU models but not all + of it (the two development GPUs differed by up to 1.8x on this metric). +- `cuda_device_bytes_ratio_vs_csr` -- device memory held, ours over the CSR's. + Deterministic, so gated tightly. This is the number the kernels exist to buy. +- `cuda_materialize_in_matmuls` -- how many of our matmuls the one-time + materialization costs. Recorded for context, not gated. + +Timing a CUDA call needs `best_gpu_time`, not `best_time`: launches are +asynchronous, so an unsynchronized timer measures the launch and not the work. + ## Adding a case Write a function returning `{metric: value}` in `cases.py`, decorated with -`@fast` (runs on every PR, keep it under a minute) or `@slow`. Add any new -gated metric to `margins` in `baselines.json`, then `--record`. +`@fast` (runs on every PR, keep it under a minute), `@slow`, or `@cuda`. Add any +new gated metric to `margins` in `baselines.json`, then `--record`. + +Record CUDA ceilings on the *slowest* device you expect to run them on, so a +faster one cannot trip a gate it should pass. ## Memory is a test, not a benchmark diff --git a/benchmarks/baselines.json b/benchmarks/baselines.json index 7ad1739..38f2500 100644 --- a/benchmarks/baselines.json +++ b/benchmarks/baselines.json @@ -10,7 +10,9 @@ "cpu_ratio_1t_parafac2": 1.5, "cpu_ratio_1t_pearson": 1.5, "cpu_ratio_1t_raw": 1.5, - "cpu_ratio_1t_scanpy": 1.5 + "cpu_ratio_1t_scanpy": 1.5, + "cuda_time_ratio_kernel_over_csr": 2.0, + "cuda_device_bytes_ratio_vs_csr": 1.1 }, "cases": { "layout_bytes_per_nonzero": { @@ -99,6 +101,38 @@ }, "misaligned_rmatmat_vs_scipy": { "time_ratio_vs_scipy": 1.0817 + }, + "cuda_cp10k_log1p_matmul_vs_csr": { + "cuda_time_ratio_kernel_over_csr": 1.6947, + "cuda_device_bytes_ratio_vs_csr": 0.6792 + }, + "cuda_parafac2_matmul_vs_csr": { + "cuda_time_ratio_kernel_over_csr": 1.4385, + "cuda_device_bytes_ratio_vs_csr": 0.6792 + }, + "cuda_pearson_matmul_vs_csr": { + "cuda_time_ratio_kernel_over_csr": 1.4353, + "cuda_device_bytes_ratio_vs_csr": 0.6792 + }, + "cuda_raw_matmul_vs_csr": { + "cuda_time_ratio_kernel_over_csr": 1.4316, + "cuda_device_bytes_ratio_vs_csr": 0.6792 + }, + "cuda_scanpy_matmul_vs_csr": { + "cuda_time_ratio_kernel_over_csr": 1.7048, + "cuda_device_bytes_ratio_vs_csr": 0.6792 + }, + "cuda_parafac2_rmatmul_vs_csr": { + "cuda_time_ratio_kernel_over_csr": 0.6483, + "cuda_device_bytes_ratio_vs_csr": 0.5575 + }, + "cuda_parafac2_matmul_misaligned_vs_csr": { + "cuda_time_ratio_kernel_over_csr": 1.6738, + "cuda_device_bytes_ratio_vs_csr": 0.5575 + }, + "cuda_parafac2_rmatmul_misaligned_vs_csr": { + "cuda_time_ratio_kernel_over_csr": 0.4632, + "cuda_device_bytes_ratio_vs_csr": 0.6792 } } } diff --git a/benchmarks/cases.py b/benchmarks/cases.py index ef2b68f..b5bb151 100644 --- a/benchmarks/cases.py +++ b/benchmarks/cases.py @@ -7,6 +7,7 @@ from benchmarks.harness import ( best_cpu_time, + best_gpu_time, best_time, integer_counts_csr, ratio_vs_scipy, @@ -14,6 +15,7 @@ FAST: dict[str, Callable[[], dict[str, float]]] = {} SLOW: dict[str, Callable[[], dict[str, float]]] = {} +CUDA: dict[str, Callable[[], dict[str, float]]] = {} def fast(fn): @@ -26,6 +28,11 @@ def slow(fn): return fn +def cuda(fn): + CUDA[fn.__name__] = fn + return fn + + # -- layout size ------------------------------------------------------------- @@ -345,4 +352,122 @@ def large_layout_and_matmul() -> dict[str, float]: } -ALL: dict[str, Callable[[], dict[str, float]]] = {**FAST, **SLOW} +# -- CUDA kernels vs materializing to a sparse matrix ------------------------ +# +# The question `vsparse._cuda` exists to answer: is walking the +# value-compressed layout on the device with our own kernels actually +# competitive with expanding it to one float per nonzero and handing the +# product to cuSPARSE -- which is what `to_cupy_sparse()` does, and the device +# twin of the `to_scipy_sparse()` path a caller would otherwise take. +# +# Two numbers matter and they pull opposite ways: +# +# `cuda_time_ratio_kernel_over_csr` -- steady-state throughput, our kernel over +# cuSPARSE on the already-built CSR. A ratio above 1 would be defensible: we +# recompute `g` on every nonzero at every call where the CSR baked it in once, +# against a heavily tuned library. In practice every direction measures below +# 1, because the kernel reads ~40% less memory (the layout is the whole point) +# and this is a bandwidth-bound problem -- and `dense @ sparse` is a shape +# cuSPARSE handles particularly poorly, where we measured 0.17x-0.32x. +# +# `cuda_device_bytes_ratio_vs_csr` -- device memory held, ours over the CSR's. +# This is what the layout buys, it is deterministic, and it is the reason the +# throughput comes out where it does, so unlike the timing it is gated tightly. +# +# The one-time materialization is reported as `cuda_materialize_in_matmuls`: +# how many of our matmuls the caller pays up front to reach the CSR baseline at +# all. These run only on a GPU runner (`--set cuda`). + +_CUDA_ROWS, _CUDA_COLS, _CUDA_WIDTH = 60_000, 2_000, 16 + + +def _cuda_setup(cls_name: str, recipe: str): + """``(gpu_view, csr, means, mat)`` for a CUDA case. + + The baseline is deliberately the *best* materialized option, not the + cheapest one to produce, so the ratio is not a strawman: CSR in every case + (including for a VCSC-backed view, whose natural materialization is CSC -- + `B @ csc` measured ~5x `B @ csr`), and index-sorted, which + `to_cupy_sparse` does by default (unsorted measured ~6x slower again). + `cuda_materialize_in_matmuls` times that same full call, so both the + conversion and the sort are counted where a caller would pay them. + """ + import vsparse + + mat = integer_counts_csr(_CUDA_ROWS, _CUDA_COLS, density=0.05) + nv = getattr(vsparse, cls_name).from_scipy(mat).normalized(recipe) + gpu = nv.to_gpu() + return gpu, gpu.to_cupy_sparse(format="csr"), gpu.means, mat + + +def _csr_device_bytes(csr) -> int: + return int(csr.data.nbytes + csr.indices.nbytes + csr.indptr.nbytes) + + +def _cuda_metrics(gpu, csr, via_kernel, via_csr) -> dict[str, float]: + """Time both sides, check they agree, and report the three metrics.""" + import cupy as cp + + cp.testing.assert_allclose(via_kernel(), via_csr(), rtol=2e-4, atol=2e-4) + kernel_time = best_gpu_time(via_kernel) + return { + "cuda_time_ratio_kernel_over_csr": kernel_time / best_gpu_time(via_csr), + "cuda_device_bytes_ratio_vs_csr": gpu.nbytes / _csr_device_bytes(csr), + "cuda_materialize_in_matmuls": best_gpu_time(lambda: gpu.to_cupy_sparse(format="csr")) + / kernel_time, + } + + +def _cuda_matmul_case(cls_name: str, recipe: str) -> Callable[[], dict[str, float]]: + """``gpu @ B`` against ``csr @ B`` plus the same rank-1 correction.""" + + def bench() -> dict[str, float]: + import cupy as cp + + gpu, csr, means, mat = _cuda_setup(cls_name, recipe) + B = cp.asarray( + np.random.default_rng(0).normal(size=(mat.shape[1], _CUDA_WIDTH)), dtype=cp.float32 + ) + return _cuda_metrics(gpu, csr, lambda: gpu @ B, lambda: csr @ B + (-means) @ B) + + direction = "matmul" if cls_name == "VCSRArray" else "matmul_misaligned" + bench.__name__ = f"cuda_{recipe}_{direction}_vs_csr" + return bench + + +def _cuda_rmatmul_case(cls_name: str, recipe: str) -> Callable[[], dict[str, float]]: + """``B @ gpu`` against ``B @ csr`` plus the same rank-1 correction.""" + + def bench() -> dict[str, float]: + import cupy as cp + + gpu, csr, means, mat = _cuda_setup(cls_name, recipe) + B = cp.asarray( + np.random.default_rng(0).normal(size=(_CUDA_WIDTH, mat.shape[0])), dtype=cp.float32 + ) + correction = B.sum(axis=1)[:, None] * (-means)[None, :] + return _cuda_metrics(gpu, csr, lambda: B @ gpu, lambda: B @ csr + correction) + + direction = "rmatmul" if cls_name == "VCSCArray" else "rmatmul_misaligned" + bench.__name__ = f"cuda_{recipe}_{direction}_vs_csr" + return bench + + +def _register_cuda_benchmarks() -> None: + from vsparse import RECIPES + + # Every recipe in the aligned direction, since the recipes differ only in + # `g` and that is a per-nonzero cost the kernel pays and the CSR does not. + for recipe in sorted(RECIPES): + cuda(_cuda_matmul_case("VCSRArray", recipe)) + # One recipe is enough for the other three directions: they exercise the + # kernel's structure (register accumulation vs atomicAdd), not `g`. + cuda(_cuda_rmatmul_case("VCSCArray", "parafac2")) + cuda(_cuda_matmul_case("VCSCArray", "parafac2")) + cuda(_cuda_rmatmul_case("VCSRArray", "parafac2")) + + +_register_cuda_benchmarks() + + +ALL: dict[str, Callable[[], dict[str, float]]] = {**FAST, **SLOW, **CUDA} diff --git a/benchmarks/harness.py b/benchmarks/harness.py index 4f2f7f9..b4d31b7 100644 --- a/benchmarks/harness.py +++ b/benchmarks/harness.py @@ -48,3 +48,26 @@ def integer_counts_csr(n_rows: int, n_cols: int, density: float, seed: int = 0) mat = sp.random_array((n_rows, n_cols), density=density, format="csr", random_state=seed) mat.data = np.round(rng.integers(1, 8, size=mat.data.shape[0])).astype(np.float64) return mat + + +def best_gpu_time(fn: Callable[[], Any], repeat: int = 7) -> float: + """Best wall-clock time over ``repeat`` runs of a CUDA call, in seconds. + + CUDA launches are asynchronous, so timing one the way :func:`best_time` + does measures the launch and not the work. This synchronizes on both sides + of each run. The warm-up run also absorbs the one-time NVRTC compile of + :mod:`vsparse._cuda`'s kernels and cuSPARSE's own first-call setup. + """ + import cupy as cp + + device = cp.cuda.Device() + fn() + device.synchronize() + best = float("inf") + for _ in range(repeat): + device.synchronize() + start = time.perf_counter() + fn() + device.synchronize() + best = min(best, time.perf_counter() - start) + return best diff --git a/benchmarks/run.py b/benchmarks/run.py index bb40711..51002da 100644 --- a/benchmarks/run.py +++ b/benchmarks/run.py @@ -1,6 +1,7 @@ """Run benchmark cases and compare them against the checked-in baselines. python -m benchmarks.run --set fast # run and compare (CI does this) + python -m benchmarks.run --set cuda # the GPU cases (needs a device) python -m benchmarks.run --set fast --record # rewrite baselines.json python -m benchmarks.run --case matvec_vs_scipy # one case @@ -56,7 +57,7 @@ def _compare(results: dict[str, dict[str, float]], baselines: dict) -> list[str] def main() -> int: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--set", choices=["fast", "slow", "all"], default="fast") + parser.add_argument("--set", choices=["fast", "slow", "cuda", "all"], default="fast") parser.add_argument("--case", help="run a single case by name") parser.add_argument("--record", action="store_true", help="rewrite baselines.json") parser.add_argument("--emit", help=argparse.SUPPRESS) # internal: run one, print JSON @@ -71,7 +72,12 @@ def main() -> int: if args.case: names = [args.case] else: - chosen = {"fast": case_module.FAST, "slow": case_module.SLOW, "all": case_module.ALL} + chosen = { + "fast": case_module.FAST, + "slow": case_module.SLOW, + "cuda": case_module.CUDA, + "all": case_module.ALL, + } names = list(chosen[args.set]) results = {} diff --git a/docs/api.md b/docs/api.md index 4b67505..afec1d9 100644 --- a/docs/api.md +++ b/docs/api.md @@ -19,6 +19,13 @@ .. autofunction:: vsparse.load_and_normalize +.. autoclass:: vsparse.CudaNormalizedView + :members: + :undoc-members: + :show-inheritance: + +.. autofunction:: vsparse.cuda_is_available + .. autoclass:: vsparse.VCSCAnnData :members: :undoc-members: diff --git a/docs/usage.md b/docs/usage.md index 954a6ba..b26039d 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -131,3 +131,64 @@ This bypasses generic `AnnData` indexing and decompression overhead: 1. Cell counts are computed in $O(n_{\text{unique}})$ time from unique-value group sizes without touching or decoding the packed `indices` byte stream. 2. The delta+varint indices are unpacked in parallel across CPU cores. 3. Row/gene filtering, compaction, and depth normalization ($(\text{cell\_scale}) \times (\text{gene\_sums})$ followed by $\log_{10}(1000x + 1)$) are executed in fused parallel passes directly into the output CSR representation. + +## CUDA + +A normalized view can be moved onto an NVIDIA GPU with `.to_gpu()`, which needs +CuPy (`pip install vsparse[cuda]`): + +```python +import numpy as np +import vsparse + +arr = vsparse.VCSRArray.from_scipy(counts) +gpu = arr.normalized("parafac2").to_gpu() + +B = np.random.default_rng(0).normal(size=(arr.shape[1], 16)) +out = gpu @ B # a float32 CuPy array +out = B.T @ gpu # the other direction, likewise +``` + +The transfer keeps the value-compressed layout: `to_gpu()` uploads the same +`major_ptr`/`values`/`value_ptr`/`indices` arrays the host holds, plus the +$O(n_{\text{rows}} + n_{\text{cols}})$ statistics, never an expanded +one-float-per-nonzero copy. That is the point on a device, where memory is the +binding constraint: the benchmarks measure the device-resident view at 0.51x to +0.62x the bytes of the equivalent CuPy CSR. + +Keeping the layout is also why `vsparse` ships its own kernels. cuSPARSE reads +CSR and CSC only, so nothing in CuPy can walk the two-level +`major_ptr -> (values, value_ptr) -> indices` structure. Both directions of the +product, in both formats, are handled by the four kernels in `vsparse._cuda`. + +### Precision + +Everything floating-point on the device is **float32**, without exception -- +the values, the statistics, the dense operand and the result. The statistics +themselves are still computed on the host in float64 and narrowed on transfer, +since the $O(nnz)$ passes that derive them run once where the matmul runs many +times. + +So a device product will not match the host's float64 one to better than +roughly `1e-6` relative. In the two misaligned directions (`VCSC @ B` and +`B @ VCSR`) the kernels also accumulate with `atomicAdd`, so repeated runs can +differ in the last bits; the other two are deterministic. + +### Materializing instead + +`CudaNormalizedView.to_cupy_sparse()` is the device twin of +`to_scipy_sparse()`: it expands the runs into a CuPy CSR/CSC of the uncentered +`Delta` term, leaving `means` to subtract externally. Useful for handing the +matrix to a library that wants real cuSPARSE input, and it is the baseline the +kernels are benchmarked against. + +Two traps it handles, both worth knowing about if you build such a matrix +yourself: + +- A VCS slice stores its indices grouped by the value they share, never + ascending, so the expanded matrix is **not** canonical. cuSPARSE measured + ~6x slower on the unsorted form, which is why `to_cupy_sparse` sorts by + default (pass `sort_indices=False` to skip it). +- cuSPARSE is much faster with a CSR operand in *both* directions -- + `B @ csc` measured ~5x `B @ csr` -- so pass `format="csr"` if you will + multiply repeatedly, even from a VCSC-backed view. diff --git a/pyproject.toml b/pyproject.toml index 7050c08..9577e8e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,11 @@ Documentation = "https://meyer-lab.github.io/vsparse/" Issues = "https://github.com/meyer-lab/vsparse/issues" [project.optional-dependencies] +# CUDA support (`vsparse._cuda`). Kept an opt-in extra because the wheel is +# ~1 GB and useless without an NVIDIA GPU; CI's CPU jobs pass `--no-extra cuda`. +cuda = [ + "cupy-cuda13x[ctk]>=13.3", +] docs = [ "sphinx>=7.0", "furo>=2024.1.29", @@ -77,6 +82,11 @@ python-version = "3.12" [tool.pytest.ini_options] testpaths = ["tests"] addopts = "-ra" +markers = [ + # tests/test_cuda.py skips itself without a device, so the marker is for + # selecting (`-m cuda`) or excluding (`-m 'not cuda'`) the GPU suite. + "cuda: requires CuPy and an NVIDIA GPU", +] [tool.coverage.run] source = ["src/vsparse"] diff --git a/src/vsparse/__init__.py b/src/vsparse/__init__.py index c91c7b0..5474f48 100644 --- a/src/vsparse/__init__.py +++ b/src/vsparse/__init__.py @@ -9,18 +9,21 @@ from vsparse._anndata import from_anndata, to_layer from vsparse._anndata_class import VCSCAnnData from vsparse._base import VCSCArray, VCSRArray +from vsparse._cuda import CudaNormalizedView, cuda_is_available from vsparse._norm_common import RECIPES, Recipe from vsparse._rapid_load import load_and_normalize, load_packed from vsparse._vcs_norm import VCSCArrayNormalized, VCSRArrayNormalized __all__ = [ "RECIPES", + "CudaNormalizedView", "Recipe", "VCSCAnnData", "VCSCArray", "VCSCArrayNormalized", "VCSRArray", "VCSRArrayNormalized", + "cuda_is_available", "from_anndata", "load_and_normalize", "load_packed", diff --git a/src/vsparse/_cuda.py b/src/vsparse/_cuda.py new file mode 100644 index 0000000..41d45f0 --- /dev/null +++ b/src/vsparse/_cuda.py @@ -0,0 +1,592 @@ +"""CUDA transfer and matmul kernels for the normalized VCS views. + +:meth:`~vsparse._norm_common.NormalizedViewBase.to_gpu` moves a normalized +view onto an NVIDIA GPU as a :class:`CudaNormalizedView`, which supports the +same ``@``/``__rmatmul__`` the CPU view does. The array keeps its +value-compressed layout on the device -- the transfer is the same five arrays +the host holds (``major_ptr``/``values``/``value_ptr``/``indices`` plus the +``O(n_rows + n_cols)`` statistics), never an expanded one-float-per-nonzero +copy. That matters more on a GPU than on a host, since device memory is the +scarcer resource: the whole reason to hold a count matrix in VCS form is that +the expanded form does not fit. + +Keeping the layout is also what forces custom kernels. cuSPARSE speaks CSR and +CSC only, so nothing in CuPy can walk the two-level +``major_ptr -> (values, value_ptr) -> indices`` structure; the four kernels +below do, and :meth:`CudaNormalizedView.to_cupy_sparse` is the expanded +alternative they are measured against (see ``benchmarks/cases.py``). + +Everything floating-point on the device is ``float32`` +-- see :data:`DEVICE_DTYPE` and the precision note on :func:`to_gpu`. + +Kernel structure +---------------- + +As on the host (:mod:`vsparse._vcs_matmul`), ``A_norm = Delta + 1 (x) (-c*s)`` +with ``Delta`` nonzero only on the stored entries, so only ``Delta @ B`` / +``B @ Delta`` need a kernel; the rank-1 term is a CuPy one-liner. + +Each of the four kernels is one warp per major slice, grid-strided. The warp's +32 lanes split two ways: ``ww`` lanes (the smallest power of two at or above +the dense width, capped at 32) cover the width, and the remaining +``groups = 32 / ww`` lane-groups walk the slice's nonzeros in stride. A +PARAFAC2-shaped width of 8 would otherwise leave 24 of 32 lanes idle on every +nonzero; splitting recovers them. + +The four are the two independent binary choices the host kernels also make: + +- *which axis the statistics hang off* -- VCSR's major slice is a row, so + ``row_scale`` is uniform across the slice and ``gene_scale``/``col_post_scale`` + vary per nonzero; VCSC is the reverse. +- *aligned or not* -- when the output's disjoint axis is the major axis + (VCSR for ``self @ B``, VCSC for ``B @ self``) each warp owns its output row + outright, so it accumulates in registers, reduces across lane-groups with + ``__shfl_down_sync``, and stores once. In the misaligned direction the output + index varies per nonzero, so the kernel ``atomicAdd``s instead. + +The host's answer to that misaligned direction -- a private per-thread +accumulator, capped by :func:`~vsparse._ops.accumulator_threads` -- has no +analogue here and needs none: ``nthreads`` on a GPU is in the tens of +thousands, so a private copy each is out of the question, while float32 +``atomicAdd`` is a single hardware instruction and contention is low when the +scattered index is a cell id. The cost is that the misaligned kernels sum in +nondeterministic order, so repeated runs can differ in the last bits. The +aligned kernels are deterministic. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import numpy as np + +from vsparse._norm_common import _G_CODES, G_IDENTITY, G_LOG1P, G_LOG1P_1000X, G_SQRT + +if TYPE_CHECKING: + from vsparse._norm_common import NormalizedViewBase + +__all__ = ["DEVICE_DTYPE", "CudaNormalizedView", "cuda_is_available", "to_gpu"] + +#: Every floating-point array on the device, without exception: values, the +#: statistics, the dense operand and the result. +DEVICE_DTYPE = np.float32 + +#: Threads per block, as (lanes, warps). One warp handles one major slice. +_WARP = 32 +_WARPS_PER_BLOCK = 8 + +#: Width accumulators each lane may hold in the aligned kernels. A warp covers +#: ``ww * _ACC_MAX`` columns of the dense operand per pass over the slice, so +#: any width up to 128 is a single pass; wider needs one pass per chunk. +_ACC_MAX = 4 + +#: Blocks to launch. Capped rather than one-per-slice so a matrix with millions +#: of cells doesn't build an enormous grid for slices that are mostly empty; +#: the kernels grid-stride, so any cap is correct. +_MAX_BLOCKS = 4096 + + +def cuda_is_available() -> bool: + """Whether CuPy is importable and a CUDA device is actually present.""" + try: + import cupy as cp + except ImportError: + return False + try: + return cp.cuda.runtime.getDeviceCount() > 0 + except Exception: # pragma: no cover - driver present but unusable + return False + + +def _require_cupy() -> Any: + try: + import cupy as cp + except ImportError as exc: # pragma: no cover - depends on the install + raise ImportError( + "vsparse CUDA support requires CuPy; install it with `pip install vsparse[cuda]`" + ) from exc + return cp + + +# -- kernel source ------------------------------------------------------------ +# +# Assembled from fragments rather than written out four times: the four kernels +# differ only in the two substitutions described in the module docstring, and +# spelling each one out in full would be ~200 lines of near-identical CUDA. + +_PREAMBLE = f""" +#define ACC_MAX {_ACC_MAX} +#define FULL_MASK 0xffffffffu + +__device__ __forceinline__ float apply_g(const float x, const int g_code) {{ + if (g_code == {G_LOG1P}) return log1pf(x); + if (g_code == {G_LOG1P_1000X}) return log10f(1.0f + 1000.0f * x); + if (g_code == {G_SQRT}) return x > 0.0f ? sqrtf(x) : 0.0f; + return x; +}} +""" + +#: The transforms ``apply_g`` implements. A ``g_code`` missing from here would +#: fall through its final ``return x`` and silently compute the identity, so +#: this asserts against the host's own set instead: adding a transform to +#: :data:`~vsparse._norm_common.RECIPES` without teaching the kernels about it +#: fails at import, not in someone's results. +_CUDA_G_CODES = frozenset({G_IDENTITY, G_LOG1P, G_LOG1P_1000X, G_SQRT}) +if _CUDA_G_CODES != _G_CODES: + raise RuntimeError( + f"vsparse._cuda's apply_g implements {sorted(_CUDA_G_CODES)} but " + f"vsparse._norm_common defines {sorted(_G_CODES)}; add the missing " + "transform to the CUDA _PREAMBLE" + ) + +#: Arguments, identical for all four kernels. ``B``/``out`` are both +#: ``(*, width)`` and C-contiguous, whichever direction the product runs in. +_SIGNATURE = """ +extern "C" __global__ void {name}( + const long long* __restrict__ major_ptr, + const float* __restrict__ values, + const long long* __restrict__ value_ptr, + const int* __restrict__ indices, + const float* __restrict__ row_scale, + const float* __restrict__ gene_scale, + const float* __restrict__ col_post_scale, + const int g_code, + const float* __restrict__ B, + float* __restrict__ out, + const long long n_major, + const int width, + const int ww, + const int groups) +""" + +#: Per-warp lane bookkeeping and the grid-stride loop header, shared by both +#: body templates. ``wl`` indexes the width, ``nl`` the lane-group. +_WARP_SETUP = """ +{ + const int lane = threadIdx.x; + const int wl = lane & (ww - 1); + const int nl = lane / ww; + const long long stride = (long long)gridDim.x * blockDim.y; + for (long long maj = (long long)blockIdx.x * blockDim.y + threadIdx.y; + maj < n_major; maj += stride) { + const long long mstart = major_ptr[maj]; + const long long mstop = major_ptr[maj + 1]; +""" + +#: VCSR: the major slice is a row, so ``row_scale`` is uniform over it and the +#: per-gene statistics are looked up per nonzero (and a zero-variance gene is +#: skipped one nonzero at a time). +_SCALES_MAJOR_IS_ROW = { + "hoist": """ + const float rs = row_scale[maj]; +""", + "delta": """ + const float gs = gene_scale[mi]; + if (!(gs > 0.0f)) continue; + const float d = col_post_scale[mi] * apply_g(v / rs / gs, g_code); +""", +} + +#: VCSC: the major slice is a gene, so the per-gene statistics are uniform over +#: it -- including the zero-variance test, which skips the whole slice -- and +#: ``row_scale`` is the per-nonzero lookup. +_SCALES_MAJOR_IS_COL = { + "hoist": """ + const float gs = gene_scale[maj]; + if (!(gs > 0.0f)) continue; + const float s = col_post_scale[maj]; +""", + "delta": """ + const float d = s * apply_g(v / row_scale[mi] / gs, g_code); +""", +} + +#: The per-nonzero walk itself, shared by both bodies so the two cannot drift +#: apart in how they stride the layout. Opens two loops -- over the slice's +#: unique values, then over the indices sharing each -- which ``_WALK_CLOSE`` +#: closes, and leaves ``mi`` (the minor index) and ``d`` (the ``Delta`` entry) +#: in scope between them. Lane-group ``nl`` starts ``nl`` into each run and +#: steps by ``groups``, so the groups split the run between them. +_WALK_OPEN = """ + for (long long u = mstart; u < mstop; ++u) {{ + const float v = values[u]; + const long long kstop = value_ptr[u + 1]; + for (long long kk = value_ptr[u] + nl; kk < kstop; kk += groups) {{ + const int mi = indices[kk]; +{delta}""" + +_WALK_CLOSE = """ + }} + }}""" + +#: Aligned: the warp owns output row ``maj``, so it accumulates the width in +#: registers, reduces across lane-groups, and stores once. The dense operand is +#: indexed by the per-nonzero minor index. +_BODY_ALIGNED = """ +{hoist} + float* const orow = out + maj * (long long)width; + for (int cbase = 0; cbase < width; cbase += ww * ACC_MAX) {{ + float acc[ACC_MAX]; + #pragma unroll + for (int a = 0; a < ACC_MAX; ++a) acc[a] = 0.0f; +{walk_open} + const float* const brow = B + (long long)mi * width; + #pragma unroll + for (int a = 0; a < ACC_MAX; ++a) {{ + const int c = cbase + a * ww + wl; + if (c < width) acc[a] += d * brow[c]; + }}{walk_close} + + // Lanes sharing a `wl` sit `ww` apart, so folding by ww, 2*ww, ... + // lands each column's total in the nl == 0 group. + #pragma unroll + for (int a = 0; a < ACC_MAX; ++a) {{ + for (int off = ww; off < {warp}; off <<= 1) + acc[a] += __shfl_down_sync(FULL_MASK, acc[a], off); + }} + if (nl == 0) {{ + #pragma unroll + for (int a = 0; a < ACC_MAX; ++a) {{ + const int c = cbase + a * ww + wl; + if (c < width) orow[c] = acc[a]; + }} + }} + }} + }} +}} +""" + +#: Misaligned: the output row varies per nonzero, so there is nothing to +#: accumulate in a register and the kernel scatters with atomicAdd. The dense +#: operand is the one indexed by ``maj`` here, so its row is loop-invariant. +_BODY_MISALIGNED = """ +{hoist} + const float* const brow = B + maj * (long long)width; +{walk_open} + float* const orow = out + (long long)mi * width; + for (int c = wl; c < width; c += ww) + atomicAdd(orow + c, d * brow[c]);{walk_close} + }} +}} +""" + +#: ``name -> (body template, scale fragments)``. The names encode the format +#: the kernel walks and the direction of the product. +_KERNELS = { + "matmul_delta_csr": (_BODY_ALIGNED, _SCALES_MAJOR_IS_ROW), + "rmatmul_delta_csc": (_BODY_ALIGNED, _SCALES_MAJOR_IS_COL), + "matmul_delta_csc": (_BODY_MISALIGNED, _SCALES_MAJOR_IS_COL), + "rmatmul_delta_csr": (_BODY_MISALIGNED, _SCALES_MAJOR_IS_ROW), +} + + +def _build_source() -> str: + """The full CUDA source: all four kernels, assembled from the fragments above.""" + parts = [_PREAMBLE] + for name, (body, scales) in _KERNELS.items(): + # The walk is spliced in before formatting, not passed to it, because + # it carries a `{delta}` of its own that the same pass has to expand. + spliced = body.replace("{walk_open}", _WALK_OPEN).replace("{walk_close}", _WALK_CLOSE) + parts.append(_SIGNATURE.format(name=name)) + parts.append(_WARP_SETUP) + parts.append(spliced.format(warp=_WARP, **scales)) + return "".join(parts) + + +_module_cache: Any = None + + +def _module() -> Any: + """The compiled :class:`cupy.RawModule`, built once per process.""" + global _module_cache + if _module_cache is None: + cp = _require_cupy() + # --use_fast_math: measured 24% off the parafac2 kernel (every recipe + # but "raw" puts a transcendental on every nonzero) while the worst + # relative error against the float64 host view moved from 2.3e-7 to + # 3.4e-7 -- nothing, next to the float32 the whole device path is + # already committed to. Its denormal flushing is if anything the safer + # behaviour here: a denormal `gene_scale` fails the `> 0.0f` guard and + # its column is skipped, rather than dividing into an infinity. + _module_cache = cp.RawModule(code=_build_source(), options=("--use_fast_math",)) + return _module_cache + + +def _launch_config(n_major: int, width: int) -> tuple[tuple[int], tuple[int, int], int, int]: + """``(grid, block, ww, groups)`` for a slice count and a dense width.""" + ww = 1 + while ww < width and ww < _WARP: + ww <<= 1 + groups = _WARP // ww + blocks = min(_MAX_BLOCKS, max(1, -(-n_major // _WARPS_PER_BLOCK))) + return (blocks,), (_WARP, _WARPS_PER_BLOCK), ww, groups + + +# -- the device-resident view ------------------------------------------------- + + +class CudaNormalizedView: + """A :class:`~vsparse._norm_common.NormalizedViewBase` resident on a CUDA device. + + Built by :meth:`~vsparse._norm_common.NormalizedViewBase.to_gpu`, not + directly. Holds the wrapped array's value-compressed structure and the + recipe's statistics as CuPy arrays, and computes ``self @ B`` / ``B @ self`` + with the kernels in this module. All floating-point state is + :data:`DEVICE_DTYPE` (``float32``). + + The object pins its device memory for as long as it is alive; drop the + reference (and, if you need the memory back immediately, + ``cupy.get_default_memory_pool().free_all_blocks()``) to release it. + """ + + # Tell numpy to defer `ndarray @ CudaNormalizedView` to our __rmatmul__ + # instead of trying to broadcast us, exactly as the host types do. + __array_ufunc__ = None + + __slots__ = ( + "_format", + "col_mean", + "col_post_scale", + "gene_scale", + "indices", + "major_ptr", + "recipe", + "row_scale", + "shape", + "value_ptr", + "values", + ) + + def __init__(self, view: NormalizedViewBase) -> None: + cp = _require_cupy() + arr = view._arr + self._format = view._format + self.shape = view.shape + self.recipe = view.recipe + + # Structure. The pointer arrays index into `values`/`indices`, which can + # exceed 2**31 entries, so they stay 64-bit; `indices` holds a row or + # column number and is narrowed to int32 for the kernels to share one + # signature (a 2**31-row matrix is far past what fits on a device). + self.major_ptr = cp.asarray(arr.major_ptr, dtype=cp.int64) + self.value_ptr = cp.asarray(arr.value_ptr, dtype=cp.int64) + self.indices = cp.asarray(arr.indices, dtype=cp.int32) + self.values = cp.asarray(arr.values, dtype=DEVICE_DTYPE) + + # Statistics: O(n_rows + n_cols), negligible next to the structure. + self.row_scale = cp.asarray(view.row_scale, dtype=DEVICE_DTYPE) + self.gene_scale = cp.asarray(view.gene_scale, dtype=DEVICE_DTYPE) + self.col_mean = cp.asarray(view.col_mean, dtype=DEVICE_DTYPE) + self.col_post_scale = cp.asarray(view.col_post_scale, dtype=DEVICE_DTYPE) + + # -- introspection -------------------------------------------------------- + + @property + def dtype(self) -> np.dtype: + return np.dtype(DEVICE_DTYPE) + + @property + def nnz(self) -> int: + return int(self.indices.shape[0]) + + @property + def nbytes(self) -> int: + """Device bytes held, structure plus statistics.""" + return sum( + int(a.nbytes) + for a in ( + self.major_ptr, + self.value_ptr, + self.indices, + self.values, + self.row_scale, + self.gene_scale, + self.col_mean, + self.col_post_scale, + ) + ) + + @property + def means(self) -> Any: + """Per-gene mean-correction vector, ``col_post_scale * col_mean``.""" + return self.col_mean * self.col_post_scale + + def __repr__(self) -> str: # pragma: no cover - cosmetic + return ( + f"<{type(self).__name__} shape={self.shape} dtype={self.dtype} " + f"recipe={self.recipe.name!r} nnz={self.nnz}>" + ) + + # -- kernels -------------------------------------------------------------- + + def _kernel_args(self, g_code: int) -> tuple[Any, ...]: + return ( + self.major_ptr, + self.values, + self.value_ptr, + self.indices, + self.row_scale, + self.gene_scale, + self.col_post_scale, + np.int32(g_code), + ) + + def _run(self, name: str, B: Any, out: Any, n_major: int) -> None: + width = int(B.shape[1]) + grid, block, ww, groups = _launch_config(n_major, width) + kernel = _module().get_function(name) + kernel( + grid, + block, + ( + *self._kernel_args(self.recipe.g_code), + B, + out, + np.int64(n_major), + np.int32(width), + np.int32(ww), + np.int32(groups), + ), + ) + + def __matmul__(self, other: Any) -> Any: + """``self @ other`` for a dense ``other``; returns a float32 CuPy array.""" + cp = _require_cupy() + n_rows, n_cols = self.shape + B, squeeze = _prep_dense(cp, other, n_cols) + out = cp.zeros((n_rows, B.shape[1]), dtype=DEVICE_DTYPE) + + if self._format == "csr": + self._run("matmul_delta_csr", B, out, n_rows) + else: + self._run("matmul_delta_csc", B, out, n_cols) + + out += (-self.means) @ B # every row's implicit-zero contribution + return out[:, 0] if squeeze else out + + def __rmatmul__(self, other: Any) -> Any: + """``other @ self`` for a dense ``other``; returns a float32 CuPy array.""" + cp = _require_cupy() + n_rows, n_cols = self.shape + B2, squeeze = _prep_dense_right(cp, other, n_rows) + # (n_rows, p), so the kernels' inner loop over the width walks one + # contiguous row -- the same reason the host kernels transpose. + Bt = cp.ascontiguousarray(B2.T) + out_t = cp.zeros((n_cols, Bt.shape[1]), dtype=DEVICE_DTYPE) + + if self._format == "csc": + self._run("rmatmul_delta_csc", Bt, out_t, n_cols) + else: + self._run("rmatmul_delta_csr", Bt, out_t, n_rows) + + out = cp.ascontiguousarray(out_t.T) + out += B2.sum(axis=1)[:, None] * (-self.means)[None, :] + return out[0, :] if squeeze else out + + # -- the expanded alternative --------------------------------------------- + + def to_cupy_sparse(self, format: str | None = None, *, sort_indices: bool = True) -> Any: + """The uncentered ``Delta`` term as a CuPy CSR/CSC matrix, one float per nonzero. + + The device-side counterpart of + :meth:`~vsparse._norm_common.NormalizedViewBase.to_scipy_sparse`, and + the baseline the kernels are benchmarked against: it hands the product + to cuSPARSE, at the cost of expanding every value-compressed run and so + giving up the layout's whole memory advantage. :attr:`means` is left to + subtract externally, exactly as on the host. + + Parameters + ---------- + format : {"csr", "csc"}, optional + Sparse format to build. Defaults to whichever one matches the + wrapped array's own major axis, which is the one that costs + nothing to assemble. Worth overriding: cuSPARSE is markedly + faster with a CSR operand in *both* directions -- measurably so + for ``dense @ sparse``, where CSC ran ~5x slower in + ``benchmarks/cases.py`` -- so a caller multiplying repeatedly + should pay the one-time conversion and pass ``"csr"``. + sort_indices : bool, default True + Whether to sort each slice's indices ascending before returning. + A VCS slice stores its indices grouped by the value they share, + never in index order, so the expanded matrix is *not* canonical + unless this sorts it -- and cuSPARSE measured ~6x slower on the + unsorted form. That penalty is why this defaults to ``True`` + where the host's + :meth:`~vsparse._norm_common.NormalizedViewBase.to_scipy_sparse`, + whose result usually goes somewhere less picky, returns the + unsorted order. Pass ``False`` to skip the sort when the result + is used once, or by something that does not care. + """ + cp = _require_cupy() + import cupyx.scipy.sparse as cps + + indptr = self.value_ptr[self.major_ptr] + raw = cp.repeat(self.values, cp.diff(self.value_ptr)) + major = cp.repeat(cp.arange(len(indptr) - 1, dtype=cp.int32), cp.diff(indptr)) + if self._format == "csr": + row, col, ctor = major, self.indices, cps.csr_matrix + else: + row, col, ctor = self.indices, major, cps.csc_matrix + + gs = self.gene_scale[col] + scaled = cp.where(gs > 0, raw / self.row_scale[row] / gs, DEVICE_DTYPE(0)) + data = (self.col_post_scale[col] * _g_cupy(cp, scaled, self.recipe.g_code)).astype( + DEVICE_DTYPE + ) + out = ctor((data, self.indices.astype(cp.int32), indptr.astype(cp.int32)), shape=self.shape) + if format not in (None, "csr", "csc"): + raise ValueError(f"format must be 'csr', 'csc' or None, got {format!r}") + if format is not None and format != self._format: + out = out.tocsr() if format == "csr" else out.tocsc() + elif sort_indices: + out.sort_indices() + return out + + +def _g_cupy(cp: Any, x: Any, g_code: int) -> Any: + """Vectorized ``g`` on a CuPy array -- the device twin of ``_g_np``.""" + if g_code == G_LOG1P: + return cp.log1p(x) + if g_code == G_LOG1P_1000X: + return cp.log10(1.0 + 1000.0 * x) + if g_code == G_SQRT: + return cp.sqrt(cp.clip(x, 0.0, None)) + return x + + +def _prep_dense(cp: Any, other: Any, expect_rows: int) -> tuple[Any, bool]: + """A C-contiguous float32 device ``(expect_rows, k)``, and whether it was 1-D.""" + B = cp.asarray(other, dtype=DEVICE_DTYPE) + squeeze = B.ndim == 1 + if squeeze: + B = B.reshape(-1, 1) + if B.ndim != 2 or B.shape[0] != expect_rows: + raise ValueError(f"shape mismatch: expected first dimension {expect_rows}, got {B.shape}") + return cp.ascontiguousarray(B), squeeze + + +def _prep_dense_right(cp: Any, other: Any, expect_cols: int) -> tuple[Any, bool]: + """A C-contiguous float32 device ``(p, expect_cols)``, and whether it was 1-D.""" + B = cp.asarray(other, dtype=DEVICE_DTYPE) + squeeze = B.ndim == 1 + if squeeze: + B = B.reshape(1, -1) + if B.ndim != 2 or B.shape[1] != expect_cols: + raise ValueError(f"shape mismatch: expected last dimension {expect_cols}, got {B.shape}") + return cp.ascontiguousarray(B), squeeze + + +def to_gpu(view: NormalizedViewBase) -> CudaNormalizedView: + """Move ``view``'s structure and statistics onto the current CUDA device. + + Backs :meth:`~vsparse._norm_common.NormalizedViewBase.to_gpu`. + + The statistics are computed on the host, in float64, and narrowed on + transfer: the ``O(nnz)`` passes that derive them run once, where the matmul + the device is here for runs many times. Everything on the device is then + float32, so a product will not agree with the host's float64 one to better + than roughly ``1e-6`` relative -- and in the misaligned directions + (``VCSC @ B``, ``B @ VCSR``) the atomic accumulation also makes it + nondeterministic in the last bits. + """ + return CudaNormalizedView(view) diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index 63bd638..2c90fdb 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -48,6 +48,10 @@ in :mod:`vsparse._vcs_matmul` walk that same already-materialized ``indices`` array directly. :class:`VCSCArrayNormalized`/:class:`VCSRArrayNormalized` each supply their own ``__matmul__``/``__rmatmul__`` wired to those kernels. + +A view can also be moved onto a CUDA device with +:meth:`NormalizedViewBase.to_gpu`; :mod:`vsparse._cuda` walks the same layout +there with its own kernels, in float32. """ from __future__ import annotations @@ -1007,6 +1011,18 @@ def to_scipy_sparse(self, dtype: npt.DTypeLike = np.float64) -> Any: indptr = arr.value_ptr[arr.major_ptr] return ctor((data, arr.indices, indptr), shape=self.shape) + def to_gpu(self) -> Any: + """This view, moved onto the current CUDA device -- see :mod:`vsparse._cuda`. + + Returns a :class:`~vsparse._cuda.CudaNormalizedView`, which keeps the + value-compressed layout on the device (no one-float-per-nonzero + expansion) and supports the same ``@``/``__rmatmul__`` this view does, + in float32. Requires CuPy: ``pip install vsparse[cuda]``. + """ + from vsparse._cuda import to_gpu + + return to_gpu(self) + def norm_sq(self) -> float: """Squared Frobenius norm of the full normalized matrix, in ``O(nnz + n_cols)``.""" arr = self._arr diff --git a/tests/test_cuda.py b/tests/test_cuda.py new file mode 100644 index 0000000..b2505ec --- /dev/null +++ b/tests/test_cuda.py @@ -0,0 +1,189 @@ +"""CUDA kernels against the CPU view they mirror. + +Skipped wholesale without CuPy and a device. The GPU is float32 throughout +(:data:`vsparse._cuda.DEVICE_DTYPE`) while the host view is float64, so every +comparison here is a float32-scale tolerance, never an equality -- and in the +misaligned directions the kernels accumulate with ``atomicAdd``, so even two +GPU runs need not agree bit for bit. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import scipy.sparse as sp + +from vsparse import RECIPES, VCSCArray, VCSRArray, cuda_is_available + +pytestmark = pytest.mark.cuda + +if not cuda_is_available(): # pragma: no cover - environment dependent + pytest.skip("no CUDA device / CuPy available", allow_module_level=True) + + +#: Relative tolerance for a float32 product against the float64 reference. +#: The kernels sum ``nnz / n`` terms per output entry, so the bound is +#: float32 eps times a modest growth factor, not eps itself. +RTOL = 2e-4 +ATOL = 2e-4 + + +def counts(rng: np.random.Generator, shape: tuple[int, int], density: float = 0.4) -> sp.csr_array: + """Integer counts with repeated values, so the layout actually dedupes.""" + dense = rng.integers(1, 8, size=shape).astype(np.float64) + dense[rng.random(shape) > density] = 0.0 + return sp.csr_array(dense) + + +def assert_close(got, want) -> None: + import cupy as cp + + got = cp.asnumpy(got) + assert got.dtype == np.float32 + assert got.shape == want.shape + # `scale` guards the empty case, where `.max()` has no identity. + scale = max(1.0, float(np.abs(want).max())) if want.size else 1.0 + np.testing.assert_allclose(got, want, rtol=RTOL, atol=ATOL * scale) + + +@pytest.fixture(scope="module") +def mat() -> sp.csr_array: + return counts(np.random.default_rng(0), (300, 90)) + + +@pytest.mark.parametrize("recipe", sorted(RECIPES)) +@pytest.mark.parametrize("cls", [VCSRArray, VCSCArray]) +@pytest.mark.parametrize("width", [1, 3, 8, 32, 64, 130]) +def test_matmul_matches_cpu(mat, recipe, cls, width): + """``gpu @ B`` against ``cpu @ B``, in both formats and across the width chunking. + + ``width`` spans the lane-splitting cases: below a warp (lanes split across + nonzeros), exactly a warp, and past ``ww * _ACC_MAX`` (more than one pass). + """ + nv = cls.from_scipy(mat).normalized(recipe) + B = np.random.default_rng(1).normal(size=(mat.shape[1], width)) + assert_close(nv.to_gpu() @ B, nv @ B) + + +@pytest.mark.parametrize("recipe", sorted(RECIPES)) +@pytest.mark.parametrize("cls", [VCSRArray, VCSCArray]) +@pytest.mark.parametrize("width", [1, 3, 8, 32, 64, 130]) +def test_rmatmul_matches_cpu(mat, recipe, cls, width): + """``B @ gpu`` against ``B @ cpu``, in both formats and across the width chunking.""" + nv = cls.from_scipy(mat).normalized(recipe) + B = np.random.default_rng(2).normal(size=(width, mat.shape[0])) + assert_close(B @ nv.to_gpu(), B @ nv) + + +@pytest.mark.parametrize("cls", [VCSRArray, VCSCArray]) +def test_one_dimensional_operands_squeeze(mat, cls): + """A 1-D operand comes back 1-D, on both sides, as it does on the host.""" + nv = cls.from_scipy(mat).normalized("parafac2") + gpu = nv.to_gpu() + rng = np.random.default_rng(3) + + x = rng.normal(size=mat.shape[1]) + got = gpu @ x + assert got.ndim == 1 + assert_close(got, nv @ x) + + y = rng.normal(size=mat.shape[0]) + got = y @ gpu + assert got.ndim == 1 + assert_close(got, y @ nv) + + +@pytest.mark.parametrize("cls", [VCSRArray, VCSCArray]) +def test_device_operand_is_accepted(mat, cls): + """A CuPy operand is used as-is, with no host round-trip.""" + import cupy as cp + + nv = cls.from_scipy(mat).normalized("scanpy") + B = np.random.default_rng(4).normal(size=(mat.shape[1], 8)) + assert_close(nv.to_gpu() @ cp.asarray(B, dtype=cp.float32), nv @ B) + + +@pytest.mark.parametrize("cls", [VCSRArray, VCSCArray]) +def test_shape_mismatch_raises(mat, cls): + gpu = cls.from_scipy(mat).normalized("raw").to_gpu() + with pytest.raises(ValueError, match="shape mismatch"): + gpu @ np.zeros((mat.shape[1] + 1, 4)) + with pytest.raises(ValueError, match="shape mismatch"): + np.zeros((4, mat.shape[0] + 1)) @ gpu + + +@pytest.mark.parametrize("recipe", sorted(RECIPES)) +@pytest.mark.parametrize("cls", [VCSRArray, VCSCArray]) +def test_to_cupy_sparse_reproduces_the_view(mat, recipe, cls): + """The expanded ``Delta`` plus ``means`` is the same matrix the kernels imply. + + This is the baseline ``benchmarks/cases.py`` times the kernels against, so + it has to compute the same thing rather than merely run. + """ + import cupy as cp + + nv = cls.from_scipy(mat).normalized(recipe) + gpu = nv.to_gpu() + dense = cp.asnumpy(gpu.to_cupy_sparse().toarray()) - cp.asnumpy(gpu.means)[None, :] + np.testing.assert_allclose(dense, nv.toarray(), rtol=RTOL, atol=ATOL) + + +@pytest.mark.parametrize("cls", [VCSRArray, VCSCArray]) +def test_transfer_keeps_the_compressed_layout(mat, cls): + """``to_gpu`` uploads the VCS structure, not an expanded float-per-nonzero copy. + + The point of the device view: ``values`` stays ``n_unique``-long, so a + matrix that only fits in device memory compressed still fits. + """ + arr = cls.from_scipy(mat) + gpu = arr.normalized("parafac2").to_gpu() + assert gpu.values.shape[0] == arr.n_unique + assert arr.n_unique < arr.nnz # the fixture does dedupe, so this is a real check + assert gpu.nnz == arr.nnz + # An expanded CSR would need 4 bytes of data + 4 of indices per nonzero, + # before either pointer array. + assert gpu.nbytes < 8 * arr.nnz + + +def test_stats_are_float32_on_device(mat): + """Nothing floating-point survives the transfer at float64.""" + import cupy as cp + + gpu = VCSRArray.from_scipy(mat).normalized("pearson").to_gpu() + for name in ("values", "row_scale", "gene_scale", "col_mean", "col_post_scale"): + assert getattr(gpu, name).dtype == cp.float32, name + assert ( + VCSRArray.from_scipy(mat).normalized("raw").to_gpu() @ np.zeros(mat.shape[1]) + ).dtype == (cp.float32) + + +def test_empty_and_degenerate_columns_are_handled(): + """A zero column gives ``gene_scale == 0``; the kernels must skip, not divide.""" + dense = np.zeros((20, 6)) + dense[:, 1] = np.arange(20) % 4 + dense[:, 4] = 3.0 + nv = VCSRArray.from_scipy(sp.csr_array(dense)).normalized("parafac2") + B = np.random.default_rng(5).normal(size=(6, 8)) + got = nv.to_gpu() @ B + assert np.isfinite(np.asarray(got.get())).all() + assert_close(got, nv @ B) + + +@pytest.mark.parametrize("shape", [(1, 1), (1, 7), (5, 1), (3, 3)]) +@pytest.mark.parametrize("cls", [VCSRArray, VCSCArray]) +@pytest.mark.parametrize("width", [0, 1, 5]) +def test_degenerate_shapes_and_widths(shape, cls, width): + """Shapes and operand widths that leave a kernel with nothing to do. + + A zero width makes the aligned kernels' chunk loop trip zero times, and a + single-row or single-column matrix leaves the grid-stride loop with one + slice for thousands of warps. Both should produce a correctly shaped, + correct result rather than a launch error. + """ + dense = np.zeros(shape) + dense[0, 0] = 2.0 + nv = cls.from_scipy(sp.csr_array(dense)).normalized("parafac2") + gpu = nv.to_gpu() + + assert_close(gpu @ np.ones((shape[1], width)), nv @ np.ones((shape[1], width))) + assert_close(np.ones((width, shape[0])) @ gpu, np.ones((width, shape[0])) @ nv) diff --git a/uv.lock b/uv.lock index 5593e93..a4dbdab 100644 --- a/uv.lock +++ b/uv.lock @@ -350,6 +350,83 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b1/5a/234e8fadf85c3cc48cb31c247b9e8e0c7f06ece80f5b29f9b8c241f9da4c/coverage-7.16.0-py3-none-any.whl", hash = "sha256:245f7de6d023a5bba375dbec9f2e0869bfa26ac0cc639bbb7b4c814884000b73", size = 214977, upload-time = "2026-08-28T21:54:35.189Z" }, ] +[[package]] +name = "cuda-pathfinder" +version = "1.8.2" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/98/59/239c7259e669c46ddbcac0aa60e3a0ef00bfeaaa687f905b24dd6a7a10fe/cuda_pathfinder-1.8.2-py3-none-any.whl", hash = "sha256:4e65059febdb4d19d5cbc4798677e19db2b582f2f702f457b609e571690d357e", size = 62551, upload-time = "2026-09-17T22:19:09.396Z" }, +] + +[[package]] +name = "cuda-toolkit" +version = "13.4.2" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/80/02/ba4b5eaec47fb1ae01138ab0d88c735e14e594a7c1bd8146cdd4dfe15d89/cuda_toolkit-13.4.2-py2.py3-none-any.whl", hash = "sha256:2e79d99df4f3c5b3102fa4a5a69eb4c37a06ebf06b9aef92c18078c26e8db004", size = 2693, upload-time = "2026-09-16T20:57:11.438Z" }, +] + +[package.optional-dependencies] +cublas = [ + { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +cudart = [ + { name = "nvidia-cuda-runtime", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +cufft = [ + { name = "nvidia-cufft", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +curand = [ + { name = "nvidia-curand", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +cusolver = [ + { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-cusolver", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +cusparse = [ + { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +nvrtc = [ + { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] + +[[package]] +name = "cupy-cuda13x" +version = "14.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-pathfinder" }, + { name = "numpy" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/dd/c8/68e07ff959f6a9f93169e37a64a4c24223487102b8625191e4a566cab465/cupy_cuda13x-14.2.0-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:0a9b6538079c5151dcd99cf44b94cc3f1242315b7de0498b4619cef9b2a7c2c8", size = 73382974, upload-time = "2026-08-20T02:41:13.648Z" }, + { url = "https://files.pythonhosted.org/packages/7e/bf/c4c25e97552320451afdcd459d4b6e0b715459594e998b23051c06db64c3/cupy_cuda13x-14.2.0-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:9aafadbe29db3d4285026044aafeb28f2a3de4bd0df70718ec6892b7d27a4ded", size = 69439247, upload-time = "2026-08-20T02:41:17.375Z" }, + { url = "https://files.pythonhosted.org/packages/5e/88/0a4d2d075485c4a82bfcd4725300ff87d17de36dce27b36f720d286f0f70/cupy_cuda13x-14.2.0-cp312-cp312-win_amd64.whl", hash = "sha256:9e8bc8a7693c5c212cd1450d728f47e1f05aa07ae003de30415893d0daf7d243", size = 35986921, upload-time = "2026-08-20T02:41:20.253Z" }, + { url = "https://files.pythonhosted.org/packages/a4/ab/8a756f553397b8509906d4a1d8083d92d3d21cbde115504e7a5db97de61f/cupy_cuda13x-14.2.0-cp312-cp312-win_arm64.whl", hash = "sha256:0dae6d788752c987c1cd87adb903640620820feb895339e8bc51a01fba0fb804", size = 35049803, upload-time = "2026-09-01T06:41:43.442Z" }, + { url = "https://files.pythonhosted.org/packages/3c/02/3cb39774c71a16275c50b5d1e92ddcefa948ff1e69f464249b70a79b0abe/cupy_cuda13x-14.2.0-cp313-cp313-manylinux2014_aarch64.whl", hash = "sha256:449af7f608dc70ce64e8c45cf5a6ea0bbef51623e7c806b1a43445ae3c76e1e0", size = 72969322, upload-time = "2026-08-20T02:41:23.311Z" }, + { url = "https://files.pythonhosted.org/packages/4e/27/216370ae8c39fa9835645be3c27513ab72f52cf4f51faf58d799fba981b6/cupy_cuda13x-14.2.0-cp313-cp313-manylinux2014_x86_64.whl", hash = "sha256:5b79a606f6639d9bed74ef7072f36756a95a804c86d7adb367b39307ac771600", size = 69048238, upload-time = "2026-08-20T02:41:26.831Z" }, + { url = "https://files.pythonhosted.org/packages/5d/75/9df3c545baad4ff2e384db1ff6c8005fa4040492c4ee5672ca8b8711d29d/cupy_cuda13x-14.2.0-cp313-cp313-win_amd64.whl", hash = "sha256:f9a2a55abf889f6df68b5f62cd5f3ec4a2f305ac57ed09becd52e16201edcd86", size = 35965490, upload-time = "2026-08-20T02:41:29.943Z" }, + { url = "https://files.pythonhosted.org/packages/5c/3d/ee85ce39c2f32e8561c81c507df458a914b45d9a0eb1a54e5396c2cfd632/cupy_cuda13x-14.2.0-cp313-cp313-win_arm64.whl", hash = "sha256:5ddaf38291f60ec20e8c7889b28e339bf17ad1b4df768b8bb5e6c9f9e445efbb", size = 35029087, upload-time = "2026-09-01T06:41:48.593Z" }, + { url = "https://files.pythonhosted.org/packages/e5/70/8c80c9e7010193afed8ccb7110d6173aa23bf3a6ce2fc2450244b0fa50a1/cupy_cuda13x-14.2.0-cp314-cp314-manylinux2014_aarch64.whl", hash = "sha256:7ac57c4e16c62f265b5ecc2b45d6bbafe9dc70a99b849ed6cfb8d16f32204589", size = 72840193, upload-time = "2026-08-20T02:41:33.146Z" }, + { url = "https://files.pythonhosted.org/packages/c2/b8/4f4c4f34fc31ab8d136ed965a919505297974629d9852c50e15bfd616281/cupy_cuda13x-14.2.0-cp314-cp314-manylinux2014_x86_64.whl", hash = "sha256:ff0bdebd1b43c0c6db53095784c787c4e4eae671356cf521c4b7482ed78a1a7e", size = 68422828, upload-time = "2026-08-20T02:41:36.527Z" }, + { url = "https://files.pythonhosted.org/packages/7d/5c/58941ece5ed0b2cc8304f8ba81b5e8334759d409306475d53c52ac99809b/cupy_cuda13x-14.2.0-cp314-cp314-win_amd64.whl", hash = "sha256:b99a8bf9d5391954c0c08675a8cb39219923fe2ee249faa3546ebbf75783a360", size = 36110882, upload-time = "2026-08-20T02:41:39.837Z" }, + { url = "https://files.pythonhosted.org/packages/e2/c5/cc7d1a9b764b3f828e09b48a73d0fdeaad15ecd2c7e50507f6d5fcf5fd1a/cupy_cuda13x-14.2.0-cp314-cp314-win_arm64.whl", hash = "sha256:c61c59ba665948c0a855926997ee96026420b628816b2e406998ec797caddb22", size = 35184806, upload-time = "2026-09-01T06:41:58.317Z" }, + { url = "https://files.pythonhosted.org/packages/6f/54/a582ac70e9e75d3fad7ae073051bcd50dd128a1f66dd275291dc5f48d997/cupy_cuda13x-14.2.0-cp314-cp314t-manylinux2014_aarch64.whl", hash = "sha256:e8bb453551deced08ec5c4576da61ebf1e2b4a45aeff48e09f47ec332f7cad3d", size = 73336756, upload-time = "2026-08-20T02:41:42.718Z" }, + { url = "https://files.pythonhosted.org/packages/f9/aa/222151ba5061d7a01efe4009efad9f9a7164e15bb2a494d9aa86efc97d4f/cupy_cuda13x-14.2.0-cp314-cp314t-manylinux2014_x86_64.whl", hash = "sha256:313500c5d415f65d388e84a9d2eb7f07f97c23c85dc5785f986e75d46d9c9693", size = 68792663, upload-time = "2026-08-20T02:41:46.079Z" }, + { url = "https://files.pythonhosted.org/packages/d4/20/13dd37dbeaec879c8348fde2bfca67d96daa56a6e63e4a5083ff61660e53/cupy_cuda13x-14.2.0-cp314-cp314t-win_amd64.whl", hash = "sha256:c517fc55502a7e3d1a8b9310abad95cec61dded5a38d328fe0220bcf7b036268", size = 37004845, upload-time = "2026-08-20T02:41:48.963Z" }, + { url = "https://files.pythonhosted.org/packages/d1/5b/b1367f79abd90a259dbb7f0603dd162c03c3a76ec1d02c2527b0c79d46ba/cupy_cuda13x-14.2.0-cp314-cp314t-win_arm64.whl", hash = "sha256:83885db963eb0b95a0757b4907a37115a1d2aebfa9cf43b99ebac8db35713d37", size = 35485280, upload-time = "2026-09-01T06:41:53.618Z" }, +] + +[package.optional-dependencies] +ctk = [ + { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "curand", "cusolver", "cusparse", "nvrtc"] }, +] + [[package]] name = "docutils" version = "0.22.4" @@ -947,6 +1024,108 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/0b/b12a2df5d1b774bd9007a6fdff9381145b6223d37f11afc9c37ab0efd9a1/numpy-2.5.3-cp315-cp315t-win_arm64.whl", hash = "sha256:befa1ae5bd6030b3f512b43ff3fa5290bbed6b84411a44244b14adf835f5b89d", size = 10850807, upload-time = "2026-09-06T16:27:43.868Z" }, ] +[[package]] +name = "nvidia-cublas" +version = "13.8.0.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cuda-nvrtc" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c9/9e/0c73b14c14cd32e8b3a028133371f4bf1d9d14b811c515c5fd1201a3c38d/nvidia_cublas-13.8.0.4-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e22b25be18f8b7d267dbbe2a8ec78e2c69c05762536d299dd0afa2c847115740", size = 552729449, upload-time = "2026-09-16T20:40:37.357Z" }, + { url = "https://files.pythonhosted.org/packages/7a/38/bdd540bf511d2c9b6f9efc71a81c60cb88e295be0b9312b61d19bbed2212/nvidia_cublas-13.8.0.4-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:9f17797dfcc048694461f4e47de17d2e3c25adf172ef723d2db0a07cd8744b89", size = 439317144, upload-time = "2026-09-16T20:41:11.569Z" }, + { url = "https://files.pythonhosted.org/packages/a3/df/f1246959833e2c437db8be3e5b477f66b87f8817821ed40de6c7561c9a36/nvidia_cublas-13.8.0.4-py3-none-win_amd64.whl", hash = "sha256:8c5494423bb8a46822cb6b0cb95d7fa4be2d7b96a31155dff083839ec8297910", size = 423266897, upload-time = "2026-09-16T20:48:16.789Z" }, + { url = "https://files.pythonhosted.org/packages/96/38/723c97681824250b261c91761de3f1a1899111a9b1dda65eec5b2daa1d58/nvidia_cublas-13.8.0.4-py3-none-win_arm64.whl", hash = "sha256:a2ffda7a27d8315e6b75c1ced526d5a68c7a2a50adfe866235401c9592d3c9e4", size = 153856980, upload-time = "2026-09-16T20:54:58.249Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc" +version = "13.4.92" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/9c/1342ebb460ce2afd014ec5a002adc99108e3c53a6b70f399cf64dd1f267d/nvidia_cuda_nvrtc-13.4.92-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:5ce8c97b00b232c4f50c8c4b5a3b68cafee08bdb82ea86f2052ff01d03194f4a", size = 53301733, upload-time = "2026-09-16T20:39:06.686Z" }, + { url = "https://files.pythonhosted.org/packages/4f/73/cf76dc0083d68bd1e29309bc51fc95ada28d4b64d341e945eb3f8ec8edec/nvidia_cuda_nvrtc-13.4.92-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:24b9f5eccc6a5a19779038cf468aecb7cecfa8716269beaef5640bd989d21c28", size = 50927348, upload-time = "2026-09-16T20:38:55.548Z" }, + { url = "https://files.pythonhosted.org/packages/64/5b/aa91896f64444eff1dd7f189f4db05046c7f0ba08c0f531283c1c1aac581/nvidia_cuda_nvrtc-13.4.92-py3-none-win_amd64.whl", hash = "sha256:6af7ac5372920f6a7a560d0699348560fe44cb52c1d722d6afd3119f8af948c4", size = 46972283, upload-time = "2026-09-16T20:47:13.295Z" }, + { url = "https://files.pythonhosted.org/packages/25/22/af15f5ef51ba0f90d5439487b8efe482d7b2cdafa8ff6a01e13791622b98/nvidia_cuda_nvrtc-13.4.92-py3-none-win_arm64.whl", hash = "sha256:1620066e967e93119d67338628935cbb1196b53a1474b07f14b7b77aba477284", size = 42709926, upload-time = "2026-09-16T20:54:16.396Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime" +version = "13.4.92" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/51/3b/2c6e9d88efa7b3585572760bca5e7da83fa9e06098a9386c7b29c81c46b2/nvidia_cuda_runtime-13.4.92-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:bef071788589550ab02846fcbc732a2b88c2b16f2cd8a0f45d98e68e8c5fb0c8", size = 2509438, upload-time = "2026-09-16T20:37:16.36Z" }, + { url = "https://files.pythonhosted.org/packages/98/8a/3431271f6344874b8f1ac03f16b3d679c91493f8da63f716160403e6d0a0/nvidia_cuda_runtime-13.4.92-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9641f797da20ce1dd8e779b6e96d08cf9ba564cec8e8225458811ee26423f3a5", size = 2494438, upload-time = "2026-09-16T20:37:21.89Z" }, + { url = "https://files.pythonhosted.org/packages/86/00/d5436004268f049214193659ebc36550b5ef3925c3d13b4cc980e13be6f5/nvidia_cuda_runtime-13.4.92-py3-none-win_amd64.whl", hash = "sha256:08dca5e4aba480c2fd5b55075c0fa71b84ef9dcf0521f2d58baa14a803a7311c", size = 2778543, upload-time = "2026-09-16T20:46:42.459Z" }, + { url = "https://files.pythonhosted.org/packages/b5/bc/9a141b73d65b9f09e548fcb9b7f4415da22f35e0fe78c60e3c0a9b458efd/nvidia_cuda_runtime-13.4.92-py3-none-win_arm64.whl", hash = "sha256:43972819798ca06ad6354f6cbdb82a0a113f88b94e9ea90256d660027b0b3b60", size = 2770284, upload-time = "2026-09-16T20:53:45.718Z" }, +] + +[[package]] +name = "nvidia-cufft" +version = "12.4.0.43" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/41/f1/bb8cc4e1ffee345a51fe4cf656352bc705274431dcbd6c23ed9f00585ec7/nvidia_cufft-12.4.0.43-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3938644a5b594e06d396e02d6e52bdd84589b3c1f3a8a127a2704211c2abc441", size = 161750525, upload-time = "2026-09-16T20:41:56.892Z" }, + { url = "https://files.pythonhosted.org/packages/76/bf/3fea3d1c6262bded26ae00e3106432d63235954965d7785f5051ac146651/nvidia_cufft-12.4.0.43-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0e8385013596b112d29c9ce8c63dc575b308d77636c7169104e18714f03961a8", size = 161750089, upload-time = "2026-09-16T20:42:15.469Z" }, + { url = "https://files.pythonhosted.org/packages/7b/cc/be7fe31058127336a66c88414a2ecc6beecf680baf85994a43eaf292b202/nvidia_cufft-12.4.0.43-py3-none-win_amd64.whl", hash = "sha256:4ff7075f2d0b5f69291f70938d37a86ec632cbe5747184c74ba1f50f17accacc", size = 160953139, upload-time = "2026-09-16T20:48:43.428Z" }, + { url = "https://files.pythonhosted.org/packages/0b/ee/9217a02eae5b0b183b62b39b65f1e64ee8ffb6d6da67a79fd7a2f295ed3a/nvidia_cufft-12.4.0.43-py3-none-win_arm64.whl", hash = "sha256:4e8d551542bd661aef431422acddc5f06f793ac9b1dc7af53de63cbfde5781f9", size = 161850334, upload-time = "2026-09-16T20:55:16.698Z" }, +] + +[[package]] +name = "nvidia-curand" +version = "10.4.4.72" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/9d/42faff77e90e5498eecb9b099894d5deea744b3d6b35e700dafd0381e07a/nvidia_curand-10.4.4.72-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:53bef256d4362eb3d70c9c048a8d31bcf003f9dcade17720c73d3a494a7dcaa3", size = 63771251, upload-time = "2026-09-16T20:42:40.644Z" }, + { url = "https://files.pythonhosted.org/packages/07/73/3ee8e5b4cb891401e603ffd3a59b35c6afe785fd2de123afe7c7029603dc/nvidia_curand-10.4.4.72-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:25c3457ae7a224fdd484dab90b0fc5dc0e842fab5db3012afa4a5bd2af4eb7e5", size = 61498332, upload-time = "2026-09-16T20:42:51.685Z" }, + { url = "https://files.pythonhosted.org/packages/e2/d4/f59ab342b82ff05ff27d674acd470ba625d788bf71338f7da6c89669d6b9/nvidia_curand-10.4.4.72-py3-none-win_amd64.whl", hash = "sha256:e0bce83e083ef25976ee74f59e8f067c15149a74538f1be6c2462e286e7c9c68", size = 56731059, upload-time = "2026-09-16T20:48:53.631Z" }, + { url = "https://files.pythonhosted.org/packages/8f/60/6734d25ee55e688a88d4ea5461d71f2cd9d8f0e8da50d5d10f3e6408f855/nvidia_curand-10.4.4.72-py3-none-win_arm64.whl", hash = "sha256:4635b2c8a727f51b585614f590509cb5c50fb2fc55fa9350f01c22c7f8c0fcdf", size = 65140169, upload-time = "2026-09-16T20:55:30.897Z" }, +] + +[[package]] +name = "nvidia-cusolver" +version = "12.3.4.7" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas" }, + { name = "nvidia-cusparse" }, + { name = "nvidia-nvjitlink" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/77/f1/8c8a667c59cb4ea0dd1543ed99c6552d818e2a4bb922d0ee98f0a49de207/nvidia_cusolver-12.3.4.7-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:4a38d88a1ea3f7b656e52001caf943b625822a89df3a8b63d895f56ef777b1a4", size = 291070656, upload-time = "2026-09-16T20:43:15.438Z" }, + { url = "https://files.pythonhosted.org/packages/3f/43/e1a29568cc9989c95bcf31519f69baccfaebdf5ef5e624773726cc9740e2/nvidia_cusolver-12.3.4.7-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:225dd543c7b93ca22e62a35b5f4d8b4f9caa515733793925c05cdcf58da5cf02", size = 246439468, upload-time = "2026-09-16T20:43:34.367Z" }, + { url = "https://files.pythonhosted.org/packages/65/0b/1d751c92e1429ef3261021eb4b6cf27076b547d7512c6dc4b9aebeb9c49a/nvidia_cusolver-12.3.4.7-py3-none-win_amd64.whl", hash = "sha256:7ed56898cd98abe36d8727eaf1408622c67557f2acb3baca5a27e24e595a1ccb", size = 261421900, upload-time = "2026-09-16T20:49:15.844Z" }, + { url = "https://files.pythonhosted.org/packages/51/e7/e001caaf3101d01d80583753253a041b47db91e67fe71fb6320ff7afd7c7/nvidia_cusolver-12.3.4.7-py3-none-win_arm64.whl", hash = "sha256:87891ab21de591da154070bd58d5fd3ab52ccaeb900f741a7ead0c8a859e44ac", size = 71318077, upload-time = "2026-09-16T20:55:41.128Z" }, +] + +[[package]] +name = "nvidia-cusparse" +version = "12.8.6.72" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/2e/bf04fbd787d6b227da62abbe3674de5d89460dacebbf5cc97b5717b89a49/nvidia_cusparse-12.8.6.72-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c3917c86fd419cbd42229b5cccb8bbe5200ed5b7dec3aa86d28e25c90a256e1b", size = 188672939, upload-time = "2026-09-16T20:43:52.751Z" }, + { url = "https://files.pythonhosted.org/packages/80/cb/409df72613bf1c6cd9b545a257ab5489ecfd98c88cdf7ee9d7fc434e94a1/nvidia_cusparse-12.8.6.72-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a739f6ff51ea2a8a2990b267b9b4d3063a5d94890efb8ccf393bfcab9dd487aa", size = 170468146, upload-time = "2026-09-16T20:44:11.113Z" }, + { url = "https://files.pythonhosted.org/packages/85/f7/e0d0edefab227c9967df12484fd50e5586b9af86439f38354779c7086e8b/nvidia_cusparse-12.8.6.72-py3-none-win_amd64.whl", hash = "sha256:013d3f83316431dd3338cf15141dd5420b24dffc6663fc083c053c8409fa3cb7", size = 168202645, upload-time = "2026-09-16T20:49:32.639Z" }, + { url = "https://files.pythonhosted.org/packages/23/ff/bedb3859cedf60ba14c761181f91bcfe926fb072bf686deb5f93c6cb7b10/nvidia_cusparse-12.8.6.72-py3-none-win_arm64.whl", hash = "sha256:7efe07b54c505f3eeec0469078c992a31b15ab716d40f07a8109f7cbcc4fd7fe", size = 184995442, upload-time = "2026-09-16T20:55:58.013Z" }, +] + +[[package]] +name = "nvidia-nvjitlink" +version = "13.4.92" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1d/6b/eef7a9e32872b8f41e145bf10cddc9af26e153c338852811fe9a9baddf9e/nvidia_nvjitlink-13.4.92-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:e0391f24ed94ec879b84e3da4d4ec320c879aff681f2c7a638462f7199284323", size = 42452378, upload-time = "2026-09-16T20:45:29.042Z" }, + { url = "https://files.pythonhosted.org/packages/1f/a8/1cbd4014898af8b419e69b0d7dbc63da2121ee92d92b47d59f4fe9075349/nvidia_nvjitlink-13.4.92-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:25f74fad0d654271c921ac4dca614bd6258bc21791242fc7b2289dad7ae9c099", size = 40420120, upload-time = "2026-09-16T20:45:19.163Z" }, + { url = "https://files.pythonhosted.org/packages/72/7c/c44bd7277afafed9de262a23baec4dc8dcf7ecbb1b9ecb228821442d1e9b/nvidia_nvjitlink-13.4.92-py3-none-win_amd64.whl", hash = "sha256:b286f3a4f227a9363efdec263c7b91788cef1478d2b8a5fa8bab7f3e82ff82fd", size = 39053950, upload-time = "2026-09-16T20:52:54.244Z" }, + { url = "https://files.pythonhosted.org/packages/3a/8e/afaa7687fac05a507d10765f6097937e99735f835c5b262276d1e2d2cad5/nvidia_nvjitlink-13.4.92-py3-none-win_arm64.whl", hash = "sha256:9e4a7ff4f0cafa8c624917055b863dc11f5c2912ead23c462889e166f3b0e57d", size = 35981621, upload-time = "2026-09-16T20:56:30.819Z" }, +] + [[package]] name = "packaging" version = "26.3" @@ -1653,6 +1832,9 @@ dependencies = [ ] [package.optional-dependencies] +cuda = [ + { name = "cupy-cuda13x", extra = ["ctk"] }, +] docs = [ { name = "furo" }, { name = "myst-parser" }, @@ -1677,6 +1859,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "anndata", specifier = ">=0.13" }, + { name = "cupy-cuda13x", extras = ["ctk"], marker = "extra == 'cuda'", specifier = ">=13.3" }, { name = "furo", marker = "extra == 'docs'", specifier = ">=2024.1.29" }, { name = "hdf5plugin", specifier = ">=4.0" }, { name = "myst-parser", marker = "extra == 'docs'", specifier = ">=3.0" }, @@ -1686,7 +1869,7 @@ requires-dist = [ { name = "sphinx", marker = "extra == 'docs'", specifier = ">=7.0" }, { name = "sphinx-autodoc-typehints", marker = "extra == 'docs'", specifier = ">=2.0" }, ] -provides-extras = ["docs"] +provides-extras = ["cuda", "docs"] [package.metadata.requires-dev] dev = [ From 34fe8e9f459f02ac9439b157219a5e045fff143d Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sun, 20 Sep 2026 07:10:11 -0700 Subject: [PATCH 2/2] Run the CUDA job on any self-hosted runner, not a labelled one Extra labels are another thing to register on each machine and keep in sync with the workflow, and getting them wrong leaves the job queued forever rather than failing. Plain `self-hosted` needs no setup beyond registering a runner. The trade is that the job can now land on a machine with no usable GPU. The "Assert a CUDA device is visible" step already caught that -- it has to exist regardless, since tests/test_cuda.py skips itself when CuPy cannot see a device and the job would otherwise pass green having tested nothing -- so it just becomes the job's only GPU gate rather than a second one. `nvidia-smi` becomes non-fatal so that assert prints the actionable message instead of the step before it dying on "not found". Co-Authored-By: Claude Opus 5 --- .github/workflows/ci.yml | 25 ++++++++++++++++--------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 10fcebb..2e19e1c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -93,13 +93,15 @@ jobs: cuda: name: CUDA Tests and Benchmarks - # A self-hosted runner with an NVIDIA GPU and a driver new enough for the - # `cupy-cuda13x` wheel. Labelled rather than named so more than one machine - # can serve the job. Forked PRs never see this runner, and `pull_request` - # from a fork cannot reach a self-hosted runner anyway -- but the guard is - # explicit so an untrusted PR can never run code on our own hardware. + # Any self-hosted runner, with no extra labels to register or keep in sync. + # That means the job can land on a machine without a usable GPU, which the + # "Assert a CUDA device is visible" step below turns into a loud failure + # rather than a silent skip -- the same check that has to be there anyway. + # Forked PRs never see this runner, and `pull_request` from a fork cannot + # reach a self-hosted runner regardless -- but the guard is explicit so an + # untrusted PR can never run code on our own hardware. if: github.event_name == 'push' || github.event.pull_request.head.repo.full_name == github.repository - runs-on: [self-hosted, linux, x64, gpu, cuda] + runs-on: self-hosted steps: - name: Checkout code uses: actions/checkout@v7 @@ -110,15 +112,20 @@ jobs: enable-cache: true python-version: "3.12" + # Informational, and deliberately non-fatal: on a runner with no GPU at + # all this would otherwise fail first, with a bare "nvidia-smi: not + # found" instead of the actionable message the assert step below prints. - name: Show the device the job landed on + continue-on-error: true run: nvidia-smi --query-gpu=name,driver_version,compute_cap,memory.total --format=csv - name: Install dependencies run: uv sync --all-extras --dev - # tests/test_cuda.py skips itself if CuPy cannot see a device, which on - # this runner would mean the job silently passed without testing - # anything. Assert the device is really there first. + # Load-bearing, not belt-and-braces: `runs-on: self-hosted` carries no + # promise of a GPU, and tests/test_cuda.py skips itself when CuPy cannot + # see one -- so without this the job would pass green having tested + # nothing at all. - name: Assert a CUDA device is visible run: uv run python -c "import vsparse, sys; sys.exit(0 if vsparse.cuda_is_available() else 'no CUDA device visible to CuPy')"