diff --git a/docs/release-notes/0.17.0.md b/docs/release-notes/0.17.0.md index 35a6daff..ef41eebe 100644 --- a/docs/release-notes/0.17.0.md +++ b/docs/release-notes/0.17.0.md @@ -9,6 +9,7 @@ * Fix {func}`~rapids_singlecell.pp.neighbors` with ``algorithm="all_neighbors"`` losing the neighbors that cross cluster boundaries on multiple GPUs; ``overlap_factor`` now scales with the cluster count {pr}`769` {smaller}`S Dicks` * Fix {func}`~rapids_singlecell.pp.neighbors` with ``metric="inner_product"`` failing in the connectivity step for every ``algorithm``. The metric is no longer forwarded to ``cuml``'s ``fuzzy_simplicial_set``, which rejects the metrics it does not know and ignores it anyway when the neighbors are precomputed {pr}`769` {smaller}`S Dicks` * Fix {func}`~rapids_singlecell.pp.neighbors` with ``algorithm="all_neighbors"`` and ``algorithm_kwds={"n_clusters": 2}`` always raising; the derived ``overlap_factor`` is now capped at ``n_clusters - 1`` {pr}`769` {smaller}`S Dicks` +* Fix batched ``all_neighbors`` returning incomplete or inaccurate results on multi-GPU CUDA 12 systems. {pr}`782` {smaller}`S Dicks` ```{rubric} Features ``` diff --git a/src/rapids_singlecell/preprocessing/_neighbors/_algorithms/_all_neighbors.py b/src/rapids_singlecell/preprocessing/_neighbors/_algorithms/_all_neighbors.py index 7a1b7bf1..a9d2ffda 100644 --- a/src/rapids_singlecell/preprocessing/_neighbors/_algorithms/_all_neighbors.py +++ b/src/rapids_singlecell/preprocessing/_neighbors/_algorithms/_all_neighbors.py @@ -4,7 +4,9 @@ from typing import TYPE_CHECKING import cupy as cp +import cuvs import numpy as np +from packaging.version import parse as parse_version from rapids_singlecell.preprocessing._neighbors._helper import _compute_nlist @@ -18,6 +20,9 @@ from rapids_singlecell.preprocessing._neighbors import _Metrics +_CUVS_HOST_OUTPUT_MIN_VERSION = parse_version("26.08") + + def _default_overlap_factor(n_clusters: int) -> int: """Overlap needed to hold recall as the dataset is split into more clusters.""" if n_clusters <= 1: @@ -66,14 +71,18 @@ def _all_neighbors_knn( ) algo = algorithm_kwds.get("algo", "nn_descent") n_clusters, overlap_factor = _all_neighbors_batching(algorithm_kwds) - if n_clusters == 1: - from cuvs.common import Resources - - res = Resources() - else: + use_host_output = ( + n_clusters > 1 + and parse_version(cuvs.__version__) >= _CUVS_HOST_OUTPUT_MIN_VERSION + ) + if use_host_output: from cuvs.common import MultiGpuResources res = MultiGpuResources() + else: + from cuvs.common import Resources + + res = Resources() cuvs_metric = "sqeuclidean" if metric == "euclidean" else metric if algo == "ivf_pq" or algo == "ivfpq": from cuvs.neighbors import ivf_pq @@ -111,8 +120,11 @@ def _all_neighbors_knn( ivf_pq_params=ivf_pq_params, nn_descent_params=nn_descent_params, ) - neighbors = cp.zeros([X.shape[0], k], dtype=np.int64) - distances = cp.zeros([X.shape[0], k], dtype=np.float32) + # Host outputs use cuVS's synchronized merge path. Older releases only accept + # device outputs, so keep batching on one GPU to avoid concurrent writes. + output_module = np if use_host_output else cp + neighbors = output_module.zeros([X.shape[0], k], dtype=np.int64) + distances = output_module.zeros([X.shape[0], k], dtype=np.float32) all_neighbors.build( dataset=X, @@ -122,7 +134,8 @@ def _all_neighbors_knn( distances=distances, resources=res, ) - neighbors = neighbors.astype(np.int32) + neighbors = cp.asarray(neighbors, dtype=np.int32) + distances = cp.asarray(distances) if metric == "euclidean": distances = cp.sqrt(distances) return neighbors, distances