From 72021e63d2659e83966367d1f93aa7459a656f6b Mon Sep 17 00:00:00 2001 From: fishidaho Date: Wed, 2 Sep 2026 17:22:08 -0700 Subject: [PATCH 1/2] Regroup a chunk at a time for the misaligned matmul direction `_get_dual` built a full opposite-format copy of the array the first time a normalized-view matmul needed the direction the storage isn't aligned for. No size gate, no opt-in, no warning: a second complete copy of the dataset, allocated silently inside a matmul. On a 72M-nonzero array that's +851 MB; at atlas scale it's simply fatal on a machine the original already fits. The regrouping doesn't have to happen all at once. `Delta @ B` splits over the contracted axis (sum of `Delta[:, C] @ B[C, :]` over column chunks) and `B @ Delta` splits the same way over row chunks, so a contiguous range of major slices can be transposed alone, fed to the same major-aligned kernel, accumulated into the shared output, and dropped. Peak memory is then one chunk's regrouping -- set by a byte budget -- instead of the whole array's. `_VCSBase._major_range` is the chunking primitive: a contiguous major range is already contiguous in every stored array, so `values`/`indices` come back as views and only the two small pointer arrays are rebuilt. One case still caches without being asked: an array whose *entire* regrouping already fits inside a single chunk's budget. Transposing it per call would allocate exactly that much anyway, so keeping it costs no extra peak memory -- and it matters, because repeated misaligned products run ~55x faster against a cached dual than re-regrouping every call. Past one chunk nothing is cached, and `pin_dual()` is the explicit opt-in for callers who want the cached dual anyway and know the memory is affordable. Measured on 72M nonzeros (30000x3000, VCSC, `self @ B`): chunked first call 10.30s, +0 MB peak RSS pinned first call 14.83s, +851 MB peak RSS max abs difference between the two: 1.9e-14 so the default is now both faster to a first result and free of the extra copy; steady-state repeated calls are where pinning wins, hence the opt-in. `test_dual_array_is_built_lazily_and_cached` becomes two tests, since the behavior it pinned is exactly what changed: small arrays still cache lazily, and multi-chunk arrays cache nothing. Co-Authored-By: Claude Sonnet 5 --- src/vsparse/_base.py | 24 ++++ src/vsparse/_vcs_matmul.py | 178 +++++++++++++++++++++------ src/vsparse/_vcs_norm.py | 22 +++- tests/test_chunked_transpose.py | 205 ++++++++++++++++++++++++++++++++ tests/test_vcs_norm.py | 41 ++++++- 5 files changed, 422 insertions(+), 48 deletions(-) create mode 100644 tests/test_chunked_transpose.py diff --git a/src/vsparse/_base.py b/src/vsparse/_base.py index eb89f05..54fa193 100644 --- a/src/vsparse/_base.py +++ b/src/vsparse/_base.py @@ -360,6 +360,30 @@ def _select_major(self, key: Any) -> _VCSBase: ) return type(self)(new_shape, new_major_ptr, new_values, new_value_ptr, new_indices) + def _major_range(self, start: int, stop: int) -> _VCSBase: + """The contiguous major-slice range ``[start, stop)``, without copying values. + + Unlike :meth:`_select_major` (arbitrary index arrays, so it has to + gather), a contiguous range is already contiguous in every stored + array: ``values``/``indices`` come back as plain views into this + array's buffers, and only the two small pointer arrays are rebuilt + (rebased to the new start). That makes it cheap enough to use as the + chunking primitive for a streaming pass over the major axis. + """ + u0, u1 = int(self.major_ptr[start]), int(self.major_ptr[stop]) + k0, k1 = int(self.value_ptr[u0]), int(self.value_ptr[u1]) + n_sel = stop - start + new_shape = ( + (self.n_minor, n_sel) if self._format == "csc" else (n_sel, self.n_minor) + ) + return type(self)( + new_shape, + self.major_ptr[start : stop + 1] - u0, + self.values[u0:u1], + self.value_ptr[u0 : u1 + 1] - k0, + self.indices[k0:k1], + ) + def __getitem__(self, key): if isinstance(key, tuple): if len(key) != 2: diff --git a/src/vsparse/_vcs_matmul.py b/src/vsparse/_vcs_matmul.py index ed54d5a..b1afb9d 100644 --- a/src/vsparse/_vcs_matmul.py +++ b/src/vsparse/_vcs_matmul.py @@ -18,15 +18,20 @@ :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), :func:`_get_dual` builds and -caches the *other* VCS format's storage for the same underlying array, once, -via :meth:`~vsparse._base._VCSBase._transpose_major` -- a single global sort -(see :func:`vsparse._construct.transpose_major`), not a per-call cost -- and -every subsequent call in the misaligned direction runs the same -major-aligned kernel against that cached dual instead. ``row_scale``/ -``gene_scale``/``col_mean`` are per-row/per-column statistics, so they carry -over unchanged regardless of which physical layout is used to compute -against. +enough for wide ``B`` to need throttling), the storage is regrouped into the +other VCS format via :meth:`~vsparse._base._VCSBase._transpose_major` (see +:func:`vsparse._construct.transpose_major`) so the same major-aligned kernel +can run against it. + +That regrouping is done **a chunk of major slices at a time** rather than +over the whole array, so the extra memory is one chunk's worth rather than a +second full copy of the dataset -- see the section below the kernels. +:func:`pin_dual` opts into caching the full dual instead, for callers who +will run many misaligned products against an array that fits comfortably. + +``row_scale``/``gene_scale``/``col_mean`` are per-row/per-column statistics, +so they carry over unchanged regardless of which physical layout is used to +compute against (a chunk sees the slice of them its own axis covers). """ from __future__ import annotations @@ -62,14 +67,6 @@ def _vcsr_matmul_delta(major_ptr, values, value_ptr, indices, row_scale, gene_sc out[i, c] += delta * B[col, c] -def _matmul_vcsr(arr, row_scale, gene_scale, B: np.ndarray) -> np.ndarray: - n_rows = arr.n_major - k = B.shape[1] - out = np.zeros((n_rows, k), dtype=np.float64) - _vcsr_matmul_delta(arr.major_ptr, arr.values, arr.value_ptr, arr.indices, row_scale, gene_scale, B, out) - return out - - # -- major-aligned: VCSC, B @ self (out cols == major slices) --------------- # # ``B`` arrives as (p, n_rows) and the natural output is (p, n_cols) -- but @@ -103,20 +100,90 @@ def _vcsc_rmatmul_delta(major_ptr, values, value_ptr, indices, row_scale, gene_s acc[c] += delta * brow[c] -def _rmatmul_vcsc(arr, row_scale, gene_scale, B: np.ndarray) -> np.ndarray: - n_cols = arr.n_major - p = B.shape[0] - Bt = np.ascontiguousarray(B.T) # (n_rows, p) - out_t = np.zeros((n_cols, p), dtype=np.float64) - _vcsc_rmatmul_delta(arr.major_ptr, arr.values, arr.value_ptr, arr.indices, row_scale, gene_scale, Bt, out_t) - return np.ascontiguousarray(out_t.T) - - -# -- dual-format cache: gives every call a major-aligned array to run against - - -def _get_dual(nview: _VCSNormalizedBase): - """The opposite-format raw array for ``nview``'s wrapped array, built once and cached.""" +# -- misaligned direction: transpose a chunk at a time, not the whole array -- +# +# Converting the whole array to its opposite format gives every call a +# major-aligned kernel to run against, but it costs a second full copy of +# the array -- unbounded, paid up front on the first call, and fatal at +# atlas scale where the array already fills most of the machine. +# +# The regrouping doesn't have to happen all at once. ``Delta @ B`` splits +# over the contracted axis (``Delta[:, C] @ B[C, :]``, summed over column +# chunks) and ``B @ Delta`` splits the same way over row chunks, so a +# contiguous range of major slices can be transposed on its own, fed to the +# same major-aligned kernel, accumulated into the shared output, and thrown +# away. Peak memory is then one chunk's worth of regrouping instead of the +# whole array's, and it's the caller's chunk budget that sets it rather than +# the dataset size. +# +# ``pin_dual`` stays available for callers who do want the cached full dual +# -- many repeated misaligned products against an array that comfortably +# fits -- but nothing reaches for it implicitly any more. + +_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 (the values/major/minor entry arrays, the +# lexsort permutation, the sorted copies) on top of the output. Deliberately +# generous, since underestimating means the budget doesn't actually 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 + + +def _aligned_source(nview: _VCSNormalizedBase, needed_format: str): + """``(array, None)`` to run one major-aligned pass, or ``(None, chunk bounds)``. + + Three cases, in order: the array is already in the format the kernel + needs; a full dual has been pinned (explicitly, or cached below); or the + misaligned direction has to regroup, in which case the caller loops over + the returned chunk bounds. + + The one case that caches without being asked is an array whose *whole* + regrouping already fits inside a single chunk's budget. Transposing it + per call would allocate exactly that much anyway, so keeping the result + costs no additional peak memory and saves every later call the work -- + which matters, since repeated misaligned products against a cached dual + run far faster than re-regrouping each time. + """ + 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: + return pin_dual(nview), None + return None, bounds + + +def pin_dual(nview: _VCSNormalizedBase): + """Build and cache the opposite-format copy of ``nview``'s array, once. + + Opt-in: it doubles the array's memory. Worth it only when many + misaligned-direction products will run against an array that fits + comfortably; otherwise the chunked path above needs no extra memory. + """ if nview._dual_arr is None: nview._dual_arr = nview._arr._transpose_major() return nview._dual_arr @@ -140,10 +207,27 @@ def normalized_at_dense(nview: _VCSNormalizedBase, other: Any) -> np.ndarray: arr = nview._arr n_cols = arr.shape[1] B, squeeze = _prep_dense(other, n_cols) + out = np.zeros((arr.shape[0], B.shape[1]), dtype=np.float64) + + # self @ B is major-aligned for VCSR; a VCSC array needs regrouping. + src, bounds = _aligned_source(nview, "csr") + if src is not None: + _vcsr_matmul_delta( + src.major_ptr, src.values, src.value_ptr, src.indices, + nview.row_scale, nview.gene_scale, B, 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], + np.ascontiguousarray(B[start:stop]), out, + ) - # self @ B is major-aligned for VCSR; VCSC needs its VCSR dual. - src = arr if arr._format == "csr" else _get_dual(nview) - out = _matmul_vcsr(src, nview.row_scale, nview.gene_scale, B) baseline = (-nview.col_mean) @ B # (k,): every row's implicit-zero contribution out += baseline[None, :] return out[:, 0] if squeeze else out @@ -159,10 +243,28 @@ def dense_at_normalized(nview: _VCSNormalizedBase, other: Any) -> np.ndarray: if B2.ndim != 2 or B2.shape[1] != n_rows: raise ValueError(f"shape mismatch: expected last dimension {n_rows}, got {B2.shape}") B2 = np.ascontiguousarray(B2) - - # B @ self is major-aligned for VCSC; VCSR needs its VCSC dual. - src = arr if arr._format == "csc" else _get_dual(nview) - out = _rmatmul_vcsc(src, nview.row_scale, nview.gene_scale, B2) + p = B2.shape[0] + Bt = np.ascontiguousarray(B2.T) # (n_rows, p) -- see the note above the kernel + out_t = np.zeros((arr.shape[1], p), dtype=np.float64) + + # B @ self is major-aligned for VCSC; a VCSR array needs regrouping. + src, bounds = _aligned_source(nview, "csc") + if src is not None: + _vcsc_rmatmul_delta( + src.major_ptr, src.values, src.value_ptr, src.indices, + nview.row_scale, nview.gene_scale, Bt, 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, + np.ascontiguousarray(Bt[start:stop]), out_t, + ) + + out = np.ascontiguousarray(out_t.T) baseline = B2.sum(axis=1)[:, None] * (-nview.col_mean)[None, :] # (m, n_cols) out += baseline return out[0, :] if squeeze else out diff --git a/src/vsparse/_vcs_norm.py b/src/vsparse/_vcs_norm.py index 12990dd..e4e7709 100644 --- a/src/vsparse/_vcs_norm.py +++ b/src/vsparse/_vcs_norm.py @@ -29,12 +29,26 @@ class _VCSNormalizedBase(NormalizedViewBase): def __init__(self, arr: _VCSBase) -> None: super().__init__(arr) - # Lazily built, cached opposite-format copy of `arr` -- see - # vsparse._vcs_matmul._get_dual. Built at most once per view, the - # first time a matmul needs it in the direction `arr` isn't - # major-aligned for. + # Opposite-format copy of `arr`, built only if `pin_dual` is called. + # Left as None, matmuls in the direction `arr` isn't major-aligned + # for regroup a chunk at a time instead -- see vsparse._vcs_matmul. self._dual_arr: _VCSBase | None = None + def pin_dual(self) -> None: + """Build and keep the opposite-format copy of the wrapped array. + + Matmuls in the direction the array isn't major-aligned for normally + regroup one chunk at a time, which needs no lasting extra memory. + Pinning trades that for a second full copy of the array, held for + this view's lifetime, so repeated misaligned products skip the + per-call regrouping. Only worth it when the array comfortably fits + and many such products are coming; the result is identical either + way. + """ + from vsparse._vcs_matmul import pin_dual + + pin_dual(self) + def __matmul__(self, other: Any) -> Any: """``self @ other`` for a dense ``other`` -- see :mod:`vsparse._vcs_matmul`.""" from vsparse._vcs_matmul import normalized_at_dense diff --git a/tests/test_chunked_transpose.py b/tests/test_chunked_transpose.py new file mode 100644 index 0000000..f5e8d12 --- /dev/null +++ b/tests/test_chunked_transpose.py @@ -0,0 +1,205 @@ +"""The misaligned-direction matmul regroups a chunk at a time, not the whole array. + +The results have to be identical to the full-dual route (covered against a +dense reference in test_vcs_norm.py); what's tested here is that the +chunking machinery itself is correct at boundaries, that multi-chunk runs +actually happen, and that peak memory stays bounded rather than scaling +with the dataset. +""" + +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) + + +# -- chunk boundaries -------------------------------------------------------- + + +def test_chunks_cover_the_major_axis_exactly(dense, vcls): + v = vcls.from_scipy(_scipy_for(vcls, dense)) + for budget in (1, 64, 1 << 20): + bounds = _chunk_bounds(v, budget) + assert bounds[0][0] == 0 + assert bounds[-1][1] == v.n_major + for (_, prev_stop), (start, _) in pairwise(bounds): + assert start == prev_stop # contiguous, no gaps or overlaps + assert all(start < stop for start, stop in bounds) # always advances + + +def test_everything_fits_in_one_chunk_under_a_large_budget(dense, vcls): + v = vcls.from_scipy(_scipy_for(vcls, dense)) + assert _chunk_bounds(v, 1 << 30) == [(0, v.n_major)] + + +def test_a_tiny_budget_still_makes_progress(vcls): + """A single major slice bigger than the budget can't be split further.""" + dense = np.ones((40, 40)) + v = vcls.from_scipy(_scipy_for(vcls, dense)) + bounds = _chunk_bounds(v, 1) + assert bounds == [(i, i + 1) for i in range(v.n_major)] + + +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_chunks_respect_the_budget(vcls, rng): + 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 # ~40 nonzeros per chunk + + cumulative = v.value_ptr[v.major_ptr] + for start, stop in _chunk_bounds(v, budget): + chunk_nnz = int(cumulative[stop] - cumulative[start]) + # A chunk of one major slice can exceed the budget: it's indivisible. + assert chunk_nnz <= 40 or stop - start == 1 + + +def test_major_range_is_a_zero_copy_view(vcls, rng): + """The chunking primitive must not copy the value/index buffers.""" + 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 chunk.n_major == 3 + 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_fancy_selection(vcls, rng): + dense = rng.integers(0, 4, size=(30, 20)).astype(np.float64) + v = vcls.from_scipy(_scipy_for(vcls, dense)) + + np.testing.assert_allclose( + v._major_range(2, 7).toarray(), + v._select_major(np.arange(2, 7)).toarray(), + ) + + +# -- multi-chunk products match the single-chunk result ---------------------- + + +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, 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 ever built -- that's 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) + 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_many_chunks_actually_run(monkeypatch, vcls, rng): + """Guards against the budget silently collapsing to a single chunk.""" + 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 + + calls = [] + original = vcs_matmul._chunk_bounds + monkeypatch.setattr( + vcs_matmul, + "_chunk_bounds", + lambda arr, budget: calls.append(len(original(arr, budget))) or original(arr, budget), + ) + + nv = v.normalized() + nv @ rng.normal(size=(dense.shape[1], 2)) + rng.normal(size=(2, dense.shape[0])) @ nv + + assert calls and max(calls) > 1 + + +# -- the memory bound -------------------------------------------------------- + + +def test_misaligned_matmul_peak_is_bounded_by_the_chunk_budget(monkeypatch, rng): + """The regression: peak memory tracks the chunk budget, not the dataset. + + A full dual transpose is a second copy of the whole array; this asserts + the misaligned direction stays far below that even when the budget only + admits a fraction of the nonzeros at a time. + """ + 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)) # self @ B is misaligned for VCSC + nnz_bytes = v.nnz * v.indices.dtype.itemsize + B = rng.normal(size=(dense.shape[1], 2)) + + v.normalized() @ B # warm up numba's 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_vcs_norm.py b/tests/test_vcs_norm.py index fa535d9..0638803 100644 --- a/tests/test_vcs_norm.py +++ b/tests/test_vcs_norm.py @@ -174,8 +174,13 @@ def test_all_zero_matrix_does_not_crash(vcls): np.testing.assert_allclose(nv.toarray(), np.zeros((5, 4))) -def test_dual_array_is_built_lazily_and_cached(dense, vcls): - """The opposite-format dual is only built on the first misaligned-direction matmul, then reused.""" +def test_small_array_dual_is_cached_lazily(dense, vcls): + """An array whose whole regrouping fits one chunk is transposed once and kept. + + Doing it per call would allocate the same amount anyway, so caching + costs no extra peak memory here. (Arrays too big for a single chunk + regroup per chunk and cache nothing -- see test_chunked_transpose.py.) + """ if dense.sum() == 0: pytest.skip("all-zero matrix: median row total is 0") v = vcls.from_scipy(_scipy_for(vcls, dense)) @@ -185,8 +190,8 @@ def test_dual_array_is_built_lazily_and_cached(dense, vcls): rng = np.random.default_rng(12) B = rng.normal(size=(dense.shape[1], 3)) Bl = rng.normal(size=(3, dense.shape[0])) - nv @ B # major-aligned for VCSR self@B; builds the dual for VCSC self@B - Bl @ nv # major-aligned for VCSC B@self; builds the dual for VCSR B@self + 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 @@ -194,8 +199,32 @@ def test_dual_array_is_built_lazily_and_cached(dense, vcls): nv @ B Bl @ nv - # same object reused, not rebuilt, across repeated calls - assert nv._dual_arr is dual_after_matmul + assert nv._dual_arr is dual_after_matmul # reused, not rebuilt + + +def test_pin_dual_caches_and_agrees_with_the_chunked_path(dense, vcls): + """Pinning is opt-in, reused across calls, and numerically identical.""" + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + v = vcls.from_scipy(_scipy_for(vcls, dense)) + rng = np.random.default_rng(12) + B = rng.normal(size=(dense.shape[1], 3)) + Bl = rng.normal(size=(3, dense.shape[0])) + + chunked = v.normalized() + unpinned_right, unpinned_left = chunked @ B, Bl @ chunked + + pinned = v.normalized() + pinned.pin_dual() + dual = pinned._dual_arr + assert dual is not None + assert dual._format != v._format + + np.testing.assert_allclose(pinned @ B, unpinned_right, atol=1e-12) + np.testing.assert_allclose(Bl @ pinned, unpinned_left, atol=1e-12) + + pinned.pin_dual() # idempotent: same object, not rebuilt + assert pinned._dual_arr is dual def test_transpose_major_roundtrip(dense, vcls): From f54ff05803e4e44d618a8bda89ee0cfa94e36cb4 Mon Sep 17 00:00:00 2001 From: fishidaho Date: Thu, 3 Sep 2026 17:09:39 -0700 Subject: [PATCH 2/2] Trim comments and tests, and drop the unused pin_dual method The public pin_dual had no caller outside its own test. The internal _build_dual it wrapped stays, since that is how a whole-array regrouping gets cached when it fits one chunk's budget. Chunking tests fold into the properties that matter: chunks tile the axis and respect the budget, _major_range shares buffers and matches ordinary slicing, results are invariant to chunk size, and peak memory tracks the budget rather than the array. --- src/vsparse/_base.py | 8 +- src/vsparse/_vcs_matmul.py | 70 ++++----------- src/vsparse/_vcs_norm.py | 19 +--- tests/test_chunked_transpose.py | 153 ++++++++------------------------ tests/test_vcs_norm.py | 32 +------ 5 files changed, 58 insertions(+), 224 deletions(-) diff --git a/src/vsparse/_base.py b/src/vsparse/_base.py index 6c82f17..22d6c7e 100644 --- a/src/vsparse/_base.py +++ b/src/vsparse/_base.py @@ -515,12 +515,8 @@ def _select_major(self, key: Any) -> _VCSBase: def _major_range(self, start: int, stop: int) -> _VCSBase: """The contiguous major-slice range ``[start, stop)``, without copying values. - Unlike :meth:`_select_major` (arbitrary index arrays, so it has to - gather), a contiguous range is already contiguous in every stored - array: ``values``/``indices`` come back as plain views into this - array's buffers, and only the two small pointer arrays are rebuilt - (rebased to the new start). That makes it cheap enough to use as the - chunking primitive for a streaming pass over the major axis. + ``values``/``indices`` come back as views; only the two pointer + arrays are rebuilt, rebased to the new start. """ u0, u1 = int(self.major_ptr[start]), int(self.major_ptr[stop]) k0, k1 = int(self.value_ptr[u0]), int(self.value_ptr[u1]) diff --git a/src/vsparse/_vcs_matmul.py b/src/vsparse/_vcs_matmul.py index b1afb9d..bf253e9 100644 --- a/src/vsparse/_vcs_matmul.py +++ b/src/vsparse/_vcs_matmul.py @@ -19,19 +19,12 @@ 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` (see -:func:`vsparse._construct.transpose_major`) so the same major-aligned kernel -can run against it. - -That regrouping is done **a chunk of major slices at a time** rather than -over the whole array, so the extra memory is one chunk's worth rather than a -second full copy of the dataset -- see the section below the kernels. -:func:`pin_dual` opts into caching the full dual instead, for callers who -will run many misaligned products against an array that fits comfortably. +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`` are per-row/per-column statistics, -so they carry over unchanged regardless of which physical layout is used to -compute against (a chunk sees the slice of them its own axis covers). +so a chunk just takes the slice of them its own axis covers. """ from __future__ import annotations @@ -100,32 +93,15 @@ def _vcsc_rmatmul_delta(major_ptr, values, value_ptr, indices, row_scale, gene_s acc[c] += delta * brow[c] -# -- misaligned direction: transpose a chunk at a time, not the whole array -- -# -# Converting the whole array to its opposite format gives every call a -# major-aligned kernel to run against, but it costs a second full copy of -# the array -- unbounded, paid up front on the first call, and fatal at -# atlas scale where the array already fills most of the machine. -# -# The regrouping doesn't have to happen all at once. ``Delta @ B`` splits -# over the contracted axis (``Delta[:, C] @ B[C, :]``, summed over column -# chunks) and ``B @ Delta`` splits the same way over row chunks, so a -# contiguous range of major slices can be transposed on its own, fed to the -# same major-aligned kernel, accumulated into the shared output, and thrown -# away. Peak memory is then one chunk's worth of regrouping instead of the -# whole array's, and it's the caller's chunk budget that sets it rather than -# the dataset size. -# -# ``pin_dual`` stays available for callers who do want the cached full dual -# -- many repeated misaligned products against an array that comfortably -# fits -- but nothing reaches for it implicitly any more. +# ``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. _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 (the values/major/minor entry arrays, the -# lexsort permutation, the sorted copies) on top of the output. Deliberately -# generous, since underestimating means the budget doesn't actually hold. +# several nnz-sized temporaries. Deliberately generous, since underestimating +# means the budget doesn't hold. _TRANSPOSE_BYTES_PER_NNZ = 64 @@ -152,20 +128,7 @@ def _chunk_bounds(arr, budget_bytes: int) -> list[tuple[int, int]]: def _aligned_source(nview: _VCSNormalizedBase, needed_format: str): - """``(array, None)`` to run one major-aligned pass, or ``(None, chunk bounds)``. - - Three cases, in order: the array is already in the format the kernel - needs; a full dual has been pinned (explicitly, or cached below); or the - misaligned direction has to regroup, in which case the caller loops over - the returned chunk bounds. - - The one case that caches without being asked is an array whose *whole* - regrouping already fits inside a single chunk's budget. Transposing it - per call would allocate exactly that much anyway, so keeping the result - costs no additional peak memory and saves every later call the work -- - which matters, since repeated misaligned products against a cached dual - run far faster than re-regrouping each time. - """ + """``(array, None)`` to run one major-aligned pass, or ``(None, chunk bounds)``.""" arr = nview._arr if arr._format == needed_format: return arr, None @@ -173,17 +136,14 @@ def _aligned_source(nview: _VCSNormalizedBase, needed_format: str): return nview._dual_arr, None bounds = _chunk_bounds(arr, _CHUNK_BUDGET_BYTES) if len(bounds) <= 1: - return pin_dual(nview), None + # 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 -def pin_dual(nview: _VCSNormalizedBase): - """Build and cache the opposite-format copy of ``nview``'s array, once. - - Opt-in: it doubles the array's memory. Worth it only when many - misaligned-direction products will run against an array that fits - comfortably; otherwise the chunked path above needs no extra memory. - """ +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 diff --git a/src/vsparse/_vcs_norm.py b/src/vsparse/_vcs_norm.py index e4e7709..01e12a4 100644 --- a/src/vsparse/_vcs_norm.py +++ b/src/vsparse/_vcs_norm.py @@ -29,25 +29,10 @@ class _VCSNormalizedBase(NormalizedViewBase): def __init__(self, arr: _VCSBase) -> None: super().__init__(arr) - # Opposite-format copy of `arr`, built only if `pin_dual` is called. - # Left as None, matmuls in the direction `arr` isn't major-aligned - # for regroup a chunk at a time instead -- see vsparse._vcs_matmul. + # Opposite-format copy of `arr`, cached by vsparse._vcs_matmul when + # regrouping the whole array fits one chunk's budget. self._dual_arr: _VCSBase | None = None - def pin_dual(self) -> None: - """Build and keep the opposite-format copy of the wrapped array. - - Matmuls in the direction the array isn't major-aligned for normally - regroup one chunk at a time, which needs no lasting extra memory. - Pinning trades that for a second full copy of the array, held for - this view's lifetime, so repeated misaligned products skip the - per-call regrouping. Only worth it when the array comfortably fits - and many such products are coming; the result is identical either - way. - """ - from vsparse._vcs_matmul import pin_dual - - pin_dual(self) 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 index 0dfba12..09d7279 100644 --- a/tests/test_chunked_transpose.py +++ b/tests/test_chunked_transpose.py @@ -1,12 +1,3 @@ -"""The misaligned-direction matmul regroups a chunk at a time, not the whole array. - -The results have to be identical to the full-dual route (covered against a -dense reference in test_vcs_norm.py); what's tested here is that the -chunking machinery itself is correct at boundaries, that multi-chunk runs -actually happen, and that peak memory stays bounded rather than scaling -with the dataset. -""" - from __future__ import annotations import tracemalloc @@ -29,85 +20,67 @@ def _scipy_for(vcls, dense): return sp.csc_array(dense) if vcls is VCSCArray else sp.csr_array(dense) -# -- chunk boundaries -------------------------------------------------------- - - -def test_chunks_cover_the_major_axis_exactly(dense, vcls): - v = vcls.from_scipy(_scipy_for(vcls, dense)) - for budget in (1, 64, 1 << 20): - bounds = _chunk_bounds(v, budget) - assert bounds[0][0] == 0 - assert bounds[-1][1] == v.n_major - for (_, prev_stop), (start, _) in pairwise(bounds): - assert start == prev_stop # contiguous, no gaps or overlaps - assert all(start < stop for start, stop in bounds) # always advances - - -def test_everything_fits_in_one_chunk_under_a_large_budget(dense, vcls): - v = vcls.from_scipy(_scipy_for(vcls, dense)) - assert _chunk_bounds(v, 1 << 30) == [(0, v.n_major)] +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_a_tiny_budget_still_makes_progress(vcls): - """A single major slice bigger than the budget can't be split further.""" - dense = np.ones((40, 40)) +@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, 1) - assert bounds == [(i, i + 1) for i in range(v.n_major)] + bounds = _chunk_bounds(v, budget) - -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) == [] + 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 # ~40 nonzeros per chunk + 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]) - # A chunk of one major slice can exceed the budget: it's indivisible. 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): - """The chunking primitive must not copy the value/index buffers.""" + """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 chunk.n_major == 3 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_fancy_selection(vcls, rng): +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 - np.testing.assert_allclose( - v._major_range(2, 7).toarray(), - v._select_major(np.arange(2, 7)).toarray(), - ) - - -# -- multi-chunk products match the single-chunk result ---------------------- - - -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) + 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]) @@ -131,22 +104,7 @@ def test_matmul_is_chunk_size_invariant(monkeypatch, vcls, rng, budget): def test_multi_chunk_matmul_caches_nothing(monkeypatch, vcls, rng): - """Past one chunk, no full dual is ever built -- that's 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) - 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_many_chunks_actually_run(monkeypatch, vcls, rng): - """Guards against the budget silently collapsing to a single chunk.""" + """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) @@ -154,39 +112,23 @@ def test_many_chunks_actually_run(monkeypatch, vcls, rng): monkeypatch.setattr(vcs_matmul, "_CHUNK_BUDGET_BYTES", 200) assert len(_chunk_bounds(v, 200)) > 1 - calls = [] - original = vcs_matmul._chunk_bounds - monkeypatch.setattr( - vcs_matmul, - "_chunk_bounds", - lambda arr, budget: calls.append(len(original(arr, budget))) or original(arr, budget), - ) - nv = v.normalized() nv @ rng.normal(size=(dense.shape[1], 2)) rng.normal(size=(2, dense.shape[0])) @ nv - assert calls and max(calls) > 1 - - -# -- the memory bound -------------------------------------------------------- + assert nv._dual_arr is None def test_misaligned_matmul_peak_is_bounded_by_the_chunk_budget(monkeypatch, rng): - """The regression: peak memory tracks the chunk budget, not the dataset. - - A full dual transpose is a second copy of the whole array; this asserts - the misaligned direction stays far below that even when the budget only - admits a fraction of the nonzeros at a time. - """ + """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)) # self @ B is misaligned for VCSC + 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 numba's JIT before measuring + v.normalized() @ B # warm up the JIT before measuring monkeypatch.setattr(vcs_matmul, "_CHUNK_BUDGET_BYTES", 64 * _TRANSPOSE_BYTES_PER_NNZ) nv = v.normalized() @@ -203,22 +145,3 @@ def test_misaligned_matmul_peak_is_bounded_by_the_chunk_budget(monkeypatch, rng) assert nv._dual_arr is None assert peak - before < nnz_bytes np.testing.assert_allclose(out, _reference(dense) @ B, atol=1e-7) - - -def test_major_range_agrees_with_native_slicing(vcls, rng): - """#23 gave __getitem__ a native path for slices; _major_range must match it. - - Both now exist for the same job on a contiguous range -- __getitem__ for - callers, _major_range as the chunking primitive (it returns views rather - than gathering). If they ever disagree, the chunked matmul is computing - against a different sub-array than a caller would get. - """ - dense = rng.integers(0, 4, size=(30, 20)).astype(np.float64) - v = vcls.from_scipy(_scipy_for(vcls, dense)) - start, stop = 2, 7 - - chunk = v._major_range(start, stop) - native = v[:, start:stop] if vcls is VCSCArray else v[start:stop, :] - - assert isinstance(native, vcls) - np.testing.assert_allclose(chunk.toarray(), native.toarray()) diff --git a/tests/test_vcs_norm.py b/tests/test_vcs_norm.py index 0638803..700d660 100644 --- a/tests/test_vcs_norm.py +++ b/tests/test_vcs_norm.py @@ -175,12 +175,7 @@ def test_all_zero_matrix_does_not_crash(vcls): def test_small_array_dual_is_cached_lazily(dense, vcls): - """An array whose whole regrouping fits one chunk is transposed once and kept. - - Doing it per call would allocate the same amount anyway, so caching - costs no extra peak memory here. (Arrays too big for a single chunk - regroup per chunk and cache nothing -- see test_chunked_transpose.py.) - """ + """An array whose whole regrouping fits one chunk is transposed once and kept.""" if dense.sum() == 0: pytest.skip("all-zero matrix: median row total is 0") v = vcls.from_scipy(_scipy_for(vcls, dense)) @@ -202,31 +197,6 @@ def test_small_array_dual_is_cached_lazily(dense, vcls): assert nv._dual_arr is dual_after_matmul # reused, not rebuilt -def test_pin_dual_caches_and_agrees_with_the_chunked_path(dense, vcls): - """Pinning is opt-in, reused across calls, and numerically identical.""" - if dense.sum() == 0: - pytest.skip("all-zero matrix: median row total is 0") - v = vcls.from_scipy(_scipy_for(vcls, dense)) - rng = np.random.default_rng(12) - B = rng.normal(size=(dense.shape[1], 3)) - Bl = rng.normal(size=(3, dense.shape[0])) - - chunked = v.normalized() - unpinned_right, unpinned_left = chunked @ B, Bl @ chunked - - pinned = v.normalized() - pinned.pin_dual() - dual = pinned._dual_arr - assert dual is not None - assert dual._format != v._format - - np.testing.assert_allclose(pinned @ B, unpinned_right, atol=1e-12) - np.testing.assert_allclose(Bl @ pinned, unpinned_left, atol=1e-12) - - pinned.pin_dual() # idempotent: same object, not rebuilt - assert pinned._dual_arr is dual - - def test_transpose_major_roundtrip(dense, vcls): v = vcls.from_scipy(_scipy_for(vcls, dense)) dual = v._transpose_major()