Multi-GPU support with dask

docs/conf.py

@@ -62,7 +62,7 @@
 
 autosummary_generate = True
 autodoc_member_order = "bysource"
-autodoc_mock_imports = ["cudf", "cuml", "cugraph", "cupy", "cupyx", "pylibraft"]
+autodoc_mock_imports = ["cudf", "cuml", "cugraph", "cupy", "cupyx", "pylibraft", "dask"]
 default_role = "literal"
 napoleon_google_docstring = False
 napoleon_numpy_docstring = True
@@ -108,6 +108,7 @@
     "rmm": ("https://docs.rapids.ai/api/rmm/stable/", None),
     "statsmodels": ("https://www.statsmodels.org/stable/", None),
     "omnipath": ("https://omnipath.readthedocs.io/en/latest/", None),
+    "dask": ("https://docs.dask.org/en/stable/", None),
 }
 
 # List of patterns, relative to source directory, that match files and

pyproject.toml

@@ -37,6 +37,7 @@ doc = [
     "scanpydoc[typehints,theme]>=0.9.4",
     "readthedocs-sphinx-ext",
     "sphinx_copybutton",
+    "dask",
 ]
 test = [
     "pytest",

src/rapids_singlecell/_compat.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import cupy as cp
+from cupyx.scipy.sparse import csr_matrix
+
+try:
+    from dask.array import Array as DaskArray
+except ImportError:
+
+    class DaskArray:
+        pass
+
+
+try:
+    from dask.distributed import Client as DaskClient
+except ImportError:
+
+    class DaskClient:
+        pass
+
+
+def _get_dask_client(client=None):
+    from dask.distributed import default_client
+
+    if client is None or not isinstance(client, DaskClient):
+        return default_client()
+    return client
+
+
+def _meta_dense(dtype):
+    return cp.zeros([0], dtype=dtype)
+
+
+def _meta_sparse(dtype):
+    return csr_matrix(cp.array((1.0,), dtype=dtype))

src/rapids_singlecell/preprocessing/__init__.py

@@ -5,11 +5,11 @@
 from ._neighbors import neighbors
 from ._normalize import log1p, normalize_pearson_residuals, normalize_total
 from ._pca import pca
+from ._qc import calculate_qc_metrics
 from ._regress_out import regress_out
 from ._scale import scale
 from ._scrublet import scrublet, scrublet_simulate_doublets
 from ._simple import (
-    calculate_qc_metrics,
     filter_cells,
     filter_genes,
     filter_highly_variable,

src/rapids_singlecell/preprocessing/_hvg.py

@@ -9,10 +9,12 @@
 import cupy as cp
 import numpy as np
 import pandas as pd
-from cupyx.scipy.sparse import issparse, isspmatrix_csc
+from cupyx.scipy.sparse import csr_matrix, issparse, isspmatrix_csc
 from scanpy.get import _get_obs_rep
 
-from ._simple import calculate_qc_metrics
+from rapids_singlecell._compat import DaskArray, DaskClient, _meta_dense, _meta_sparse
+
+from ._qc import calculate_qc_metrics
 from ._utils import _check_gpu_X, _check_nonnegative_integers, _get_mean_var
 
 if TYPE_CHECKING:
@@ -47,6 +49,7 @@ def highly_variable_genes(
     chunksize: int = 1000,
     n_samples: int = 10000,
     batch_key: str = None,
+    client: DaskClient | None = None,
 ) -> None:
     """\
     Annotate highly variable genes.
@@ -116,6 +119,8 @@ def highly_variable_genes(
             of enrichment of zeros for each gene (only for `flavor='poisson_gene_selection'`).
         batch_key
             If specified, highly-variable genes are selected within each batch separately and merged.
+        client
+            Dask client to use for computation. If `None`, the default client is used. Only used if `X` is a Dask array.
 
     Returns
     -------
@@ -188,7 +193,12 @@ def highly_variable_genes(
 
         if batch_key is None:
             df = _highly_variable_genes_single_batch(
-                adata, layer=layer, cutoff=cutoff, n_bins=n_bins, flavor=flavor
+                adata,
+                layer=layer,
+                cutoff=cutoff,
+                n_bins=n_bins,
+                flavor=flavor,
+                client=client,
             )
         else:
             df = _highly_variable_genes_batched(
@@ -198,6 +208,7 @@ def highly_variable_genes(
                 cutoff=cutoff,
                 n_bins=n_bins,
                 flavor=flavor,
+                client=client,
             )
 
         adata.uns["hvg"] = {"flavor": flavor}
@@ -267,6 +278,7 @@ def _highly_variable_genes_single_batch(
     cutoff: _Cutoffs | int,
     n_bins: int = 20,
     flavor: Literal["seurat", "cell_ranger"] = "seurat",
+    client: DaskClient | None = None,
 ) -> pd.DataFrame:
     """\
     See `highly_variable_genes`.
@@ -277,18 +289,24 @@ def _highly_variable_genes_single_batch(
     `highly_variable`, `means`, `dispersions`, and `dispersions_norm`.
     """
     X = _get_obs_rep(adata, layer=layer)
-
+    _check_gpu_X(X, allow_dask=True)
     if hasattr(X, "_view_args"):  # AnnData array view
-        # For compatibility with anndata<0.9
-        X = X.copy()  # Doesn't actually copy memory, just removes View class wrapper
+        X = X.copy()
 
     if flavor == "seurat":
-        X = X.copy()
-        if issparse(X):
-            X = X.expm1()
+        if isinstance(X, DaskArray):
+            if isinstance(X._meta, cp.ndarray):
+                X = X.map_blocks(lambda X: cp.expm1(X), meta=_meta_dense(X.dtype))
+            elif isinstance(X._meta, csr_matrix):
+                X = X.map_blocks(lambda X: X.expm1(), meta=_meta_sparse(X.dtype))
         else:
-            X = cp.expm1(X)
-    mean, var = _get_mean_var(X, axis=0)
+            X = X.copy()
+            if issparse(X):
+                X = X.expm1()
+            else:
+                X = cp.expm1(X)
+
+    mean, var = _get_mean_var(X, axis=0, client=client)
     mean[mean == 0] = 1e-12
     disp = var / mean
     if flavor == "seurat":  # logarithmized mean as in Seurat
@@ -407,20 +425,26 @@ def _highly_variable_genes_batched(
     n_bins: int,
     flavor: Literal["seurat", "cell_ranger"],
     cutoff: _Cutoffs | int,
+    client: DaskClient | None = None,
 ) -> pd.DataFrame:
     adata._sanitize()
     batches = adata.obs[batch_key].cat.categories
     dfs = []
     gene_list = adata.var_names
     for batch in batches:
-        adata_subset = adata[adata.obs[batch_key] == batch]
+        adata_subset = adata[adata.obs[batch_key] == batch].copy()
 
-        calculate_qc_metrics(adata_subset, layer=layer)
+        calculate_qc_metrics(adata_subset, layer=layer, client=client)
         filt = adata_subset.var["n_cells_by_counts"].to_numpy() > 0
-        adata_subset = adata_subset[:, filt]
+        adata_subset = adata_subset[:, filt].copy()
 
         hvg = _highly_variable_genes_single_batch(
-            adata_subset, layer=layer, cutoff=cutoff, n_bins=n_bins, flavor=flavor
+            adata_subset,
+            layer=layer,
+            cutoff=cutoff,
+            n_bins=n_bins,
+            flavor=flavor,
+            client=client,
         )
         hvg.reset_index(drop=False, inplace=True, names=["gene"])
 

src/rapids_singlecell/preprocessing/_kernels/_scale_kernel.py

@@ -20,7 +20,7 @@
 """
 
 _csr_scale_diff_kernel = r"""
-(const int *indptr, const int *indices, {0} *data, const double * std, const int *mask, {0} clipper,int nrows) {
+(const int *indptr, const int *indices, {0} *data, const {0} * std, const int *mask, {0} clipper,int nrows) {
         int row = blockIdx.x;
 
         if(row >= nrows){