From 675272c147ea9c3c75149dea4dc0231f4c70bb3b Mon Sep 17 00:00:00 2001 From: fishidaho Date: Thu, 10 Sep 2026 16:26:35 -0700 Subject: [PATCH] Parallelize the misaligned direction of the core matmul path `_ops._major_matvec`/`_major_matmat` -- the scatter direction, reached by `VCSC @ B` and `B @ VCSR` -- ran single-threaded with no strategy at all, losing to scipy by up to 4x: VCSC @ x 8.7 ms vs scipy 6.7 ms 1.30x VCSC @ B (k=8) 20.1 ms vs scipy 21.0 ms 0.96x x @ VCSR 20.5 ms vs scipy 5.1 ms 3.99x B @ VCSR (p=8) 28.5 ms vs scipy 14.6 ms 1.96x They cannot be parallelized over the major axis directly, since two major slices can collide on the same output index. Regrouping a chunk into the opposite format and running an aligned kernel -- the strategy `_vcs_matmul` uses for the normalized view -- was tried first and is much worse here. That path caches its transposed copy and amortizes it across an iterative algorithm's many products; a bare `A @ B` pays the regroup once per call, which turned the 8.7 ms product into 700 ms. Reverted. Thread-local accumulators instead -- the shape `minor_sums`/`minor_extrema` already use -- reduced across threads afterwards. The thread count is the whole design problem: the scatter is `nnz * width` work split across threads while the reduction is `nthreads * n_minor * width`, so more threads is not better. Measured on 6M nonzeros over 60k x 2k, letting the existing 64 MiB accumulator budget pick 48 threads for a width-8 product ran it in 56.6 ms, *slower than the serial kernel's* 23.3 ms; 4 threads and 15 MB ran it in 11.6 ms. `scatter_threads()` takes the smaller of two bounds. Setting the derivative of scatter-plus-reduction to zero puts the optimum at `sqrt(nnz / n_minor)`, which is 10 for that array against a measured best of 16/8/4 at widths 1/4/8 -- right in magnitude, but blind to width. A 16 MiB byte cap supplies the missing width-dependence, admitting fewer threads exactly as the accumulator grows. Together they track the measured optimum across widths: VCSC @ x 2.4 ms vs scipy 6.7 ms 0.35x (was 1.30x) VCSC @ B (k=8) 12.0 ms vs scipy 21.6 ms 0.56x (was 0.96x) x @ VCSR 1.8 ms vs scipy 5.2 ms 0.36x (was 3.99x) B @ VCSR (p=8) 5.2 ms vs scipy 14.3 ms 0.36x (was 1.96x) Peak allocation stays bounded by the cap rather than by `nnz`: 5.3 MB at width 1, 19.2 MB at width 8, 23.0 MB at width 16. Single-thread cases fall back to the serial kernels rather than pay an allocation and a reduction pass for one partial. Three benchmark cases now cover this direction; the existing `matvec_vs_scipy`/`matmat_vs_scipy` only ever ran VCSR aligned. Rebased from the pre-0.4.0 branch onto main. Two adaptations were needed, both from #44's dtype-promotion fix, which landed after this work was written: - The wrappers now call `_promote()` before dispatching, so the serial fallback (nthreads <= 1) does not bypass it. - The kernels accumulate in `values.dtype` rather than a hardcoded float64, and reduce the thread partials into an explicitly-typed output. `partial.sum(axis=0)` widens a narrow dtype (uint16 -> uint64), which would have made the result dtype depend on the thread count the machine happened to pick. Verified that `uint16 @ uint16` still wraps exactly as numpy and scipy do, and that the mixed `uint16 @ float64` case matches numpy to 3.7e-12 relative -- tighter than scipy's own agreement with numpy on the same input. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_019ZwuJMcTRUacxqARwjMngE --- benchmarks/baselines.json | 11 ++++ benchmarks/cases.py | 56 +++++++++++++++++++ src/vsparse/_ops.py | 112 +++++++++++++++++++++++++++++++++++++- 3 files changed, 177 insertions(+), 2 deletions(-) diff --git a/benchmarks/baselines.json b/benchmarks/baselines.json index 251c274..9209740 100644 --- a/benchmarks/baselines.json +++ b/benchmarks/baselines.json @@ -157,6 +157,17 @@ "cpu_ratio_1t_pearson": 2.85, "cpu_ratio_1t_raw": 2.445, "cpu_ratio_1t_scanpy": 7.845 + }, + "misaligned_matvec_vs_scipy": { + "time_ratio_vs_scipy": 1.2909 + }, + "misaligned_matmat_vs_scipy": { + "time_ratio_vs_scipy": 2.0436, + "peak_alloc_mb": 38.4009 + }, + "misaligned_rmatmat_vs_scipy": { + "time_ratio_vs_scipy": 1.0817, + "peak_alloc_mb": 20.2251 } } } diff --git a/benchmarks/cases.py b/benchmarks/cases.py index f813ab0..62a15bd 100644 --- a/benchmarks/cases.py +++ b/benchmarks/cases.py @@ -351,6 +351,62 @@ def via_sparse(delta=delta, offset=offset): return out +# -- misaligned direction of the *core* matmul path -------------------------- +# +# `matvec_vs_scipy`/`matmat_vs_scipy` above run VCSR in its aligned direction. +# These run the other one -- the scatter path in `_ops` -- which is where the +# array's layout works against the product and where the serial kernel used to +# lose to scipy by up to 4x. + + +@fast +def misaligned_matvec_vs_scipy() -> dict[str, float]: + """``VCSC @ x``: iterate columns, scatter into rows.""" + import scipy.sparse as sp + + from vsparse import VCSCArray + + mat = integer_counts_csr(60_000, 2_000, density=0.05) + v = VCSCArray.from_scipy(mat) + csc = sp.csc_array(mat) + x = np.random.default_rng(0).normal(size=mat.shape[1]) + return {"time_ratio_vs_scipy": ratio_vs_scipy(lambda: v @ x, lambda: csc @ x)} + + +@fast +def misaligned_matmat_vs_scipy() -> dict[str, float]: + """``VCSC @ B``, width 8 -- the case where the accumulator is widest.""" + import scipy.sparse as sp + + from vsparse import VCSCArray + + mat = integer_counts_csr(60_000, 2_000, density=0.05) + v = VCSCArray.from_scipy(mat) + csc = sp.csc_array(mat) + B = np.random.default_rng(0).normal(size=(mat.shape[1], 8)) + return { + "time_ratio_vs_scipy": ratio_vs_scipy(lambda: v @ B, lambda: csc @ B), + "peak_alloc_mb": peak_alloc_mb(lambda: v @ B), + } + + +@fast +def misaligned_rmatmat_vs_scipy() -> dict[str, float]: + """``B @ VCSR``: same scatter, reached from the other side.""" + import scipy.sparse as sp + + from vsparse import VCSRArray + + mat = integer_counts_csr(60_000, 2_000, density=0.05) + v = VCSRArray.from_scipy(mat) + csr = sp.csr_array(mat) + B = np.random.default_rng(0).normal(size=(8, mat.shape[0])) + return { + "time_ratio_vs_scipy": ratio_vs_scipy(lambda: B @ v, lambda: B @ csr), + "peak_alloc_mb": peak_alloc_mb(lambda: B @ v), + } + + # -- larger, for the scheduled job ------------------------------------------- diff --git a/src/vsparse/_ops.py b/src/vsparse/_ops.py index 594fcbc..6a686c2 100644 --- a/src/vsparse/_ops.py +++ b/src/vsparse/_ops.py @@ -2,6 +2,8 @@ from __future__ import annotations +import math + import numba import numpy as np @@ -32,6 +34,30 @@ def accumulator_threads(n_minor: int, bytes_per_element: int = 8) -> int: return int(min(numba.get_num_threads(), affordable)) +#: Tighter than `_ACCUMULATOR_BUDGET_BYTES`, because this scatter also pays to +#: reduce its partials afterwards: a budget that admits enough threads to make +#: that reduction dominate is slower than not parallelizing at all. +_SCATTER_ACCUMULATOR_BUDGET_BYTES = 16 << 20 # 16 MiB + + +def scatter_threads(nnz: int, n_minor: int, width: int = 1) -> int: + """Threads for a major-axis scatter of ``nnz`` values into an ``n_minor`` output. + + ``width`` is the number of output columns, 1 for a matvec. + """ + if n_minor <= 0 or nnz <= 0: + return 1 + width = max(1, width) + # The scatter is `nnz * width` work split across threads, but reducing the + # partials afterwards costs `nthreads * n_minor * width`, so more threads + # is not better: their sum is minimized at `sqrt(nnz / n_minor)`. The byte + # cap supplies the width-dependence that estimate lacks, admitting fewer + # threads as the accumulator grows. + affordable = _SCATTER_ACCUMULATOR_BUDGET_BYTES // (n_minor * width * 8) + balanced = math.isqrt(max(1, nnz // n_minor)) + return int(max(1, min(numba.get_num_threads(), affordable, balanced))) + + @numba.njit(cache=True) def _major_matvec(major_ptr, values, value_ptr, indices, x, n_major, n_minor): """y = A @ x where A's major axis (columns for VCSC) has length n_major. @@ -176,9 +202,78 @@ def _promote(values, other): return values.astype(out_dtype, copy=False), other.astype(out_dtype, copy=False) +# -- misaligned direction, parallelized over thread-local accumulators ------- +# +# The serial kernels above scatter into the output along the major axis, so +# they cannot be parallelized over it: two major slices can collide on the +# same output index. Each thread accumulates into a private partial instead -- +# the shape `minor_sums`/`minor_extrema` already use -- reduced afterwards. + + +@numba.njit(cache=True, parallel=True) +def _major_matvec_par(major_ptr, values, value_ptr, indices, x, n_major, n_minor, nthreads): + partial = np.zeros((nthreads, n_minor), dtype=values.dtype) + span = (n_major + nthreads - 1) // nthreads + for t in numba.prange(nthreads): # ty: ignore[not-iterable] + start = t * span + end = min(n_major, start + span) + local = partial[t] + for j in range(start, end): + xj = x[j] + if xj == 0: + continue + for u in range(major_ptr[j], major_ptr[j + 1]): + val = values[u] * xj + for k in range(value_ptr[u], value_ptr[u + 1]): + local[indices[k]] += val + # `partial.sum(axis=0)` would widen a narrow dtype (uint16 -> uint64) and + # so return a different dtype than the serial kernel for the same input. + out = np.zeros(n_minor, dtype=values.dtype) + for t in range(nthreads): + out += partial[t] + return out + + +@numba.njit(cache=True, parallel=True) +def _major_matmat_par(major_ptr, values, value_ptr, indices, b, n_major, n_minor, nthreads): + width = b.shape[1] + partial = np.zeros((nthreads, n_minor, width), dtype=values.dtype) + span = (n_major + nthreads - 1) // nthreads + for t in numba.prange(nthreads): # ty: ignore[not-iterable] + start = t * span + end = min(n_major, start + span) + local = partial[t] + for j in range(start, end): + brow = b[j] + for u in range(major_ptr[j], major_ptr[j + 1]): + val = values[u] + for k in range(value_ptr[u], value_ptr[u + 1]): + acc = local[indices[k]] + for c in range(width): + acc[c] += val * brow[c] + out = np.zeros((n_minor, width), dtype=values.dtype) + for t in range(nthreads): + out += partial[t] + return out + + def major_matvec(major_ptr, values, value_ptr, indices, x, n_major, n_minor): values, x = _promote(values, np.asarray(x)) - return _major_matvec(major_ptr, values, value_ptr, indices, x, n_major, n_minor) + nthreads = scatter_threads(int(value_ptr[-1]) if value_ptr.shape[0] else 0, n_minor) + if nthreads <= 1: + # One thread's worth of accumulator is the serial kernel with an extra + # allocation and a reduction pass, so skip both. + return _major_matvec(major_ptr, values, value_ptr, indices, x, n_major, n_minor) + return _major_matvec_par( + major_ptr, + values, + value_ptr, + indices, + np.ascontiguousarray(x), + n_major, + n_minor, + nthreads, + ) def minor_matvec(major_ptr, values, value_ptr, indices, x, n_major): @@ -188,7 +283,20 @@ def minor_matvec(major_ptr, values, value_ptr, indices, x, n_major): def major_matmat(major_ptr, values, value_ptr, indices, b, n_major, n_minor): values, b = _promote(values, np.ascontiguousarray(b)) - return _major_matmat(major_ptr, values, value_ptr, indices, b, n_major, n_minor) + width = b.shape[1] if b.ndim == 2 else 1 + nthreads = scatter_threads(int(value_ptr[-1]) if value_ptr.shape[0] else 0, n_minor, width) + if nthreads <= 1: + return _major_matmat(major_ptr, values, value_ptr, indices, b, n_major, n_minor) + return _major_matmat_par( + major_ptr, + values, + value_ptr, + indices, + b, + n_major, + n_minor, + nthreads, + ) def minor_matmat(major_ptr, values, value_ptr, indices, b, n_major):