diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 34f61ba..2e19e1c 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,52 @@ 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 + # 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 + 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" + + # 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 + + # 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')" + + - 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 = [