From ffc95b2494a0eb3db4ad7b8dadfce4be89dd4671 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sun, 13 Sep 2026 11:42:45 -0700 Subject: [PATCH] Replace chunked-transpose matmul fallback with a native minor-axis kernel The misaligned direction of a normalized view's matmul (self@B for VCSC, B@self for VCSR) previously regrouped the array into the other VCS format a chunk of major slices at a time (_chunk_bounds/_transpose_major), caching a full opposite-format copy (_dual_arr) when the whole array fit one chunk's budget. At scale (~2.3B nonzeros, 1.3M x 12.2K), this needed ~1,150 chunks per pass, repeated across every power iteration and every rank/trial in a BiCV sweep -- tens of thousands of variably-sized alloc/free cycles that fragmented the allocator and drove RSS to ~140 GB on a shared machine, even though no single chunk's live memory was large. Replace it with a direct minor-axis kernel that walks the array's own storage as-is: each thread gets a private, full-output-sized accumulator and scatters into it while owning a disjoint range of major slices, summed across threads at the end -- the same pattern _ops.py already uses for minor_sums/minor_counts/minor_extrema. accumulator_threads caps the thread count to a fixed byte budget, so the accumulator block is always nthreads * output_dim * width * 8 bytes: bounded, independent of nnz, and allocated once per call instead of thousands of times. For a huge output axis (e.g. millions of cells) that caps down to a single thread, trading parallelism for a hard memory bound rather than the previous unbounded chunk churn. Removes _chunk_bounds/_aligned_source/_build_dual/_dual_arr and the chunked-transpose test suite; adds tests/test_minor_axis_matmul.py covering correctness against a dense reference, thread-count invariance, that no second copy of the array is ever built, and that peak memory stays bounded by the accumulator budget rather than nnz. Co-Authored-By: Claude Sonnet 5 --- src/vsparse/_norm_common.py | 12 +- src/vsparse/_vcs_matmul.py | 263 +++++++++++++++++++------------- src/vsparse/_vcs_norm.py | 22 +-- tests/test_chunked_transpose.py | 147 ------------------ tests/test_minor_axis_matmul.py | 125 +++++++++++++++ tests/test_vcs_norm.py | 19 ++- tests/test_vcs_norm_recipes.py | 2 +- 7 files changed, 300 insertions(+), 290 deletions(-) delete mode 100644 tests/test_chunked_transpose.py create mode 100644 tests/test_minor_axis_matmul.py diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index 9b6d406..1bb4835 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -138,12 +138,10 @@ def resolve_recipe(view: str | Recipe) -> Recipe: class _NormCache: """Bounded LRU of computed normalizations for one array, keyed by :class:`Recipe`. - Holds each view's *statistics* strongly and the view itself only weakly. - That split matters: a view can carry an ``O(nnz)`` ``_dual_arr`` (a whole - opposite-format copy of the array, cached by the matmul kernels), and a - cache that kept views alive would pin one of those per recipe for as long - as the array lived, even after the caller had dropped every reference. - The statistics are ``O(n_rows + n_cols)``, so retaining those is cheap. + Holds each view's *statistics* strongly and the view itself only weakly, + so a cache hit never pins a view alive for longer than the caller's own + references would. The statistics are ``O(n_rows + n_cols)``, so retaining + those is cheap. A hit on a still-live view hands back that exact object, so ``recalculate=False`` is identity-stable for as long as the caller holds @@ -786,7 +784,7 @@ def _from_internal_stats( return self def _init_extra(self) -> None: - """Hook for subclasses with extra per-instance state (e.g. ``_dual_arr``). + """Hook for subclasses with extra per-instance state to initialize. ``__init__`` normally initializes that state itself; :meth:`from_stats` builds an instance via ``object.__new__`` instead, bypassing it, so it diff --git a/src/vsparse/_vcs_matmul.py b/src/vsparse/_vcs_matmul.py index a6f16b7..f00043b 100644 --- a/src/vsparse/_vcs_matmul.py +++ b/src/vsparse/_vcs_matmul.py @@ -18,17 +18,34 @@ ``self @ B``, :class:`~vsparse.VCSCArray` for ``B @ self`` -- see :func:`_vcsr_matmul_delta`/:func:`_vcsc_rmatmul_delta`). The other direction on a given array (:class:`~vsparse.VCSCArray` for ``self @ B``, -:class:`~vsparse.VCSRArray` for ``B @ self``) doesn't have that alignment -- -rather than run a scatter kernel there (thread-local output-shaped -accumulators, reduced across threads: cache-unfriendly, and memory-hungry -enough for wide ``B`` to need throttling), the storage is regrouped into the -other VCS format via :meth:`~vsparse._base._VCSBase._transpose_major`, a -chunk of major slices at a time so the extra memory is one chunk's worth -rather than a second copy of the array. - -``row_scale``/``gene_scale``/``col_mean``/``col_post_scale`` are per-row/ -per-column statistics, so a chunk just takes the slice of them its own axis -covers. +:class:`~vsparse.VCSRArray` for ``B @ self``) doesn't have that alignment: +parallelizing naively over major slices there would have different threads +scatter-add into the *same* output row/column, a data race. + +That misaligned direction is handled the same way :mod:`vsparse._ops` +already handles the minor-axis reductions (``minor_sums``/``minor_counts``/ +``minor_extrema``): each thread gets its own private, full-output-sized +accumulator and walks a disjoint contiguous range of major slices, scattering +into that private copy; the per-thread copies are summed once at the end +(:func:`_vcsc_matmul_delta_minor`/:func:`_vcsr_rmatmul_delta_minor`). No +second copy of the sparse array's structure is ever built, and no chunk of it +is ever regrouped into the other VCS format -- the kernel walks the array's +own ``indices`` exactly as stored, so the cost of a call is ``O(nnz * width)`` +work plus one fixed-size accumulator allocation, never a sequence of +allocate/free cycles that scales with array size. + +The one thing this can't avoid is that the accumulator block itself costs +``nthreads * out_major_dim * width * 8`` bytes. +:func:`~vsparse._ops.accumulator_threads` (shared with the reduction +kernels) caps ``nthreads`` to keep that block under a fixed budget -- for a +huge output axis (e.g. millions of cells) that caps down to a single thread, +trading parallelism for a hard memory bound. That is the right trade at +that scale: a small, constant footprint and a correct answer, rather than +either an unbounded thread count or a chunked-transpose fallback whose total +allocation churn (thousands of variably-sized chunk buffers, over many +power iterations and many ranks/trials in a BiCV sweep) is what previously +fragmented the allocator and drove RSS up to ~140 GB on a shared machine +(see the project issue this replaced). """ from __future__ import annotations @@ -39,6 +56,7 @@ import numpy as np from vsparse._norm_common import _g +from vsparse._ops import accumulator_threads if TYPE_CHECKING: from vsparse._vcs_norm import _VCSNormalizedBase @@ -104,60 +122,97 @@ def _vcsc_rmatmul_delta( acc[c] += delta * brow[c] -# ``Delta @ B`` splits over the contracted axis and ``B @ Delta`` over rows, -# so a contiguous range of major slices can be regrouped on its own, -# accumulated into the shared output, and dropped. +# -- misaligned direction: per-thread private accumulators, no regrouping ---- -_CHUNK_BUDGET_BYTES = 128 << 20 # 128 MiB of transient regrouping per chunk - -# transpose_major sorts globally over the chunk's nonzeros, so its peak is -# several nnz-sized temporaries. Deliberately generous, since underestimating -# means the budget doesn't hold. -_TRANSPOSE_BYTES_PER_NNZ = 64 - - -def _chunk_bounds(arr, budget_bytes: int) -> list[tuple[int, int]]: - """Contiguous ``[start, stop)`` major-slice ranges, each within the byte budget.""" - n_major = arr.n_major - if n_major == 0: - return [] - max_nnz = max(1, budget_bytes // _TRANSPOSE_BYTES_PER_NNZ) - if arr.nnz <= max_nnz: - return [(0, n_major)] - - # nnz of major slices [0, j) -- value_ptr indexed by the group boundary. - cumulative = arr.value_ptr[arr.major_ptr] - bounds = [] - start = 0 - while start < n_major: - # Furthest stop whose chunk stays under budget; always advance by >= 1. - stop = int(np.searchsorted(cumulative, cumulative[start] + max_nnz, side="right")) - 1 - stop = min(max(stop, start + 1), n_major) - bounds.append((start, stop)) - start = stop - return bounds +@numba.njit(cache=True, parallel=True) +def _vcsc_matmul_delta_minor( + major_ptr, + values, + value_ptr, + indices, + row_scale, + gene_scale, + col_post_scale, + g_code, + B, + nthreads, + n_rows, +): + """``Delta @ B`` for a VCSC array (major=columns), walked directly (no VCSR regroup). -def _aligned_source(nview: _VCSNormalizedBase, needed_format: str): - """``(array, None)`` to run one major-aligned pass, or ``(None, chunk bounds)``.""" - arr = nview._arr - if arr._format == needed_format: - return arr, None - if nview._dual_arr is not None: - return nview._dual_arr, None - bounds = _chunk_bounds(arr, _CHUNK_BUDGET_BYTES) - if len(bounds) <= 1: - # Regrouping the whole array already fits the per-call budget, so - # keeping it costs no extra peak memory and saves every later call. - return _build_dual(nview), None - return None, bounds + ``B`` is ``(n_cols, k)``; returns ``(n_rows, k)``. Parallel-safe because + each thread scatters into its own private ``(n_rows, k)`` accumulator + while owning a disjoint, contiguous range of columns. + """ + n_major = major_ptr.shape[0] - 1 # n_cols + k = B.shape[1] + chunk = (n_major + nthreads - 1) // nthreads + partial = np.zeros((nthreads, n_rows, k), dtype=np.float64) + for t in numba.prange(nthreads): # ty: ignore[not-iterable] + start = t * chunk + stop = min(n_major, start + chunk) + local = partial[t] + for j in range(start, stop): + gs = gene_scale[j] + if gs == 0.0: + continue + s = col_post_scale[j] + brow = B[j] + for u in range(major_ptr[j], major_ptr[j + 1]): + v = values[u] + for kk in range(value_ptr[u], value_ptr[u + 1]): + row = indices[kk] + delta = s * _g(v / row_scale[row] / gs, g_code) + acc = local[row] + for c in range(k): + acc[c] += delta * brow[c] + return partial.sum(axis=0) -def _build_dual(nview: _VCSNormalizedBase): - """Build and cache the opposite-format copy of ``nview``'s array, once.""" - if nview._dual_arr is None: - nview._dual_arr = nview._arr._transpose_major() - return nview._dual_arr +@numba.njit(cache=True, parallel=True) +def _vcsr_rmatmul_delta_minor( + major_ptr, + values, + value_ptr, + indices, + row_scale, + gene_scale, + col_post_scale, + g_code, + Bt, + nthreads, + n_cols, +): + """``B @ Delta`` for a VCSR array (major=rows), walked directly (no VCSC regroup). + + ``Bt`` is ``(n_rows, p)`` (``B`` transposed); returns ``(n_cols, p)``. + Parallel-safe because each thread scatters into its own private + ``(n_cols, p)`` accumulator while owning a disjoint, contiguous range of + rows. + """ + n_major = major_ptr.shape[0] - 1 # n_rows + p = Bt.shape[1] + chunk = (n_major + nthreads - 1) // nthreads + partial = np.zeros((nthreads, n_cols, p), dtype=np.float64) + for t in numba.prange(nthreads): # ty: ignore[not-iterable] + start = t * chunk + stop = min(n_major, start + chunk) + local = partial[t] + for i in range(start, stop): + rs = row_scale[i] + brow = Bt[i] + for u in range(major_ptr[i], major_ptr[i + 1]): + v = values[u] + for kk in range(value_ptr[u], value_ptr[u + 1]): + col = indices[kk] + gs = gene_scale[col] + if gs > 0.0: + delta = col_post_scale[col] * _g(v / rs / gs, g_code) + acc = local[col] + for c in range(p): + acc[c] += delta * brow[c] + return partial.sum(axis=0) # -- public entry points: dense correction + sparse delta ------------------- @@ -182,14 +237,13 @@ def normalized_at_dense(nview: _VCSNormalizedBase, other: Any) -> np.ndarray: g_code = nview.recipe.g_code - # self @ B is major-aligned for VCSR; a VCSC array needs regrouping. - src, bounds = _aligned_source(nview, "csr") - if src is not None: + if arr._format == "csr": + # self @ B is major-aligned for VCSR. _vcsr_matmul_delta( - src.major_ptr, - src.values, - src.value_ptr, - src.indices, + arr.major_ptr, + arr.values, + arr.value_ptr, + arr.indices, nview.row_scale, nview.gene_scale, nview.col_post_scale, @@ -198,23 +252,21 @@ def normalized_at_dense(nview: _VCSNormalizedBase, other: Any) -> np.ndarray: out, ) else: - # Column chunk at a time: Delta @ B == sum over chunks of - # Delta[:, chunk] @ B[chunk, :], so each chunk's contribution - # accumulates into the same output and is then discarded. - for start, stop in bounds: - chunk = arr._major_range(start, stop)._transpose_major() - _vcsr_matmul_delta( - chunk.major_ptr, - chunk.values, - chunk.value_ptr, - chunk.indices, - nview.row_scale, - nview.gene_scale[start:stop], - nview.col_post_scale[start:stop], - g_code, - np.ascontiguousarray(B[start:stop]), - out, - ) + # arr is VCSC: self @ B is the misaligned direction for this format. + nthreads = accumulator_threads(arr.shape[0], bytes_per_element=8 * B.shape[1]) + out += _vcsc_matmul_delta_minor( + arr.major_ptr, + arr.values, + arr.value_ptr, + arr.indices, + nview.row_scale, + nview.gene_scale, + nview.col_post_scale, + g_code, + B, + nthreads, + arr.shape[0], + ) offset = nview.col_mean * nview.col_post_scale baseline = (-offset) @ B # (k,): every row's implicit-zero contribution @@ -238,14 +290,13 @@ def dense_at_normalized(nview: _VCSNormalizedBase, other: Any) -> np.ndarray: g_code = nview.recipe.g_code - # B @ self is major-aligned for VCSC; a VCSR array needs regrouping. - src, bounds = _aligned_source(nview, "csc") - if src is not None: + if arr._format == "csc": + # B @ self is major-aligned for VCSC. _vcsc_rmatmul_delta( - src.major_ptr, - src.values, - src.value_ptr, - src.indices, + arr.major_ptr, + arr.values, + arr.value_ptr, + arr.indices, nview.row_scale, nview.gene_scale, nview.col_post_scale, @@ -254,21 +305,21 @@ def dense_at_normalized(nview: _VCSNormalizedBase, other: Any) -> np.ndarray: out_t, ) else: - # Row chunk at a time, accumulating into the same output. - for start, stop in bounds: - chunk = arr._major_range(start, stop)._transpose_major() - _vcsc_rmatmul_delta( - chunk.major_ptr, - chunk.values, - chunk.value_ptr, - chunk.indices, - nview.row_scale[start:stop], - nview.gene_scale, - nview.col_post_scale, - g_code, - np.ascontiguousarray(Bt[start:stop]), - out_t, - ) + # arr is VCSR: B @ self is the misaligned direction for this format. + nthreads = accumulator_threads(arr.shape[1], bytes_per_element=8 * p) + out_t += _vcsr_rmatmul_delta_minor( + arr.major_ptr, + arr.values, + arr.value_ptr, + arr.indices, + nview.row_scale, + nview.gene_scale, + nview.col_post_scale, + g_code, + Bt, + nthreads, + arr.shape[1], + ) out = np.ascontiguousarray(out_t.T) offset = nview.col_mean * nview.col_post_scale diff --git a/src/vsparse/_vcs_norm.py b/src/vsparse/_vcs_norm.py index bd9fdbf..eaab3cd 100644 --- a/src/vsparse/_vcs_norm.py +++ b/src/vsparse/_vcs_norm.py @@ -13,12 +13,9 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import Any -from vsparse._norm_common import DEFAULT_RECIPE, NormalizedViewBase, Recipe - -if TYPE_CHECKING: - from vsparse._base import _VCSBase +from vsparse._norm_common import NormalizedViewBase __all__ = ["VCSCArrayNormalized", "VCSRArrayNormalized"] @@ -26,20 +23,7 @@ class _VCSNormalizedBase(NormalizedViewBase): """Shared implementation for :class:`VCSCArrayNormalized`/:class:`VCSRArrayNormalized`.""" - __slots__ = ("_dual_arr",) - - _dual_arr: _VCSBase | None - - def __init__( - self, arr: _VCSBase, recipe: str | Recipe = DEFAULT_RECIPE, *, stale: bool = False - ) -> None: - super().__init__(arr, recipe, stale=stale) - self._init_extra() - - def _init_extra(self) -> None: - # Opposite-format copy of `arr`, cached by vsparse._vcs_matmul when - # regrouping the whole array fits one chunk's budget. - self._dual_arr = None + __slots__ = () def __matmul__(self, other: Any) -> Any: """``self @ other`` for a dense ``other`` -- see :mod:`vsparse._vcs_matmul`.""" diff --git a/tests/test_chunked_transpose.py b/tests/test_chunked_transpose.py deleted file mode 100644 index 09d7279..0000000 --- a/tests/test_chunked_transpose.py +++ /dev/null @@ -1,147 +0,0 @@ -from __future__ import annotations - -import tracemalloc -from itertools import pairwise - -import numpy as np -import pytest -import scipy.sparse as sp - -from vsparse import VCSCArray, VCSRArray -from vsparse._vcs_matmul import _TRANSPOSE_BYTES_PER_NNZ, _chunk_bounds - - -@pytest.fixture(params=[VCSCArray, VCSRArray]) -def vcls(request): - return request.param - - -def _scipy_for(vcls, dense): - return sp.csc_array(dense) if vcls is VCSCArray else sp.csr_array(dense) - - -def _reference(dense: np.ndarray) -> np.ndarray: - row_totals = dense.sum(axis=1) - row_scale = row_totals / np.median(row_totals) - row_scale[row_scale == 0.0] = 1.0 - scaled = dense / row_scale[:, None] - gene_scale = scaled.sum(axis=0) - with np.errstate(divide="ignore", invalid="ignore"): - normalized = np.where(gene_scale > 0, scaled / gene_scale[None, :], 0.0) - transformed = np.log10(1.0 + 1000.0 * normalized) - return transformed - transformed.mean(axis=0, keepdims=True) - - -@pytest.mark.parametrize("budget", [1, 200, 1 << 30]) -def test_chunks_tile_the_major_axis(dense, vcls, budget): - """Chunks must cover every major slice exactly once, at any budget.""" - v = vcls.from_scipy(_scipy_for(vcls, dense)) - bounds = _chunk_bounds(v, budget) - - assert bounds[0][0] == 0 - assert bounds[-1][1] == v.n_major - assert all(start < stop for start, stop in bounds) - for (_, prev_stop), (start, _) in pairwise(bounds): - assert start == prev_stop - - -def test_chunks_respect_the_budget(vcls, rng): - """A chunk may only exceed the budget when it is a single, indivisible slice.""" - dense = rng.integers(0, 4, size=(60, 50)).astype(np.float64) - v = vcls.from_scipy(_scipy_for(vcls, dense)) - budget = 40 * _TRANSPOSE_BYTES_PER_NNZ - - cumulative = v.value_ptr[v.major_ptr] - for start, stop in _chunk_bounds(v, budget): - chunk_nnz = int(cumulative[stop] - cumulative[start]) - assert chunk_nnz <= 40 or stop - start == 1 - - -def test_empty_array_has_no_chunks(vcls): - shape = (0, 4) if vcls is VCSRArray else (4, 0) - v = vcls.from_scipy(_scipy_for(vcls, np.zeros(shape))) - assert _chunk_bounds(v, 1 << 20) == [] - - -def test_major_range_is_a_zero_copy_view(vcls, rng): - """Chunking is only affordable because the buffers are shared, not gathered.""" - dense = rng.integers(0, 4, size=(30, 20)).astype(np.float64) - v = vcls.from_scipy(_scipy_for(vcls, dense)) - - chunk = v._major_range(1, 4) - assert np.shares_memory(chunk.values, v.values) - assert chunk.nnz == 0 or np.shares_memory(chunk.indices, v.indices) - - -def test_major_range_matches_ordinary_slicing(vcls, rng): - """A chunk must hold the same sub-array a caller would get by slicing.""" - dense = rng.integers(0, 4, size=(30, 20)).astype(np.float64) - v = vcls.from_scipy(_scipy_for(vcls, dense)) - start, stop = 2, 7 - - native = v[:, start:stop] if vcls is VCSCArray else v[start:stop, :] - np.testing.assert_allclose(v._major_range(start, stop).toarray(), native.toarray()) - - -@pytest.mark.parametrize("budget", [1, 200, 4000, 1 << 30]) -def test_matmul_is_chunk_size_invariant(monkeypatch, vcls, rng, budget): - """Same answer whether it runs in one chunk or one major slice at a time.""" - import vsparse._vcs_matmul as vcs_matmul - - dense = rng.integers(0, 5, size=(40, 24)).astype(np.float64) - v = vcls.from_scipy(_scipy_for(vcls, dense)) - ref = _reference(dense) - B = rng.normal(size=(dense.shape[1], 3)) - Bl = rng.normal(size=(2, dense.shape[0])) - - monkeypatch.setattr(vcs_matmul, "_CHUNK_BUDGET_BYTES", budget) - nv = v.normalized() - - np.testing.assert_allclose(nv @ B, ref @ B, atol=1e-7) - np.testing.assert_allclose(Bl @ nv, Bl @ ref, atol=1e-7) - np.testing.assert_allclose(nv @ B[:, 0], ref @ B[:, 0], atol=1e-7) - np.testing.assert_allclose(Bl[0] @ nv, Bl[0] @ ref, atol=1e-7) - - -def test_multi_chunk_matmul_caches_nothing(monkeypatch, vcls, rng): - """Past one chunk no full dual is built, which is the memory guarantee.""" - import vsparse._vcs_matmul as vcs_matmul - - dense = rng.integers(1, 5, size=(40, 24)).astype(np.float64) - v = vcls.from_scipy(_scipy_for(vcls, dense)) - monkeypatch.setattr(vcs_matmul, "_CHUNK_BUDGET_BYTES", 200) - assert len(_chunk_bounds(v, 200)) > 1 - - nv = v.normalized() - nv @ rng.normal(size=(dense.shape[1], 2)) - rng.normal(size=(2, dense.shape[0])) @ nv - - assert nv._dual_arr is None - - -def test_misaligned_matmul_peak_is_bounded_by_the_chunk_budget(monkeypatch, rng): - """Peak memory has to track the budget, not the size of the array.""" - import vsparse._vcs_matmul as vcs_matmul - - dense = rng.integers(1, 5, size=(1200, 400)).astype(np.float64) - v = VCSCArray.from_scipy(sp.csc_array(dense)) - nnz_bytes = v.nnz * v.indices.dtype.itemsize - B = rng.normal(size=(dense.shape[1], 2)) - - v.normalized() @ B # warm up the JIT before measuring - - monkeypatch.setattr(vcs_matmul, "_CHUNK_BUDGET_BYTES", 64 * _TRANSPOSE_BYTES_PER_NNZ) - nv = v.normalized() - - tracemalloc.start() - try: - before = tracemalloc.get_traced_memory()[0] - tracemalloc.reset_peak() - out = nv @ B - peak = tracemalloc.get_traced_memory()[1] - finally: - tracemalloc.stop() - - assert nv._dual_arr is None - assert peak - before < nnz_bytes - np.testing.assert_allclose(out, _reference(dense) @ B, atol=1e-7) diff --git a/tests/test_minor_axis_matmul.py b/tests/test_minor_axis_matmul.py new file mode 100644 index 0000000..b558cb4 --- /dev/null +++ b/tests/test_minor_axis_matmul.py @@ -0,0 +1,125 @@ +"""Misaligned-direction normalized-view matmul: native minor-axis kernel. + +Replaces the old chunked-transpose/dual-copy approach (see +``vsparse._vcs_matmul``): the misaligned direction now walks the array's own +storage directly with per-thread private accumulators, never building a +second copy of the sparse structure or regrouping any part of it into the +other VCS format. +""" + +from __future__ import annotations + +import tracemalloc + +import numpy as np +import pytest +import scipy.sparse as sp + +from vsparse import VCSCArray, VCSRArray +from vsparse._ops import accumulator_threads + + +@pytest.fixture(params=[VCSCArray, VCSRArray]) +def vcls(request): + return request.param + + +def _scipy_for(vcls, dense): + return sp.csc_array(dense) if vcls is VCSCArray else sp.csr_array(dense) + + +def _reference(dense: np.ndarray) -> np.ndarray: + row_totals = dense.sum(axis=1) + row_scale = row_totals / np.median(row_totals) + row_scale[row_scale == 0.0] = 1.0 + scaled = dense / row_scale[:, None] + gene_scale = scaled.sum(axis=0) + with np.errstate(divide="ignore", invalid="ignore"): + normalized = np.where(gene_scale > 0, scaled / gene_scale[None, :], 0.0) + transformed = np.log10(1.0 + 1000.0 * normalized) + return transformed - transformed.mean(axis=0, keepdims=True) + + +def test_misaligned_matmul_matches_reference(dense, vcls): + """Both matmul directions match a dense reference, whichever direction is misaligned.""" + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + v = vcls.from_scipy(_scipy_for(vcls, dense)) + nv = v.normalized() + ref = _reference(dense) + + rng = np.random.default_rng(7) + B = rng.normal(size=(dense.shape[1], 3)) + Bl = rng.normal(size=(2, dense.shape[0])) + + np.testing.assert_allclose(nv @ B, ref @ B, atol=1e-7) + np.testing.assert_allclose(Bl @ nv, Bl @ ref, atol=1e-7) + np.testing.assert_allclose(nv @ B[:, 0], ref @ B[:, 0], atol=1e-7) + np.testing.assert_allclose(Bl[0] @ nv, Bl[0] @ ref, atol=1e-7) + + +def test_matmul_result_is_accumulator_thread_count_invariant(monkeypatch, vcls, rng): + """Same answer regardless of how many private accumulators the kernel splits work across.""" + import vsparse._vcs_matmul as vcs_matmul + + dense = rng.integers(0, 5, size=(40, 24)).astype(np.float64) + v = vcls.from_scipy(_scipy_for(vcls, dense)) + ref = _reference(dense) + B = rng.normal(size=(dense.shape[1], 3)) + Bl = rng.normal(size=(2, dense.shape[0])) + nv = v.normalized() + + for forced_threads in (1, 2, 8): + monkeypatch.setattr( + vcs_matmul, "accumulator_threads", lambda *a, _n=forced_threads, **k: _n + ) + np.testing.assert_allclose(nv @ B, ref @ B, atol=1e-7) + np.testing.assert_allclose(Bl @ nv, Bl @ ref, atol=1e-7) + + +def test_no_second_copy_of_the_array_is_built(vcls, rng): + """The normalized view carries nothing beyond the original array and O(n_rows+n_cols) stats.""" + dense = rng.integers(1, 5, size=(30, 20)).astype(np.float64) + v = vcls.from_scipy(_scipy_for(vcls, dense)) + nv = v.normalized() + + B = rng.normal(size=(dense.shape[1], 2)) + Bl = rng.normal(size=(2, dense.shape[0])) + nv @ B + Bl @ nv + + assert not hasattr(nv, "_dual_arr") + # The wrapped array's own buffers are untouched/unreplaced by any matmul call. + assert nv._arr is v + + +def test_misaligned_matmul_peak_is_bounded_by_the_accumulator_budget(rng): + """Peak memory tracks the (fixed) accumulator budget, not the size of the array.""" + dense = rng.integers(1, 5, size=(1200, 400)).astype(np.float64) + v = VCSCArray.from_scipy(sp.csc_array(dense)) + nnz_bytes = v.nnz * v.indices.dtype.itemsize + B = rng.normal(size=(dense.shape[1], 2)) + + v.normalized() @ B # warm up the JIT before measuring + + nv = v.normalized() + tracemalloc.start() + try: + before = tracemalloc.get_traced_memory()[0] + tracemalloc.reset_peak() + out = nv @ B + peak = tracemalloc.get_traced_memory()[1] + finally: + tracemalloc.stop() + + # Accumulator block is nthreads * n_rows * width * 8 bytes -- independent + # of nnz, so it stays far below a per-nonzero cost for this shape. + assert peak - before < nnz_bytes + np.testing.assert_allclose(out, _reference(dense) @ B, atol=1e-7) + + +def test_accumulator_threads_degrades_to_one_for_a_huge_output_axis(): + """A huge output axis (e.g. millions of cells) must not blow the accumulator budget.""" + huge_axis = 2_000_000 + wide_b = 200 # e.g. rank + oversampling in a randomized SVD + assert accumulator_threads(huge_axis, bytes_per_element=8 * wide_b) == 1 diff --git a/tests/test_vcs_norm.py b/tests/test_vcs_norm.py index 09f60d1..16858c3 100644 --- a/tests/test_vcs_norm.py +++ b/tests/test_vcs_norm.py @@ -121,13 +121,17 @@ def test_all_zero_matrix_does_not_crash(vcls): np.testing.assert_allclose(nv.toarray(), np.zeros((5, 4))) -def test_small_array_dual_is_cached_lazily(dense, vcls): - """An array whose whole regrouping fits one chunk is transposed once and kept.""" +def test_matmul_never_builds_a_second_copy_of_the_array(dense, vcls): + """Neither matmul direction builds an opposite-format copy of the array. + + See ``tests/test_minor_axis_matmul.py`` for the misaligned direction's own + correctness/memory-bound coverage; this just confirms the view stays thin. + """ if dense.sum() == 0: pytest.skip("all-zero matrix: median row total is 0") v = vcls.from_scipy(_scipy_for(vcls, dense)) nv = v.normalized() - assert nv._dual_arr is None # not built at construction + assert nv._arr is v rng = np.random.default_rng(12) B = rng.normal(size=(dense.shape[1], 3)) @@ -135,13 +139,8 @@ def test_small_array_dual_is_cached_lazily(dense, vcls): nv @ B # major-aligned for VCSR self@B; misaligned for VCSC Bl @ nv # major-aligned for VCSC B@self; misaligned for VCSR - dual_after_matmul = nv._dual_arr - assert dual_after_matmul is not None - assert dual_after_matmul._format != v._format - - nv @ B - Bl @ nv - assert nv._dual_arr is dual_after_matmul # reused, not rebuilt + assert nv._arr is v # still the original array, not replaced by a dual + assert not hasattr(nv, "_dual_arr") def test_transpose_major_roundtrip(dense, vcls): diff --git a/tests/test_vcs_norm_recipes.py b/tests/test_vcs_norm_recipes.py index 4d5a7e0..e9a4a3e 100644 --- a/tests/test_vcs_norm_recipes.py +++ b/tests/test_vcs_norm_recipes.py @@ -296,7 +296,7 @@ def test_recipe_with_an_unknown_g_code_is_rejected(vcls, dense): def test_cache_does_not_pin_a_dropped_view(vcls, dense): - """The cache holds views weakly, so it can't keep an O(nnz) _dual_arr alive.""" + """The cache holds views weakly, so a dropped view is actually collected.""" if dense.sum() == 0: pytest.skip("all-zero matrix: median row total is 0") v = vcls.from_scipy(_scipy_for(vcls, dense))