Benchmark suite ci

.github/workflows/ci.yml

@@ -55,3 +55,26 @@ jobs:
 
       - name: Run pytest
         run: uv run pytest --cov
+
+  benchmarks:
+    name: Benchmark Regression Gate
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v5
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v10.0.1
+        with:
+          enable-cache: true
+          python-version: "3.12"
+
+      - name: Install dependencies
+        run: uv sync --all-extras --dev
+
+      # Fails if a gated metric exceeds its ceiling in benchmarks/baselines.json.
+      # Memory and layout numbers are deterministic and gated tightly; timing is
+      # recorded as a ratio against scipy in the same process, and gated loosely,
+      # 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

benchmarks/README.md

@@ -0,0 +1,57 @@
+# Benchmarks
+
+A small suite CI runs as a regression gate, plus larger cases for running by
+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 --case matvec_vs_scipy
+uv run python -m benchmarks.run --set fast --record   # rewrite baselines.json
+```
+
+Exits nonzero if a gated metric exceeds its ceiling.
+
+## Metrics
+
+Correctness tests do not catch cost regressions, so the suite measures two
+things that move silently:
+
+**Layout size and memory allocated by an operation.** Both deterministic and
+comparable across machines, so they are gated tightly. Memory uses
+`tracemalloc` rather than `ru_maxrss`, which is a process-lifetime high-water
+mark and reports zero for an operation staying under the peak set while
+building its input.
+
+**Throughput relative to scipy**, never absolute seconds. The same work is
+timed through scipy in the same process and the ratio recorded, which cancels
+most of the difference between machines. Still the noisiest metric, so its
+gate is much looser.
+
+Each case runs in its own subprocess, since measurement state and JIT warm-up
+leak between them otherwise.
+
+## Baselines
+
+`baselines.json` holds a ceiling per gated metric, regenerated with
+`--record` by multiplying a fresh measurement by that metric's margin. Only
+metrics named in `margins` are gated; anything else a case returns is
+recorded for context.
+
+The checked-in ceilings were recorded before the memory fixes landed, so the
+memory ones are deliberately generous and should be re-recorded as those
+merge. On a 4M-nonzero array, for results of length `n_minor`:
+
+| metric | recorded | with the fix |
+|---|---|---|
+| `minor_sum_peak_mb` | 66 MB | 0.8 MB |
+| `minor_extrema_peak_mb` | 66 MB | 1.6 MB |
+| `minor_getnnz_peak_mb` | 32 MB | 0.8 MB |
+| `minor_selection_peak_mb` | 62 MB | 4.8 MB |
+| `misaligned_matmul_peak_mb` | 204 MB | 115 MB |
+
+## 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`.

benchmarks/baselines.json

@@ -0,0 +1,38 @@
+{
+  "comment": "Ceilings a metric must stay under, regenerated with `python -m benchmarks.run --set fast --record`. Only metrics listed in `margins` are gated; the rest are recorded by cases for context. Margins are the slack applied to a measured value when recording: tight for deterministic layout/memory numbers, loose for timing ratios, which vary with the runner.",
+  "margins": {
+    "bytes_per_nonzero": 1.1,
+    "indices_bytes_per_nonzero": 1.1,
+    "vs_scipy_ratio": 1.1,
+    "peak_alloc_mb": 2.0,
+    "time_ratio_vs_scipy": 4.0
+  },
+  "cases": {
+    "layout_bytes_per_nonzero": {
+      "bytes_per_nonzero": 5.72,
+      "indices_bytes_per_nonzero": 4.4,
+      "vs_scipy_ratio": 0.4751
+    },
+    "minor_sum_peak_mb": {
+      "peak_alloc_mb": 132.5131
+    },
+    "misaligned_matmul_peak_mb": {
+      "peak_alloc_mb": 407.2383
+    },
+    "matvec_vs_scipy": {
+      "time_ratio_vs_scipy": 0.9768
+    },
+    "matmat_vs_scipy": {
+      "time_ratio_vs_scipy": 0.6333
+    },
+    "minor_extrema_peak_mb": {
+      "peak_alloc_mb": 132.5453
+    },
+    "minor_getnnz_peak_mb": {
+      "peak_alloc_mb": 64.0324
+    },
+    "minor_selection_peak_mb": {
+      "peak_alloc_mb": 124.5457
+    }
+  }
+}

