diff --git a/src/vsparse/_base.py b/src/vsparse/_base.py index 9a3ed29..22d6c7e 100644 --- a/src/vsparse/_base.py +++ b/src/vsparse/_base.py @@ -512,6 +512,26 @@ 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. + + ``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]) + 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 _select_minor(self, key: Any) -> _VCSBase: """Select along the minor axis (rows for VCSC, columns for VCSR). diff --git a/src/vsparse/_vcs_matmul.py b/src/vsparse/_vcs_matmul.py index ed54d5a..bf253e9 100644 --- a/src/vsparse/_vcs_matmul.py +++ b/src/vsparse/_vcs_matmul.py @@ -18,15 +18,13 @@ :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`, 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 a chunk just takes the slice of them its own axis covers. """ from __future__ import annotations @@ -62,14 +60,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 +93,57 @@ 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) +# ``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. Deliberately generous, since underestimating +# means the budget doesn't hold. +_TRANSPOSE_BYTES_PER_NNZ = 64 -# -- dual-format cache: gives every call a major-aligned array to run against +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)] -def _get_dual(nview: _VCSNormalizedBase): - """The opposite-format raw array for ``nview``'s wrapped array, built once and cached.""" + # 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)``.""" + 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 + + +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 @@ -140,10 +167,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 +203,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..01e12a4 100644 --- a/src/vsparse/_vcs_norm.py +++ b/src/vsparse/_vcs_norm.py @@ -29,12 +29,11 @@ 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`, cached by vsparse._vcs_matmul when + # regrouping the whole array fits one chunk's budget. self._dual_arr: _VCSBase | None = None + 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..09d7279 --- /dev/null +++ b/tests/test_chunked_transpose.py @@ -0,0 +1,147 @@ +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_vcs_norm.py b/tests/test_vcs_norm.py index fa535d9..700d660 100644 --- a/tests/test_vcs_norm.py +++ b/tests/test_vcs_norm.py @@ -174,8 +174,8 @@ 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.""" if dense.sum() == 0: pytest.skip("all-zero matrix: median row total is 0") v = vcls.from_scipy(_scipy_for(vcls, dense)) @@ -185,8 +185,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 +194,7 @@ 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_transpose_major_roundtrip(dense, vcls):