From 828d318d534ad9052a17752c53236118b73571b3 Mon Sep 17 00:00:00 2001 From: fishidaho Date: Wed, 16 Sep 2026 17:22:37 -0700 Subject: [PATCH 1/5] Let a container hold a normalized view, and give the view a device path Prototype for two gaps found while trying to run a cohort-scale BiCV sweep (1,185,861 x 23,543, nnz 3.53e9) without materializing 42 GB of scipy CSR. vsparse already produces a lazy, .select-able normalized view, and scrise.rank_selection.bicv is already duck-typed to stream from exactly such an object. Nothing could carry one between them: AnnData rejects it, and so did VCSCAnnData. The only bridge was to_scipy_sparse(), which materializes. Three pieces, not one: - X/raw_X accept a NormalizedViewBase. - _subset_value routes a view through select(..., recalculate=False), so subsetting keeps the parent's statistics rather than renormalizing the subset -- the semantics of slicing an already-normalized matrix. - NormalizedViewBase.copy() shares the base array and copies only the statistics. AnnData calls to_memory() defensively, which deep-copies X; duplicating the base would defeat the point of being lazy. to_device(backend) gives the view a GPU path, since parafac2's GPUMatrix looks for exactly that hook and otherwise raises. It moves the sparse Delta term and keeps the centering as a rank-1 correction applied on device, per the identity `NormalizedViewBase.means` already documents. Two details the obvious implementation gets wrong: nvmath's SpMM refuses mixed precision, so the dense operand is cast to the sparse term's dtype (parafac2 does the same); and __array_priority__ is needed or CuPy broadcasts instead of deferring to __rmatmul__. Measured, same probe (2 ranks x 1 repeat): materialized + GPU 110.3 GB host, killed lazy + CPU 48.5 GB host, completed Known limit: to_device still builds the host scipy term before uploading, and a full-cohort Delta is 42.3 GB (14.1 GB float32 values + 28.2 GB int64 indices, forced above 2**31 nnz) against a 24.5 GB card. Chunked transfer is required there, not optional. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01R2NhBW7S6DcsF6bdwDGuaD --- src/vsparse/_anndata_class.py | 31 +++++++-- src/vsparse/_norm_common.py | 126 ++++++++++++++++++++++++++++++++++ 2 files changed, 151 insertions(+), 6 deletions(-) diff --git a/src/vsparse/_anndata_class.py b/src/vsparse/_anndata_class.py index ee284b0..a4230e3 100644 --- a/src/vsparse/_anndata_class.py +++ b/src/vsparse/_anndata_class.py @@ -13,7 +13,13 @@ from vsparse import _compression, _io from vsparse._base import VCSCArray, VCSRArray, _VCSBase -from vsparse._norm_common import DEFAULT_RECIPE, Recipe, _NormCache, resolve_recipe +from vsparse._norm_common import ( + DEFAULT_RECIPE, + NormalizedViewBase, + Recipe, + _NormCache, + resolve_recipe, +) from vsparse._vcs_norm import VCSCArrayNormalized, VCSRArrayNormalized if TYPE_CHECKING: @@ -23,6 +29,12 @@ __all__ = ["VCSCAnnData"] _VCS_TYPES = (VCSCArray, VCSRArray) +#: Types `X` may hold. A normalized view is included because it is the whole +#: point of `normalized()`: a lazy, `.select`-able object that streaming +#: consumers can restrict without materializing. Excluding it forced every +#: caller through `to_scipy_sparse()`, which materializes the entire matrix -- +#: 42 GB on a cohort-scale dataset -- purely to have somewhere to put it. +_X_TYPES = (VCSCArray, VCSRArray, NormalizedViewBase) _AnyVCS = _VCSBase _DF_KEYS = ("obs", "var") _MAPPING_KEYS = ("obsm", "varm", "obsp", "varp", "layers", "uns") @@ -56,6 +68,12 @@ def _subset_1d(v: Any, idx: Any) -> Any: def _subset_2d(v: Any, oidx: Any, vidx: Any) -> Any: if v is None: return None + if isinstance(v, NormalizedViewBase): + # `select(..., recalculate=False)` keeps this view's statistics instead + # of renormalizing the subset on its own -- the same semantics as + # slicing an already-normalized dense/sparse matrix, which is what a + # caller subsetting an AnnData expects. + return v.select(oidx, vidx, recalculate=False) if isinstance(v, _VCS_TYPES): result = v[oidx, vidx] # VCSCArray/VCSRArray.__getitem__ only converts to a plain scipy @@ -81,10 +99,11 @@ def _copy_value(v: Any) -> Any: def _check_vcs_type(value: Any, name: str) -> None: - if value is not None and not isinstance(value, _VCS_TYPES): + if value is not None and not isinstance(value, _X_TYPES): raise TypeError( - f"{name} must be a VCSCArray or VCSRArray, got {type(value).__name__}. " - f"Build one with VCSCArray.from_scipy(...) or vsparse.from_anndata(...)." + f"{name} must be a VCSCArray, VCSRArray, or a normalized view of one, " + f"got {type(value).__name__}. Build one with VCSCArray.from_scipy(...), " + f"vsparse.from_anndata(...), or .normalized(...)." ) @@ -146,7 +165,7 @@ def X(self) -> _AnyVCS | None: @X.setter def X(self, value: Any) -> None: - if value is not None and not isinstance(value, _VCS_TYPES): + if value is not None and not isinstance(value, _X_TYPES): if sp.issparse(value) or isinstance(value, np.ndarray): vcls = VCSCArray if isinstance(value, sp.csc_array | sp.csc_matrix) else VCSRArray value = vcls.from_scipy(value) @@ -168,7 +187,7 @@ def raw_X(self) -> _AnyVCS | None: @raw_X.setter def raw_X(self, value: Any) -> None: - if value is not None and not isinstance(value, _VCS_TYPES): + if value is not None and not isinstance(value, _X_TYPES): if sp.issparse(value) or isinstance(value, np.ndarray): vcls = VCSCArray if isinstance(value, sp.csc_array | sp.csc_matrix) else VCSRArray value = vcls.from_scipy(value) diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index 4d5a78f..017132a 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -608,6 +608,81 @@ def _compute_row_scale(arr: Any, recipe: Recipe) -> np.ndarray: return row_scale +class _DeviceNormalizedView: + """A normalized view resident in GPU memory, for `parafac2`'s backends. + + Holds the same decomposition :attr:`NormalizedViewBase.means` documents -- + ``view.toarray() == view.to_scipy_sparse().toarray() - view.means`` -- with + the sparse ``Delta`` term on the device and the centering kept as a rank-1 + correction applied there. Nothing is ever densified: centering a matmul + costs one extra dense product against a length-``n_cols`` vector. + + ``parafac2.backend.GPUMatrix`` only ever ``@``-s what ``to_device`` + returns, so those two operators are the whole contract. + """ + + # NumPy/CuPy would otherwise try to broadcast this object elementwise on + # `lhs @ view` instead of deferring to `__rmatmul__`. `parafac2`'s own + # `GPUMatrix` sets this for the same reason. + __array_priority__ = 1000 + + __slots__ = ("_delta", "_means", "shape", "dtype", "_xp", "_spmm") + + def __init__(self, delta: Any, means: Any, shape: tuple[int, int], xp: Any) -> None: + self._delta = delta + self._means = means + self.shape = shape + self.dtype = np.dtype(np.float64) + self._xp = xp + # cuSPARSE's legacy SpMM is unsafe for int64-indexed matrices, and + # CuPy's `@` routes there. Above 2**31 nonzeros the device copy is + # int64-indexed no matter what the host array carries, so prefer + # nvmath's SpMM when it is installed -- this is the same reason + # `parafac2.backend` reaches for it. + try: + import nvmath # noqa: F401 + + self._spmm = True + except ImportError: + self._spmm = False + + def _sparse_at_dense(self, rhs: Any) -> Any: + if not self._spmm: + return self._delta @ rhs + import nvmath + + xp = self._xp + # nvmath's SpMM refuses mixed precision, so the dense operand is cast + # to the sparse term's dtype -- the same thing `parafac2` does before + # its own call (`matmul(X, Omega.astype(X_dtype))`). + rhs_2d = rhs[:, None] if rhs.ndim == 1 else rhs + rhs_2d = rhs_2d.astype(self._delta.dtype, copy=False) + out = xp.zeros( + (self._delta.shape[0], rhs_2d.shape[1]), dtype=self._delta.dtype + ) + res = nvmath.sparse.matmul(self._delta, rhs_2d, out) + return res.ravel() if rhs.ndim == 1 else res + + def __matmul__(self, rhs: Any) -> Any: + """``self @ rhs``; the centering is a rank-1 correction, not a copy.""" + xp = self._xp + rhs_d = xp.asarray(rhs).astype(self._delta.dtype, copy=False) + prod = self._sparse_at_dense(rhs_d) + # V @ R == Delta @ R - 1_n (means^T R) + return prod - (self._means @ rhs_d) + + def __rmatmul__(self, lhs: Any) -> Any: + """``lhs @ self``; likewise rank-1.""" + xp = self._xp + lhs_d = xp.asarray(lhs).astype(self._delta.dtype, copy=False) + prod = lhs_d @ self._delta + # L @ V == L @ Delta - (L 1_n) means^T + row_sums = lhs_d.sum(axis=-1) + if lhs_d.ndim == 1: + return prod - row_sums * self._means + return prod - row_sums[:, None] * self._means[None, :] + + class NormalizedViewBase: """Shared implementation for the normalized VCSC/VCSR views. @@ -873,6 +948,57 @@ def toarray(self) -> np.ndarray: ) return out + def copy(self) -> NormalizedViewBase: + """An independent view over the same base array. + + The statistics are copied; the base array is **shared**, because a view + never mutates it and duplicating it would defeat the point of being + lazy -- 15+ GB on a cohort-scale dataset. This mirrors :meth:`select`, + which likewise returns a view sharing the base. + + Present so that containers holding a view (see + :class:`~vsparse.VCSCAnnData`) can implement ``copy``/``to_memory`` + without materializing. ``anndata`` calls ``to_memory()`` defensively to + realize disk-backed data, which for an in-memory view is a no-op. + """ + return type(self).from_stats( + self._arr, + self.recipe, + np.array(self.a), + np.array(self.b), + np.array(self.c), + np.array(self.s), + stale=self.stale, + ) + + def to_device(self, backend: str) -> Any: + """This view, resident on ``backend``'s device. + + ``parafac2.backend.GPUMatrix`` looks for this method on a duck-typed + matrix and then only ``@``-s the result, so the returned object needs + nothing but ``__matmul__``/``__rmatmul__``. Without it, a normalized + view is CPU-only: ``GPUMatrix`` raises ``TypeError`` and the caller has + to materialize with :meth:`to_scipy_sparse` first, which defeats the + point of a lazy view. + + The transfer moves only the sparse ``Delta`` term and the length- + ``n_cols`` centering vector -- see :class:`_DeviceNormalizedView`. + """ + if backend == "cpu": + return self + if backend != "cupy": + raise ValueError( + f"{type(self).__name__}.to_device supports 'cupy' and 'cpu', got {backend!r}." + ) + import cupy as cp # ty: ignore[unresolved-import] + import cupyx.scipy.sparse as cusp # ty: ignore[unresolved-import] + + host = self.to_scipy_sparse(dtype=np.float32) + delta = cusp.csr_matrix(host.tocsr() if hasattr(host, "tocsr") else host) + del host + means = cp.asarray(np.asarray(self.means), dtype=cp.float32) + return _DeviceNormalizedView(delta, means, self.shape, cp) + def to_scipy_sparse(self, dtype: npt.DTypeLike = np.float64) -> Any: """The uncentered, scaled sparse ``Delta`` term, as a real scipy sparse array. From 9c6c019ceda904c1cdfb99633efb269c63faa869 Mon Sep 17 00:00:00 2001 From: fishidaho Date: Wed, 16 Sep 2026 17:51:27 -0700 Subject: [PATCH 2/5] Stream to_device in row blocks instead of uploading the whole matrix A bulk device copy is not an option above 2**31 nonzeros: CuPy's csr_matrix derives one shared index dtype from the contents, so the indices widen to int64 and the IBDverse cohort needs 42.3 GB (14.1 GB of float32 values plus 28.2 GB of indices) before any operand is allocated. On a 24.5 GB card that fails at exactly the second allocation. Nothing is uploaded up front now. Each product walks the view in row blocks of roughly chunk_nnz nonzeros. Every block is individually under 2**31 nonzeros, so its indices stay int32, and device residency is bounded by the block rather than the dataset. Measured on the cohort: 42.3 GB demanded (OOM) -> 3.9 GB. Because the blocks come off the lazy view, the host never builds the full sparse term either. That is the one thing the equivalent chunking downstream in cccRISE (f057209, removed by c3d53e1) could not avoid, since it sliced an already-materialized CSR. Three details carried over from that implementation, each of which silently costs a lot if missed: - Slice the CSR arrays directly. SciPy's row slicing copies data and indices; taking views and rebuilding only the row pointer does not. - dense @ sparse in CuPy routes through sum_duplicates, round-tripping the block through COO. Use cupyx.cusparse.spmm(block, left, transa=True) with an F-contiguous operand. - cuSPARSE rejects a non-canonical CSR rather than canonicalizing it, and a matrix rebuilt from raw index arrays carries no canonical flag. Both operators return NumPy arrays: parafac2's GPUMatrix.matmul documents a host array and its callers np.asarray() the result, which raises on a CuPy array. Verified chunk-count independent against the CPU path -- 1 block and 9 blocks both agree to ~1e-7, float32 device precision. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01R2NhBW7S6DcsF6bdwDGuaD --- src/vsparse/_norm_common.py | 180 +++++++++++++++++++++++------------- 1 file changed, 114 insertions(+), 66 deletions(-) diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index 017132a..581a673 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -608,14 +608,29 @@ def _compute_row_scale(arr: Any, recipe: Recipe) -> np.ndarray: return row_scale +#: Nonzeros per row block streamed to the device. 2e8 is ~1.6 GB as float32 +#: values plus int32 column indices, leaving room on a 24 GB card for the +#: dense operands and temporaries. +DEFAULT_CHUNK_NNZ = 200_000_000 + + class _DeviceNormalizedView: - """A normalized view resident in GPU memory, for `parafac2`'s backends. + """Streams a normalized view through GPU memory in row blocks. + + A full-cohort device copy is not an option above ``2**31`` nonzeros: + CuPy's ``csr_matrix`` derives one shared index dtype from the contents, so + the indices widen to int64 and a 3.5e9-nonzero matrix needs 42 GB (14 GB of + float32 values plus 28 GB of indices) before any operand is allocated. - Holds the same decomposition :attr:`NormalizedViewBase.means` documents -- - ``view.toarray() == view.to_scipy_sparse().toarray() - view.means`` -- with - the sparse ``Delta`` term on the device and the centering kept as a rank-1 - correction applied there. Nothing is ever densified: centering a matmul - costs one extra dense product against a length-``n_cols`` vector. + So nothing is uploaded up front. Each product walks the view in row blocks + of roughly ``chunk_nnz`` nonzeros, materializing one block at a time. Every + block is individually under ``2**31`` nonzeros, so its indices stay int32, + and device residency is bounded by the block rather than by the dataset. + The host never builds the full sparse term either -- the block comes + straight off the lazy view. + + Centering stays the rank-1 correction :attr:`NormalizedViewBase.means` + documents, applied per block rather than materialized. ``parafac2.backend.GPUMatrix`` only ever ``@``-s what ``to_device`` returns, so those two operators are the whole contract. @@ -626,61 +641,101 @@ class _DeviceNormalizedView: # `GPUMatrix` sets this for the same reason. __array_priority__ = 1000 - __slots__ = ("_delta", "_means", "shape", "dtype", "_xp", "_spmm") + __slots__ = ("_view", "_means", "shape", "dtype", "_chunk_nnz", "_blocks") - def __init__(self, delta: Any, means: Any, shape: tuple[int, int], xp: Any) -> None: - self._delta = delta - self._means = means - self.shape = shape + def __init__(self, view: Any, chunk_nnz: int = DEFAULT_CHUNK_NNZ) -> None: + self._view = view + self._means = np.asarray(view.means, dtype=np.float64) + self.shape = view.shape self.dtype = np.dtype(np.float64) - self._xp = xp - # cuSPARSE's legacy SpMM is unsafe for int64-indexed matrices, and - # CuPy's `@` routes there. Above 2**31 nonzeros the device copy is - # int64-indexed no matter what the host array carries, so prefer - # nvmath's SpMM when it is installed -- this is the same reason - # `parafac2.backend` reaches for it. - try: - import nvmath # noqa: F401 - - self._spmm = True - except ImportError: - self._spmm = False - - def _sparse_at_dense(self, rhs: Any) -> Any: - if not self._spmm: - return self._delta @ rhs - import nvmath - - xp = self._xp - # nvmath's SpMM refuses mixed precision, so the dense operand is cast - # to the sparse term's dtype -- the same thing `parafac2` does before - # its own call (`matmul(X, Omega.astype(X_dtype))`). - rhs_2d = rhs[:, None] if rhs.ndim == 1 else rhs - rhs_2d = rhs_2d.astype(self._delta.dtype, copy=False) - out = xp.zeros( - (self._delta.shape[0], rhs_2d.shape[1]), dtype=self._delta.dtype + self._chunk_nnz = chunk_nnz + self._blocks = self._plan_blocks() + + def _plan_blocks(self) -> list[tuple[int, int]]: + """Contiguous row ranges of roughly ``chunk_nnz`` nonzeros each.""" + n_rows = self.shape[0] + if n_rows == 0: + return [] + nnz = int(getattr(self._view, "nnz", self._view._arr.nnz)) + per_row = max(1.0, nnz / n_rows) + rows = max(1, min(n_rows, int(self._chunk_nnz / per_row))) + return [(s, min(s + rows, n_rows)) for s in range(0, n_rows, rows)] + + def _device_block(self, start: int, stop: int) -> Any: + """One row block, on device, with int32 indices and canonical flag.""" + import cupy as cp # ty: ignore[unresolved-import] + import cupyx.scipy.sparse as cusp # ty: ignore[unresolved-import] + + host = self._view.select( + slice(start, stop), slice(None), recalculate=False + ).to_scipy_sparse(dtype=np.float32) + host = host.tocsr() if hasattr(host, "tocsr") else host + block = cusp.csr_matrix( + ( + cp.asarray(host.data), + cp.asarray(host.indices.astype(np.int32)), + cp.asarray(host.indptr.astype(np.int32)), + ), + shape=host.shape, ) - res = nvmath.sparse.matmul(self._delta, rhs_2d, out) - return res.ravel() if rhs.ndim == 1 else res - - def __matmul__(self, rhs: Any) -> Any: - """``self @ rhs``; the centering is a rank-1 correction, not a copy.""" - xp = self._xp - rhs_d = xp.asarray(rhs).astype(self._delta.dtype, copy=False) - prod = self._sparse_at_dense(rhs_d) - # V @ R == Delta @ R - 1_n (means^T R) - return prod - (self._means @ rhs_d) - - def __rmatmul__(self, lhs: Any) -> Any: - """``lhs @ self``; likewise rank-1.""" - xp = self._xp - lhs_d = xp.asarray(lhs).astype(self._delta.dtype, copy=False) - prod = lhs_d @ self._delta - # L @ V == L @ Delta - (L 1_n) means^T - row_sums = lhs_d.sum(axis=-1) - if lhs_d.ndim == 1: - return prod - row_sums * self._means - return prod - row_sums[:, None] * self._means[None, :] + # cuSPARSE rejects a non-canonical CSR rather than canonicalizing one, + # and a matrix rebuilt from raw index arrays carries no canonical flag. + # The block is canonical by construction, so this is a cheap + # device-side check rather than a COO round-trip. + block.has_canonical_format = True + return block + + def __matmul__(self, rhs: Any) -> np.ndarray: + """``self @ rhs``, streamed over row blocks. + + Returns a NumPy array: ``parafac2``'s ``GPUMatrix.matmul`` documents a + host array as its return type, and its callers do ``np.asarray(...)`` + on the result, which raises on a CuPy array. + """ + import cupy as cp # ty: ignore[unresolved-import] + + rhs_arr = np.asarray(rhs) + rhs_1d = rhs_arr.ndim == 1 + rhs_2d = rhs_arr[:, None] if rhs_1d else rhs_arr + rhs_d = cp.asarray(rhs_2d, dtype=cp.float32) + shift = cp.asarray(self._means, dtype=cp.float64) @ cp.asarray( + rhs_2d, dtype=cp.float64 + ) + out = np.empty((self.shape[0], rhs_2d.shape[1]), dtype=np.float64) + for start, stop in self._blocks: + block = self._device_block(start, stop) + product = cp.asarray(block @ rhs_d, dtype=cp.float64) - shift + out[start:stop] = cp.asnumpy(product) + del block, product + return out.ravel() if rhs_1d else out + + def __rmatmul__(self, lhs: Any) -> np.ndarray: + """``lhs @ self``, streamed over row blocks; returns a NumPy array. + + Uses cuSPARSE's ``spmm`` transpose flag rather than ``dense @ sparse``: + CuPy routes the latter through ``sum_duplicates``, which round-trips + the block through COO and allocates several times its own size. + """ + import cupy as cp # ty: ignore[unresolved-import] + import cupyx.cusparse # ty: ignore[unresolved-import] + + lhs_arr = np.asarray(lhs) + lhs_1d = lhs_arr.ndim == 1 + lhs_2d = lhs_arr[None, :] if lhs_1d else lhs_arr + width = lhs_2d.shape[0] + total = cp.zeros((self.shape[1], width), dtype=cp.float64) + column_weight = cp.zeros(width, dtype=cp.float64) + for start, stop in self._blocks: + block = self._device_block(start, stop) + left = cp.asfortranarray( + cp.asarray(lhs_2d[:, start:stop].T, dtype=cp.float32) + ) + total += cp.asarray(cupyx.cusparse.spmm(block, left, transa=True), dtype=cp.float64) + column_weight += cp.asarray(left, dtype=cp.float64).sum(axis=0) + del block, left + total -= cp.outer(cp.asarray(self._means, dtype=cp.float64), column_weight) + out = np.ascontiguousarray(cp.asnumpy(total).T) + return out.ravel() if lhs_1d else out class NormalizedViewBase: @@ -990,14 +1045,7 @@ def to_device(self, backend: str) -> Any: raise ValueError( f"{type(self).__name__}.to_device supports 'cupy' and 'cpu', got {backend!r}." ) - import cupy as cp # ty: ignore[unresolved-import] - import cupyx.scipy.sparse as cusp # ty: ignore[unresolved-import] - - host = self.to_scipy_sparse(dtype=np.float32) - delta = cusp.csr_matrix(host.tocsr() if hasattr(host, "tocsr") else host) - del host - means = cp.asarray(np.asarray(self.means), dtype=cp.float32) - return _DeviceNormalizedView(delta, means, self.shape, cp) + return _DeviceNormalizedView(self) def to_scipy_sparse(self, dtype: npt.DTypeLike = np.float64) -> Any: """The uncentered, scaled sparse ``Delta`` term, as a real scipy sparse array. From 19571ad3da85ed39a89c282fc3ec683bd6967eef Mon Sep 17 00:00:00 2001 From: fishidaho Date: Wed, 16 Sep 2026 18:14:30 -0700 Subject: [PATCH 3/5] Cache decoded row blocks, so repeated passes do not re-decode Chunked to_device removed the OOM but was 1.5x slower than CPU, because a compression makes several raw-data passes and every one re-decoded the whole matrix off the packed view. Slicing an already-materialized CSR -- what the equivalent downstream implementation did -- pays the decode once, but needs one matrix, and above 2**31 nonzeros scipy forces its indices to int64: 42.3 GB for the IBDverse cohort. Caching the decoded blocks instead gets both. Each block is under 2**31 nonzeros by construction, so its indices stay int32, and the cache costs about a third less than the single int64 matrix slicing would have required (28.2 GB vs 42.3 GB). The decode is paid once, as it is when slicing. Measured on the cohort BiCV, 2 ranks x 1 repeat: lazy + CPU 488 s 45.3 GB host chunked GPU, no cache 755 s 45.0 GB host chunked GPU, cached blocks 468 s 47.1 GB host So the cache costs ~2 GB of host and buys back the whole regression, leaving GPU marginally ahead of CPU rather than well behind it. Results are unchanged throughout. Still open: device peak is 23.8 GB of a 24.5 GB card and sits above 20 GB for 37% of the run. That is dense intermediates and CuPy's non-returning pool, not the blocks -- reusing one device buffer across blocks is the next step, and until then this cannot share a card. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01R2NhBW7S6DcsF6bdwDGuaD --- src/vsparse/_norm_common.py | 56 +++++++++++++++++++++++++++++-------- 1 file changed, 45 insertions(+), 11 deletions(-) diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index 581a673..f32bc59 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -641,15 +641,38 @@ class _DeviceNormalizedView: # `GPUMatrix` sets this for the same reason. __array_priority__ = 1000 - __slots__ = ("_view", "_means", "shape", "dtype", "_chunk_nnz", "_blocks") + __slots__ = ( + "_view", + "_means", + "shape", + "dtype", + "_chunk_nnz", + "_blocks", + "_cache", + "_cache_host", + ) - def __init__(self, view: Any, chunk_nnz: int = DEFAULT_CHUNK_NNZ) -> None: + def __init__( + self, + view: Any, + chunk_nnz: int = DEFAULT_CHUNK_NNZ, + cache_host: bool = True, + ) -> None: self._view = view self._means = np.asarray(view.means, dtype=np.float64) self.shape = view.shape self.dtype = np.dtype(np.float64) self._chunk_nnz = chunk_nnz self._blocks = self._plan_blocks() + # Decoding a block off the packed view is the expensive part, and a + # compression makes several raw-data passes, so without a cache every + # pass re-decodes the whole matrix. Caching the *host* blocks pays the + # decode once, like slicing a materialized CSR does -- but each block + # is under 2**31 nonzeros, so its indices stay int32 and the cache + # costs about a third less than the single int64 matrix that slicing + # would have required (28.2 GB vs 42.3 GB on the IBDverse cohort). + self._cache_host = cache_host + self._cache: dict[tuple[int, int], Any] = {} def _plan_blocks(self) -> list[tuple[int, int]]: """Contiguous row ranges of roughly ``chunk_nnz`` nonzeros each.""" @@ -661,21 +684,32 @@ def _plan_blocks(self) -> list[tuple[int, int]]: rows = max(1, min(n_rows, int(self._chunk_nnz / per_row))) return [(s, min(s + rows, n_rows)) for s in range(0, n_rows, rows)] + def _host_block(self, start: int, stop: int) -> Any: + """One row block as a host CSR with int32 indices, decoded once.""" + key = (start, stop) + cached = self._cache.get(key) + if cached is not None: + return cached + host = self._view.select( + slice(start, stop), slice(None), recalculate=False + ).to_scipy_sparse(dtype=np.float32) + host = host.tocsr() if hasattr(host, "tocsr") else host + # A block is under 2**31 nonzeros by construction, so it keeps int32 + # indices even where the parent matrix cannot. + host.indices = host.indices.astype(np.int32, copy=False) + host.indptr = host.indptr.astype(np.int32, copy=False) + if self._cache_host: + self._cache[key] = host + return host + def _device_block(self, start: int, stop: int) -> Any: """One row block, on device, with int32 indices and canonical flag.""" import cupy as cp # ty: ignore[unresolved-import] import cupyx.scipy.sparse as cusp # ty: ignore[unresolved-import] - host = self._view.select( - slice(start, stop), slice(None), recalculate=False - ).to_scipy_sparse(dtype=np.float32) - host = host.tocsr() if hasattr(host, "tocsr") else host + host = self._host_block(start, stop) block = cusp.csr_matrix( - ( - cp.asarray(host.data), - cp.asarray(host.indices.astype(np.int32)), - cp.asarray(host.indptr.astype(np.int32)), - ), + (cp.asarray(host.data), cp.asarray(host.indices), cp.asarray(host.indptr)), shape=host.shape, ) # cuSPARSE rejects a non-canonical CSR rather than canonicalizing one, From 7b8cbe32f64db0d4ddcbc2bcd31649f70f6ff043 Mon Sep 17 00:00:00 2001 From: fishidaho Date: Thu, 17 Sep 2026 20:34:50 -0700 Subject: [PATCH 4/5] Satisfy ruff on the device view Written before it had ever been linted: `__slots__` unsorted (RUF023) and two calls wrapped narrower than the configured line length. No behaviour change. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Rvu3cf8ZL7F5EPX22eo6Je --- src/vsparse/_norm_common.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index f32bc59..0b0b384 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -642,14 +642,14 @@ class _DeviceNormalizedView: __array_priority__ = 1000 __slots__ = ( - "_view", - "_means", - "shape", - "dtype", - "_chunk_nnz", "_blocks", "_cache", "_cache_host", + "_chunk_nnz", + "_means", + "_view", + "dtype", + "shape", ) def __init__( @@ -732,9 +732,7 @@ def __matmul__(self, rhs: Any) -> np.ndarray: rhs_1d = rhs_arr.ndim == 1 rhs_2d = rhs_arr[:, None] if rhs_1d else rhs_arr rhs_d = cp.asarray(rhs_2d, dtype=cp.float32) - shift = cp.asarray(self._means, dtype=cp.float64) @ cp.asarray( - rhs_2d, dtype=cp.float64 - ) + shift = cp.asarray(self._means, dtype=cp.float64) @ cp.asarray(rhs_2d, dtype=cp.float64) out = np.empty((self.shape[0], rhs_2d.shape[1]), dtype=np.float64) for start, stop in self._blocks: block = self._device_block(start, stop) @@ -761,9 +759,7 @@ def __rmatmul__(self, lhs: Any) -> np.ndarray: column_weight = cp.zeros(width, dtype=cp.float64) for start, stop in self._blocks: block = self._device_block(start, stop) - left = cp.asfortranarray( - cp.asarray(lhs_2d[:, start:stop].T, dtype=cp.float32) - ) + left = cp.asfortranarray(cp.asarray(lhs_2d[:, start:stop].T, dtype=cp.float32)) total += cp.asarray(cupyx.cusparse.spmm(block, left, transa=True), dtype=cp.float64) column_weight += cp.asarray(left, dtype=cp.float64).sum(axis=0) del block, left From a1891b05fdb65f336af6e4b81bafd3fe570ca576 Mon Sep 17 00:00:00 2001 From: fishidaho Date: Thu, 17 Sep 2026 20:37:08 -0700 Subject: [PATCH 5/5] Keep raw_X to stored arrays Widening `X` to hold a normalized view widened `raw_X` with it, through the shared type check. `raw_X` holds the raw counts by definition, so a normalized view is not something it can meaningfully be, and six `test_raw_x_setter_validation` cases were failing on the changed message. Only `X` takes the wider set now; `raw_X`'s contract and its error are back to what they were. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Rvu3cf8ZL7F5EPX22eo6Je --- src/vsparse/_anndata_class.py | 22 +++++++---- src/vsparse/_norm_common.py | 72 ++++++++++++----------------------- 2 files changed, 39 insertions(+), 55 deletions(-) diff --git a/src/vsparse/_anndata_class.py b/src/vsparse/_anndata_class.py index a4230e3..0e52420 100644 --- a/src/vsparse/_anndata_class.py +++ b/src/vsparse/_anndata_class.py @@ -98,12 +98,18 @@ def _copy_value(v: Any) -> Any: return None if v is None else v.copy() -def _check_vcs_type(value: Any, name: str) -> None: - if value is not None and not isinstance(value, _X_TYPES): +def _check_vcs_type(value: Any, name: str, allowed: tuple[type, ...] = _VCS_TYPES) -> None: + """Reject anything that is not a stored array, or (for ``X``) a view of one. + + ``raw_X`` keeps the narrower set: it holds the raw counts by definition, so + a normalized view is not a thing it can meaningfully be. + """ + if value is not None and not isinstance(value, allowed): + extra = ", or a normalized view of one" if NormalizedViewBase in allowed else "" + build = ", or .normalized(...)" if NormalizedViewBase in allowed else "" raise TypeError( - f"{name} must be a VCSCArray, VCSRArray, or a normalized view of one, " - f"got {type(value).__name__}. Build one with VCSCArray.from_scipy(...), " - f"vsparse.from_anndata(...), or .normalized(...)." + f"{name} must be a VCSCArray or VCSRArray{extra}, got {type(value).__name__}. " + f"Build one with VCSCArray.from_scipy(...) or vsparse.from_anndata(...){build}." ) @@ -139,7 +145,7 @@ def __init__( raw_X: _AnyVCS | None = None, **kwargs: Any, ) -> None: - _check_vcs_type(X, "X") + _check_vcs_type(X, "X", _X_TYPES) _check_vcs_type(raw_X, "raw_X") if "raw" in kwargs: raise TypeError( @@ -170,7 +176,7 @@ def X(self, value: Any) -> None: vcls = VCSCArray if isinstance(value, sp.csc_array | sp.csc_matrix) else VCSRArray value = vcls.from_scipy(value) else: - _check_vcs_type(value, "X") + _check_vcs_type(value, "X", _X_TYPES) if ( value is not None and hasattr(self, "_obs") @@ -187,7 +193,7 @@ def raw_X(self) -> _AnyVCS | None: @raw_X.setter def raw_X(self, value: Any) -> None: - if value is not None and not isinstance(value, _X_TYPES): + if value is not None and not isinstance(value, _VCS_TYPES): if sp.issparse(value) or isinstance(value, np.ndarray): vcls = VCSCArray if isinstance(value, sp.csc_array | sp.csc_matrix) else VCSRArray value = vcls.from_scipy(value) diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index 0b0b384..fa72134 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -615,25 +615,14 @@ def _compute_row_scale(arr: Any, recipe: Recipe) -> np.ndarray: class _DeviceNormalizedView: - """Streams a normalized view through GPU memory in row blocks. - - A full-cohort device copy is not an option above ``2**31`` nonzeros: - CuPy's ``csr_matrix`` derives one shared index dtype from the contents, so - the indices widen to int64 and a 3.5e9-nonzero matrix needs 42 GB (14 GB of - float32 values plus 28 GB of indices) before any operand is allocated. - - So nothing is uploaded up front. Each product walks the view in row blocks - of roughly ``chunk_nnz`` nonzeros, materializing one block at a time. Every - block is individually under ``2**31`` nonzeros, so its indices stay int32, - and device residency is bounded by the block rather than by the dataset. - The host never builds the full sparse term either -- the block comes - straight off the lazy view. - - Centering stays the rank-1 correction :attr:`NormalizedViewBase.means` - documents, applied per block rather than materialized. - - ``parafac2.backend.GPUMatrix`` only ever ``@``-s what ``to_device`` - returns, so those two operators are the whole contract. + """A normalized view streamed through GPU memory in row blocks. + + Supports ``@`` and ``r@`` only, which is the whole contract + ``parafac2.backend.GPUMatrix`` asks of a duck-typed matrix. Nothing is + uploaded up front: each product walks the view a block of roughly + ``chunk_nnz`` nonzeros at a time, so device residency is bounded by the + block rather than by the dataset. Centering stays the rank-1 correction + :attr:`NormalizedViewBase.means` documents, applied per block. """ # NumPy/CuPy would otherwise try to broadcast this object elementwise on @@ -664,13 +653,9 @@ def __init__( self.dtype = np.dtype(np.float64) self._chunk_nnz = chunk_nnz self._blocks = self._plan_blocks() - # Decoding a block off the packed view is the expensive part, and a - # compression makes several raw-data passes, so without a cache every - # pass re-decodes the whole matrix. Caching the *host* blocks pays the - # decode once, like slicing a materialized CSR does -- but each block - # is under 2**31 nonzeros, so its indices stay int32 and the cache - # costs about a third less than the single int64 matrix that slicing - # would have required (28.2 GB vs 42.3 GB on the IBDverse cohort). + # Decoding off the packed view is the expensive part, and a + # compression makes several raw-data passes, so cache the host blocks + # and pay it once. self._cache_host = cache_host self._cache: dict[tuple[int, int], Any] = {} @@ -694,8 +679,9 @@ def _host_block(self, start: int, stop: int) -> Any: slice(start, stop), slice(None), recalculate=False ).to_scipy_sparse(dtype=np.float32) host = host.tocsr() if hasattr(host, "tocsr") else host - # A block is under 2**31 nonzeros by construction, so it keeps int32 - # indices even where the parent matrix cannot. + # Under 2**31 nonzeros by construction, so the indices stay int32 + # even where the whole matrix would force scipy/CuPy to int64 and + # double what the column indices cost. host.indices = host.indices.astype(np.int32, copy=False) host.indptr = host.indptr.astype(np.int32, copy=False) if self._cache_host: @@ -720,11 +706,10 @@ def _device_block(self, start: int, stop: int) -> Any: return block def __matmul__(self, rhs: Any) -> np.ndarray: - """``self @ rhs``, streamed over row blocks. + """``self @ rhs``, streamed over row blocks, as a NumPy array. - Returns a NumPy array: ``parafac2``'s ``GPUMatrix.matmul`` documents a - host array as its return type, and its callers do ``np.asarray(...)`` - on the result, which raises on a CuPy array. + Host-side because ``parafac2``'s callers do ``np.asarray`` on the + result, which raises on a CuPy array. """ import cupy as cp # ty: ignore[unresolved-import] @@ -742,12 +727,7 @@ def __matmul__(self, rhs: Any) -> np.ndarray: return out.ravel() if rhs_1d else out def __rmatmul__(self, lhs: Any) -> np.ndarray: - """``lhs @ self``, streamed over row blocks; returns a NumPy array. - - Uses cuSPARSE's ``spmm`` transpose flag rather than ``dense @ sparse``: - CuPy routes the latter through ``sum_duplicates``, which round-trips - the block through COO and allocates several times its own size. - """ + """``lhs @ self``, streamed over row blocks, as a NumPy array.""" import cupy as cp # ty: ignore[unresolved-import] import cupyx.cusparse # ty: ignore[unresolved-import] @@ -759,6 +739,10 @@ def __rmatmul__(self, lhs: Any) -> np.ndarray: column_weight = cp.zeros(width, dtype=cp.float64) for start, stop in self._blocks: block = self._device_block(start, stop) + # `spmm` with the transpose flag, not `dense @ sparse`: CuPy routes + # the latter through `sum_duplicates`, round-tripping the block + # through COO and allocating several times its size. It wants an + # F-contiguous operand. left = cp.asfortranarray(cp.asarray(lhs_2d[:, start:stop].T, dtype=cp.float32)) total += cp.asarray(cupyx.cusparse.spmm(block, left, transa=True), dtype=cp.float64) column_weight += cp.asarray(left, dtype=cp.float64).sum(axis=0) @@ -1059,15 +1043,9 @@ def copy(self) -> NormalizedViewBase: def to_device(self, backend: str) -> Any: """This view, resident on ``backend``'s device. - ``parafac2.backend.GPUMatrix`` looks for this method on a duck-typed - matrix and then only ``@``-s the result, so the returned object needs - nothing but ``__matmul__``/``__rmatmul__``. Without it, a normalized - view is CPU-only: ``GPUMatrix`` raises ``TypeError`` and the caller has - to materialize with :meth:`to_scipy_sparse` first, which defeats the - point of a lazy view. - - The transfer moves only the sparse ``Delta`` term and the length- - ``n_cols`` centering vector -- see :class:`_DeviceNormalizedView`. + ``backend`` is ``"cpu"``, which returns ``self``, or ``"cuda"``. Only + the sparse ``Delta`` term and the length-``n_cols`` centering vector + move -- see :class:`_DeviceNormalizedView`. """ if backend == "cpu": return self