benchmarks/cases.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import numpy as np
+
+from benchmarks.harness import (
+    integer_counts_csr,
+    peak_alloc_mb,
+    ratio_vs_scipy,
+)
+
+FAST: dict[str, Callable[[], dict[str, float]]] = {}
+SLOW: dict[str, Callable[[], dict[str, float]]] = {}
+
+
+def fast(fn):
+    FAST[fn.__name__] = fn
+    return fn
+
+
+def slow(fn):
+    SLOW[fn.__name__] = fn
+    return fn
+
+
+# -- layout size -------------------------------------------------------------
+
+
+@fast
+def layout_bytes_per_nonzero() -> dict[str, float]:
+    """Bytes the layout holds per stored nonzero, and the same for ``indices`` alone."""
+    from vsparse import VCSRArray
+
+    mat = integer_counts_csr(20_000, 2_000, density=0.05)
+    v = VCSRArray.from_scipy(mat)
+    stored = v.major_ptr.nbytes + v.values.nbytes + v.value_ptr.nbytes + v.indices.nbytes
+    scipy_bytes = mat.indptr.nbytes + mat.indices.nbytes + mat.data.nbytes
+    return {
+        "bytes_per_nonzero": stored / v.nnz,
+        "indices_bytes_per_nonzero": v.indices.nbytes / v.nnz,
+        "vs_scipy_ratio": stored / scipy_bytes,
+    }
+
+
+# -- memory ceilings ---------------------------------------------------------
+
+
+@fast
+def minor_sum_peak_mb() -> dict[str, float]:
+    """Memory allocated by a minor-axis sum."""
+    from vsparse import VCSRArray
+
+    v = VCSRArray.from_scipy(integer_counts_csr(40_000, 2_000, density=0.05))
+    nnz_mb = v.nnz * 8 / 1e6
+    return {
+        "peak_alloc_mb": peak_alloc_mb(lambda: v.sum(axis=0)),
+        "expanded_nnz_mb": nnz_mb,  # what a per-nonzero temporary would cost
+    }
+
+
+@fast
+def misaligned_matmul_peak_mb() -> dict[str, float]:
+    """Memory allocated by the matmul direction the storage isn't aligned for."""
+    from vsparse import VCSCArray
+
+    v = VCSCArray.from_scipy(integer_counts_csr(40_000, 2_000, density=0.05))
+    rng = np.random.default_rng(0)
+    B = rng.normal(size=(v.shape[1], 4))
+    array_mb = (v.values.nbytes + v.value_ptr.nbytes + v.indices.nbytes) / 1e6
+
+    return {
+        "peak_alloc_mb": peak_alloc_mb(lambda: v.normalized() @ B),
+        "array_mb": array_mb,  # what a full second copy would cost
+    }
+
+
+@fast
+def minor_extrema_peak_mb() -> dict[str, float]:
+    """Memory allocated by a minor-axis max/min."""
+    from vsparse import VCSRArray
+
+    v = VCSRArray.from_scipy(integer_counts_csr(40_000, 2_000, density=0.05))
+    return {
+        "peak_alloc_mb": peak_alloc_mb(lambda: v.max(axis=0)),
+        "expanded_nnz_mb": v.nnz * 8 / 1e6,
+    }
+
+
+@fast
+def minor_getnnz_peak_mb() -> dict[str, float]:
+    """Memory allocated by a per-minor-index stored-element count."""
+    from vsparse import VCSRArray
+
+    v = VCSRArray.from_scipy(integer_counts_csr(40_000, 2_000, density=0.05))
+    return {
+        "peak_alloc_mb": peak_alloc_mb(lambda: v.getnnz(axis=0)),
+        "indices_nnz_mb": v.nnz * 8 / 1e6,
+    }
+
+
+@fast
+def minor_selection_peak_mb() -> dict[str, float]:
+    """Memory allocated by a minor-axis selection."""
+    from vsparse import VCSRArray
+
+    v = VCSRArray.from_scipy(integer_counts_csr(40_000, 2_000, density=0.05))
+    cols = np.arange(0, v.shape[1], 2)
+    return {
+        "peak_alloc_mb": peak_alloc_mb(lambda: v[:, cols]),
+        "indices_nnz_mb": v.nnz * 8 / 1e6,
+    }
+
+
+# -- throughput, relative to scipy -------------------------------------------
+
+
+@fast
+def matvec_vs_scipy() -> dict[str, float]:
+    """Aligned-direction matrix-vector product, against scipy's CSR."""
+    from vsparse import VCSRArray
+
+    mat = integer_counts_csr(20_000, 2_000, density=0.05)
+    v = VCSRArray.from_scipy(mat)
+    rng = np.random.default_rng(0)
+    x = rng.normal(size=mat.shape[1])
+    return {"time_ratio_vs_scipy": ratio_vs_scipy(lambda: v @ x, lambda: mat @ x)}
+
+
+@fast
+def matmat_vs_scipy() -> dict[str, float]:
+    """Aligned-direction matrix-matrix product, against scipy's CSR."""
+    from vsparse import VCSRArray
+
+    mat = integer_counts_csr(20_000, 2_000, density=0.05)
+    v = VCSRArray.from_scipy(mat)
+    rng = np.random.default_rng(0)
+    B = rng.normal(size=(mat.shape[1], 8))
+    return {"time_ratio_vs_scipy": ratio_vs_scipy(lambda: v @ B, lambda: mat @ B)}
+
+
+# -- larger, for the scheduled job -------------------------------------------
+
+
+@slow
+def large_layout_and_matmul() -> dict[str, float]:
+    """The same measurements an order of magnitude up."""
+    from vsparse import VCSRArray
+
+    mat = integer_counts_csr(200_000, 3_000, density=0.02)
+    v = VCSRArray.from_scipy(mat)
+    rng = np.random.default_rng(0)
+    B = rng.normal(size=(mat.shape[1], 8))
+    stored = v.major_ptr.nbytes + v.values.nbytes + v.value_ptr.nbytes + v.indices.nbytes
+    v @ B
+    return {
+        "bytes_per_nonzero": stored / v.nnz,
+        "time_ratio_vs_scipy": ratio_vs_scipy(lambda: v @ B, lambda: mat @ B),
+        "matmul_peak_alloc_mb": peak_alloc_mb(lambda: v @ B),
+    }
+
+
+ALL: dict[str, Callable[[], dict[str, float]]] = {**FAST, **SLOW}

