diff --git a/src/vsparse/_base.py b/src/vsparse/_base.py index 428ea6b..055b003 100644 --- a/src/vsparse/_base.py +++ b/src/vsparse/_base.py @@ -227,10 +227,11 @@ def _major_sums(self) -> np.ndarray: return np.bincount(group_of_major, weights=weighted, minlength=self.n_major) def _minor_sums(self) -> np.ndarray: - """Per-minor-index totals -- a scatter-add over every nonzero.""" - group_sizes = np.diff(self.value_ptr) - expanded = np.repeat(self.values.astype(np.float64), group_sizes) - return np.bincount(self.indices, weights=expanded, minlength=self.n_minor) + """A parallel scatter-add over every nonzero to get per-minor-index totals.""" + return _ops.minor_sums( + self.values, self.value_ptr, self.indices, self.n_minor, + _ops.accumulator_threads(self.n_minor), + ) def sum(self, axis: int | None = None) -> np.ndarray | float: """Sum of (structural) values along ``axis`` (0=rows, 1=columns), or overall if ``None``.""" @@ -252,8 +253,11 @@ def _major_nnz(self) -> np.ndarray: ) def _minor_nnz(self) -> np.ndarray: - """Per-minor-index stored-element counts -- a scatter-add over every nonzero.""" - return np.bincount(self.indices, minlength=self.n_minor).astype(np.int64) + """A parallel scatter over every nonzero to get per-minor-index counts.""" + return _ops.minor_counts( + self.value_ptr, self.indices, self.n_minor, + _ops.accumulator_threads(self.n_minor), + ) def getnnz(self, axis: int | None = None) -> np.ndarray | int: """Count of stored elements along ``axis``, or overall if ``None``.""" @@ -298,11 +302,11 @@ def _major_reduce(self, ufunc: np.ufunc, initial: Any) -> np.ndarray: def _minor_reduce(self, ufunc: np.ufunc, initial: Any) -> np.ndarray: """Per-minor-index max/min, accounting for implicit zeros in sparse slices.""" - group_sizes = np.diff(self.value_ptr) - expanded = np.repeat(self.values, group_sizes) - out = np.full(self.n_minor, initial, dtype=self.values.dtype) - ufunc.at(out, self.indices, expanded) - counts = np.bincount(self.indices, minlength=self.n_minor) + out, counts = _ops.minor_extrema( + self.values, self.value_ptr, self.indices, self.n_minor, initial, ufunc is np.maximum + ) + # A minor index missing from some major slice has a structural zero + # competing in the reduction. not_dense = counts < self.n_major out[not_dense] = ufunc(out[not_dense], 0) return out diff --git a/src/vsparse/_ops.py b/src/vsparse/_ops.py index 20f8d26..e81fd70 100644 --- a/src/vsparse/_ops.py +++ b/src/vsparse/_ops.py @@ -6,14 +6,31 @@ import numpy as np __all__ = [ + "accumulator_threads", "major_matmat", "major_matvec", + "minor_counts", + "minor_extrema", "minor_matmat", "minor_matvec", "minor_select_counts", "minor_select_fill", + "minor_sums", ] +# Thread-local accumulators for the minor-axis scatters below cost +# ``nthreads * n_minor * bytes_per_element``. The thread count is capped to +# keep that block under this budget. +_ACCUMULATOR_BUDGET_BYTES = 64 << 20 # 64 MiB + + +def accumulator_threads(n_minor: int, bytes_per_element: int = 8) -> int: + """Threads to run a minor-axis scatter with, capped by accumulator size.""" + if n_minor <= 0: + return 1 + affordable = max(1, _ACCUMULATOR_BUDGET_BYTES // (n_minor * max(1, bytes_per_element))) + return int(min(numba.get_num_threads(), affordable)) + @numba.njit(cache=True) def _major_matvec(major_ptr, values, value_ptr, indices, x, n_major, n_minor): @@ -82,6 +99,76 @@ def _minor_matmat(major_ptr, values, value_ptr, indices, b, n_major): return y +@numba.njit(cache=True, parallel=True) +def minor_sums(values, value_ptr, indices, n_minor, nthreads): + n_groups = values.shape[0] + chunk = (n_groups + nthreads - 1) // nthreads + partial = np.zeros((nthreads, n_minor), dtype=np.float64) + for t in numba.prange(nthreads): # ty: ignore[not-iterable] + start = t * chunk + end = min(n_groups, start + chunk) + local = partial[t] + for u in range(start, end): + v = np.float64(values[u]) + for k in range(value_ptr[u], value_ptr[u + 1]): + local[indices[k]] += v + return partial.sum(axis=0) + + +@numba.njit(cache=True, parallel=True) +def _minor_extrema_kernel(values, value_ptr, indices, n_minor, nthreads, initial, is_max): + # The count comes along for free and tells the caller which minor indices + # have an implicit zero to fold in. + n_groups = values.shape[0] + chunk = (n_groups + nthreads - 1) // nthreads + part_val = np.full((nthreads, n_minor), initial, dtype=values.dtype) + part_cnt = np.zeros((nthreads, n_minor), dtype=np.int64) + for t in numba.prange(nthreads): # ty: ignore[not-iterable] + start = t * chunk + end = min(n_groups, start + chunk) + local_val = part_val[t] + local_cnt = part_cnt[t] + for u in range(start, end): + v = values[u] + for k in range(value_ptr[u], value_ptr[u + 1]): + idx = indices[k] + local_cnt[idx] += 1 + if is_max: + if v > local_val[idx]: + local_val[idx] = v + else: + if v < local_val[idx]: + local_val[idx] = v + return part_val, part_cnt + + +def minor_extrema(values, value_ptr, indices, n_minor, initial, is_max): + """Extremum over the stored values, and the stored count, per minor index.""" + nthreads = accumulator_threads(n_minor, values.dtype.itemsize + 8) + part_val, part_cnt = _minor_extrema_kernel( + values, value_ptr, indices, n_minor, nthreads, initial, is_max + ) + extrema = part_val.max(axis=0) if is_max else part_val.min(axis=0) + return extrema, part_cnt.sum(axis=0) + + +@numba.njit(cache=True, parallel=True) +def minor_counts(value_ptr, indices, n_minor, nthreads): + n_groups = value_ptr.shape[0] - 1 + chunk = (n_groups + nthreads - 1) // nthreads + partial = np.zeros((nthreads, n_minor), dtype=np.int64) + for t in numba.prange(nthreads): # ty: ignore[not-iterable] + start = t * chunk + end = min(n_groups, start + chunk) + local = partial[t] + for u in range(start, end): + for k in range(value_ptr[u], value_ptr[u + 1]): + local[indices[k]] += 1 + return partial.sum(axis=0) + + + + def major_matvec(major_ptr, values, value_ptr, indices, x, n_major, n_minor): return _major_matvec(major_ptr, values, value_ptr, indices, np.asarray(x), n_major, n_minor) diff --git a/src/vsparse/_rapid_load.py b/src/vsparse/_rapid_load.py index e09eaa8..3bc2864 100644 --- a/src/vsparse/_rapid_load.py +++ b/src/vsparse/_rapid_load.py @@ -62,6 +62,7 @@ from vsparse import _ivcsc from vsparse._anndata_class import VCSCAnnData from vsparse._indexutils import smallest_index_dtype +from vsparse._ops import accumulator_threads if TYPE_CHECKING: from os import PathLike @@ -469,11 +470,11 @@ def load_and_normalize( data = _build_data(values, value_ptr, indices.shape[0]) row_indptr = value_ptr[major_ptr] - gene_totals_raw = _weighted_bincount(indices, data, n_genes, numba.get_num_threads()) + gene_totals_raw = _weighted_bincount(indices, data, n_genes, accumulator_threads(n_genes)) gene_mask = gene_totals_raw > (gene_threshold * n_cells) if min_cells is not None: gene_detection_counts = _gene_detection_counts( - indices, data, n_genes, numba.get_num_threads() + indices, data, n_genes, accumulator_threads(n_genes) ) gene_mask &= gene_detection_counts >= min_cells metadata_cell_mask = cell_mask @@ -500,11 +501,11 @@ def load_and_normalize( ) del packed - gene_totals_raw = _weighted_bincount(indices, data, n_genes, numba.get_num_threads()) + gene_totals_raw = _weighted_bincount(indices, data, n_genes, accumulator_threads(n_genes)) gene_mask = gene_totals_raw > (gene_threshold * selected_rows.shape[0]) if min_cells is not None: gene_detection_counts = _gene_detection_counts( - indices, data, n_genes, numba.get_num_threads() + indices, data, n_genes, accumulator_threads(n_genes) ) gene_mask &= gene_detection_counts >= min_cells diff --git a/tests/test_reduction_memory.py b/tests/test_reduction_memory.py new file mode 100644 index 0000000..52a0239 --- /dev/null +++ b/tests/test_reduction_memory.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +import tracemalloc + +import numba +import numpy as np +import pytest +import scipy.sparse as sp + +from vsparse import VCSCArray, VCSRArray +from vsparse._ops import _ACCUMULATOR_BUDGET_BYTES, 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 _minor_axis(vcls): + return 1 if vcls is VCSCArray else 0 + + +@pytest.mark.parametrize("axis", [None, 0, 1]) +def test_sum_matches_dense(dense, vcls, axis): + """Totals along each axis, and overall.""" + v = vcls.from_scipy(_scipy_for(vcls, dense)) + np.testing.assert_allclose(v.sum(axis=axis), dense.sum(axis=axis)) + + +@pytest.mark.parametrize("axis", [0, 1]) +@pytest.mark.parametrize("kind", ["max", "min"]) +def test_extrema_match_dense(dense, vcls, axis, kind): + """Extrema have to fold in the implicit zeros the layout never stores.""" + v = vcls.from_scipy(_scipy_for(vcls, dense)) + np.testing.assert_allclose(getattr(v, kind)(axis=axis), getattr(dense, kind)(axis=axis)) + + +@pytest.mark.parametrize("axis", [0, 1]) +def test_getnnz_matches_dense(dense, vcls, axis): + v = vcls.from_scipy(_scipy_for(vcls, dense)) + np.testing.assert_array_equal(v.getnnz(axis=axis), (dense != 0).sum(axis=axis)) + + +def test_reductions_on_an_all_zero_array(vcls): + """Every entry is an implicit zero, so the stored values are empty.""" + v = vcls.from_scipy(_scipy_for(vcls, np.zeros((4, 3)))) + np.testing.assert_allclose(v.max(axis=0), np.zeros(3)) + np.testing.assert_allclose(v.min(axis=0), np.zeros(3)) + np.testing.assert_allclose(v.sum(axis=0), np.zeros(3)) + np.testing.assert_array_equal(v.getnnz(axis=0), np.zeros(3, dtype=np.int64)) + + +def test_integer_values_use_integer_sentinels(vcls): + """A float sentinel would make an integer max come back wrong or upcast.""" + dense = np.array([[7, 0, 2], [0, 3, 9]], dtype=np.int32) + v = vcls.from_scipy(_scipy_for(vcls, dense.astype(np.float64))) + np.testing.assert_array_equal(v.max(axis=0), dense.max(axis=0)) + np.testing.assert_array_equal(v.min(axis=0), dense.min(axis=0)) + + +def test_sums_accumulate_in_float64(vcls): + """Accumulating in the stored dtype would overflow or lose precision.""" + v = vcls.from_scipy(_scipy_for(vcls, np.array([[1.0, 0.0, 2.0], [3.0, 4.0, 0.0]]))) + assert v._minor_sums().dtype == np.float64 + + +@pytest.mark.parametrize("n_minor", [0, 1_000, 10**8]) +@pytest.mark.parametrize("bytes_per_element", [8, 16]) +def test_accumulator_block_stays_within_budget(n_minor, bytes_per_element): + """Thread-local accumulators must not become the new unbounded allocation.""" + nthreads = accumulator_threads(n_minor, bytes_per_element) + assert 1 <= nthreads <= numba.get_num_threads() + if nthreads > 1: + assert nthreads * n_minor * bytes_per_element <= _ACCUMULATOR_BUDGET_BYTES + + +@pytest.mark.parametrize( + ("label", "call"), + [ + ("sum", lambda v: v.sum(axis=0)), + ("max", lambda v: v.max(axis=0)), + ("getnnz", lambda v: v.getnnz(axis=0)), + ], +) +def test_minor_axis_reductions_allocate_nothing_nnz_sized(label, call): + """An n_minor-sized result must not cost nnz-sized scratch.""" + rng = np.random.default_rng(0) + n_rows, n_cols = 2_000, 500 + dense = rng.integers(1, 5, size=(n_rows, n_cols)).astype(np.float64) + v = VCSRArray.from_scipy(sp.csr_array(dense)) + + call(v) # warm up the JIT before measuring + + tracemalloc.start() + try: + before = tracemalloc.get_traced_memory()[0] + tracemalloc.reset_peak() + out = call(v) + peak = tracemalloc.get_traced_memory()[1] + finally: + tracemalloc.stop() + + assert peak - before < 1 << 20, f"{label} allocated {(peak - before) / 1e6:.1f} MB" + assert out.shape == (n_cols,)