Chunked dual transpose

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).
 

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

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

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)

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,17 +185,16 @@ 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
     assert dual_after_matmul._format != v._format
 
     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):