benchmarks/harness.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import time
+import tracemalloc
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np
+import scipy.sparse as sp
+
+
+def peak_alloc_mb(fn: Callable[[], Any]) -> float:
+    """Peak memory allocated during ``fn``, in MB.
+
+    Not RSS, which is a process-lifetime high-water mark and so reports zero
+    for anything staying under the peak set while building its input. numpy
+    allocations are traced, numba's internal ones are not.
+    """
+    fn()  # JIT compile / warm caches outside the measurement
+    tracemalloc.start()
+    try:
+        before = tracemalloc.get_traced_memory()[0]
+        tracemalloc.reset_peak()
+        fn()
+        peak = tracemalloc.get_traced_memory()[1]
+    finally:
+        tracemalloc.stop()
+    return max(0.0, (peak - before) / 1e6)
+
+
+def best_time(fn: Callable[[], Any], repeat: int = 7) -> float:
+    """Best wall-clock time over ``repeat`` runs, in seconds. Warms up first."""
+    fn()  # JIT compile / allocate caches outside the measurement
+    best = float("inf")
+    for _ in range(repeat):
+        start = time.perf_counter()
+        fn()
+        best = min(best, time.perf_counter() - start)
+    return best
+
+
+def ratio_vs_scipy(ours: Callable[[], Any], theirs: Callable[[], Any], repeat: int = 7) -> float:
+    """``our time / scipy's time`` for the same work."""
+    return best_time(ours, repeat) / best_time(theirs, repeat)
+
+
+def integer_counts_csr(n_rows: int, n_cols: int, density: float, seed: int = 0) -> sp.csr_array:
+    """Integer-valued sparse counts, repeated enough for the layout to dedupe."""
+    rng = np.random.default_rng(seed)
+    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