From 7a1377588b9b40ee5cb9899865d1eeb22f99f5a5 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Mon, 24 Aug 2026 12:56:39 +0200 Subject: [PATCH 1/2] first draft Signed-off-by: Intron7 --- src/rapids_singlecell/_utils/__init__.py | 10 + src/rapids_singlecell/_utils/_multi_gpu.py | 209 +++++++++++ .../pertpy_gpu/_metrics/_base_metric.py | 22 +- .../pertpy_gpu/_metrics/_edistance.py | 325 +++++++++++++----- .../pertpy_gpu/_metrics/_wasserstein.py | 86 +++-- .../squidpy_gpu/_autocorr.py | 81 +++-- src/rapids_singlecell/squidpy_gpu/_co_oc.py | 75 ++-- src/rapids_singlecell/squidpy_gpu/_gearysc.py | 38 +- src/rapids_singlecell/squidpy_gpu/_moransi.py | 38 +- .../_rank_genes_groups/_stream_multi_gpu.py | 27 +- .../_rank_genes_groups/_wilcoxon_host.py | 23 +- tests/pertpy/test_distances.py | 48 +++ tests/test_multi_gpu_utils.py | 191 +++++++++- tests/test_rank_genes_groups_wilcoxon.py | 50 +++ 14 files changed, 1025 insertions(+), 198 deletions(-) diff --git a/src/rapids_singlecell/_utils/__init__.py b/src/rapids_singlecell/_utils/__init__.py index 0d7677953..0c5487079 100644 --- a/src/rapids_singlecell/_utils/__init__.py +++ b/src/rapids_singlecell/_utils/__init__.py @@ -8,19 +8,29 @@ from dask.array import Array as DaskArray from ._multi_gpu import ( + MultiGPUFallbackWarning, _calculate_blocks_per_pair, + _copy_to_device_via_host, _create_category_index_mapping, _get_device_attrs, _split_pairs, parse_device_ids, + peer_copy_verified, + peer_copy_works, + validate_multi_gpu, ) __all__ = [ + "MultiGPUFallbackWarning", "_calculate_blocks_per_pair", + "_copy_to_device_via_host", "_create_category_index_mapping", "_get_device_attrs", "_split_pairs", "parse_device_ids", + "peer_copy_verified", + "peer_copy_works", + "validate_multi_gpu", ] diff --git a/src/rapids_singlecell/_utils/_multi_gpu.py b/src/rapids_singlecell/_utils/_multi_gpu.py index baafd75c6..7799cf42e 100644 --- a/src/rapids_singlecell/_utils/_multi_gpu.py +++ b/src/rapids_singlecell/_utils/_multi_gpu.py @@ -12,12 +12,221 @@ from __future__ import annotations +import warnings +from functools import cache + import cupy as cp +import numpy as np # Cache for device attributes per device (lazy initialization) _DEVICE_ATTRS_CACHE: dict[int, dict] = {} +CUDA_ERROR_PEER_ACCESS_UNSUPPORTED = 217 +CUDA_ERROR_PEER_ACCESS_ALREADY_ENABLED = 704 +CUDA_ERROR_PEER_ACCESS_NOT_ENABLED = 705 +CUDA_ERROR_TOO_MANY_PEERS = 711 + +_SAFE_PEER_FALLBACK_ERRORS = { + CUDA_ERROR_PEER_ACCESS_UNSUPPORTED, + CUDA_ERROR_PEER_ACCESS_NOT_ENABLED, + CUDA_ERROR_TOO_MANY_PEERS, +} +_P2P_CANARY_EXPECTED = np.arange(1, 9, dtype=np.float64) +_P2P_CANARY_POISON = np.full(8, -7.0, dtype=np.float64) +_WARNED_P2P_FAILURES: set[tuple[int, tuple[tuple[int, int], ...]]] = set() + + +class MultiGPUFallbackWarning(RuntimeWarning): + """Warning emitted when unsafe P2P makes an operation use one GPU.""" + + +def _copy_to_device_via_host(array: cp.ndarray, device_id: int) -> cp.ndarray: + """Copy a device array without touching a peer link. + + This is intentionally reserved for small control arrays whose owner may + differ from the main input's device. Capturing the source before changing + devices is essential: calling ``cp.asarray`` first would already attempt + the peer transfer this helper exists to avoid. + """ + source_device = array.device.id + if source_device == device_id: + return array + + with cp.cuda.Device(source_device): + host = array.get() + with cp.cuda.Device(device_id): + copied = cp.asarray(host) + cp.cuda.get_current_stream().synchronize() + return copied + + +def _run_peer_copy_canary(destination: int, source: int) -> bool: + """Return whether a small peer copy arrives intact. + + The explicit synchronizations make the canary independent of whichever + CuPy stream was current when the multi-GPU operation was entered. + """ + with cp.cuda.Device(source): + expected = cp.asarray(_P2P_CANARY_EXPECTED) + cp.cuda.get_current_stream().synchronize() + + with cp.cuda.Device(destination): + probe = cp.asarray(_P2P_CANARY_POISON) + cp.cuda.get_current_stream().synchronize() + try: + cp.cuda.runtime.memcpyPeer( + probe.data.ptr, + destination, + expected.data.ptr, + source, + expected.nbytes, + ) + cp.cuda.runtime.deviceSynchronize() + arrived = cp.asnumpy(probe) + except cp.cuda.runtime.CUDARuntimeError as error: + if error.status in _SAFE_PEER_FALLBACK_ERRORS: + return False + # memcpy/readback can surface an earlier asynchronous failure. In + # particular, an illegal address poisons the CUDA context and + # cannot be repaired by switching to one GPU in this process. + raise + + return bool(np.array_equal(arrived, _P2P_CANARY_EXPECTED)) + + +@cache +def peer_copy_works(destination: int, source: int) -> bool: + """Return whether ``source -> destination`` P2P transfers are usable. + + ``deviceCanAccessPeer`` is only a capability query. Some affected systems + report support and return success from ``memcpyPeer`` while leaving the + destination unchanged, so this function verifies the transferred bytes. + The result is cached per ordered device pair for the lifetime of the + process. + """ + if destination == source: + return True + + with cp.cuda.Device(destination): + if not cp.cuda.runtime.deviceCanAccessPeer(destination, source): + return False + try: + cp.cuda.runtime.deviceEnablePeerAccess(source) + except cp.cuda.runtime.CUDARuntimeError as error: + if error.status == CUDA_ERROR_PEER_ACCESS_ALREADY_ENABLED: + pass + elif error.status in _SAFE_PEER_FALLBACK_ERRORS: + return False + else: + raise + + return _run_peer_copy_canary(destination, source) + + +# Compatibility name used by the container-level P2P diagnostics. +peer_copy_verified = peer_copy_works + + +def validate_multi_gpu( + device_ids: list[int], + *, + source_device: int | None = None, + gather_device: int | None = None, +) -> list[int]: + """Return usable devices, falling back to one GPU when P2P is unsafe. + + Multi-GPU implementations in this package fan input buffers out from a + source GPU and gather results onto a selected GPU. This validates exactly + those ordered transfers before sharding or worker threads start. + If a required link is unavailable or silently corrupts the canary, the + operation falls back to the source GPU and emits one warning per failed + link set. Unexpected CUDA errors propagate because they may indicate an + already-poisoned context. + + Parameters + ---------- + device_ids + Requested execution devices. + source_device + Device holding shared input buffers. Defaults to the current device. + gather_device + Device receiving worker results. Defaults to the first requested + device. + + Returns + ------- + list[int] + The requested devices when all required P2P directions work, otherwise + a one-element list containing ``source_device``. + """ + if not device_ids: + raise ValueError("device_ids must contain at least one device") + + device_ids = list(dict.fromkeys(device_ids)) + if source_device is None: + source_device = cp.cuda.Device().id + + n_available = cp.cuda.runtime.getDeviceCount() + invalid_ids = [ + device_id + for device_id in device_ids + if device_id < 0 or device_id >= n_available + ] + if invalid_ids: + raise ValueError( + f"Invalid GPU device ID(s): {invalid_ids}. " + f"Available devices: {list(range(n_available))}" + ) + if source_device < 0 or source_device >= n_available: + raise ValueError( + f"Invalid source GPU device ID {source_device}. " + f"Available devices: {list(range(n_available))}" + ) + + if gather_device is None: + gather_device = device_ids[0] + if gather_device < 0 or gather_device >= n_available: + raise ValueError( + f"Invalid gather GPU device ID {gather_device}. " + f"Available devices: {list(range(n_available))}" + ) + + required_pairs: set[tuple[int, int]] = set() + for device_id in device_ids: + if device_id != source_device: + required_pairs.add((device_id, source_device)) + if device_id != gather_device: + required_pairs.add((gather_device, device_id)) + + failed_pairs = () + for pair in sorted(required_pairs): + if not peer_copy_works(*pair): + # Once fallback is required, do not exercise any more suspect + # links in the caller's CUDA context. + failed_pairs = (pair,) + break + if not failed_pairs: + return device_ids + + warning_key = (source_device, failed_pairs) + if warning_key not in _WARNED_P2P_FAILURES: + transfers = ", ".join( + f"GPU {source} -> GPU {destination}" for destination, source in failed_pairs + ) + warnings.warn( + "Multi-GPU execution was disabled because the required P2P " + f"transfer(s) failed validation: {transfers}. Falling back to " + f"GPU {source_device} for this operation. This avoids silent " + "result corruption but may be slower.", + MultiGPUFallbackWarning, + stacklevel=2, + ) + _WARNED_P2P_FAILURES.add(warning_key) + + return [source_device] + + def parse_device_ids(*, multi_gpu: bool | list[int] | str | None) -> list[int]: """Parse multi_gpu parameter into a list of device IDs. diff --git a/src/rapids_singlecell/pertpy_gpu/_metrics/_base_metric.py b/src/rapids_singlecell/pertpy_gpu/_metrics/_base_metric.py index 0c1fa1f05..caf332ea8 100644 --- a/src/rapids_singlecell/pertpy_gpu/_metrics/_base_metric.py +++ b/src/rapids_singlecell/pertpy_gpu/_metrics/_base_metric.py @@ -7,7 +7,11 @@ import numpy as np from rapids_singlecell._keys import _preset_obsm_names, _resolve_obsm_key -from rapids_singlecell._utils import _create_category_index_mapping, parse_device_ids +from rapids_singlecell._utils import ( + _copy_to_device_via_host, + _create_category_index_mapping, + parse_device_ids, +) from rapids_singlecell.squidpy_gpu._utils import _assert_categorical_obs if TYPE_CHECKING: @@ -119,13 +123,21 @@ def _subset_to_groups( Ordered group names matching the category indices. """ embedding_raw = self._get_embedding(adata) + source_device = ( + embedding_raw.device.id + if isinstance(embedding_raw, cp.ndarray) + else cp.cuda.Device().id + ) mask, cat_offsets, cell_indices, groups_list = self._subset_indices( adata, groupby, needed_groups ) - if mask is not None: - embedding = cp.asarray(embedding_raw[mask]) - else: - embedding = cp.asarray(embedding_raw) + with cp.cuda.Device(source_device): + if mask is not None: + embedding = cp.asarray(embedding_raw[mask]) + else: + embedding = cp.asarray(embedding_raw) + cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) + cell_indices = _copy_to_device_via_host(cell_indices, source_device) return embedding, cat_offsets, cell_indices, groups_list def _subset_indices( diff --git a/src/rapids_singlecell/pertpy_gpu/_metrics/_edistance.py b/src/rapids_singlecell/pertpy_gpu/_metrics/_edistance.py index 91d6e0b28..4e6736b9d 100644 --- a/src/rapids_singlecell/pertpy_gpu/_metrics/_edistance.py +++ b/src/rapids_singlecell/pertpy_gpu/_metrics/_edistance.py @@ -14,7 +14,9 @@ from rapids_singlecell._cuda import _edistance_cuda as _ed from rapids_singlecell._utils import ( _calculate_blocks_per_pair, + _copy_to_device_via_host, _split_pairs, + validate_multi_gpu, ) from rapids_singlecell.squidpy_gpu._utils import _assert_categorical_obs @@ -105,11 +107,19 @@ def _materialize_source(embedding_raw, selector): ``selector`` is a boolean mask, an integer row-index array, or ``None``. """ - if _is_sparse(embedding_raw): - return _build_csr_source(embedding_raw, selector) - if selector is None: - return cp.asarray(embedding_raw) - return cp.asarray(embedding_raw[selector]) + source_device = ( + embedding_raw.device.id + if isinstance(embedding_raw, cp.ndarray) + else embedding_raw.data.device.id + if cpsp.issparse(embedding_raw) + else cp.cuda.Device().id + ) + with cp.cuda.Device(source_device): + if _is_sparse(embedding_raw): + return _build_csr_source(embedding_raw, selector) + if selector is None: + return cp.asarray(embedding_raw) + return cp.asarray(embedding_raw[selector]) class EDistanceMetric(BaseMetric): @@ -162,6 +172,11 @@ def _load_source( adata, groupby, needed_groups ) source = _materialize_source(embedding_raw, mask) + source_device = ( + source.data.device.id if isinstance(source, _CSRData) else source.device.id + ) + cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) + cell_indices = _copy_to_device_via_host(cell_indices, source_device) return source, cat_offsets, cell_indices, groups_list def pairwise( @@ -305,6 +320,11 @@ def onesided_distances( embedding, cat_offsets, cell_indices, groups_list = self._load_source( adata, groupby, needed ) + source_device = ( + embedding.data.device.id + if isinstance(embedding, _CSRData) + else embedding.device.id + ) k = len(groups_list) group_map = {v: i for i, v in enumerate(groups_list)} selected_indices = [group_map[sg] for sg in selected_groups] @@ -327,14 +347,15 @@ def onesided_distances( # e[s,b] = 2*d[s,b] - d[s,s] - d[b,b] ed_cols = {} var_cols = {} - for i, (sg, si) in enumerate(zip(selected_groups, selected_indices)): - ed_row = 2 * cross_mean[i, :] - diag_mean[si] - diag_mean - ed_row[si] = 0.0 - ed_cols[sg] = ed_row.get() + with cp.cuda.Device(source_device): + for i, (sg, si) in enumerate(zip(selected_groups, selected_indices)): + ed_row = 2 * cross_mean[i, :] - diag_mean[si] - diag_mean + ed_row[si] = 0.0 + ed_cols[sg] = ed_row.get() - var_row = 4 * cross_var[i, :] + diag_var[si] + diag_var - var_row[si] = 0.0 - var_cols[sg] = var_row.get() + var_row = 4 * cross_var[i, :] + diag_var[si] + diag_var + var_row[si] = 0.0 + var_cols[sg] = var_row.get() distances = pd.DataFrame(ed_cols, index=groups_list) distances.index.name = groupby @@ -364,10 +385,11 @@ def onesided_distances( # cross_means[i, j] = mean dist from selected[i] to group j # diag_means[j] = mean within-group dist for group j ed_cols = {} - for i, (sg, si) in enumerate(zip(selected_groups, selected_indices)): - ed_row = 2 * cross_means[i, :] - diag_means[si] - diag_means - ed_row[si] = 0.0 - ed_cols[sg] = ed_row.get() + with cp.cuda.Device(source_device): + for i, (sg, si) in enumerate(zip(selected_groups, selected_indices)): + ed_row = 2 * cross_means[i, :] - diag_means[si] - diag_means + ed_row[si] = 0.0 + ed_cols[sg] = ed_row.get() df = pd.DataFrame(ed_cols, index=groups_list) df.index.name = groupby @@ -526,14 +548,22 @@ def contrast_distances( embedding = _materialize_source(embedding_raw, original_indices) cell_indices = cp.arange(len(original_indices), dtype=cp.int32) elif len(original_indices) < int(len(embedding_raw) * 0.7): - embedding = cp.asarray(embedding_raw[original_indices]) + embedding = _materialize_source(embedding_raw, original_indices) cell_indices = cp.arange(len(original_indices), dtype=cp.int32) else: - embedding = cp.asarray(embedding_raw) + embedding = _materialize_source(embedding_raw, None) cell_indices = cp.array(original_indices, dtype=cp.int32) - group_sizes = cp.diff(cat_offsets).astype(cp.int64) - group_sizes_cpu = group_sizes.get() + source_device = ( + embedding.data.device.id + if isinstance(embedding, _CSRData) + else embedding.device.id + ) + cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) + cell_indices = _copy_to_device_via_host(cell_indices, source_device) + with cp.cuda.Device(source_device): + group_sizes = cp.diff(cat_offsets).astype(cp.int64) + group_sizes_cpu = group_sizes.get() # Build deduplicated pairs pair_to_flat: dict[tuple[int, int], int] = {} for idx_a, idx_b in contrast_pairs: @@ -553,8 +583,9 @@ def contrast_distances( return result pairs = sorted(pair_to_flat.keys(), key=lambda p: pair_to_flat[p]) - pair_left = cp.array([p[0] for p in pairs], dtype=cp.int32) - pair_right = cp.array([p[1] for p in pairs], dtype=cp.int32) + with cp.cuda.Device(source_device): + pair_left = cp.array([p[0] for p in pairs], dtype=cp.int32) + pair_right = cp.array([p[1] for p in pairs], dtype=cp.int32) flat_sums = self._launch_distance_kernel( embedding, @@ -565,17 +596,18 @@ def contrast_distances( device_ids=device_ids, ) - # Vectorized normalization - is_diag = pair_left == pair_right - sizes_l = group_sizes[pair_left.astype(cp.intp)] - sizes_r = group_sizes[pair_right.astype(cp.intp)] - flat_norms = cp.where( - is_diag, - cp.maximum(sizes_l * (sizes_l - 1) // 2, 1), - sizes_l * sizes_r, - ).astype(embedding.dtype) - flat_means = flat_sums / flat_norms - flat_means_cpu = flat_means.get() + with cp.cuda.Device(source_device): + # Vectorized normalization + is_diag = pair_left == pair_right + sizes_l = group_sizes[pair_left.astype(cp.intp)] + sizes_r = group_sizes[pair_right.astype(cp.intp)] + flat_norms = cp.where( + is_diag, + cp.maximum(sizes_l * (sizes_l - 1) // 2, 1), + sizes_l * sizes_r, + ).astype(embedding.dtype) + flat_means = flat_sums / flat_norms + flat_means_cpu = flat_means.get() # Extract edistances edistances = np.empty(len(contrast_pairs), dtype=np.float64) @@ -739,6 +771,36 @@ def _launch_distance_kernel( pair_left: cp.ndarray, pair_right: cp.ndarray, device_ids: list[int], + ) -> cp.ndarray: + """Run distribution, launch, and gather from the input's device.""" + source_device = ( + embedding.data.device.id + if isinstance(embedding, _CSRData) + else embedding.device.id + ) + cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) + cell_indices = _copy_to_device_via_host(cell_indices, source_device) + pair_left = _copy_to_device_via_host(pair_left, source_device) + pair_right = _copy_to_device_via_host(pair_right, source_device) + with cp.cuda.Device(source_device): + return self._launch_distance_kernel_on_source( + embedding, + cat_offsets, + cell_indices, + pair_left=pair_left, + pair_right=pair_right, + device_ids=device_ids, + ) + + def _launch_distance_kernel_on_source( + self, + embedding: cp.ndarray, + cat_offsets: cp.ndarray, + cell_indices: cp.ndarray, + *, + pair_left: cp.ndarray, + pair_right: cp.ndarray, + device_ids: list[int], ) -> cp.ndarray: """Launch the edistance kernel across GPUs and return raw flat sums. @@ -765,6 +827,27 @@ def _launch_distance_kernel( cp.ndarray Raw distance sums of shape (n_pairs,), NOT normalized. """ + source_device = ( + embedding.data.device.id + if isinstance(embedding, _CSRData) + else embedding.device.id + ) + # Control arrays may have been created on the caller's current device + # even when a device-resident embedding lives elsewhere. Stage these + # small arrays through host memory so fallback never depends on P2P. + cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) + cell_indices = _copy_to_device_via_host(cell_indices, source_device) + pair_left = _copy_to_device_via_host(pair_left, source_device) + pair_right = _copy_to_device_via_host(pair_right, source_device) + device_ids = list(dict.fromkeys(device_ids)) + device_ids.sort(key=lambda device_id: device_id != source_device) + device_ids = validate_multi_gpu( + device_ids, + source_device=source_device, + gather_device=source_device, + ) + with cp.cuda.Device(source_device): + cp.cuda.get_current_stream().synchronize() n_devices = len(device_ids) n_total_pairs = len(pair_left) _, n_features = embedding.shape @@ -852,7 +935,7 @@ def _launch_distance_kernel( feat_tile, block_size, shared_mem, - cp.cuda.get_current_stream().ptr, + streams[device_id].ptr, ) else: _ed.compute_distances( @@ -869,17 +952,17 @@ def _launch_distance_kernel( feat_tile, block_size, shared_mem, - cp.cuda.get_current_stream().ptr, + streams[device_id].ptr, ) # Phase 3: Synchronize all devices for data in device_data: if data is not None: with cp.cuda.Device(data["device_id"]): - cp.cuda.Stream.null.synchronize() + streams[data["device_id"]].synchronize() - # Phase 4: Aggregate on GPU 0 - with cp.cuda.Device(device_ids[0]): + # Phase 4: Aggregate where the input lives. + with cp.cuda.Device(source_device): total_sums = cp.zeros(n_total_pairs, dtype=embedding.dtype) for i, data in enumerate(device_data): if data is not None: @@ -896,6 +979,27 @@ def _pairwise_means( cell_indices: cp.ndarray, k: int, device_ids: list[int], + ) -> cp.ndarray: + """Run pairwise reconstruction on the embedding's owning device.""" + source_device = ( + embedding.data.device.id + if isinstance(embedding, _CSRData) + else embedding.device.id + ) + cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) + cell_indices = _copy_to_device_via_host(cell_indices, source_device) + with cp.cuda.Device(source_device): + return self._pairwise_means_on_source( + embedding, cat_offsets, cell_indices, k, device_ids + ) + + def _pairwise_means_on_source( + self, + embedding: cp.ndarray, + cat_offsets: cp.ndarray, + cell_indices: cp.ndarray, + k: int, + device_ids: list[int], ) -> cp.ndarray: """Compute between-group mean distances for all group pairs. @@ -969,6 +1073,34 @@ def _onesided_means( *, selected_indices: list[int], device_ids: list[int], + ) -> tuple[cp.ndarray, cp.ndarray]: + """Run one-sided reconstruction on the embedding's owning device.""" + source_device = ( + embedding.data.device.id + if isinstance(embedding, _CSRData) + else embedding.device.id + ) + cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) + cell_indices = _copy_to_device_via_host(cell_indices, source_device) + with cp.cuda.Device(source_device): + return self._onesided_means_on_source( + embedding, + cat_offsets, + cell_indices, + k, + selected_indices=selected_indices, + device_ids=device_ids, + ) + + def _onesided_means_on_source( + self, + embedding: cp.ndarray, + cat_offsets: cp.ndarray, + cell_indices: cp.ndarray, + k: int, + *, + selected_indices: list[int], + device_ids: list[int], ) -> tuple[cp.ndarray, cp.ndarray]: """Compute mean distances from selected group(s) to all groups. @@ -1109,19 +1241,27 @@ def _pairwise_means_bootstrap( tuple (means, variances) matrices (k x k each) """ - # Get group sizes for bootstrap sampling (on GPU 0) - group_sizes = cp.diff(cat_offsets) + source_device = ( + embedding.data.device.id + if isinstance(embedding, _CSRData) + else embedding.device.id + ) + cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) + cell_indices = _copy_to_device_via_host(cell_indices, source_device) + with cp.cuda.Device(source_device): + group_sizes = cp.diff(cat_offsets) # Run bootstrap iterations - each uses all GPUs for pairwise computation all_results = [] for i in range(n_bootstrap): # Generate bootstrap sample on GPU 0 - boot_cat_offsets, boot_cell_indices = self._bootstrap_sample_cells( - cat_offsets=cat_offsets, - cell_indices=cell_indices, - group_sizes_gpu=group_sizes, - seed=random_state + i, - ) + with cp.cuda.Device(source_device): + boot_cat_offsets, boot_cell_indices = self._bootstrap_sample_cells( + cat_offsets=cat_offsets, + cell_indices=cell_indices, + group_sizes_gpu=group_sizes, + seed=random_state + i, + ) # Compute pairwise means using all GPUs pairwise_means = self._pairwise_means( @@ -1133,8 +1273,8 @@ def _pairwise_means_bootstrap( ) all_results.append(pairwise_means.get()) - # Compute statistics on first GPU - with cp.cuda.Device(device_ids[0]): + # Keep the returned statistics with the embedding. + with cp.cuda.Device(source_device): bootstrap_stack = cp.array(all_results) # [n_bootstrap, k, k] means = cp.mean(bootstrap_stack, axis=0) variances = cp.var(bootstrap_stack, axis=0) @@ -1187,20 +1327,28 @@ def _onesided_means_bootstrap( diag_var Variance of bootstrap diag_means, shape (k,) """ - # Get group sizes for bootstrap sampling (on GPU 0) - group_sizes = cp.diff(cat_offsets) + source_device = ( + embedding.data.device.id + if isinstance(embedding, _CSRData) + else embedding.device.id + ) + cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) + cell_indices = _copy_to_device_via_host(cell_indices, source_device) + with cp.cuda.Device(source_device): + group_sizes = cp.diff(cat_offsets) # Run bootstrap iterations - each uses all GPUs for onesided computation all_cross = [] all_diag = [] for i in range(n_bootstrap): # Generate bootstrap sample on GPU 0 - boot_cat_offsets, boot_cell_indices = self._bootstrap_sample_cells( - cat_offsets=cat_offsets, - cell_indices=cell_indices, - group_sizes_gpu=group_sizes, - seed=random_state + i, - ) + with cp.cuda.Device(source_device): + boot_cat_offsets, boot_cell_indices = self._bootstrap_sample_cells( + cat_offsets=cat_offsets, + cell_indices=cell_indices, + group_sizes_gpu=group_sizes, + seed=random_state + i, + ) # Compute onesided means using all GPUs cross_means, diag_means = self._onesided_means( @@ -1214,8 +1362,8 @@ def _onesided_means_bootstrap( all_cross.append(cross_means.get()) all_diag.append(diag_means.get()) - # Compute statistics on first GPU - with cp.cuda.Device(device_ids[0]): + # Keep the returned statistics with the embedding. + with cp.cuda.Device(source_device): cross_stack = cp.array(all_cross) diag_stack = cp.array(all_diag) cross_mean = cp.mean(cross_stack, axis=0) @@ -1306,32 +1454,36 @@ def _prepare_edistance_df_bootstrap( random_state=random_state, device_ids=device_ids, ) - - # Vectorized edistance: e[a,b] = 2*d[a,b] - d[a,a] - d[b,b] - diag_means = cp.diag(pairwise_means_boot) - edistance_means = ( - 2 * pairwise_means_boot - diag_means[:, None] - diag_means[None, :] + source_device = ( + embedding.data.device.id + if isinstance(embedding, _CSRData) + else embedding.device.id ) - cp.fill_diagonal(edistance_means, 0) - # Vectorized variance computation (delta method approximation): - # var[a, b] = 4 * var[a, b] + var[a, a] + var[b, b] - diag_vars = cp.diag(pairwise_vars_boot) - edistance_vars = ( - 4 * pairwise_vars_boot + diag_vars[:, None] + diag_vars[None, :] - ) - cp.fill_diagonal(edistance_vars, 0) + with cp.cuda.Device(source_device): + # Vectorized edistance: e[a,b] = 2*d[a,b] - d[a,a] - d[b,b] + diag_means = cp.diag(pairwise_means_boot) + edistance_means = ( + 2 * pairwise_means_boot - diag_means[:, None] - diag_means[None, :] + ) + cp.fill_diagonal(edistance_means, 0) - df_mean = pd.DataFrame( - edistance_means.get(), index=groups_list, columns=groups_list - ) + # Vectorized variance computation (delta method approximation): + # var[a, b] = 4 * var[a, b] + var[a, a] + var[b, b] + diag_vars = cp.diag(pairwise_vars_boot) + edistance_vars = ( + 4 * pairwise_vars_boot + diag_vars[:, None] + diag_vars[None, :] + ) + cp.fill_diagonal(edistance_vars, 0) + means_host = edistance_means.get() + vars_host = edistance_vars.get() + + df_mean = pd.DataFrame(means_host, index=groups_list, columns=groups_list) df_mean.index.name = groupby df_mean.columns.name = groupby df_mean.name = "pairwise edistance" - df_var = pd.DataFrame( - edistance_vars.get(), index=groups_list, columns=groups_list - ) + df_var = pd.DataFrame(vars_host, index=groups_list, columns=groups_list) df_var.index.name = groupby df_var.columns.name = groupby df_var.name = "pairwise edistance variance" @@ -1354,15 +1506,20 @@ def _prepare_edistance_df( pairwise_means = self._pairwise_means( embedding, cat_offsets, cell_indices, k, device_ids ) + source_device = ( + embedding.data.device.id + if isinstance(embedding, _CSRData) + else embedding.device.id + ) - # Vectorized edistance: e[a,b] = 2*d[a,b] - d[a,a] - d[b,b] - diag = cp.diag(pairwise_means) - edistance_matrix = 2 * pairwise_means - diag[:, None] - diag[None, :] - cp.fill_diagonal(edistance_matrix, 0) # Self-distance is 0 + with cp.cuda.Device(source_device): + # Vectorized edistance: e[a,b] = 2*d[a,b] - d[a,a] - d[b,b] + diag = cp.diag(pairwise_means) + edistance_matrix = 2 * pairwise_means - diag[:, None] - diag[None, :] + cp.fill_diagonal(edistance_matrix, 0) # Self-distance is 0 + edistance_host = edistance_matrix.get() - df = pd.DataFrame( - edistance_matrix.get(), index=groups_list, columns=groups_list - ) + df = pd.DataFrame(edistance_host, index=groups_list, columns=groups_list) df.index.name = groupby df.columns.name = groupby df.name = "pairwise edistance" diff --git a/src/rapids_singlecell/pertpy_gpu/_metrics/_wasserstein.py b/src/rapids_singlecell/pertpy_gpu/_metrics/_wasserstein.py index a9c856c2f..dccd3a66f 100644 --- a/src/rapids_singlecell/pertpy_gpu/_metrics/_wasserstein.py +++ b/src/rapids_singlecell/pertpy_gpu/_metrics/_wasserstein.py @@ -14,6 +14,7 @@ import pandas as pd from rapids_singlecell._cuda import _sinkhorn_cuda as _sk +from rapids_singlecell._utils import _copy_to_device_via_host, validate_multi_gpu from rapids_singlecell.squidpy_gpu._utils import _assert_categorical_obs from ._base_metric import BaseMetric, parse_device_ids @@ -327,14 +328,27 @@ def _solve_pairs( device_ids = [0] n_pairs = len(pair_left) if n_pairs == 0: - return cp.zeros(0, dtype=dtype) + with cp.cuda.Device(embedding.device.id): + return cp.zeros(0, dtype=dtype) + + source_device = embedding.device.id + cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) + cell_indices = _copy_to_device_via_host(cell_indices, source_device) + device_ids = list(dict.fromkeys(device_ids)) + device_ids.sort(key=lambda device_id: device_id != source_device) + device_ids = validate_multi_gpu( + device_ids, + source_device=source_device, + gather_device=source_device, + ) pl_host = np.asarray(pair_left, dtype=np.int32) pr_host = np.asarray(pair_right, dtype=np.int32) - group_sizes = (cat_offsets[1:] - cat_offsets[:-1]).astype(cp.int32) - # Group sizes to host (one sync) -> all batch planning is host-only. - n_left = group_sizes[cp.asarray(pl_host)].get() - n_right = group_sizes[cp.asarray(pr_host)].get() + with cp.cuda.Device(source_device): + group_sizes = (cat_offsets[1:] - cat_offsets[:-1]).astype(cp.int32) + # Group sizes to host (one sync) -> all batch planning is host-only. + n_left = group_sizes[cp.asarray(pl_host)].get() + n_right = group_sizes[cp.asarray(pr_host)].get() # Orient larger group as columns. swap = n_left > n_right rows = np.where(swap, pr_host, pl_host) @@ -345,8 +359,10 @@ def _solve_pairs( plans = _plan_device_batches(n_row, n_col, itemsize, len(device_ids)) - out = cp.empty(n_pairs, dtype=dtype) - home = device_ids[0] + home = source_device + with cp.cuda.Device(home): + cp.cuda.get_current_stream().synchronize() + out = cp.empty(n_pairs, dtype=dtype) # Move the shared inputs to each participating device once. streams: dict[int, cp.cuda.Stream] = {} @@ -474,8 +490,19 @@ def _bootstrap_solve( raise ValueError(f"n_bootstrap must be >= 1, got {n_bootstrap}") n_pairs = len(pair_left) if n_pairs == 0: - empty = cp.zeros(0, dtype=dtype) - return empty, empty + with cp.cuda.Device(embedding.device.id): + empty = cp.zeros(0, dtype=dtype) + return empty, empty + source_device = embedding.device.id + cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) + cell_indices = _copy_to_device_via_host(cell_indices, source_device) + device = validate_multi_gpu( + [device], + source_device=source_device, + gather_device=source_device, + )[0] + with cp.cuda.Device(source_device): + cp.cuda.get_current_stream().synchronize() with cp.cuda.Device(device): emb = cp.ascontiguousarray(cp.asarray(embedding)) offs = cp.asarray(cat_offsets) @@ -531,7 +558,12 @@ def _bootstrap_solve( if not converged: self._warn_not_converged() reg = reg.reshape(n_pairs, n_bootstrap) - return reg.mean(axis=1), reg.var(axis=1) + mean = reg.mean(axis=1) + var = reg.var(axis=1) + return ( + _copy_to_device_via_host(mean, source_device), + _copy_to_device_via_host(var, source_device), + ) def bootstrap_arrays( self, @@ -631,13 +663,15 @@ def pairwise( pair_right.append(j) def _to_matrix(flat: cp.ndarray, name: str) -> pd.DataFrame: - mat = cp.zeros((k, k), dtype=embedding.dtype) - if pair_left: - il = cp.asarray(pair_left, dtype=cp.intp) - jr = cp.asarray(pair_right, dtype=cp.intp) - mat[il, jr] = flat - mat[jr, il] = flat - df = pd.DataFrame(mat.get(), index=groups_list, columns=groups_list) + with cp.cuda.Device(embedding.device.id): + mat = cp.zeros((k, k), dtype=embedding.dtype) + if pair_left: + il = cp.asarray(pair_left, dtype=cp.intp) + jr = cp.asarray(pair_right, dtype=cp.intp) + mat[il, jr] = flat + mat[jr, il] = flat + mat_host = mat.get() + df = pd.DataFrame(mat_host, index=groups_list, columns=groups_list) df.index.name = groupby df.columns.name = groupby df.name = name @@ -710,7 +744,8 @@ def onesided_distances( pair_right.append(j) def _to_df(flat: cp.ndarray) -> pd.DataFrame: - flat_cpu = flat.get() + with cp.cuda.Device(embedding.device.id): + flat_cpu = flat.get() ed_cols: dict[str, np.ndarray] = {} cursor = 0 for ii, si in enumerate(selected_indices): @@ -767,6 +802,11 @@ def contrast_distances( device_ids = parse_device_ids(multi_gpu=multi_gpu) groupby, split_by = self._parse_contrasts(adata, contrasts) embedding_raw = self._get_embedding(adata) + source_device = ( + embedding_raw.device.id + if isinstance(embedding_raw, cp.ndarray) + else cp.cuda.Device().id + ) all_cols = [groupby, *split_by] grouped = adata.obs.groupby(all_cols, observed=True) @@ -806,9 +846,10 @@ def contrast_distances( flat_cell_idx = ( np.concatenate(all_cells) if all_cells else np.array([], dtype=np.int64) ) - cat_offsets = cp.asarray(offsets_host, dtype=cp.int32) - cell_indices = cp.asarray(flat_cell_idx, dtype=cp.int32) - embedding = cp.asarray(embedding_raw) + with cp.cuda.Device(source_device): + cat_offsets = cp.asarray(offsets_host, dtype=cp.int32) + cell_indices = cp.asarray(flat_cell_idx, dtype=cp.int32) + embedding = cp.asarray(embedding_raw) dtype = embedding.dtype # Deduplicate canonical pairs (i, j) with i < j @@ -833,7 +874,8 @@ def contrast_distances( dtype=dtype, device_ids=device_ids, ) - flat_cpu = flat.get() + with cp.cuda.Device(embedding.device.id): + flat_cpu = flat.get() distances = np.empty(len(contrast_pairs), dtype=np.float64) for n, (idx_a, idx_b) in enumerate(contrast_pairs): diff --git a/src/rapids_singlecell/squidpy_gpu/_autocorr.py b/src/rapids_singlecell/squidpy_gpu/_autocorr.py index 76b2a3f3a..eff2a7f88 100644 --- a/src/rapids_singlecell/squidpy_gpu/_autocorr.py +++ b/src/rapids_singlecell/squidpy_gpu/_autocorr.py @@ -31,21 +31,28 @@ def _to_cupy(vals, *, use_sparse: bool, dtype): Dense input is always returned as a dense CuPy array. """ is_sparse = sparse.issparse(vals) or sparse_gpu.isspmatrix(vals) - - # Dense input - use_sparse is ignored - if not is_sparse: - return cp.array(vals, dtype=dtype, order="C") - - # Sparse input - respect use_sparse parameter - if use_sparse: - if sparse_gpu.isspmatrix(vals): - return vals.tocsr().astype(dtype) - return sparse_gpu.csr_matrix(vals.tocsr(), dtype=dtype) - - # Sparse input but use_sparse=False - convert to dense - if not sparse_gpu.isspmatrix(vals): - vals = sparse_gpu.csr_matrix(vals.tocsr(), dtype=dtype) - return _sparse_to_dense(vals, order="C") + source_device = ( + vals.device.id + if isinstance(vals, cp.ndarray) + else vals.data.device.id + if sparse_gpu.isspmatrix(vals) + else cp.cuda.Device().id + ) + with cp.cuda.Device(source_device): + # Dense input - use_sparse is ignored + if not is_sparse: + return cp.array(vals, dtype=dtype, order="C") + + # Sparse input - respect use_sparse parameter + if use_sparse: + if sparse_gpu.isspmatrix(vals): + return vals.tocsr().astype(dtype) + return sparse_gpu.csr_matrix(vals.tocsr(), dtype=dtype) + + # Sparse input but use_sparse=False - convert to dense + if not sparse_gpu.isspmatrix(vals): + vals = sparse_gpu.csr_matrix(vals.tocsr(), dtype=dtype) + return _sparse_to_dense(vals, order="C") def spatial_autocorr( @@ -143,15 +150,23 @@ def spatial_autocorr( if compute_dtype not in (np.float32, np.float64): compute_dtype = np.float32 - # create Adj-Matrix - adj_matrix = adata.obsp[connectivity_key] - adj_matrix_cupy = sparse_gpu.csr_matrix(adj_matrix, dtype=compute_dtype) - - if transformation: # row-normalize - row_sums = adj_matrix_cupy.sum(axis=1).reshape(-1, 1) - non_zero_rows = row_sums != 0 - row_sums[non_zero_rows] = 1.0 / row_sums[non_zero_rows] - adj_matrix_cupy = adj_matrix_cupy.multiply(sparse_gpu.csr_matrix(row_sums)) + source_device = ( + vals.device.id + if isinstance(vals, cp.ndarray) + else vals.data.device.id + if sparse_gpu.isspmatrix(vals) + else cp.cuda.Device().id + ) + with cp.cuda.Device(source_device): + # create Adj-Matrix alongside the data + adj_matrix = adata.obsp[connectivity_key] + adj_matrix_cupy = sparse_gpu.csr_matrix(adj_matrix, dtype=compute_dtype) + + if transformation: # row-normalize + row_sums = adj_matrix_cupy.sum(axis=1).reshape(-1, 1) + non_zero_rows = row_sums != 0 + row_sums[non_zero_rows] = 1.0 / row_sums[non_zero_rows] + adj_matrix_cupy = adj_matrix_cupy.multiply(sparse_gpu.csr_matrix(row_sums)) params = {"two_tailed": two_tailed} @@ -168,10 +183,13 @@ def _run_autocorr(data, adj_matrix_cupy, mode, n_perms, multi_gpu): else: raise ValueError(f"Invalid mode: {mode}") - data = _to_cupy(vals, use_sparse=use_sparse, dtype=compute_dtype) + with cp.cuda.Device(source_device): + data = _to_cupy(vals, use_sparse=use_sparse, dtype=compute_dtype) - # Run full computation - score, score_perms = _run_autocorr(data, adj_matrix_cupy, mode, n_perms, multi_gpu) + # Run full computation + score, score_perms = _run_autocorr( + data, adj_matrix_cupy, mode, n_perms, multi_gpu + ) # Set mode-specific params if mode == "moran": @@ -185,10 +203,11 @@ def _run_autocorr(data, adj_matrix_cupy, mode, n_perms, multi_gpu): params["ascending"] = True params["mode"] = "gearyC" - g = sparse.csr_matrix(adj_matrix_cupy.get()) - score = score.get() - if n_perms is not None: - score_perms = score_perms.get() + with cp.cuda.Device(source_device): + g = sparse.csr_matrix(adj_matrix_cupy.get()) + score = score.get() + if n_perms is not None: + score_perms = score_perms.get() with np.errstate(divide="ignore"): pval_results = _p_value_calc(score, sims=score_perms, weights=g, params=params) diff --git a/src/rapids_singlecell/squidpy_gpu/_co_oc.py b/src/rapids_singlecell/squidpy_gpu/_co_oc.py index 7f1b3462e..b4505ad9d 100644 --- a/src/rapids_singlecell/squidpy_gpu/_co_oc.py +++ b/src/rapids_singlecell/squidpy_gpu/_co_oc.py @@ -9,9 +9,11 @@ from rapids_singlecell._cuda import _cooc_cuda as _co from rapids_singlecell._utils import ( _calculate_blocks_per_pair, + _copy_to_device_via_host, _create_category_index_mapping, _split_pairs, parse_device_ids, + validate_multi_gpu, ) from ._utils import _assert_categorical_obs, _assert_spatial_basis @@ -67,26 +69,37 @@ def co_occurrence( _assert_categorical_obs(adata, key=cluster_key) _assert_spatial_basis(adata, key=spatial_key) - spatial = cp.array(adata.obsm[spatial_key]).astype(np.float32) - original_clust = adata.obs[cluster_key] - clust_map = {v: i for i, v in enumerate(original_clust.cat.categories.values)} - labs = cp.array([clust_map[c] for c in original_clust], dtype=np.int32) - # create intervals thresholds - if isinstance(interval, int): - thresh_min, thresh_max = _find_min_max(spatial) - interval = cp.linspace(thresh_min, thresh_max, num=interval, dtype=np.float32) - else: - interval = cp.array(sorted(interval), dtype=np.float32, copy=True) - if len(interval) <= 1: - raise ValueError( - f"Expected interval to be of length `>= 2`, found `{len(interval)}`." - ) - - device_ids = parse_device_ids(multi_gpu=multi_gpu) - out = _co_occurrence_helper( - spatial, interval, labs, fast=True, device_ids=device_ids + spatial_input = adata.obsm[spatial_key] + source_device = ( + spatial_input.device.id + if isinstance(spatial_input, cp.ndarray) + else cp.cuda.Device().id ) - out, interval = out.get(), interval.get() + with cp.cuda.Device(source_device): + spatial = cp.array(spatial_input).astype(np.float32) + original_clust = adata.obs[cluster_key] + clust_map = {v: i for i, v in enumerate(original_clust.cat.categories.values)} + labs = cp.array([clust_map[c] for c in original_clust], dtype=np.int32) + # create intervals thresholds + if isinstance(interval, int): + thresh_min, thresh_max = _find_min_max(spatial) + interval = cp.linspace( + thresh_min, thresh_max, num=interval, dtype=np.float32 + ) + else: + if isinstance(interval, cp.ndarray): + interval = _copy_to_device_via_host(interval, source_device) + interval = cp.array(sorted(interval), dtype=np.float32, copy=True) + if len(interval) <= 1: + raise ValueError( + f"Expected interval to be of length `>= 2`, found `{len(interval)}`." + ) + + device_ids = parse_device_ids(multi_gpu=multi_gpu) + out = _co_occurrence_helper( + spatial, interval, labs, fast=True, device_ids=device_ids + ) + out, interval = out.get(), interval.get() if copy: return out, interval @@ -292,9 +305,23 @@ def _co_occurrence_gpu( if not valid_device_ids: return cp.zeros((k, k, l_val), dtype=cp.uint64), False - device_ids = valid_device_ids - n_devices = len(device_ids) source_device_id = spatial.device.id + thresholds = _copy_to_device_via_host(thresholds, source_device_id) + cat_offsets = _copy_to_device_via_host(cat_offsets, source_device_id) + cell_indices = _copy_to_device_via_host(cell_indices, source_device_id) + pair_left = _copy_to_device_via_host(pair_left, source_device_id) + pair_right = _copy_to_device_via_host(pair_right, source_device_id) + valid_device_ids.sort(key=lambda device_id: device_id != source_device_id) + device_ids = validate_multi_gpu( + valid_device_ids, + source_device=source_device_id, + gather_device=source_device_id, + ) + if any(device_id not in kernel_configs for device_id in device_ids): + return cp.zeros((k, k, l_val), dtype=cp.uint64), False + with cp.cuda.Device(source_device_id): + cp.cuda.get_current_stream().synchronize() + n_devices = len(device_ids) # Split pairs across devices with load balancing group_sizes = cp.diff(cat_offsets).astype(cp.int64) @@ -378,7 +405,7 @@ def _co_occurrence_gpu( cell_tile=cell_tile, block_size=block_size, shared_mem=shared_mem, - stream=cp.cuda.get_current_stream().ptr, + stream=streams[device_id].ptr, ) # Phase 3: Synchronize all devices (wait for kernels to complete) @@ -387,8 +414,8 @@ def _co_occurrence_gpu( with cp.cuda.Device(data["device_id"]): streams[data["device_id"]].synchronize() - # Phase 4: Aggregate counts on first device - with cp.cuda.Device(device_ids[0]): + # Phase 4: Aggregate counts where the input lives. + with cp.cuda.Device(source_device_id): counts = cp.zeros((k, k, l_val), dtype=cp.uint64) for data in device_data: if data is not None: diff --git a/src/rapids_singlecell/squidpy_gpu/_gearysc.py b/src/rapids_singlecell/squidpy_gpu/_gearysc.py index 118ef754a..49d41d142 100644 --- a/src/rapids_singlecell/squidpy_gpu/_gearysc.py +++ b/src/rapids_singlecell/squidpy_gpu/_gearysc.py @@ -6,7 +6,7 @@ from cupyx.scipy import sparse from rapids_singlecell._cuda import _autocorr_cuda as _ac -from rapids_singlecell._utils import parse_device_ids +from rapids_singlecell._utils import parse_device_ids, validate_multi_gpu from ._utils import _check_precision_issues @@ -80,6 +80,16 @@ def _run_permutations_dense( if device_ids is None: device_ids = [0] + source_device = data.device.id + device_ids = list(dict.fromkeys(device_ids)) + device_ids.sort(key=lambda device_id: device_id != source_device) + device_ids = validate_multi_gpu( + device_ids, + source_device=source_device, + gather_device=source_device, + ) + with cp.cuda.Device(source_device): + cp.cuda.get_current_stream().synchronize() n_devices = len(device_ids) streams: dict[int, cp.cuda.Stream] = {} device_data: list[dict] = [] @@ -94,7 +104,7 @@ def _run_permutations_dense( with streams[device_id]: # Copy data to this device - if device_id == device_ids[0]: + if device_id == source_device: dev_data = data dev_adj = adj_matrix_cupy dev_den = den @@ -127,7 +137,7 @@ def _run_permutations_dense( for p in range(perms_per_device): for dd in device_data: device_id = dd["device_id"] - with cp.cuda.Device(device_id): + with cp.cuda.Device(device_id), streams[device_id]: streams[device_id].synchronize() num_permuted = cp.zeros(n_features, dtype=dtype) @@ -151,8 +161,8 @@ def _run_permutations_dense( with cp.cuda.Device(dd["device_id"]): streams[dd["device_id"]].synchronize() - # Phase 3: Gather results on first device and cut to exact size - with cp.cuda.Device(device_ids[0]): + # Phase 3: Gather results where the input lives and cut to exact size. + with cp.cuda.Device(source_device): all_perms = [cp.asarray(dd["perms"]) for dd in device_data] gearys_C_permutations = cp.concatenate(all_perms, axis=0)[:n_permutations] @@ -240,6 +250,16 @@ def _run_permutations_sparse( if device_ids is None: device_ids = [0] + source_device = data.data.device.id + device_ids = list(dict.fromkeys(device_ids)) + device_ids.sort(key=lambda device_id: device_id != source_device) + device_ids = validate_multi_gpu( + device_ids, + source_device=source_device, + gather_device=source_device, + ) + with cp.cuda.Device(source_device): + cp.cuda.get_current_stream().synchronize() n_devices = len(device_ids) streams: dict[int, cp.cuda.Stream] = {} device_data: list[dict] = [] @@ -254,7 +274,7 @@ def _run_permutations_sparse( with streams[device_id]: # Copy data to this device - if device_id == device_ids[0]: + if device_id == source_device: dev_data = data dev_adj = adj_matrix_cupy dev_den = den @@ -294,7 +314,7 @@ def _run_permutations_sparse( for p in range(perms_per_device): for dd in device_data: device_id = dd["device_id"] - with cp.cuda.Device(device_id): + with cp.cuda.Device(device_id), streams[device_id]: streams[device_id].synchronize() num_permuted = cp.zeros(n_features, dtype=dtype) @@ -320,8 +340,8 @@ def _run_permutations_sparse( with cp.cuda.Device(dd["device_id"]): streams[dd["device_id"]].synchronize() - # Phase 3: Gather results on first device and cut to exact size - with cp.cuda.Device(device_ids[0]): + # Phase 3: Gather results where the input lives and cut to exact size. + with cp.cuda.Device(source_device): all_perms = [cp.asarray(dd["perms"]) for dd in device_data] gearys_C_permutations = cp.concatenate(all_perms, axis=0)[:n_permutations] diff --git a/src/rapids_singlecell/squidpy_gpu/_moransi.py b/src/rapids_singlecell/squidpy_gpu/_moransi.py index 5c3e6d96a..c1c69a004 100644 --- a/src/rapids_singlecell/squidpy_gpu/_moransi.py +++ b/src/rapids_singlecell/squidpy_gpu/_moransi.py @@ -6,7 +6,7 @@ from cupyx.scipy import sparse from rapids_singlecell._cuda import _autocorr_cuda as _ac -from rapids_singlecell._utils import parse_device_ids +from rapids_singlecell._utils import parse_device_ids, validate_multi_gpu from ._utils import _check_precision_issues @@ -79,6 +79,16 @@ def _run_permutations_dense( if device_ids is None: device_ids = [0] + source_device = data_centered_cupy.device.id + device_ids = list(dict.fromkeys(device_ids)) + device_ids.sort(key=lambda device_id: device_id != source_device) + device_ids = validate_multi_gpu( + device_ids, + source_device=source_device, + gather_device=source_device, + ) + with cp.cuda.Device(source_device): + cp.cuda.get_current_stream().synchronize() n_devices = len(device_ids) streams: dict[int, cp.cuda.Stream] = {} device_data: list[dict] = [] @@ -93,7 +103,7 @@ def _run_permutations_dense( with streams[device_id]: # Copy data to this device - if device_id == device_ids[0]: + if device_id == source_device: dev_data = data_centered_cupy dev_adj = adj_matrix_cupy dev_den = den @@ -126,7 +136,7 @@ def _run_permutations_dense( for p in range(perms_per_device): for dd in device_data: device_id = dd["device_id"] - with cp.cuda.Device(device_id): + with cp.cuda.Device(device_id), streams[device_id]: streams[device_id].synchronize() num_permuted = cp.zeros(n_features, dtype=dtype) @@ -150,8 +160,8 @@ def _run_permutations_dense( with cp.cuda.Device(dd["device_id"]): streams[dd["device_id"]].synchronize() - # Phase 3: Gather results on first device and cut to exact size - with cp.cuda.Device(device_ids[0]): + # Phase 3: Gather results where the input lives and cut to exact size. + with cp.cuda.Device(source_device): all_perms = [cp.asarray(dd["perms"]) for dd in device_data] morans_I_permutations = cp.concatenate(all_perms, axis=0)[:n_permutations] @@ -240,6 +250,16 @@ def _run_permutations_sparse( if device_ids is None: device_ids = [0] + source_device = data.data.device.id + device_ids = list(dict.fromkeys(device_ids)) + device_ids.sort(key=lambda device_id: device_id != source_device) + device_ids = validate_multi_gpu( + device_ids, + source_device=source_device, + gather_device=source_device, + ) + with cp.cuda.Device(source_device): + cp.cuda.get_current_stream().synchronize() n_devices = len(device_ids) streams: dict[int, cp.cuda.Stream] = {} device_data: list[dict] = [] @@ -254,7 +274,7 @@ def _run_permutations_sparse( with streams[device_id]: # Copy data to this device - if device_id == device_ids[0]: + if device_id == source_device: dev_data = data dev_adj = adj_matrix_cupy dev_means = means @@ -297,7 +317,7 @@ def _run_permutations_sparse( for p in range(perms_per_device): for dd in device_data: device_id = dd["device_id"] - with cp.cuda.Device(device_id): + with cp.cuda.Device(device_id), streams[device_id]: streams[device_id].synchronize() num_permuted = cp.zeros(n_features, dtype=dtype) @@ -324,8 +344,8 @@ def _run_permutations_sparse( with cp.cuda.Device(dd["device_id"]): streams[dd["device_id"]].synchronize() - # Phase 3: Gather results on first device and cut to exact size - with cp.cuda.Device(device_ids[0]): + # Phase 3: Gather results where the input lives and cut to exact size. + with cp.cuda.Device(source_device): all_perms = [cp.asarray(dd["perms"]) for dd in device_data] morans_I_permutations = cp.concatenate(all_perms, axis=0)[:n_permutations] diff --git a/src/rapids_singlecell/tools/_rank_genes_groups/_stream_multi_gpu.py b/src/rapids_singlecell/tools/_rank_genes_groups/_stream_multi_gpu.py index 6445dc187..0c13918d8 100644 --- a/src/rapids_singlecell/tools/_rank_genes_groups/_stream_multi_gpu.py +++ b/src/rapids_singlecell/tools/_rank_genes_groups/_stream_multi_gpu.py @@ -17,7 +17,7 @@ import scipy.sparse as sp from rapids_singlecell._cuda import _rank_stream_cuda as _rss -from rapids_singlecell._utils import parse_device_ids +from rapids_singlecell._utils import parse_device_ids, validate_multi_gpu from ._wilcoxon_host import _copy_gpu_array_to_device @@ -152,8 +152,10 @@ def stream_planes_multi( comp_pts = rg.comp_pts col_shard = _is_col_shard(X) axis_len = n_genes if col_shard else n_cells + home = cp.cuda.Device().id + device_ids = device_ids[: min(len(device_ids), axis_len)] + device_ids = validate_multi_gpu(device_ids, source_device=home, gather_device=home) bands = _bands(axis_len, len(device_ids)) - device_ids = device_ids[: len(bands)] def run_shard(index: int): device_id = device_ids[index] @@ -168,11 +170,14 @@ def run_shard(index: int): cp.cuda.runtime.deviceSynchronize() return out - with ThreadPoolExecutor(max_workers=len(device_ids)) as executor: - shards = list(executor.map(run_shard, range(len(device_ids)))) + if len(device_ids) == 1: + shards = [run_shard(0)] + else: + with ThreadPoolExecutor(max_workers=len(device_ids)) as executor: + shards = list(executor.map(run_shard, range(len(device_ids)))) # Gather onto the caller's device so downstream stats math stays local. - dev0 = cp.cuda.Device().id + dev0 = home if col_shard: sums = _concat_to_device([s[0] for s in shards], dev0, axis=1) sqsums = _concat_to_device([s[1] for s in shards], dev0, axis=1) @@ -220,8 +225,10 @@ def run_binned_hist_multi( raise ValueError("invalid multi-GPU histogram gene window") col_shard = _is_col_shard(X) axis_len = stop - start if col_shard else n_cells + home = cp.cuda.Device().id + device_ids = device_ids[: min(len(device_ids), axis_len)] + device_ids = validate_multi_gpu(device_ids, source_device=home, gather_device=home) bands = _bands(axis_len, len(device_ids)) - device_ids = device_ids[: len(bands)] def run_shard(index: int): device_id = device_ids[index] @@ -266,10 +273,12 @@ def run_shard(index: int): cp.cuda.runtime.deviceSynchronize() return hist, gsum, gnnz - with ThreadPoolExecutor(max_workers=len(device_ids)) as executor: - shards = list(executor.map(run_shard, range(len(device_ids)))) + if len(device_ids) == 1: + shards = [run_shard(0)] + else: + with ThreadPoolExecutor(max_workers=len(device_ids)) as executor: + shards = list(executor.map(run_shard, range(len(device_ids)))) - home = cp.cuda.Device().id if col_shard: hist = _concat_to_device([s[0] for s in shards], home, axis=0) gsum = ( diff --git a/src/rapids_singlecell/tools/_rank_genes_groups/_wilcoxon_host.py b/src/rapids_singlecell/tools/_rank_genes_groups/_wilcoxon_host.py index d654de97d..ad37b7bbb 100644 --- a/src/rapids_singlecell/tools/_rank_genes_groups/_wilcoxon_host.py +++ b/src/rapids_singlecell/tools/_rank_genes_groups/_wilcoxon_host.py @@ -13,7 +13,7 @@ from rapids_singlecell._cuda import _wilcoxon_cuda as _wc from rapids_singlecell._cuda import _wilcoxon_sparse_cuda as _wcs -from rapids_singlecell._utils import parse_device_ids +from rapids_singlecell._utils import parse_device_ids, validate_multi_gpu if TYPE_CHECKING: from collections.abc import Callable @@ -400,7 +400,7 @@ def _prepare_device_column_shards(X, ranges, device_ids): def _concat_gpu_shards(arrays: list[cp.ndarray], device_id: int) -> cp.ndarray: if len(arrays) == 1: - return arrays[0] + return _copy_gpu_array_to_device(arrays[0], device_id) local_arrays = [_copy_gpu_array_to_device(array, device_id) for array in arrays] result = cp.concatenate(local_arrays, axis=1) # Keep peer-copy buffers alive until concatenation finishes. @@ -429,6 +429,9 @@ def _run_sharded_wilcoxon( if cpsp.issparse(X) else None ) + transfer_source = ( + source_device if source_device is not None else cp.cuda.Device().id + ) auto_single_device = multi_gpu is None and is_device_input and rg.ireference is None if multi_gpu is False or auto_single_device: device_ids = [ @@ -436,13 +439,25 @@ def _run_sharded_wilcoxon( ] else: device_ids = list(dict.fromkeys(parse_device_ids(multi_gpu=multi_gpu))) - device_ids.sort(key=lambda device_id: device_id != source_device) + device_ids.sort(key=lambda device_id: device_id != transfer_source) ranges = _split_gene_ranges( X, n_devices=len(device_ids), dense_fallback=rg._sparse_negative_fallback, ) device_ids = device_ids[: len(ranges)] + validated_device_ids = validate_multi_gpu( + device_ids, + source_device=transfer_source, + gather_device=transfer_source, + ) + if validated_device_ids != device_ids: + device_ids = validated_device_ids + ranges = _split_gene_ranges( + X, + n_devices=1, + dense_fallback=rg._sparse_negative_fallback, + ) ovo_host_context = ( _build_ovo_host_context(rg) if rg.ireference is not None else None ) @@ -548,7 +563,7 @@ def run_shard( ): msg = "Inconsistent sharded Wilcoxon group ordering." raise RuntimeError(msg) - result_device = device_ids[0] + result_device = transfer_source with cp.cuda.Device(result_device): scores = _concat_gpu_shards( [result[1] for result in complete_gpu_results], result_device diff --git a/tests/pertpy/test_distances.py b/tests/pertpy/test_distances.py index 19dddb432..ceeaa9998 100644 --- a/tests/pertpy/test_distances.py +++ b/tests/pertpy/test_distances.py @@ -7,6 +7,7 @@ from anndata import AnnData from scipy.spatial.distance import cdist +from rapids_singlecell._utils import MultiGPUFallbackWarning, _multi_gpu from rapids_singlecell.pertpy_gpu import Distance, MeanVar @@ -1909,6 +1910,53 @@ def test_multi_gpu_pairwise_matches_single_gpu() -> None: ) +@pytest.mark.skipif(not _has_multiple_gpus(), reason="Requires 2+ GPUs") +def test_multi_gpu_failed_preflight_falls_back_on_input_device(monkeypatch) -> None: + caller_device = cp.cuda.Device().id + source_device = next( + device_id + for device_id in range(cp.cuda.runtime.getDeviceCount()) + if device_id != caller_device + ) + rng = np.random.default_rng(42) + cpu_embedding = rng.normal(size=(60, 8)).astype(np.float32) + obs = pd.DataFrame( + { + "group": pd.Categorical( + [f"g{index}" for index in range(3) for _ in range(20)] + ) + } + ) + adata = AnnData(cpu_embedding.copy(), obs=obs) + with cp.cuda.Device(source_device): + adata.obsm["X_pca"] = cp.asarray(cpu_embedding) + + distance = Distance(metric="edistance") + expected = distance.pairwise( + adata, + groupby="group", + multi_gpu=[source_device], + ) + + monkeypatch.setattr(_multi_gpu, "peer_copy_works", lambda *_: False) + _multi_gpu._WARNED_P2P_FAILURES.clear() + try: + with pytest.warns( + MultiGPUFallbackWarning, + match=rf"Falling back to GPU {source_device}", + ): + actual = distance.pairwise( + adata, + groupby="group", + multi_gpu=[caller_device, source_device], + ) + finally: + _multi_gpu._WARNED_P2P_FAILURES.clear() + + assert cp.cuda.Device().id == caller_device + np.testing.assert_allclose(actual.values, expected.values, rtol=1e-7, atol=1e-7) + + @pytest.mark.skipif(not _has_multiple_gpus(), reason="Requires 2+ GPUs") def test_multi_gpu_onesided_matches_single_gpu() -> None: """Test that multi-GPU onesided_distances produces same results as single-GPU.""" diff --git a/tests/test_multi_gpu_utils.py b/tests/test_multi_gpu_utils.py index 0f85c77b2..e6470677a 100644 --- a/tests/test_multi_gpu_utils.py +++ b/tests/test_multi_gpu_utils.py @@ -2,9 +2,19 @@ from __future__ import annotations +import warnings +from contextlib import nullcontext + import cupy as cp +import pytest -from rapids_singlecell._utils import _split_pairs, parse_device_ids +from rapids_singlecell._utils import ( + MultiGPUFallbackWarning, + _multi_gpu, + _split_pairs, + parse_device_ids, + validate_multi_gpu, +) class TestSplitPairs: @@ -274,3 +284,182 @@ def test_single_gpu_always_device_0(self): assert parse_device_ids(multi_gpu=None) == [0] assert parse_device_ids(multi_gpu=True) == [0] assert parse_device_ids(multi_gpu=False) == [0] + + +class TestValidateMultiGPU: + @pytest.fixture(autouse=True) + def _clear_validation_state(self): + cached_peer_copy_works = _multi_gpu.peer_copy_works + cached_peer_copy_works.cache_clear() + _multi_gpu._WARNED_P2P_FAILURES.clear() + yield + cached_peer_copy_works.cache_clear() + _multi_gpu._WARNED_P2P_FAILURES.clear() + + def test_local_device_skips_peer_check(self, monkeypatch): + monkeypatch.setattr(cp.cuda.runtime, "getDeviceCount", lambda: 2) + + def unexpected_peer_check(*_): + raise AssertionError("local execution must not run a P2P canary") + + monkeypatch.setattr(_multi_gpu, "peer_copy_works", unexpected_peer_check) + + assert validate_multi_gpu([1], source_device=1) == [1] + + def test_invalid_device_fails_before_peer_check(self, monkeypatch): + monkeypatch.setattr(cp.cuda.runtime, "getDeviceCount", lambda: 2) + + def unexpected_peer_check(*_): + raise AssertionError("invalid IDs must fail before P2P validation") + + monkeypatch.setattr(_multi_gpu, "peer_copy_works", unexpected_peer_check) + + with pytest.raises(ValueError, match=r"Invalid GPU device ID.*2"): + validate_multi_gpu([0, 2], source_device=0) + + def test_validates_fanout_and_gather_directions(self, monkeypatch): + monkeypatch.setattr(cp.cuda.runtime, "getDeviceCount", lambda: 3) + checked = [] + + def peer_copy_works(destination, source): + checked.append((destination, source)) + return True + + monkeypatch.setattr(_multi_gpu, "peer_copy_works", peer_copy_works) + + assert validate_multi_gpu([0, 1, 2], source_device=1) == [0, 1, 2] + assert set(checked) == {(0, 1), (0, 2), (2, 1)} + + @pytest.mark.parametrize("failed_pair", [(0, 1), (1, 0)]) + def test_failed_direction_warns_and_falls_back_once(self, monkeypatch, failed_pair): + monkeypatch.setattr(cp.cuda.runtime, "getDeviceCount", lambda: 2) + monkeypatch.setattr( + _multi_gpu, + "peer_copy_works", + lambda destination, source: (destination, source) != failed_pair, + ) + + with pytest.warns(MultiGPUFallbackWarning, match="Falling back to GPU 0"): + assert validate_multi_gpu([0, 1], source_device=0) == [0] + + with warnings.catch_warnings(record=True) as warnings_record: + warnings.simplefilter("always") + assert validate_multi_gpu([0, 1], source_device=0) == [0] + assert not warnings_record + + def test_single_remote_target_is_not_assumed_safe(self, monkeypatch): + monkeypatch.setattr(cp.cuda.runtime, "getDeviceCount", lambda: 2) + monkeypatch.setattr(_multi_gpu, "peer_copy_works", lambda *_: False) + + with pytest.warns(MultiGPUFallbackWarning): + assert validate_multi_gpu([1], source_device=0) == [0] + + def test_stops_after_first_failed_pair(self, monkeypatch): + monkeypatch.setattr(cp.cuda.runtime, "getDeviceCount", lambda: 3) + checked = [] + + def fail_first(destination, source): + checked.append((destination, source)) + return False + + monkeypatch.setattr(_multi_gpu, "peer_copy_works", fail_first) + + with pytest.warns(MultiGPUFallbackWarning): + assert validate_multi_gpu([0, 1, 2], source_device=0) == [0] + assert checked == [(0, 1)] + + def test_unexpected_cuda_error_propagates(self, monkeypatch): + monkeypatch.setattr(cp.cuda.runtime, "getDeviceCount", lambda: 2) + + def raise_illegal_address(*_): + raise cp.cuda.runtime.CUDARuntimeError(700) + + monkeypatch.setattr(_multi_gpu, "peer_copy_works", raise_illegal_address) + + with pytest.raises(cp.cuda.runtime.CUDARuntimeError) as error: + validate_multi_gpu([0, 1], source_device=0) + assert error.value.status == 700 + + +class TestPeerCopyWorks: + @pytest.fixture(autouse=True) + def _clear_peer_cache(self): + cached_peer_copy_works = _multi_gpu.peer_copy_works + cached_peer_copy_works.cache_clear() + yield + cached_peer_copy_works.cache_clear() + + def test_canary_rejects_silent_success(self, monkeypatch): + monkeypatch.setattr(cp.cuda.runtime, "memcpyPeer", lambda *_: None) + + assert not _multi_gpu._run_peer_copy_canary(0, 0) + + def test_capability_false_skips_enable_and_canary(self, monkeypatch): + monkeypatch.setattr(cp.cuda, "Device", lambda *_: nullcontext()) + monkeypatch.setattr(cp.cuda.runtime, "deviceCanAccessPeer", lambda *_: False) + + def unexpected(*_): + raise AssertionError("unsupported links must not be exercised") + + monkeypatch.setattr(cp.cuda.runtime, "deviceEnablePeerAccess", unexpected) + monkeypatch.setattr(_multi_gpu, "_run_peer_copy_canary", unexpected) + + assert not _multi_gpu.peer_copy_works(1, 0) + + @pytest.mark.parametrize("status", [217, 705, 711]) + def test_expected_enable_error_falls_back(self, monkeypatch, status): + monkeypatch.setattr(cp.cuda, "Device", lambda *_: nullcontext()) + monkeypatch.setattr(cp.cuda.runtime, "deviceCanAccessPeer", lambda *_: True) + + def raise_expected(*_): + raise cp.cuda.runtime.CUDARuntimeError(status) + + monkeypatch.setattr(cp.cuda.runtime, "deviceEnablePeerAccess", raise_expected) + monkeypatch.setattr( + _multi_gpu, + "_run_peer_copy_canary", + lambda *_: pytest.fail("safe enable failures must skip the canary"), + ) + + assert not _multi_gpu.peer_copy_works(1, 0) + + def test_already_enabled_runs_canary_and_caches_result(self, monkeypatch): + monkeypatch.setattr(cp.cuda, "Device", lambda *_: nullcontext()) + monkeypatch.setattr(cp.cuda.runtime, "deviceCanAccessPeer", lambda *_: True) + + def raise_already_enabled(*_): + raise cp.cuda.runtime.CUDARuntimeError(704) + + monkeypatch.setattr( + cp.cuda.runtime, + "deviceEnablePeerAccess", + raise_already_enabled, + ) + canary_calls = [] + + def successful_canary(destination, source): + canary_calls.append((destination, source)) + return True + + monkeypatch.setattr(_multi_gpu, "_run_peer_copy_canary", successful_canary) + + assert _multi_gpu.peer_copy_works(1, 0) + assert _multi_gpu.peer_copy_works(1, 0) + assert canary_calls == [(1, 0)] + + def test_unexpected_enable_error_propagates(self, monkeypatch): + monkeypatch.setattr(cp.cuda, "Device", lambda *_: nullcontext()) + monkeypatch.setattr(cp.cuda.runtime, "deviceCanAccessPeer", lambda *_: True) + + def raise_illegal_address(*_): + raise cp.cuda.runtime.CUDARuntimeError(700) + + monkeypatch.setattr( + cp.cuda.runtime, + "deviceEnablePeerAccess", + raise_illegal_address, + ) + + with pytest.raises(cp.cuda.runtime.CUDARuntimeError) as error: + _multi_gpu.peer_copy_works(1, 0) + assert error.value.status == 700 diff --git a/tests/test_rank_genes_groups_wilcoxon.py b/tests/test_rank_genes_groups_wilcoxon.py index ae93b1584..60dc9be34 100644 --- a/tests/test_rank_genes_groups_wilcoxon.py +++ b/tests/test_rank_genes_groups_wilcoxon.py @@ -2387,6 +2387,56 @@ def parse_device_ids_spy(*, multi_gpu): assert "scores" in adata.uns["rank_genes_groups"] +def test_wilcoxon_failed_preflight_repartitions_and_stays_serial(monkeypatch): + source_device = cp.cuda.Device().id + fake_peer = source_device + 1 + single = _make_multi_gpu_wilcoxon_adata("cupy_dense", source_device=source_device) + fallback = _make_multi_gpu_wilcoxon_adata("cupy_dense", source_device=source_device) + kwargs = { + "method": "wilcoxon", + "use_raw": False, + "reference": "rest", + "n_genes": single.n_vars, + } + rsc.tl.rank_genes_groups(single, "group", multi_gpu=False, **kwargs) + + monkeypatch.setattr( + _wilcoxon_host, + "parse_device_ids", + lambda *, multi_gpu: [source_device, fake_peer], + ) + + def force_fallback(device_ids, *, source_device, gather_device): + assert device_ids == [source_device, fake_peer] + assert gather_device == source_device + return [source_device] + + monkeypatch.setattr(_wilcoxon_host, "validate_multi_gpu", force_fallback) + real_split = _wilcoxon_host._split_gene_ranges + split_device_counts = [] + + def split_spy(X, *, n_devices, dense_fallback): + split_device_counts.append(n_devices) + return real_split( + X, + n_devices=n_devices, + dense_fallback=dense_fallback, + ) + + monkeypatch.setattr(_wilcoxon_host, "_split_gene_ranges", split_spy) + + class UnexpectedExecutor: + def __init__(self, *_args, **_kwargs): + raise AssertionError("single-GPU fallback must not start a thread pool") + + monkeypatch.setattr(_wilcoxon_host, "ThreadPoolExecutor", UnexpectedExecutor) + + rsc.tl.rank_genes_groups(fallback, "group", multi_gpu=True, **kwargs) + + assert split_device_counts == [2, 1] + _assert_multi_gpu_wilcoxon_equal(fallback, single) + + @pytest.mark.skipif(not MULTI_GPU_AVAILABLE, reason="requires at least two GPUs") @pytest.mark.parametrize("fmt", ["cupy_dense", "cupy_csr", "cupy_csc"]) @pytest.mark.parametrize("reference", ["rest", "1"]) From d42242b3e633eac69a391eec9f84ea730351dedf Mon Sep 17 00:00:00 2001 From: Intron7 Date: Wed, 26 Aug 2026 18:42:06 +0200 Subject: [PATCH 2/2] make this easy Signed-off-by: Intron7 --- docs/release-notes/0.17.0.md | 1 + src/rapids_singlecell/_utils/__init__.py | 12 +- src/rapids_singlecell/_utils/_multi_gpu.py | 237 +++--------- .../pertpy_gpu/_metrics/_base_metric.py | 22 +- .../pertpy_gpu/_metrics/_edistance.py | 345 +++++------------- .../pertpy_gpu/_metrics/_wasserstein.py | 106 ++---- .../squidpy_gpu/_autocorr.py | 81 ++-- src/rapids_singlecell/squidpy_gpu/_co_oc.py | 89 ++--- src/rapids_singlecell/squidpy_gpu/_gearysc.py | 60 ++- src/rapids_singlecell/squidpy_gpu/_moransi.py | 62 ++-- .../_rank_genes_groups/_stream_multi_gpu.py | 35 +- .../_rank_genes_groups/_wilcoxon_host.py | 83 +---- tests/pertpy/test_distances.py | 48 --- tests/test_multi_gpu_utils.py | 219 +++-------- tests/test_rank_genes_groups_wilcoxon.py | 116 ------ 15 files changed, 361 insertions(+), 1155 deletions(-) diff --git a/docs/release-notes/0.17.0.md b/docs/release-notes/0.17.0.md index 79123e59b..35a6daff2 100644 --- a/docs/release-notes/0.17.0.md +++ b/docs/release-notes/0.17.0.md @@ -12,6 +12,7 @@ ```{rubric} Features ``` +* Route multi-GPU copies through host memory when P2P validation fails. {pr}`771` {smaller}`S Dicks` * Split {func}`~rapids_singlecell.gr.calculate_niche` into {func}`~rapids_singlecell.gr.calculate_niche_neighborhood`, {func}`~rapids_singlecell.gr.calculate_niche_utag` and {func}`~rapids_singlecell.gr.calculate_niche_cellcharter`, with ``mask``, ``library_key`` and cross-flavor ``min_niche_size``, following {mod}`squidpy` {pr}`758` {smaller}`S Dicks` * Speed up {func}`~rapids_singlecell.pp.harmony_integrate` and make it reproducible by seeding k-means from a deterministic `float64` fit on a bounded, batch-stratified random subsample instead of a non-reproducible `float32` fit over all cells. `dtype` now defaults to `numpy.float32` {pr}`756` {smaller}`S Dicks` * Derive unset {func}`~rapids_singlecell.pp.harmony_integrate` stopping rules from `flavor`: `harmony2` follows Harmony2 defaults, `harmony1` still follows harmony-pytorch {pr}`756` {smaller}`S Dicks` diff --git a/src/rapids_singlecell/_utils/__init__.py b/src/rapids_singlecell/_utils/__init__.py index 0c5487079..5c50d76b6 100644 --- a/src/rapids_singlecell/_utils/__init__.py +++ b/src/rapids_singlecell/_utils/__init__.py @@ -8,29 +8,21 @@ from dask.array import Array as DaskArray from ._multi_gpu import ( - MultiGPUFallbackWarning, _calculate_blocks_per_pair, - _copy_to_device_via_host, + _copy_to_device, _create_category_index_mapping, _get_device_attrs, _split_pairs, parse_device_ids, - peer_copy_verified, - peer_copy_works, - validate_multi_gpu, ) __all__ = [ - "MultiGPUFallbackWarning", "_calculate_blocks_per_pair", - "_copy_to_device_via_host", + "_copy_to_device", "_create_category_index_mapping", "_get_device_attrs", "_split_pairs", "parse_device_ids", - "peer_copy_verified", - "peer_copy_works", - "validate_multi_gpu", ] diff --git a/src/rapids_singlecell/_utils/_multi_gpu.py b/src/rapids_singlecell/_utils/_multi_gpu.py index 7799cf42e..6e26929b2 100644 --- a/src/rapids_singlecell/_utils/_multi_gpu.py +++ b/src/rapids_singlecell/_utils/_multi_gpu.py @@ -12,7 +12,6 @@ from __future__ import annotations -import warnings from functools import cache import cupy as cp @@ -21,210 +20,64 @@ # Cache for device attributes per device (lazy initialization) _DEVICE_ATTRS_CACHE: dict[int, dict] = {} - -CUDA_ERROR_PEER_ACCESS_UNSUPPORTED = 217 -CUDA_ERROR_PEER_ACCESS_ALREADY_ENABLED = 704 -CUDA_ERROR_PEER_ACCESS_NOT_ENABLED = 705 -CUDA_ERROR_TOO_MANY_PEERS = 711 - -_SAFE_PEER_FALLBACK_ERRORS = { - CUDA_ERROR_PEER_ACCESS_UNSUPPORTED, - CUDA_ERROR_PEER_ACCESS_NOT_ENABLED, - CUDA_ERROR_TOO_MANY_PEERS, +_CANARY = np.arange(1, 9, dtype=np.float64) +_CANARY_POISON = -_CANARY +_CUDA_ERROR_PEER_ACCESS_UNSUPPORTED = 217 +_CUDA_ERROR_PEER_ACCESS_NOT_ENABLED = 705 +_CUDA_ERROR_TOO_MANY_PEERS = 711 +_PEER_ERRORS = { + _CUDA_ERROR_PEER_ACCESS_UNSUPPORTED, + _CUDA_ERROR_PEER_ACCESS_NOT_ENABLED, + _CUDA_ERROR_TOO_MANY_PEERS, } -_P2P_CANARY_EXPECTED = np.arange(1, 9, dtype=np.float64) -_P2P_CANARY_POISON = np.full(8, -7.0, dtype=np.float64) -_WARNED_P2P_FAILURES: set[tuple[int, tuple[tuple[int, int], ...]]] = set() - - -class MultiGPUFallbackWarning(RuntimeWarning): - """Warning emitted when unsafe P2P makes an operation use one GPU.""" - - -def _copy_to_device_via_host(array: cp.ndarray, device_id: int) -> cp.ndarray: - """Copy a device array without touching a peer link. - - This is intentionally reserved for small control arrays whose owner may - differ from the main input's device. Capturing the source before changing - devices is essential: calling ``cp.asarray`` first would already attempt - the peer transfer this helper exists to avoid. - """ - source_device = array.device.id - if source_device == device_id: - return array - - with cp.cuda.Device(source_device): - host = array.get() - with cp.cuda.Device(device_id): - copied = cp.asarray(host) - cp.cuda.get_current_stream().synchronize() - return copied - - -def _run_peer_copy_canary(destination: int, source: int) -> bool: - """Return whether a small peer copy arrives intact. - - The explicit synchronizations make the canary independent of whichever - CuPy stream was current when the multi-GPU operation was entered. - """ - with cp.cuda.Device(source): - expected = cp.asarray(_P2P_CANARY_EXPECTED) - cp.cuda.get_current_stream().synchronize() - - with cp.cuda.Device(destination): - probe = cp.asarray(_P2P_CANARY_POISON) - cp.cuda.get_current_stream().synchronize() - try: - cp.cuda.runtime.memcpyPeer( - probe.data.ptr, - destination, - expected.data.ptr, - source, - expected.nbytes, - ) - cp.cuda.runtime.deviceSynchronize() - arrived = cp.asnumpy(probe) - except cp.cuda.runtime.CUDARuntimeError as error: - if error.status in _SAFE_PEER_FALLBACK_ERRORS: - return False - # memcpy/readback can surface an earlier asynchronous failure. In - # particular, an illegal address poisons the CUDA context and - # cannot be repaired by switching to one GPU in this process. - raise - - return bool(np.array_equal(arrived, _P2P_CANARY_EXPECTED)) @cache -def peer_copy_works(destination: int, source: int) -> bool: - """Return whether ``source -> destination`` P2P transfers are usable. - - ``deviceCanAccessPeer`` is only a capability query. Some affected systems - report support and return success from ``memcpyPeer`` while leaving the - destination unchanged, so this function verifies the transferred bytes. - The result is cached per ordered device pair for the lifetime of the - process. - """ +def _peer_copy_works(destination: int, source: int) -> bool: + """Return whether a peer copy arrives intact.""" if destination == source: return True - - with cp.cuda.Device(destination): - if not cp.cuda.runtime.deviceCanAccessPeer(destination, source): + if not cp.cuda.runtime.deviceCanAccessPeer(destination, source): + return False + + try: + with cp.cuda.Device(source): + expected = cp.asarray(_CANARY, blocking=True) + with cp.cuda.Device(destination): + actual = cp.asarray(_CANARY_POISON, blocking=True) + with cp.cuda.Stream(non_blocking=True) as stream: + cp.copyto(actual, expected) + stream.synchronize() + actual = cp.asnumpy(actual) + except cp.cuda.runtime.CUDARuntimeError as error: + if error.status in _PEER_ERRORS: return False - try: - cp.cuda.runtime.deviceEnablePeerAccess(source) - except cp.cuda.runtime.CUDARuntimeError as error: - if error.status == CUDA_ERROR_PEER_ACCESS_ALREADY_ENABLED: - pass - elif error.status in _SAFE_PEER_FALLBACK_ERRORS: - return False - else: - raise - - return _run_peer_copy_canary(destination, source) - - -# Compatibility name used by the container-level P2P diagnostics. -peer_copy_verified = peer_copy_works - - -def validate_multi_gpu( - device_ids: list[int], - *, - source_device: int | None = None, - gather_device: int | None = None, -) -> list[int]: - """Return usable devices, falling back to one GPU when P2P is unsafe. + raise + return bool(np.array_equal(actual, _CANARY)) - Multi-GPU implementations in this package fan input buffers out from a - source GPU and gather results onto a selected GPU. This validates exactly - those ordered transfers before sharding or worker threads start. - If a required link is unavailable or silently corrupts the canary, the - operation falls back to the source GPU and emits one warning per failed - link set. Unexpected CUDA errors propagate because they may indicate an - already-poisoned context. - Parameters - ---------- - device_ids - Requested execution devices. - source_device - Device holding shared input buffers. Defaults to the current device. - gather_device - Device receiving worker results. Defaults to the first requested - device. - - Returns - ------- - list[int] - The requested devices when all required P2P directions work, otherwise - a one-element list containing ``source_device``. - """ - if not device_ids: - raise ValueError("device_ids must contain at least one device") +def _copy_to_device_p2p(array: cp.ndarray, destination: int) -> cp.ndarray: + """Copy an array directly to another GPU.""" + with cp.cuda.Device(destination): + return cp.asarray(array) - device_ids = list(dict.fromkeys(device_ids)) - if source_device is None: - source_device = cp.cuda.Device().id - n_available = cp.cuda.runtime.getDeviceCount() - invalid_ids = [ - device_id - for device_id in device_ids - if device_id < 0 or device_id >= n_available - ] - if invalid_ids: - raise ValueError( - f"Invalid GPU device ID(s): {invalid_ids}. " - f"Available devices: {list(range(n_available))}" - ) - if source_device < 0 or source_device >= n_available: - raise ValueError( - f"Invalid source GPU device ID {source_device}. " - f"Available devices: {list(range(n_available))}" - ) - - if gather_device is None: - gather_device = device_ids[0] - if gather_device < 0 or gather_device >= n_available: - raise ValueError( - f"Invalid gather GPU device ID {gather_device}. " - f"Available devices: {list(range(n_available))}" - ) +def _copy_to_device_via_host(array: cp.ndarray, destination: int) -> cp.ndarray: + """Copy an array to another GPU through host memory.""" + with cp.cuda.Device(array.device.id): + host = array.get(order="A") + with cp.cuda.Device(destination): + return cp.asarray(host, blocking=True) - required_pairs: set[tuple[int, int]] = set() - for device_id in device_ids: - if device_id != source_device: - required_pairs.add((device_id, source_device)) - if device_id != gather_device: - required_pairs.add((gather_device, device_id)) - - failed_pairs = () - for pair in sorted(required_pairs): - if not peer_copy_works(*pair): - # Once fallback is required, do not exercise any more suspect - # links in the caller's CUDA context. - failed_pairs = (pair,) - break - if not failed_pairs: - return device_ids - - warning_key = (source_device, failed_pairs) - if warning_key not in _WARNED_P2P_FAILURES: - transfers = ", ".join( - f"GPU {source} -> GPU {destination}" for destination, source in failed_pairs - ) - warnings.warn( - "Multi-GPU execution was disabled because the required P2P " - f"transfer(s) failed validation: {transfers}. Falling back to " - f"GPU {source_device} for this operation. This avoids silent " - "result corruption but may be slower.", - MultiGPUFallbackWarning, - stacklevel=2, - ) - _WARNED_P2P_FAILURES.add(warning_key) - return [source_device] +def _copy_to_device(array: cp.ndarray, destination: int) -> cp.ndarray: + """Copy an array using P2P when it works, otherwise through the host.""" + source = array.device.id + if source == destination: + return array + if _peer_copy_works(destination, source): + return _copy_to_device_p2p(array, destination) + return _copy_to_device_via_host(array, destination) def parse_device_ids(*, multi_gpu: bool | list[int] | str | None) -> list[int]: diff --git a/src/rapids_singlecell/pertpy_gpu/_metrics/_base_metric.py b/src/rapids_singlecell/pertpy_gpu/_metrics/_base_metric.py index caf332ea8..0c1fa1f05 100644 --- a/src/rapids_singlecell/pertpy_gpu/_metrics/_base_metric.py +++ b/src/rapids_singlecell/pertpy_gpu/_metrics/_base_metric.py @@ -7,11 +7,7 @@ import numpy as np from rapids_singlecell._keys import _preset_obsm_names, _resolve_obsm_key -from rapids_singlecell._utils import ( - _copy_to_device_via_host, - _create_category_index_mapping, - parse_device_ids, -) +from rapids_singlecell._utils import _create_category_index_mapping, parse_device_ids from rapids_singlecell.squidpy_gpu._utils import _assert_categorical_obs if TYPE_CHECKING: @@ -123,21 +119,13 @@ def _subset_to_groups( Ordered group names matching the category indices. """ embedding_raw = self._get_embedding(adata) - source_device = ( - embedding_raw.device.id - if isinstance(embedding_raw, cp.ndarray) - else cp.cuda.Device().id - ) mask, cat_offsets, cell_indices, groups_list = self._subset_indices( adata, groupby, needed_groups ) - with cp.cuda.Device(source_device): - if mask is not None: - embedding = cp.asarray(embedding_raw[mask]) - else: - embedding = cp.asarray(embedding_raw) - cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) - cell_indices = _copy_to_device_via_host(cell_indices, source_device) + if mask is not None: + embedding = cp.asarray(embedding_raw[mask]) + else: + embedding = cp.asarray(embedding_raw) return embedding, cat_offsets, cell_indices, groups_list def _subset_indices( diff --git a/src/rapids_singlecell/pertpy_gpu/_metrics/_edistance.py b/src/rapids_singlecell/pertpy_gpu/_metrics/_edistance.py index 4e6736b9d..e34edfb20 100644 --- a/src/rapids_singlecell/pertpy_gpu/_metrics/_edistance.py +++ b/src/rapids_singlecell/pertpy_gpu/_metrics/_edistance.py @@ -14,9 +14,8 @@ from rapids_singlecell._cuda import _edistance_cuda as _ed from rapids_singlecell._utils import ( _calculate_blocks_per_pair, - _copy_to_device_via_host, + _copy_to_device, _split_pairs, - validate_multi_gpu, ) from rapids_singlecell.squidpy_gpu._utils import _assert_categorical_obs @@ -107,19 +106,11 @@ def _materialize_source(embedding_raw, selector): ``selector`` is a boolean mask, an integer row-index array, or ``None``. """ - source_device = ( - embedding_raw.device.id - if isinstance(embedding_raw, cp.ndarray) - else embedding_raw.data.device.id - if cpsp.issparse(embedding_raw) - else cp.cuda.Device().id - ) - with cp.cuda.Device(source_device): - if _is_sparse(embedding_raw): - return _build_csr_source(embedding_raw, selector) - if selector is None: - return cp.asarray(embedding_raw) - return cp.asarray(embedding_raw[selector]) + if _is_sparse(embedding_raw): + return _build_csr_source(embedding_raw, selector) + if selector is None: + return cp.asarray(embedding_raw) + return cp.asarray(embedding_raw[selector]) class EDistanceMetric(BaseMetric): @@ -172,11 +163,6 @@ def _load_source( adata, groupby, needed_groups ) source = _materialize_source(embedding_raw, mask) - source_device = ( - source.data.device.id if isinstance(source, _CSRData) else source.device.id - ) - cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) - cell_indices = _copy_to_device_via_host(cell_indices, source_device) return source, cat_offsets, cell_indices, groups_list def pairwise( @@ -320,11 +306,6 @@ def onesided_distances( embedding, cat_offsets, cell_indices, groups_list = self._load_source( adata, groupby, needed ) - source_device = ( - embedding.data.device.id - if isinstance(embedding, _CSRData) - else embedding.device.id - ) k = len(groups_list) group_map = {v: i for i, v in enumerate(groups_list)} selected_indices = [group_map[sg] for sg in selected_groups] @@ -347,15 +328,14 @@ def onesided_distances( # e[s,b] = 2*d[s,b] - d[s,s] - d[b,b] ed_cols = {} var_cols = {} - with cp.cuda.Device(source_device): - for i, (sg, si) in enumerate(zip(selected_groups, selected_indices)): - ed_row = 2 * cross_mean[i, :] - diag_mean[si] - diag_mean - ed_row[si] = 0.0 - ed_cols[sg] = ed_row.get() + for i, (sg, si) in enumerate(zip(selected_groups, selected_indices)): + ed_row = 2 * cross_mean[i, :] - diag_mean[si] - diag_mean + ed_row[si] = 0.0 + ed_cols[sg] = ed_row.get() - var_row = 4 * cross_var[i, :] + diag_var[si] + diag_var - var_row[si] = 0.0 - var_cols[sg] = var_row.get() + var_row = 4 * cross_var[i, :] + diag_var[si] + diag_var + var_row[si] = 0.0 + var_cols[sg] = var_row.get() distances = pd.DataFrame(ed_cols, index=groups_list) distances.index.name = groupby @@ -385,11 +365,10 @@ def onesided_distances( # cross_means[i, j] = mean dist from selected[i] to group j # diag_means[j] = mean within-group dist for group j ed_cols = {} - with cp.cuda.Device(source_device): - for i, (sg, si) in enumerate(zip(selected_groups, selected_indices)): - ed_row = 2 * cross_means[i, :] - diag_means[si] - diag_means - ed_row[si] = 0.0 - ed_cols[sg] = ed_row.get() + for i, (sg, si) in enumerate(zip(selected_groups, selected_indices)): + ed_row = 2 * cross_means[i, :] - diag_means[si] - diag_means + ed_row[si] = 0.0 + ed_cols[sg] = ed_row.get() df = pd.DataFrame(ed_cols, index=groups_list) df.index.name = groupby @@ -548,22 +527,14 @@ def contrast_distances( embedding = _materialize_source(embedding_raw, original_indices) cell_indices = cp.arange(len(original_indices), dtype=cp.int32) elif len(original_indices) < int(len(embedding_raw) * 0.7): - embedding = _materialize_source(embedding_raw, original_indices) + embedding = cp.asarray(embedding_raw[original_indices]) cell_indices = cp.arange(len(original_indices), dtype=cp.int32) else: - embedding = _materialize_source(embedding_raw, None) + embedding = cp.asarray(embedding_raw) cell_indices = cp.array(original_indices, dtype=cp.int32) - source_device = ( - embedding.data.device.id - if isinstance(embedding, _CSRData) - else embedding.device.id - ) - cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) - cell_indices = _copy_to_device_via_host(cell_indices, source_device) - with cp.cuda.Device(source_device): - group_sizes = cp.diff(cat_offsets).astype(cp.int64) - group_sizes_cpu = group_sizes.get() + group_sizes = cp.diff(cat_offsets).astype(cp.int64) + group_sizes_cpu = group_sizes.get() # Build deduplicated pairs pair_to_flat: dict[tuple[int, int], int] = {} for idx_a, idx_b in contrast_pairs: @@ -583,9 +554,8 @@ def contrast_distances( return result pairs = sorted(pair_to_flat.keys(), key=lambda p: pair_to_flat[p]) - with cp.cuda.Device(source_device): - pair_left = cp.array([p[0] for p in pairs], dtype=cp.int32) - pair_right = cp.array([p[1] for p in pairs], dtype=cp.int32) + pair_left = cp.array([p[0] for p in pairs], dtype=cp.int32) + pair_right = cp.array([p[1] for p in pairs], dtype=cp.int32) flat_sums = self._launch_distance_kernel( embedding, @@ -596,18 +566,17 @@ def contrast_distances( device_ids=device_ids, ) - with cp.cuda.Device(source_device): - # Vectorized normalization - is_diag = pair_left == pair_right - sizes_l = group_sizes[pair_left.astype(cp.intp)] - sizes_r = group_sizes[pair_right.astype(cp.intp)] - flat_norms = cp.where( - is_diag, - cp.maximum(sizes_l * (sizes_l - 1) // 2, 1), - sizes_l * sizes_r, - ).astype(embedding.dtype) - flat_means = flat_sums / flat_norms - flat_means_cpu = flat_means.get() + # Vectorized normalization + is_diag = pair_left == pair_right + sizes_l = group_sizes[pair_left.astype(cp.intp)] + sizes_r = group_sizes[pair_right.astype(cp.intp)] + flat_norms = cp.where( + is_diag, + cp.maximum(sizes_l * (sizes_l - 1) // 2, 1), + sizes_l * sizes_r, + ).astype(embedding.dtype) + flat_means = flat_sums / flat_norms + flat_means_cpu = flat_means.get() # Extract edistances edistances = np.empty(len(contrast_pairs), dtype=np.float64) @@ -771,36 +740,6 @@ def _launch_distance_kernel( pair_left: cp.ndarray, pair_right: cp.ndarray, device_ids: list[int], - ) -> cp.ndarray: - """Run distribution, launch, and gather from the input's device.""" - source_device = ( - embedding.data.device.id - if isinstance(embedding, _CSRData) - else embedding.device.id - ) - cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) - cell_indices = _copy_to_device_via_host(cell_indices, source_device) - pair_left = _copy_to_device_via_host(pair_left, source_device) - pair_right = _copy_to_device_via_host(pair_right, source_device) - with cp.cuda.Device(source_device): - return self._launch_distance_kernel_on_source( - embedding, - cat_offsets, - cell_indices, - pair_left=pair_left, - pair_right=pair_right, - device_ids=device_ids, - ) - - def _launch_distance_kernel_on_source( - self, - embedding: cp.ndarray, - cat_offsets: cp.ndarray, - cell_indices: cp.ndarray, - *, - pair_left: cp.ndarray, - pair_right: cp.ndarray, - device_ids: list[int], ) -> cp.ndarray: """Launch the edistance kernel across GPUs and return raw flat sums. @@ -827,33 +766,13 @@ def _launch_distance_kernel_on_source( cp.ndarray Raw distance sums of shape (n_pairs,), NOT normalized. """ - source_device = ( - embedding.data.device.id - if isinstance(embedding, _CSRData) - else embedding.device.id - ) - # Control arrays may have been created on the caller's current device - # even when a device-resident embedding lives elsewhere. Stage these - # small arrays through host memory so fallback never depends on P2P. - cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) - cell_indices = _copy_to_device_via_host(cell_indices, source_device) - pair_left = _copy_to_device_via_host(pair_left, source_device) - pair_right = _copy_to_device_via_host(pair_right, source_device) - device_ids = list(dict.fromkeys(device_ids)) - device_ids.sort(key=lambda device_id: device_id != source_device) - device_ids = validate_multi_gpu( - device_ids, - source_device=source_device, - gather_device=source_device, - ) - with cp.cuda.Device(source_device): - cp.cuda.get_current_stream().synchronize() n_devices = len(device_ids) n_total_pairs = len(pair_left) _, n_features = embedding.shape group_sizes = cp.diff(cat_offsets).astype(cp.int64) is_sparse = isinstance(embedding, _CSRData) + result_device = cat_offsets.device.id # Split pairs across devices with load balancing pair_chunks = _split_pairs(pair_left, pair_right, n_devices, group_sizes) @@ -884,20 +803,20 @@ def _launch_distance_kernel_on_source( with streams[device_id]: data = { - "off": cp.asarray(cat_offsets), - "idx": cp.asarray(cell_indices), - "pair_left": cp.asarray(chunk_left), - "pair_right": cp.asarray(chunk_right), + "off": _copy_to_device(cat_offsets, device_id), + "idx": _copy_to_device(cell_indices, device_id), + "pair_left": _copy_to_device(chunk_left, device_id), + "pair_right": _copy_to_device(chunk_right, device_id), "sums": cp.zeros(n_chunk_pairs, dtype=embedding.dtype), "n_pairs": n_chunk_pairs, "device_id": device_id, } if is_sparse: - data["data"] = cp.asarray(embedding.data) - data["indices"] = cp.asarray(embedding.indices) - data["indptr"] = cp.asarray(embedding.indptr) + data["data"] = _copy_to_device(embedding.data, device_id) + data["indices"] = _copy_to_device(embedding.indices, device_id) + data["indptr"] = _copy_to_device(embedding.indptr, device_id) else: - data["emb"] = cp.asarray(embedding) + data["emb"] = _copy_to_device(embedding, device_id) device_data.append(data) # Phase 2: Synchronize data transfers, then launch kernels @@ -935,7 +854,7 @@ def _launch_distance_kernel_on_source( feat_tile, block_size, shared_mem, - streams[device_id].ptr, + cp.cuda.get_current_stream().ptr, ) else: _ed.compute_distances( @@ -952,21 +871,21 @@ def _launch_distance_kernel_on_source( feat_tile, block_size, shared_mem, - streams[device_id].ptr, + cp.cuda.get_current_stream().ptr, ) # Phase 3: Synchronize all devices for data in device_data: if data is not None: with cp.cuda.Device(data["device_id"]): - streams[data["device_id"]].synchronize() + cp.cuda.Stream.null.synchronize() - # Phase 4: Aggregate where the input lives. - with cp.cuda.Device(source_device): + # Phase 4: Aggregate on the input device + with cp.cuda.Device(result_device): total_sums = cp.zeros(n_total_pairs, dtype=embedding.dtype) for i, data in enumerate(device_data): if data is not None: - sums = cp.asarray(data["sums"]) + sums = _copy_to_device(data["sums"], result_device) start = chunk_offsets[i] total_sums[start : start + len(sums)] = sums @@ -979,27 +898,6 @@ def _pairwise_means( cell_indices: cp.ndarray, k: int, device_ids: list[int], - ) -> cp.ndarray: - """Run pairwise reconstruction on the embedding's owning device.""" - source_device = ( - embedding.data.device.id - if isinstance(embedding, _CSRData) - else embedding.device.id - ) - cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) - cell_indices = _copy_to_device_via_host(cell_indices, source_device) - with cp.cuda.Device(source_device): - return self._pairwise_means_on_source( - embedding, cat_offsets, cell_indices, k, device_ids - ) - - def _pairwise_means_on_source( - self, - embedding: cp.ndarray, - cat_offsets: cp.ndarray, - cell_indices: cp.ndarray, - k: int, - device_ids: list[int], ) -> cp.ndarray: """Compute between-group mean distances for all group pairs. @@ -1073,34 +971,6 @@ def _onesided_means( *, selected_indices: list[int], device_ids: list[int], - ) -> tuple[cp.ndarray, cp.ndarray]: - """Run one-sided reconstruction on the embedding's owning device.""" - source_device = ( - embedding.data.device.id - if isinstance(embedding, _CSRData) - else embedding.device.id - ) - cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) - cell_indices = _copy_to_device_via_host(cell_indices, source_device) - with cp.cuda.Device(source_device): - return self._onesided_means_on_source( - embedding, - cat_offsets, - cell_indices, - k, - selected_indices=selected_indices, - device_ids=device_ids, - ) - - def _onesided_means_on_source( - self, - embedding: cp.ndarray, - cat_offsets: cp.ndarray, - cell_indices: cp.ndarray, - k: int, - *, - selected_indices: list[int], - device_ids: list[int], ) -> tuple[cp.ndarray, cp.ndarray]: """Compute mean distances from selected group(s) to all groups. @@ -1241,27 +1111,19 @@ def _pairwise_means_bootstrap( tuple (means, variances) matrices (k x k each) """ - source_device = ( - embedding.data.device.id - if isinstance(embedding, _CSRData) - else embedding.device.id - ) - cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) - cell_indices = _copy_to_device_via_host(cell_indices, source_device) - with cp.cuda.Device(source_device): - group_sizes = cp.diff(cat_offsets) + # Get group sizes for bootstrap sampling (on GPU 0) + group_sizes = cp.diff(cat_offsets) # Run bootstrap iterations - each uses all GPUs for pairwise computation all_results = [] for i in range(n_bootstrap): # Generate bootstrap sample on GPU 0 - with cp.cuda.Device(source_device): - boot_cat_offsets, boot_cell_indices = self._bootstrap_sample_cells( - cat_offsets=cat_offsets, - cell_indices=cell_indices, - group_sizes_gpu=group_sizes, - seed=random_state + i, - ) + boot_cat_offsets, boot_cell_indices = self._bootstrap_sample_cells( + cat_offsets=cat_offsets, + cell_indices=cell_indices, + group_sizes_gpu=group_sizes, + seed=random_state + i, + ) # Compute pairwise means using all GPUs pairwise_means = self._pairwise_means( @@ -1273,8 +1135,8 @@ def _pairwise_means_bootstrap( ) all_results.append(pairwise_means.get()) - # Keep the returned statistics with the embedding. - with cp.cuda.Device(source_device): + # Compute statistics on the input device + with cp.cuda.Device(cat_offsets.device.id): bootstrap_stack = cp.array(all_results) # [n_bootstrap, k, k] means = cp.mean(bootstrap_stack, axis=0) variances = cp.var(bootstrap_stack, axis=0) @@ -1327,28 +1189,20 @@ def _onesided_means_bootstrap( diag_var Variance of bootstrap diag_means, shape (k,) """ - source_device = ( - embedding.data.device.id - if isinstance(embedding, _CSRData) - else embedding.device.id - ) - cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) - cell_indices = _copy_to_device_via_host(cell_indices, source_device) - with cp.cuda.Device(source_device): - group_sizes = cp.diff(cat_offsets) + # Get group sizes for bootstrap sampling (on GPU 0) + group_sizes = cp.diff(cat_offsets) # Run bootstrap iterations - each uses all GPUs for onesided computation all_cross = [] all_diag = [] for i in range(n_bootstrap): # Generate bootstrap sample on GPU 0 - with cp.cuda.Device(source_device): - boot_cat_offsets, boot_cell_indices = self._bootstrap_sample_cells( - cat_offsets=cat_offsets, - cell_indices=cell_indices, - group_sizes_gpu=group_sizes, - seed=random_state + i, - ) + boot_cat_offsets, boot_cell_indices = self._bootstrap_sample_cells( + cat_offsets=cat_offsets, + cell_indices=cell_indices, + group_sizes_gpu=group_sizes, + seed=random_state + i, + ) # Compute onesided means using all GPUs cross_means, diag_means = self._onesided_means( @@ -1362,8 +1216,8 @@ def _onesided_means_bootstrap( all_cross.append(cross_means.get()) all_diag.append(diag_means.get()) - # Keep the returned statistics with the embedding. - with cp.cuda.Device(source_device): + # Compute statistics on the input device + with cp.cuda.Device(cat_offsets.device.id): cross_stack = cp.array(all_cross) diag_stack = cp.array(all_diag) cross_mean = cp.mean(cross_stack, axis=0) @@ -1454,36 +1308,32 @@ def _prepare_edistance_df_bootstrap( random_state=random_state, device_ids=device_ids, ) - source_device = ( - embedding.data.device.id - if isinstance(embedding, _CSRData) - else embedding.device.id - ) - with cp.cuda.Device(source_device): - # Vectorized edistance: e[a,b] = 2*d[a,b] - d[a,a] - d[b,b] - diag_means = cp.diag(pairwise_means_boot) - edistance_means = ( - 2 * pairwise_means_boot - diag_means[:, None] - diag_means[None, :] - ) - cp.fill_diagonal(edistance_means, 0) + # Vectorized edistance: e[a,b] = 2*d[a,b] - d[a,a] - d[b,b] + diag_means = cp.diag(pairwise_means_boot) + edistance_means = ( + 2 * pairwise_means_boot - diag_means[:, None] - diag_means[None, :] + ) + cp.fill_diagonal(edistance_means, 0) - # Vectorized variance computation (delta method approximation): - # var[a, b] = 4 * var[a, b] + var[a, a] + var[b, b] - diag_vars = cp.diag(pairwise_vars_boot) - edistance_vars = ( - 4 * pairwise_vars_boot + diag_vars[:, None] + diag_vars[None, :] - ) - cp.fill_diagonal(edistance_vars, 0) - means_host = edistance_means.get() - vars_host = edistance_vars.get() + # Vectorized variance computation (delta method approximation): + # var[a, b] = 4 * var[a, b] + var[a, a] + var[b, b] + diag_vars = cp.diag(pairwise_vars_boot) + edistance_vars = ( + 4 * pairwise_vars_boot + diag_vars[:, None] + diag_vars[None, :] + ) + cp.fill_diagonal(edistance_vars, 0) - df_mean = pd.DataFrame(means_host, index=groups_list, columns=groups_list) + df_mean = pd.DataFrame( + edistance_means.get(), index=groups_list, columns=groups_list + ) df_mean.index.name = groupby df_mean.columns.name = groupby df_mean.name = "pairwise edistance" - df_var = pd.DataFrame(vars_host, index=groups_list, columns=groups_list) + df_var = pd.DataFrame( + edistance_vars.get(), index=groups_list, columns=groups_list + ) df_var.index.name = groupby df_var.columns.name = groupby df_var.name = "pairwise edistance variance" @@ -1506,20 +1356,15 @@ def _prepare_edistance_df( pairwise_means = self._pairwise_means( embedding, cat_offsets, cell_indices, k, device_ids ) - source_device = ( - embedding.data.device.id - if isinstance(embedding, _CSRData) - else embedding.device.id - ) - with cp.cuda.Device(source_device): - # Vectorized edistance: e[a,b] = 2*d[a,b] - d[a,a] - d[b,b] - diag = cp.diag(pairwise_means) - edistance_matrix = 2 * pairwise_means - diag[:, None] - diag[None, :] - cp.fill_diagonal(edistance_matrix, 0) # Self-distance is 0 - edistance_host = edistance_matrix.get() + # Vectorized edistance: e[a,b] = 2*d[a,b] - d[a,a] - d[b,b] + diag = cp.diag(pairwise_means) + edistance_matrix = 2 * pairwise_means - diag[:, None] - diag[None, :] + cp.fill_diagonal(edistance_matrix, 0) # Self-distance is 0 - df = pd.DataFrame(edistance_host, index=groups_list, columns=groups_list) + df = pd.DataFrame( + edistance_matrix.get(), index=groups_list, columns=groups_list + ) df.index.name = groupby df.columns.name = groupby df.name = "pairwise edistance" diff --git a/src/rapids_singlecell/pertpy_gpu/_metrics/_wasserstein.py b/src/rapids_singlecell/pertpy_gpu/_metrics/_wasserstein.py index dccd3a66f..afd9c0ecf 100644 --- a/src/rapids_singlecell/pertpy_gpu/_metrics/_wasserstein.py +++ b/src/rapids_singlecell/pertpy_gpu/_metrics/_wasserstein.py @@ -14,7 +14,7 @@ import pandas as pd from rapids_singlecell._cuda import _sinkhorn_cuda as _sk -from rapids_singlecell._utils import _copy_to_device_via_host, validate_multi_gpu +from rapids_singlecell._utils import _copy_to_device from rapids_singlecell.squidpy_gpu._utils import _assert_categorical_obs from ._base_metric import BaseMetric, parse_device_ids @@ -328,27 +328,14 @@ def _solve_pairs( device_ids = [0] n_pairs = len(pair_left) if n_pairs == 0: - with cp.cuda.Device(embedding.device.id): - return cp.zeros(0, dtype=dtype) - - source_device = embedding.device.id - cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) - cell_indices = _copy_to_device_via_host(cell_indices, source_device) - device_ids = list(dict.fromkeys(device_ids)) - device_ids.sort(key=lambda device_id: device_id != source_device) - device_ids = validate_multi_gpu( - device_ids, - source_device=source_device, - gather_device=source_device, - ) + return cp.zeros(0, dtype=dtype) pl_host = np.asarray(pair_left, dtype=np.int32) pr_host = np.asarray(pair_right, dtype=np.int32) - with cp.cuda.Device(source_device): - group_sizes = (cat_offsets[1:] - cat_offsets[:-1]).astype(cp.int32) - # Group sizes to host (one sync) -> all batch planning is host-only. - n_left = group_sizes[cp.asarray(pl_host)].get() - n_right = group_sizes[cp.asarray(pr_host)].get() + group_sizes = (cat_offsets[1:] - cat_offsets[:-1]).astype(cp.int32) + # Group sizes to host (one sync) -> all batch planning is host-only. + n_left = group_sizes[cp.asarray(pl_host)].get() + n_right = group_sizes[cp.asarray(pr_host)].get() # Orient larger group as columns. swap = n_left > n_right rows = np.where(swap, pr_host, pl_host) @@ -359,10 +346,8 @@ def _solve_pairs( plans = _plan_device_batches(n_row, n_col, itemsize, len(device_ids)) - home = source_device - with cp.cuda.Device(home): - cp.cuda.get_current_stream().synchronize() - out = cp.empty(n_pairs, dtype=dtype) + out = cp.empty(n_pairs, dtype=dtype) + output_device = out.device.id # Move the shared inputs to each participating device once. streams: dict[int, cp.cuda.Stream] = {} @@ -374,9 +359,9 @@ def _solve_pairs( streams[dev] = cp.cuda.Stream(non_blocking=True) with streams[dev]: inputs[dev] = ( - cp.ascontiguousarray(cp.asarray(embedding)), - cp.asarray(cat_offsets), - cp.asarray(cell_indices), + cp.ascontiguousarray(_copy_to_device(embedding, dev)), + _copy_to_device(cat_offsets, dev), + _copy_to_device(cell_indices, dev), ) # Grow-only per-device cost buffers, reused across rounds (avoids a fresh @@ -454,9 +439,11 @@ def _solve_pairs( for u in units: with cp.cuda.Device(u["dev"]): u["stream"].synchronize() - with cp.cuda.Device(home): + with cp.cuda.Device(output_device): for u in units: - out[u["start"] : u["stop"]] = cp.asarray(u["reg"]) + out[u["start"] : u["stop"]] = _copy_to_device( + u["reg"], output_device + ) for u in units: with cp.cuda.Device(u["dev"]): converged = converged and bool(u["state"]["conv"].all().get()) @@ -490,23 +477,12 @@ def _bootstrap_solve( raise ValueError(f"n_bootstrap must be >= 1, got {n_bootstrap}") n_pairs = len(pair_left) if n_pairs == 0: - with cp.cuda.Device(embedding.device.id): - empty = cp.zeros(0, dtype=dtype) - return empty, empty - source_device = embedding.device.id - cat_offsets = _copy_to_device_via_host(cat_offsets, source_device) - cell_indices = _copy_to_device_via_host(cell_indices, source_device) - device = validate_multi_gpu( - [device], - source_device=source_device, - gather_device=source_device, - )[0] - with cp.cuda.Device(source_device): - cp.cuda.get_current_stream().synchronize() + empty = cp.zeros(0, dtype=dtype) + return empty, empty with cp.cuda.Device(device): - emb = cp.ascontiguousarray(cp.asarray(embedding)) - offs = cp.asarray(cat_offsets) - cidx = cp.asarray(cell_indices) + emb = cp.ascontiguousarray(_copy_to_device(embedding, device)) + offs = _copy_to_device(cat_offsets, device) + cidx = _copy_to_device(cell_indices, device) # Sizes/orientation on the host so the per-chunk build never syncs. co_h = cp.asnumpy(offs) sizes_h = np.diff(co_h) @@ -558,12 +534,7 @@ def _bootstrap_solve( if not converged: self._warn_not_converged() reg = reg.reshape(n_pairs, n_bootstrap) - mean = reg.mean(axis=1) - var = reg.var(axis=1) - return ( - _copy_to_device_via_host(mean, source_device), - _copy_to_device_via_host(var, source_device), - ) + return reg.mean(axis=1), reg.var(axis=1) def bootstrap_arrays( self, @@ -663,15 +634,14 @@ def pairwise( pair_right.append(j) def _to_matrix(flat: cp.ndarray, name: str) -> pd.DataFrame: - with cp.cuda.Device(embedding.device.id): - mat = cp.zeros((k, k), dtype=embedding.dtype) - if pair_left: - il = cp.asarray(pair_left, dtype=cp.intp) - jr = cp.asarray(pair_right, dtype=cp.intp) - mat[il, jr] = flat - mat[jr, il] = flat - mat_host = mat.get() - df = pd.DataFrame(mat_host, index=groups_list, columns=groups_list) + mat = cp.zeros((k, k), dtype=embedding.dtype) + if pair_left: + il = cp.asarray(pair_left, dtype=cp.intp) + jr = cp.asarray(pair_right, dtype=cp.intp) + flat = _copy_to_device(flat, mat.device.id) + mat[il, jr] = flat + mat[jr, il] = flat + df = pd.DataFrame(mat.get(), index=groups_list, columns=groups_list) df.index.name = groupby df.columns.name = groupby df.name = name @@ -744,8 +714,7 @@ def onesided_distances( pair_right.append(j) def _to_df(flat: cp.ndarray) -> pd.DataFrame: - with cp.cuda.Device(embedding.device.id): - flat_cpu = flat.get() + flat_cpu = flat.get() ed_cols: dict[str, np.ndarray] = {} cursor = 0 for ii, si in enumerate(selected_indices): @@ -802,11 +771,6 @@ def contrast_distances( device_ids = parse_device_ids(multi_gpu=multi_gpu) groupby, split_by = self._parse_contrasts(adata, contrasts) embedding_raw = self._get_embedding(adata) - source_device = ( - embedding_raw.device.id - if isinstance(embedding_raw, cp.ndarray) - else cp.cuda.Device().id - ) all_cols = [groupby, *split_by] grouped = adata.obs.groupby(all_cols, observed=True) @@ -846,10 +810,9 @@ def contrast_distances( flat_cell_idx = ( np.concatenate(all_cells) if all_cells else np.array([], dtype=np.int64) ) - with cp.cuda.Device(source_device): - cat_offsets = cp.asarray(offsets_host, dtype=cp.int32) - cell_indices = cp.asarray(flat_cell_idx, dtype=cp.int32) - embedding = cp.asarray(embedding_raw) + cat_offsets = cp.asarray(offsets_host, dtype=cp.int32) + cell_indices = cp.asarray(flat_cell_idx, dtype=cp.int32) + embedding = cp.asarray(embedding_raw) dtype = embedding.dtype # Deduplicate canonical pairs (i, j) with i < j @@ -874,8 +837,7 @@ def contrast_distances( dtype=dtype, device_ids=device_ids, ) - with cp.cuda.Device(embedding.device.id): - flat_cpu = flat.get() + flat_cpu = flat.get() distances = np.empty(len(contrast_pairs), dtype=np.float64) for n, (idx_a, idx_b) in enumerate(contrast_pairs): diff --git a/src/rapids_singlecell/squidpy_gpu/_autocorr.py b/src/rapids_singlecell/squidpy_gpu/_autocorr.py index eff2a7f88..76b2a3f3a 100644 --- a/src/rapids_singlecell/squidpy_gpu/_autocorr.py +++ b/src/rapids_singlecell/squidpy_gpu/_autocorr.py @@ -31,28 +31,21 @@ def _to_cupy(vals, *, use_sparse: bool, dtype): Dense input is always returned as a dense CuPy array. """ is_sparse = sparse.issparse(vals) or sparse_gpu.isspmatrix(vals) - source_device = ( - vals.device.id - if isinstance(vals, cp.ndarray) - else vals.data.device.id - if sparse_gpu.isspmatrix(vals) - else cp.cuda.Device().id - ) - with cp.cuda.Device(source_device): - # Dense input - use_sparse is ignored - if not is_sparse: - return cp.array(vals, dtype=dtype, order="C") - - # Sparse input - respect use_sparse parameter - if use_sparse: - if sparse_gpu.isspmatrix(vals): - return vals.tocsr().astype(dtype) - return sparse_gpu.csr_matrix(vals.tocsr(), dtype=dtype) - - # Sparse input but use_sparse=False - convert to dense - if not sparse_gpu.isspmatrix(vals): - vals = sparse_gpu.csr_matrix(vals.tocsr(), dtype=dtype) - return _sparse_to_dense(vals, order="C") + + # Dense input - use_sparse is ignored + if not is_sparse: + return cp.array(vals, dtype=dtype, order="C") + + # Sparse input - respect use_sparse parameter + if use_sparse: + if sparse_gpu.isspmatrix(vals): + return vals.tocsr().astype(dtype) + return sparse_gpu.csr_matrix(vals.tocsr(), dtype=dtype) + + # Sparse input but use_sparse=False - convert to dense + if not sparse_gpu.isspmatrix(vals): + vals = sparse_gpu.csr_matrix(vals.tocsr(), dtype=dtype) + return _sparse_to_dense(vals, order="C") def spatial_autocorr( @@ -150,23 +143,15 @@ def spatial_autocorr( if compute_dtype not in (np.float32, np.float64): compute_dtype = np.float32 - source_device = ( - vals.device.id - if isinstance(vals, cp.ndarray) - else vals.data.device.id - if sparse_gpu.isspmatrix(vals) - else cp.cuda.Device().id - ) - with cp.cuda.Device(source_device): - # create Adj-Matrix alongside the data - adj_matrix = adata.obsp[connectivity_key] - adj_matrix_cupy = sparse_gpu.csr_matrix(adj_matrix, dtype=compute_dtype) - - if transformation: # row-normalize - row_sums = adj_matrix_cupy.sum(axis=1).reshape(-1, 1) - non_zero_rows = row_sums != 0 - row_sums[non_zero_rows] = 1.0 / row_sums[non_zero_rows] - adj_matrix_cupy = adj_matrix_cupy.multiply(sparse_gpu.csr_matrix(row_sums)) + # create Adj-Matrix + adj_matrix = adata.obsp[connectivity_key] + adj_matrix_cupy = sparse_gpu.csr_matrix(adj_matrix, dtype=compute_dtype) + + if transformation: # row-normalize + row_sums = adj_matrix_cupy.sum(axis=1).reshape(-1, 1) + non_zero_rows = row_sums != 0 + row_sums[non_zero_rows] = 1.0 / row_sums[non_zero_rows] + adj_matrix_cupy = adj_matrix_cupy.multiply(sparse_gpu.csr_matrix(row_sums)) params = {"two_tailed": two_tailed} @@ -183,13 +168,10 @@ def _run_autocorr(data, adj_matrix_cupy, mode, n_perms, multi_gpu): else: raise ValueError(f"Invalid mode: {mode}") - with cp.cuda.Device(source_device): - data = _to_cupy(vals, use_sparse=use_sparse, dtype=compute_dtype) + data = _to_cupy(vals, use_sparse=use_sparse, dtype=compute_dtype) - # Run full computation - score, score_perms = _run_autocorr( - data, adj_matrix_cupy, mode, n_perms, multi_gpu - ) + # Run full computation + score, score_perms = _run_autocorr(data, adj_matrix_cupy, mode, n_perms, multi_gpu) # Set mode-specific params if mode == "moran": @@ -203,11 +185,10 @@ def _run_autocorr(data, adj_matrix_cupy, mode, n_perms, multi_gpu): params["ascending"] = True params["mode"] = "gearyC" - with cp.cuda.Device(source_device): - g = sparse.csr_matrix(adj_matrix_cupy.get()) - score = score.get() - if n_perms is not None: - score_perms = score_perms.get() + g = sparse.csr_matrix(adj_matrix_cupy.get()) + score = score.get() + if n_perms is not None: + score_perms = score_perms.get() with np.errstate(divide="ignore"): pval_results = _p_value_calc(score, sims=score_perms, weights=g, params=params) diff --git a/src/rapids_singlecell/squidpy_gpu/_co_oc.py b/src/rapids_singlecell/squidpy_gpu/_co_oc.py index b4505ad9d..8d6066549 100644 --- a/src/rapids_singlecell/squidpy_gpu/_co_oc.py +++ b/src/rapids_singlecell/squidpy_gpu/_co_oc.py @@ -9,11 +9,10 @@ from rapids_singlecell._cuda import _cooc_cuda as _co from rapids_singlecell._utils import ( _calculate_blocks_per_pair, - _copy_to_device_via_host, + _copy_to_device, _create_category_index_mapping, _split_pairs, parse_device_ids, - validate_multi_gpu, ) from ._utils import _assert_categorical_obs, _assert_spatial_basis @@ -69,37 +68,26 @@ def co_occurrence( _assert_categorical_obs(adata, key=cluster_key) _assert_spatial_basis(adata, key=spatial_key) - spatial_input = adata.obsm[spatial_key] - source_device = ( - spatial_input.device.id - if isinstance(spatial_input, cp.ndarray) - else cp.cuda.Device().id - ) - with cp.cuda.Device(source_device): - spatial = cp.array(spatial_input).astype(np.float32) - original_clust = adata.obs[cluster_key] - clust_map = {v: i for i, v in enumerate(original_clust.cat.categories.values)} - labs = cp.array([clust_map[c] for c in original_clust], dtype=np.int32) - # create intervals thresholds - if isinstance(interval, int): - thresh_min, thresh_max = _find_min_max(spatial) - interval = cp.linspace( - thresh_min, thresh_max, num=interval, dtype=np.float32 - ) - else: - if isinstance(interval, cp.ndarray): - interval = _copy_to_device_via_host(interval, source_device) - interval = cp.array(sorted(interval), dtype=np.float32, copy=True) - if len(interval) <= 1: - raise ValueError( - f"Expected interval to be of length `>= 2`, found `{len(interval)}`." - ) - - device_ids = parse_device_ids(multi_gpu=multi_gpu) - out = _co_occurrence_helper( - spatial, interval, labs, fast=True, device_ids=device_ids + spatial = cp.array(adata.obsm[spatial_key]).astype(np.float32) + original_clust = adata.obs[cluster_key] + clust_map = {v: i for i, v in enumerate(original_clust.cat.categories.values)} + labs = cp.array([clust_map[c] for c in original_clust], dtype=np.int32) + # create intervals thresholds + if isinstance(interval, int): + thresh_min, thresh_max = _find_min_max(spatial) + interval = cp.linspace(thresh_min, thresh_max, num=interval, dtype=np.float32) + else: + interval = cp.array(sorted(interval), dtype=np.float32, copy=True) + if len(interval) <= 1: + raise ValueError( + f"Expected interval to be of length `>= 2`, found `{len(interval)}`." ) - out, interval = out.get(), interval.get() + + device_ids = parse_device_ids(multi_gpu=multi_gpu) + out = _co_occurrence_helper( + spatial, interval, labs, fast=True, device_ids=device_ids + ) + out, interval = out.get(), interval.get() if copy: return out, interval @@ -305,23 +293,9 @@ def _co_occurrence_gpu( if not valid_device_ids: return cp.zeros((k, k, l_val), dtype=cp.uint64), False - source_device_id = spatial.device.id - thresholds = _copy_to_device_via_host(thresholds, source_device_id) - cat_offsets = _copy_to_device_via_host(cat_offsets, source_device_id) - cell_indices = _copy_to_device_via_host(cell_indices, source_device_id) - pair_left = _copy_to_device_via_host(pair_left, source_device_id) - pair_right = _copy_to_device_via_host(pair_right, source_device_id) - valid_device_ids.sort(key=lambda device_id: device_id != source_device_id) - device_ids = validate_multi_gpu( - valid_device_ids, - source_device=source_device_id, - gather_device=source_device_id, - ) - if any(device_id not in kernel_configs for device_id in device_ids): - return cp.zeros((k, k, l_val), dtype=cp.uint64), False - with cp.cuda.Device(source_device_id): - cp.cuda.get_current_stream().synchronize() + device_ids = valid_device_ids n_devices = len(device_ids) + source_device_id = spatial.device.id # Split pairs across devices with load balancing group_sizes = cp.diff(cat_offsets).astype(cp.int64) @@ -349,14 +323,14 @@ def _co_occurrence_gpu( dev_cat_offsets = cat_offsets dev_cell_indices = cell_indices else: - dev_spatial = cp.asarray(spatial) - dev_thresholds = cp.asarray(thresholds) - dev_cat_offsets = cp.asarray(cat_offsets) - dev_cell_indices = cp.asarray(cell_indices) + dev_spatial = _copy_to_device(spatial, device_id) + dev_thresholds = _copy_to_device(thresholds, device_id) + dev_cat_offsets = _copy_to_device(cat_offsets, device_id) + dev_cell_indices = _copy_to_device(cell_indices, device_id) # Copy pair indices to this device - dev_pair_left = cp.asarray(chunk_left) - dev_pair_right = cp.asarray(chunk_right) + dev_pair_left = _copy_to_device(chunk_left, device_id) + dev_pair_right = _copy_to_device(chunk_right, device_id) # Initialize local counts array dev_counts = cp.zeros((k, k, l_val), dtype=cp.uint64) @@ -405,7 +379,7 @@ def _co_occurrence_gpu( cell_tile=cell_tile, block_size=block_size, shared_mem=shared_mem, - stream=streams[device_id].ptr, + stream=cp.cuda.get_current_stream().ptr, ) # Phase 3: Synchronize all devices (wait for kernels to complete) @@ -414,12 +388,11 @@ def _co_occurrence_gpu( with cp.cuda.Device(data["device_id"]): streams[data["device_id"]].synchronize() - # Phase 4: Aggregate counts where the input lives. + # Phase 4: Aggregate counts on the input device with cp.cuda.Device(source_device_id): counts = cp.zeros((k, k, l_val), dtype=cp.uint64) for data in device_data: if data is not None: - dev0_counts = cp.asarray(data["counts"]) - counts += dev0_counts + counts += _copy_to_device(data["counts"], source_device_id) return counts, True diff --git a/src/rapids_singlecell/squidpy_gpu/_gearysc.py b/src/rapids_singlecell/squidpy_gpu/_gearysc.py index 49d41d142..8312b4a2b 100644 --- a/src/rapids_singlecell/squidpy_gpu/_gearysc.py +++ b/src/rapids_singlecell/squidpy_gpu/_gearysc.py @@ -6,7 +6,7 @@ from cupyx.scipy import sparse from rapids_singlecell._cuda import _autocorr_cuda as _ac -from rapids_singlecell._utils import parse_device_ids, validate_multi_gpu +from rapids_singlecell._utils import _copy_to_device, parse_device_ids from ._utils import _check_precision_issues @@ -80,22 +80,13 @@ def _run_permutations_dense( if device_ids is None: device_ids = [0] - source_device = data.device.id - device_ids = list(dict.fromkeys(device_ids)) - device_ids.sort(key=lambda device_id: device_id != source_device) - device_ids = validate_multi_gpu( - device_ids, - source_device=source_device, - gather_device=source_device, - ) - with cp.cuda.Device(source_device): - cp.cuda.get_current_stream().synchronize() n_devices = len(device_ids) streams: dict[int, cp.cuda.Stream] = {} device_data: list[dict] = [] # Each device runs perms_per_device iterations perms_per_device = (n_permutations + n_devices - 1) // n_devices + source_device = data.device.id # Phase 1: Create streams and transfer data to all devices for device_id in device_ids: @@ -109,16 +100,16 @@ def _run_permutations_dense( dev_adj = adj_matrix_cupy dev_den = den else: - dev_data = cp.asarray(data) + dev_data = _copy_to_device(data, device_id) dev_adj = sparse.csr_matrix( ( - cp.asarray(adj_matrix_cupy.data), - cp.asarray(adj_matrix_cupy.indices), - cp.asarray(adj_matrix_cupy.indptr), + _copy_to_device(adj_matrix_cupy.data, device_id), + _copy_to_device(adj_matrix_cupy.indices, device_id), + _copy_to_device(adj_matrix_cupy.indptr, device_id), ), shape=adj_matrix_cupy.shape, ) - dev_den = cp.asarray(den) + dev_den = _copy_to_device(den, device_id) # Allocate output array for this device dev_perms = cp.zeros((perms_per_device, n_features), dtype=dtype) @@ -137,7 +128,7 @@ def _run_permutations_dense( for p in range(perms_per_device): for dd in device_data: device_id = dd["device_id"] - with cp.cuda.Device(device_id), streams[device_id]: + with cp.cuda.Device(device_id): streams[device_id].synchronize() num_permuted = cp.zeros(n_features, dtype=dtype) @@ -161,9 +152,9 @@ def _run_permutations_dense( with cp.cuda.Device(dd["device_id"]): streams[dd["device_id"]].synchronize() - # Phase 3: Gather results where the input lives and cut to exact size. + # Phase 3: Gather results on the input device and cut to exact size with cp.cuda.Device(source_device): - all_perms = [cp.asarray(dd["perms"]) for dd in device_data] + all_perms = [_copy_to_device(dd["perms"], source_device) for dd in device_data] gearys_C_permutations = cp.concatenate(all_perms, axis=0)[:n_permutations] return gearys_C_permutations @@ -250,22 +241,13 @@ def _run_permutations_sparse( if device_ids is None: device_ids = [0] - source_device = data.data.device.id - device_ids = list(dict.fromkeys(device_ids)) - device_ids.sort(key=lambda device_id: device_id != source_device) - device_ids = validate_multi_gpu( - device_ids, - source_device=source_device, - gather_device=source_device, - ) - with cp.cuda.Device(source_device): - cp.cuda.get_current_stream().synchronize() n_devices = len(device_ids) streams: dict[int, cp.cuda.Stream] = {} device_data: list[dict] = [] # Each device runs perms_per_device iterations perms_per_device = (n_permutations + n_devices - 1) // n_devices + source_device = data.data.device.id # Phase 1: Create streams and transfer data to all devices for device_id in device_ids: @@ -281,21 +263,21 @@ def _run_permutations_sparse( else: dev_data = sparse.csr_matrix( ( - cp.asarray(data.data), - cp.asarray(data.indices), - cp.asarray(data.indptr), + _copy_to_device(data.data, device_id), + _copy_to_device(data.indices, device_id), + _copy_to_device(data.indptr, device_id), ), shape=data.shape, ) dev_adj = sparse.csr_matrix( ( - cp.asarray(adj_matrix_cupy.data), - cp.asarray(adj_matrix_cupy.indices), - cp.asarray(adj_matrix_cupy.indptr), + _copy_to_device(adj_matrix_cupy.data, device_id), + _copy_to_device(adj_matrix_cupy.indices, device_id), + _copy_to_device(adj_matrix_cupy.indptr, device_id), ), shape=adj_matrix_cupy.shape, ) - dev_den = cp.asarray(den) + dev_den = _copy_to_device(den, device_id) # Allocate output array for this device dev_perms = cp.zeros((perms_per_device, n_features), dtype=dtype) @@ -314,7 +296,7 @@ def _run_permutations_sparse( for p in range(perms_per_device): for dd in device_data: device_id = dd["device_id"] - with cp.cuda.Device(device_id), streams[device_id]: + with cp.cuda.Device(device_id): streams[device_id].synchronize() num_permuted = cp.zeros(n_features, dtype=dtype) @@ -340,9 +322,9 @@ def _run_permutations_sparse( with cp.cuda.Device(dd["device_id"]): streams[dd["device_id"]].synchronize() - # Phase 3: Gather results where the input lives and cut to exact size. + # Phase 3: Gather results on the input device and cut to exact size with cp.cuda.Device(source_device): - all_perms = [cp.asarray(dd["perms"]) for dd in device_data] + all_perms = [_copy_to_device(dd["perms"], source_device) for dd in device_data] gearys_C_permutations = cp.concatenate(all_perms, axis=0)[:n_permutations] return gearys_C_permutations diff --git a/src/rapids_singlecell/squidpy_gpu/_moransi.py b/src/rapids_singlecell/squidpy_gpu/_moransi.py index c1c69a004..c05f0b083 100644 --- a/src/rapids_singlecell/squidpy_gpu/_moransi.py +++ b/src/rapids_singlecell/squidpy_gpu/_moransi.py @@ -6,7 +6,7 @@ from cupyx.scipy import sparse from rapids_singlecell._cuda import _autocorr_cuda as _ac -from rapids_singlecell._utils import parse_device_ids, validate_multi_gpu +from rapids_singlecell._utils import _copy_to_device, parse_device_ids from ._utils import _check_precision_issues @@ -79,22 +79,13 @@ def _run_permutations_dense( if device_ids is None: device_ids = [0] - source_device = data_centered_cupy.device.id - device_ids = list(dict.fromkeys(device_ids)) - device_ids.sort(key=lambda device_id: device_id != source_device) - device_ids = validate_multi_gpu( - device_ids, - source_device=source_device, - gather_device=source_device, - ) - with cp.cuda.Device(source_device): - cp.cuda.get_current_stream().synchronize() n_devices = len(device_ids) streams: dict[int, cp.cuda.Stream] = {} device_data: list[dict] = [] # Each device runs perms_per_device iterations perms_per_device = (n_permutations + n_devices - 1) // n_devices + source_device = data_centered_cupy.device.id # Phase 1: Create streams and transfer data to all devices for device_id in device_ids: @@ -108,16 +99,16 @@ def _run_permutations_dense( dev_adj = adj_matrix_cupy dev_den = den else: - dev_data = cp.asarray(data_centered_cupy) + dev_data = _copy_to_device(data_centered_cupy, device_id) dev_adj = sparse.csr_matrix( ( - cp.asarray(adj_matrix_cupy.data), - cp.asarray(adj_matrix_cupy.indices), - cp.asarray(adj_matrix_cupy.indptr), + _copy_to_device(adj_matrix_cupy.data, device_id), + _copy_to_device(adj_matrix_cupy.indices, device_id), + _copy_to_device(adj_matrix_cupy.indptr, device_id), ), shape=adj_matrix_cupy.shape, ) - dev_den = cp.asarray(den) + dev_den = _copy_to_device(den, device_id) # Allocate output array for this device dev_perms = cp.zeros((perms_per_device, n_features), dtype=dtype) @@ -136,7 +127,7 @@ def _run_permutations_dense( for p in range(perms_per_device): for dd in device_data: device_id = dd["device_id"] - with cp.cuda.Device(device_id), streams[device_id]: + with cp.cuda.Device(device_id): streams[device_id].synchronize() num_permuted = cp.zeros(n_features, dtype=dtype) @@ -160,9 +151,9 @@ def _run_permutations_dense( with cp.cuda.Device(dd["device_id"]): streams[dd["device_id"]].synchronize() - # Phase 3: Gather results where the input lives and cut to exact size. + # Phase 3: Gather results on the input device and cut to exact size with cp.cuda.Device(source_device): - all_perms = [cp.asarray(dd["perms"]) for dd in device_data] + all_perms = [_copy_to_device(dd["perms"], source_device) for dd in device_data] morans_I_permutations = cp.concatenate(all_perms, axis=0)[:n_permutations] return morans_I_permutations @@ -250,22 +241,13 @@ def _run_permutations_sparse( if device_ids is None: device_ids = [0] - source_device = data.data.device.id - device_ids = list(dict.fromkeys(device_ids)) - device_ids.sort(key=lambda device_id: device_id != source_device) - device_ids = validate_multi_gpu( - device_ids, - source_device=source_device, - gather_device=source_device, - ) - with cp.cuda.Device(source_device): - cp.cuda.get_current_stream().synchronize() n_devices = len(device_ids) streams: dict[int, cp.cuda.Stream] = {} device_data: list[dict] = [] # Each device runs perms_per_device iterations perms_per_device = (n_permutations + n_devices - 1) // n_devices + source_device = data.data.device.id # Phase 1: Create streams and transfer data to all devices for device_id in device_ids: @@ -282,22 +264,22 @@ def _run_permutations_sparse( else: dev_data = sparse.csr_matrix( ( - cp.asarray(data.data), - cp.asarray(data.indices), - cp.asarray(data.indptr), + _copy_to_device(data.data, device_id), + _copy_to_device(data.indices, device_id), + _copy_to_device(data.indptr, device_id), ), shape=data.shape, ) dev_adj = sparse.csr_matrix( ( - cp.asarray(adj_matrix_cupy.data), - cp.asarray(adj_matrix_cupy.indices), - cp.asarray(adj_matrix_cupy.indptr), + _copy_to_device(adj_matrix_cupy.data, device_id), + _copy_to_device(adj_matrix_cupy.indices, device_id), + _copy_to_device(adj_matrix_cupy.indptr, device_id), ), shape=adj_matrix_cupy.shape, ) - dev_means = cp.asarray(means) - dev_den = cp.asarray(den) + dev_means = _copy_to_device(means, device_id) + dev_den = _copy_to_device(den, device_id) # Allocate output array for this device dev_perms = cp.zeros((perms_per_device, n_features), dtype=dtype) @@ -317,7 +299,7 @@ def _run_permutations_sparse( for p in range(perms_per_device): for dd in device_data: device_id = dd["device_id"] - with cp.cuda.Device(device_id), streams[device_id]: + with cp.cuda.Device(device_id): streams[device_id].synchronize() num_permuted = cp.zeros(n_features, dtype=dtype) @@ -344,9 +326,9 @@ def _run_permutations_sparse( with cp.cuda.Device(dd["device_id"]): streams[dd["device_id"]].synchronize() - # Phase 3: Gather results where the input lives and cut to exact size. + # Phase 3: Gather results on the input device and cut to exact size with cp.cuda.Device(source_device): - all_perms = [cp.asarray(dd["perms"]) for dd in device_data] + all_perms = [_copy_to_device(dd["perms"], source_device) for dd in device_data] morans_I_permutations = cp.concatenate(all_perms, axis=0)[:n_permutations] return morans_I_permutations diff --git a/src/rapids_singlecell/tools/_rank_genes_groups/_stream_multi_gpu.py b/src/rapids_singlecell/tools/_rank_genes_groups/_stream_multi_gpu.py index 0c13918d8..63ba6f642 100644 --- a/src/rapids_singlecell/tools/_rank_genes_groups/_stream_multi_gpu.py +++ b/src/rapids_singlecell/tools/_rank_genes_groups/_stream_multi_gpu.py @@ -17,9 +17,7 @@ import scipy.sparse as sp from rapids_singlecell._cuda import _rank_stream_cuda as _rss -from rapids_singlecell._utils import parse_device_ids, validate_multi_gpu - -from ._wilcoxon_host import _copy_gpu_array_to_device +from rapids_singlecell._utils import _copy_to_device, parse_device_ids if TYPE_CHECKING: from ._core import _RankGenes @@ -78,16 +76,16 @@ def _shard_view(X, b0: int, b1: int): def _sum_to_device(parts: list[cp.ndarray], device_id: int) -> cp.ndarray: with cp.cuda.Device(device_id): - total = _copy_gpu_array_to_device(parts[0], device_id).copy() + total = _copy_to_device(parts[0], device_id).copy() for part in parts[1:]: - total += _copy_gpu_array_to_device(part, device_id) + total += _copy_to_device(part, device_id) cp.cuda.runtime.deviceSynchronize() return total def _concat_to_device(parts: list[cp.ndarray], device_id: int, axis: int) -> cp.ndarray: with cp.cuda.Device(device_id): - local = [_copy_gpu_array_to_device(part, device_id) for part in parts] + local = [_copy_to_device(part, device_id) for part in parts] out = cp.concatenate(local, axis=axis) cp.cuda.runtime.deviceSynchronize() return out @@ -152,10 +150,8 @@ def stream_planes_multi( comp_pts = rg.comp_pts col_shard = _is_col_shard(X) axis_len = n_genes if col_shard else n_cells - home = cp.cuda.Device().id - device_ids = device_ids[: min(len(device_ids), axis_len)] - device_ids = validate_multi_gpu(device_ids, source_device=home, gather_device=home) bands = _bands(axis_len, len(device_ids)) + device_ids = device_ids[: len(bands)] def run_shard(index: int): device_id = device_ids[index] @@ -170,14 +166,11 @@ def run_shard(index: int): cp.cuda.runtime.deviceSynchronize() return out - if len(device_ids) == 1: - shards = [run_shard(0)] - else: - with ThreadPoolExecutor(max_workers=len(device_ids)) as executor: - shards = list(executor.map(run_shard, range(len(device_ids)))) + with ThreadPoolExecutor(max_workers=len(device_ids)) as executor: + shards = list(executor.map(run_shard, range(len(device_ids)))) # Gather onto the caller's device so downstream stats math stays local. - dev0 = home + dev0 = cp.cuda.Device().id if col_shard: sums = _concat_to_device([s[0] for s in shards], dev0, axis=1) sqsums = _concat_to_device([s[1] for s in shards], dev0, axis=1) @@ -225,10 +218,8 @@ def run_binned_hist_multi( raise ValueError("invalid multi-GPU histogram gene window") col_shard = _is_col_shard(X) axis_len = stop - start if col_shard else n_cells - home = cp.cuda.Device().id - device_ids = device_ids[: min(len(device_ids), axis_len)] - device_ids = validate_multi_gpu(device_ids, source_device=home, gather_device=home) bands = _bands(axis_len, len(device_ids)) + device_ids = device_ids[: len(bands)] def run_shard(index: int): device_id = device_ids[index] @@ -273,12 +264,10 @@ def run_shard(index: int): cp.cuda.runtime.deviceSynchronize() return hist, gsum, gnnz - if len(device_ids) == 1: - shards = [run_shard(0)] - else: - with ThreadPoolExecutor(max_workers=len(device_ids)) as executor: - shards = list(executor.map(run_shard, range(len(device_ids)))) + with ThreadPoolExecutor(max_workers=len(device_ids)) as executor: + shards = list(executor.map(run_shard, range(len(device_ids)))) + home = cp.cuda.Device().id if col_shard: hist = _concat_to_device([s[0] for s in shards], home, axis=0) gsum = ( diff --git a/src/rapids_singlecell/tools/_rank_genes_groups/_wilcoxon_host.py b/src/rapids_singlecell/tools/_rank_genes_groups/_wilcoxon_host.py index ad37b7bbb..49dc917f3 100644 --- a/src/rapids_singlecell/tools/_rank_genes_groups/_wilcoxon_host.py +++ b/src/rapids_singlecell/tools/_rank_genes_groups/_wilcoxon_host.py @@ -4,7 +4,6 @@ from contextlib import contextmanager from copy import copy from dataclasses import dataclass -from functools import cache from typing import TYPE_CHECKING import cupy as cp @@ -13,7 +12,7 @@ from rapids_singlecell._cuda import _wilcoxon_cuda as _wc from rapids_singlecell._cuda import _wilcoxon_sparse_cuda as _wcs -from rapids_singlecell._utils import parse_device_ids, validate_multi_gpu +from rapids_singlecell._utils import _copy_to_device, parse_device_ids if TYPE_CHECKING: from collections.abc import Callable @@ -23,9 +22,6 @@ from ._core import _RankGenes CUDA_HOST_REGISTER_PORTABLE = 1 -CUDA_ERROR_PEER_ACCESS_UNSUPPORTED = 217 -CUDA_ERROR_PEER_ACCESS_ALREADY_ENABLED = 704 -CUDA_ERROR_TOO_MANY_PEERS = 711 MIN_SPARSE_GENE_WORK = 1 MAX_SPARSE_SPLIT_SAMPLES = 10_000_000 SPARSE_SPLIT_SAMPLE_BLOCKS = 32 @@ -233,52 +229,6 @@ def _concat_shard_stat(workers: list[_RankGenes], name: str) -> NDArray | None: return np.concatenate(arrays, axis=1) -@cache -def _enable_peer_access(device_id: int, source_device: int) -> bool: - """Enable destination-to-source peer access.""" - with cp.cuda.Device(device_id): - if not cp.cuda.runtime.deviceCanAccessPeer(device_id, source_device): - return False - try: - cp.cuda.runtime.deviceEnablePeerAccess(source_device) - except cp.cuda.runtime.CUDARuntimeError as error: - if error.status == CUDA_ERROR_PEER_ACCESS_ALREADY_ENABLED: - return True - if error.status in { - CUDA_ERROR_PEER_ACCESS_UNSUPPORTED, - CUDA_ERROR_TOO_MANY_PEERS, - }: - return False - raise - return True - - -def _copy_gpu_array_to_device(array: cp.ndarray, device_id: int) -> cp.ndarray: - if array.device.id == device_id: - return array - source_device = array.device.id - with cp.cuda.Device(device_id): - copied = cp.empty_like(array) - if array.nbytes == 0: - return copied - if _enable_peer_access(device_id, source_device): - # Do not hide peer-copy errors; they may report earlier async failures. - cp.cuda.runtime.memcpyPeer( - copied.data.ptr, - device_id, - array.data.ptr, - source_device, - array.nbytes, - ) - return copied - - with cp.cuda.Device(array.device.id): - host = array.get() - with cp.cuda.Device(device_id): - copied.set(host) - return copied - - def _device_sparse_from_arrays(X, data, indices, indptr, shape): result = type(X)((data, indices, indptr), shape=shape, copy=False) # CuPyX may narrow int64 metadata; restore the validated source dtype. @@ -291,9 +241,9 @@ def _device_sparse_from_arrays(X, data, indices, indptr, shape): def _copy_device_sparse_to_device(X, device_id: int): if X.data.device.id == device_id: return X - data = _copy_gpu_array_to_device(X.data, device_id) - indices = _copy_gpu_array_to_device(X.indices, device_id) - indptr = _copy_gpu_array_to_device(X.indptr, device_id) + data = _copy_to_device(X.data, device_id) + indices = _copy_to_device(X.indices, device_id) + indptr = _copy_to_device(X.indptr, device_id) with cp.cuda.Device(device_id): return _device_sparse_from_arrays(X, data, indices, indptr, X.shape) @@ -362,7 +312,7 @@ def _device_column_shard(X, start: int, stop: int, device_id: int): local = _device_csr_column_shard(X, start, stop) result = ( - _copy_gpu_array_to_device(local, device_id) + _copy_to_device(local, device_id) if isinstance(local, cp.ndarray) else _copy_device_sparse_to_device(local, device_id) ) @@ -400,8 +350,8 @@ def _prepare_device_column_shards(X, ranges, device_ids): def _concat_gpu_shards(arrays: list[cp.ndarray], device_id: int) -> cp.ndarray: if len(arrays) == 1: - return _copy_gpu_array_to_device(arrays[0], device_id) - local_arrays = [_copy_gpu_array_to_device(array, device_id) for array in arrays] + return arrays[0] + local_arrays = [_copy_to_device(array, device_id) for array in arrays] result = cp.concatenate(local_arrays, axis=1) # Keep peer-copy buffers alive until concatenation finishes. cp.cuda.runtime.deviceSynchronize() @@ -429,9 +379,6 @@ def _run_sharded_wilcoxon( if cpsp.issparse(X) else None ) - transfer_source = ( - source_device if source_device is not None else cp.cuda.Device().id - ) auto_single_device = multi_gpu is None and is_device_input and rg.ireference is None if multi_gpu is False or auto_single_device: device_ids = [ @@ -439,25 +386,13 @@ def _run_sharded_wilcoxon( ] else: device_ids = list(dict.fromkeys(parse_device_ids(multi_gpu=multi_gpu))) - device_ids.sort(key=lambda device_id: device_id != transfer_source) + device_ids.sort(key=lambda device_id: device_id != source_device) ranges = _split_gene_ranges( X, n_devices=len(device_ids), dense_fallback=rg._sparse_negative_fallback, ) device_ids = device_ids[: len(ranges)] - validated_device_ids = validate_multi_gpu( - device_ids, - source_device=transfer_source, - gather_device=transfer_source, - ) - if validated_device_ids != device_ids: - device_ids = validated_device_ids - ranges = _split_gene_ranges( - X, - n_devices=1, - dense_fallback=rg._sparse_negative_fallback, - ) ovo_host_context = ( _build_ovo_host_context(rg) if rg.ireference is not None else None ) @@ -563,7 +498,7 @@ def run_shard( ): msg = "Inconsistent sharded Wilcoxon group ordering." raise RuntimeError(msg) - result_device = transfer_source + result_device = device_ids[0] with cp.cuda.Device(result_device): scores = _concat_gpu_shards( [result[1] for result in complete_gpu_results], result_device diff --git a/tests/pertpy/test_distances.py b/tests/pertpy/test_distances.py index ceeaa9998..19dddb432 100644 --- a/tests/pertpy/test_distances.py +++ b/tests/pertpy/test_distances.py @@ -7,7 +7,6 @@ from anndata import AnnData from scipy.spatial.distance import cdist -from rapids_singlecell._utils import MultiGPUFallbackWarning, _multi_gpu from rapids_singlecell.pertpy_gpu import Distance, MeanVar @@ -1910,53 +1909,6 @@ def test_multi_gpu_pairwise_matches_single_gpu() -> None: ) -@pytest.mark.skipif(not _has_multiple_gpus(), reason="Requires 2+ GPUs") -def test_multi_gpu_failed_preflight_falls_back_on_input_device(monkeypatch) -> None: - caller_device = cp.cuda.Device().id - source_device = next( - device_id - for device_id in range(cp.cuda.runtime.getDeviceCount()) - if device_id != caller_device - ) - rng = np.random.default_rng(42) - cpu_embedding = rng.normal(size=(60, 8)).astype(np.float32) - obs = pd.DataFrame( - { - "group": pd.Categorical( - [f"g{index}" for index in range(3) for _ in range(20)] - ) - } - ) - adata = AnnData(cpu_embedding.copy(), obs=obs) - with cp.cuda.Device(source_device): - adata.obsm["X_pca"] = cp.asarray(cpu_embedding) - - distance = Distance(metric="edistance") - expected = distance.pairwise( - adata, - groupby="group", - multi_gpu=[source_device], - ) - - monkeypatch.setattr(_multi_gpu, "peer_copy_works", lambda *_: False) - _multi_gpu._WARNED_P2P_FAILURES.clear() - try: - with pytest.warns( - MultiGPUFallbackWarning, - match=rf"Falling back to GPU {source_device}", - ): - actual = distance.pairwise( - adata, - groupby="group", - multi_gpu=[caller_device, source_device], - ) - finally: - _multi_gpu._WARNED_P2P_FAILURES.clear() - - assert cp.cuda.Device().id == caller_device - np.testing.assert_allclose(actual.values, expected.values, rtol=1e-7, atol=1e-7) - - @pytest.mark.skipif(not _has_multiple_gpus(), reason="Requires 2+ GPUs") def test_multi_gpu_onesided_matches_single_gpu() -> None: """Test that multi-GPU onesided_distances produces same results as single-GPU.""" diff --git a/tests/test_multi_gpu_utils.py b/tests/test_multi_gpu_utils.py index e6470677a..19a91604f 100644 --- a/tests/test_multi_gpu_utils.py +++ b/tests/test_multi_gpu_utils.py @@ -2,19 +2,10 @@ from __future__ import annotations -import warnings -from contextlib import nullcontext - import cupy as cp import pytest -from rapids_singlecell._utils import ( - MultiGPUFallbackWarning, - _multi_gpu, - _split_pairs, - parse_device_ids, - validate_multi_gpu, -) +from rapids_singlecell._utils import _multi_gpu, _split_pairs, parse_device_ids class TestSplitPairs: @@ -286,180 +277,76 @@ def test_single_gpu_always_device_0(self): assert parse_device_ids(multi_gpu=False) == [0] -class TestValidateMultiGPU: +class TestDeviceCopy: @pytest.fixture(autouse=True) - def _clear_validation_state(self): - cached_peer_copy_works = _multi_gpu.peer_copy_works - cached_peer_copy_works.cache_clear() - _multi_gpu._WARNED_P2P_FAILURES.clear() + def _clear_cache(self): + _multi_gpu._peer_copy_works.cache_clear() yield - cached_peer_copy_works.cache_clear() - _multi_gpu._WARNED_P2P_FAILURES.clear() - - def test_local_device_skips_peer_check(self, monkeypatch): - monkeypatch.setattr(cp.cuda.runtime, "getDeviceCount", lambda: 2) - - def unexpected_peer_check(*_): - raise AssertionError("local execution must not run a P2P canary") - - monkeypatch.setattr(_multi_gpu, "peer_copy_works", unexpected_peer_check) - - assert validate_multi_gpu([1], source_device=1) == [1] - - def test_invalid_device_fails_before_peer_check(self, monkeypatch): - monkeypatch.setattr(cp.cuda.runtime, "getDeviceCount", lambda: 2) - - def unexpected_peer_check(*_): - raise AssertionError("invalid IDs must fail before P2P validation") + _multi_gpu._peer_copy_works.cache_clear() - monkeypatch.setattr(_multi_gpu, "peer_copy_works", unexpected_peer_check) - - with pytest.raises(ValueError, match=r"Invalid GPU device ID.*2"): - validate_multi_gpu([0, 2], source_device=0) - - def test_validates_fanout_and_gather_directions(self, monkeypatch): - monkeypatch.setattr(cp.cuda.runtime, "getDeviceCount", lambda: 3) + def test_check_is_cached_per_direction(self, monkeypatch): checked = [] - def peer_copy_works(destination, source): - checked.append((destination, source)) - return True - - monkeypatch.setattr(_multi_gpu, "peer_copy_works", peer_copy_works) - - assert validate_multi_gpu([0, 1, 2], source_device=1) == [0, 1, 2] - assert set(checked) == {(0, 1), (0, 2), (2, 1)} + def cannot_access(*pair): + checked.append(pair) + return False - @pytest.mark.parametrize("failed_pair", [(0, 1), (1, 0)]) - def test_failed_direction_warns_and_falls_back_once(self, monkeypatch, failed_pair): - monkeypatch.setattr(cp.cuda.runtime, "getDeviceCount", lambda: 2) + monkeypatch.setattr(cp.cuda.runtime, "deviceCanAccessPeer", cannot_access) monkeypatch.setattr( - _multi_gpu, - "peer_copy_works", - lambda destination, source: (destination, source) != failed_pair, + cp, + "copyto", + lambda *_: pytest.fail("a non-P2P pair must not run the canary"), ) - with pytest.warns(MultiGPUFallbackWarning, match="Falling back to GPU 0"): - assert validate_multi_gpu([0, 1], source_device=0) == [0] - - with warnings.catch_warnings(record=True) as warnings_record: - warnings.simplefilter("always") - assert validate_multi_gpu([0, 1], source_device=0) == [0] - assert not warnings_record - - def test_single_remote_target_is_not_assumed_safe(self, monkeypatch): - monkeypatch.setattr(cp.cuda.runtime, "getDeviceCount", lambda: 2) - monkeypatch.setattr(_multi_gpu, "peer_copy_works", lambda *_: False) - - with pytest.warns(MultiGPUFallbackWarning): - assert validate_multi_gpu([1], source_device=0) == [0] - - def test_stops_after_first_failed_pair(self, monkeypatch): - monkeypatch.setattr(cp.cuda.runtime, "getDeviceCount", lambda: 3) - checked = [] - - def fail_first(destination, source): - checked.append((destination, source)) - return False - - monkeypatch.setattr(_multi_gpu, "peer_copy_works", fail_first) - - with pytest.warns(MultiGPUFallbackWarning): - assert validate_multi_gpu([0, 1, 2], source_device=0) == [0] - assert checked == [(0, 1)] - - def test_unexpected_cuda_error_propagates(self, monkeypatch): - monkeypatch.setattr(cp.cuda.runtime, "getDeviceCount", lambda: 2) - - def raise_illegal_address(*_): - raise cp.cuda.runtime.CUDARuntimeError(700) - - monkeypatch.setattr(_multi_gpu, "peer_copy_works", raise_illegal_address) + assert not _multi_gpu._peer_copy_works(1, 0) + assert not _multi_gpu._peer_copy_works(1, 0) + assert not _multi_gpu._peer_copy_works(0, 1) + assert checked == [(1, 0), (0, 1)] - with pytest.raises(cp.cuda.runtime.CUDARuntimeError) as error: - validate_multi_gpu([0, 1], source_device=0) - assert error.value.status == 700 - - -class TestPeerCopyWorks: - @pytest.fixture(autouse=True) - def _clear_peer_cache(self): - cached_peer_copy_works = _multi_gpu.peer_copy_works - cached_peer_copy_works.cache_clear() - yield - cached_peer_copy_works.cache_clear() - - def test_canary_rejects_silent_success(self, monkeypatch): - monkeypatch.setattr(cp.cuda.runtime, "memcpyPeer", lambda *_: None) - - assert not _multi_gpu._run_peer_copy_canary(0, 0) - - def test_capability_false_skips_enable_and_canary(self, monkeypatch): - monkeypatch.setattr(cp.cuda, "Device", lambda *_: nullcontext()) - monkeypatch.setattr(cp.cuda.runtime, "deviceCanAccessPeer", lambda *_: False) - - def unexpected(*_): - raise AssertionError("unsupported links must not be exercised") - - monkeypatch.setattr(cp.cuda.runtime, "deviceEnablePeerAccess", unexpected) - monkeypatch.setattr(_multi_gpu, "_run_peer_copy_canary", unexpected) - - assert not _multi_gpu.peer_copy_works(1, 0) - - @pytest.mark.parametrize("status", [217, 705, 711]) - def test_expected_enable_error_falls_back(self, monkeypatch, status): - monkeypatch.setattr(cp.cuda, "Device", lambda *_: nullcontext()) + def test_silent_copy_failure_is_detected(self, monkeypatch): + if cp.cuda.runtime.getDeviceCount() < 2: + pytest.skip("requires two GPUs") monkeypatch.setattr(cp.cuda.runtime, "deviceCanAccessPeer", lambda *_: True) - - def raise_expected(*_): - raise cp.cuda.runtime.CUDARuntimeError(status) - - monkeypatch.setattr(cp.cuda.runtime, "deviceEnablePeerAccess", raise_expected) + monkeypatch.setattr(cp, "copyto", lambda *_: None) + + assert not _multi_gpu._peer_copy_works(1, 0) + + @pytest.mark.parametrize( + ("works", "selected", "unused"), + [ + (True, "_copy_to_device_p2p", "_copy_to_device_via_host"), + (False, "_copy_to_device_via_host", "_copy_to_device_p2p"), + ], + ) + def test_copy_route(self, monkeypatch, works, selected, unused): + source = cp.arange(4) + copied = object() + destination = source.device.id + 1 + monkeypatch.setattr(_multi_gpu, "_peer_copy_works", lambda *_: works) + monkeypatch.setattr(_multi_gpu, selected, lambda *_: copied) monkeypatch.setattr( - _multi_gpu, - "_run_peer_copy_canary", - lambda *_: pytest.fail("safe enable failures must skip the canary"), + _multi_gpu, unused, lambda *_: pytest.fail("wrong copy route") ) - assert not _multi_gpu.peer_copy_works(1, 0) - - def test_already_enabled_runs_canary_and_caches_result(self, monkeypatch): - monkeypatch.setattr(cp.cuda, "Device", lambda *_: nullcontext()) - monkeypatch.setattr(cp.cuda.runtime, "deviceCanAccessPeer", lambda *_: True) - - def raise_already_enabled(*_): - raise cp.cuda.runtime.CUDARuntimeError(704) + assert _multi_gpu._copy_to_device(source, destination) is copied + def test_same_device_returns_original(self, monkeypatch): + source = cp.arange(4) monkeypatch.setattr( - cp.cuda.runtime, - "deviceEnablePeerAccess", - raise_already_enabled, + _multi_gpu, + "_peer_copy_works", + lambda *_: pytest.fail("same-device copies must not check P2P"), ) - canary_calls = [] - def successful_canary(destination, source): - canary_calls.append((destination, source)) - return True + assert _multi_gpu._copy_to_device(source, source.device.id) is source - monkeypatch.setattr(_multi_gpu, "_run_peer_copy_canary", successful_canary) + def test_host_copy(self): + if cp.cuda.runtime.getDeviceCount() < 2: + pytest.skip("requires two GPUs") + with cp.cuda.Device(0): + source = cp.arange(8) - assert _multi_gpu.peer_copy_works(1, 0) - assert _multi_gpu.peer_copy_works(1, 0) - assert canary_calls == [(1, 0)] - - def test_unexpected_enable_error_propagates(self, monkeypatch): - monkeypatch.setattr(cp.cuda, "Device", lambda *_: nullcontext()) - monkeypatch.setattr(cp.cuda.runtime, "deviceCanAccessPeer", lambda *_: True) - - def raise_illegal_address(*_): - raise cp.cuda.runtime.CUDARuntimeError(700) - - monkeypatch.setattr( - cp.cuda.runtime, - "deviceEnablePeerAccess", - raise_illegal_address, - ) + copied = _multi_gpu._copy_to_device_via_host(source, 1) - with pytest.raises(cp.cuda.runtime.CUDARuntimeError) as error: - _multi_gpu.peer_copy_works(1, 0) - assert error.value.status == 700 + assert copied.device.id == 1 + assert copied.get().tolist() == list(range(8)) diff --git a/tests/test_rank_genes_groups_wilcoxon.py b/tests/test_rank_genes_groups_wilcoxon.py index 60dc9be34..f5220f33a 100644 --- a/tests/test_rank_genes_groups_wilcoxon.py +++ b/tests/test_rank_genes_groups_wilcoxon.py @@ -2278,72 +2278,6 @@ def _assert_multi_gpu_wilcoxon_equal(actual, expected): ) -@pytest.mark.parametrize( - "peer_case", - [ - pytest.param((False, None, False), id="inaccessible"), - pytest.param((True, None, True), id="enabled"), - pytest.param( - (True, _wilcoxon_host.CUDA_ERROR_PEER_ACCESS_ALREADY_ENABLED, True), - id="already_enabled", - ), - pytest.param( - (True, _wilcoxon_host.CUDA_ERROR_PEER_ACCESS_UNSUPPORTED, False), - id="unsupported", - ), - pytest.param( - (True, _wilcoxon_host.CUDA_ERROR_TOO_MANY_PEERS, False), - id="too_many_peers", - ), - ], -) -def test_wilcoxon_enable_peer_access_expected_outcomes(monkeypatch, peer_case): - can_access, error_status, expected = peer_case - calls = {"can_access": 0, "enable": 0} - - def device_can_access_peer(device_id, source_device): - calls["can_access"] += 1 - return can_access - - def device_enable_peer_access(source_device): - calls["enable"] += 1 - if error_status is not None: - raise cp.cuda.runtime.CUDARuntimeError(error_status) - - monkeypatch.setattr(cp.cuda.runtime, "deviceCanAccessPeer", device_can_access_peer) - monkeypatch.setattr( - cp.cuda.runtime, "deviceEnablePeerAccess", device_enable_peer_access - ) - _wilcoxon_host._enable_peer_access.cache_clear() - try: - assert _wilcoxon_host._enable_peer_access(0, 0) is expected - assert _wilcoxon_host._enable_peer_access(0, 0) is expected - finally: - _wilcoxon_host._enable_peer_access.cache_clear() - - assert calls["can_access"] == 1 - assert calls["enable"] == int(can_access) - - -def test_wilcoxon_enable_peer_access_reraises_unexpected_error(monkeypatch): - unexpected_status = 999 - - monkeypatch.setattr(cp.cuda.runtime, "deviceCanAccessPeer", lambda *_: True) - - def raise_unexpected(source_device): - raise cp.cuda.runtime.CUDARuntimeError(unexpected_status) - - monkeypatch.setattr(cp.cuda.runtime, "deviceEnablePeerAccess", raise_unexpected) - _wilcoxon_host._enable_peer_access.cache_clear() - try: - with pytest.raises(cp.cuda.runtime.CUDARuntimeError) as error: - _wilcoxon_host._enable_peer_access(0, 0) - finally: - _wilcoxon_host._enable_peer_access.cache_clear() - - assert error.value.status == unexpected_status - - @pytest.mark.parametrize( "route_case", [ @@ -2387,56 +2321,6 @@ def parse_device_ids_spy(*, multi_gpu): assert "scores" in adata.uns["rank_genes_groups"] -def test_wilcoxon_failed_preflight_repartitions_and_stays_serial(monkeypatch): - source_device = cp.cuda.Device().id - fake_peer = source_device + 1 - single = _make_multi_gpu_wilcoxon_adata("cupy_dense", source_device=source_device) - fallback = _make_multi_gpu_wilcoxon_adata("cupy_dense", source_device=source_device) - kwargs = { - "method": "wilcoxon", - "use_raw": False, - "reference": "rest", - "n_genes": single.n_vars, - } - rsc.tl.rank_genes_groups(single, "group", multi_gpu=False, **kwargs) - - monkeypatch.setattr( - _wilcoxon_host, - "parse_device_ids", - lambda *, multi_gpu: [source_device, fake_peer], - ) - - def force_fallback(device_ids, *, source_device, gather_device): - assert device_ids == [source_device, fake_peer] - assert gather_device == source_device - return [source_device] - - monkeypatch.setattr(_wilcoxon_host, "validate_multi_gpu", force_fallback) - real_split = _wilcoxon_host._split_gene_ranges - split_device_counts = [] - - def split_spy(X, *, n_devices, dense_fallback): - split_device_counts.append(n_devices) - return real_split( - X, - n_devices=n_devices, - dense_fallback=dense_fallback, - ) - - monkeypatch.setattr(_wilcoxon_host, "_split_gene_ranges", split_spy) - - class UnexpectedExecutor: - def __init__(self, *_args, **_kwargs): - raise AssertionError("single-GPU fallback must not start a thread pool") - - monkeypatch.setattr(_wilcoxon_host, "ThreadPoolExecutor", UnexpectedExecutor) - - rsc.tl.rank_genes_groups(fallback, "group", multi_gpu=True, **kwargs) - - assert split_device_counts == [2, 1] - _assert_multi_gpu_wilcoxon_equal(fallback, single) - - @pytest.mark.skipif(not MULTI_GPU_AVAILABLE, reason="requires at least two GPUs") @pytest.mark.parametrize("fmt", ["cupy_dense", "cupy_csr", "cupy_csc"]) @pytest.mark.parametrize("reference", ["rest", "1"])