make p2p datatransfers safe

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`

src/rapids_singlecell/_utils/__init__.py

@@ -9,6 +9,7 @@
 
 from ._multi_gpu import (
     _calculate_blocks_per_pair,
+    _copy_to_device,
     _create_category_index_mapping,
     _get_device_attrs,
     _split_pairs,
@@ -17,6 +18,7 @@
 
 __all__ = [
     "_calculate_blocks_per_pair",
+    "_copy_to_device",
     "_create_category_index_mapping",
     "_get_device_attrs",
     "_split_pairs",

src/rapids_singlecell/_utils/_multi_gpu.py

@@ -12,11 +12,73 @@
 
 from __future__ import annotations
 
+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] = {}
 
+_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,
+}
+
+
+@cache
+def _peer_copy_works(destination: int, source: int) -> bool:
+    """Return whether a peer copy arrives intact."""
+    if destination == source:
+        return True
+    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
+        raise
+    return bool(np.array_equal(actual, _CANARY))
+
+
+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)
+
+
+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)
+
+
+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]:
     """Parse multi_gpu parameter into a list of device IDs.

src/rapids_singlecell/pertpy_gpu/_metrics/_edistance.py

@@ -14,6 +14,7 @@
 from rapids_singlecell._cuda import _edistance_cuda as _ed
 from rapids_singlecell._utils import (
     _calculate_blocks_per_pair,
+    _copy_to_device,
     _split_pairs,
 )
 from rapids_singlecell.squidpy_gpu._utils import _assert_categorical_obs
@@ -771,6 +772,7 @@ def _launch_distance_kernel(
         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)
@@ -801,20 +803,20 @@ def _launch_distance_kernel(
 
                 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
@@ -878,12 +880,12 @@ def _launch_distance_kernel(
                 with cp.cuda.Device(data["device_id"]):
                     cp.cuda.Stream.null.synchronize()
 
-        # Phase 4: Aggregate on GPU 0
-        with cp.cuda.Device(device_ids[0]):
+        # 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
 
@@ -1133,8 +1135,8 @@ def _pairwise_means_bootstrap(
             )
             all_results.append(pairwise_means.get())
 
-        # Compute statistics on first GPU
-        with cp.cuda.Device(device_ids[0]):
+        # 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)
@@ -1214,8 +1216,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]):
+        # 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)

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
 from rapids_singlecell.squidpy_gpu._utils import _assert_categorical_obs
 
 from ._base_metric import BaseMetric, parse_device_ids
@@ -346,7 +347,7 @@ 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]
+        output_device = out.device.id
 
         # Move the shared inputs to each participating device once.
         streams: dict[int, cp.cuda.Stream] = {}
@@ -358,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
@@ -438,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())
@@ -477,9 +480,9 @@ def _bootstrap_solve(
             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)
@@ -635,6 +638,7 @@ def _to_matrix(flat: cp.ndarray, name: str) -> pd.DataFrame:
             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)

src/rapids_singlecell/squidpy_gpu/_co_oc.py

@@ -9,6 +9,7 @@
 from rapids_singlecell._cuda import _cooc_cuda as _co
 from rapids_singlecell._utils import (
     _calculate_blocks_per_pair,
+    _copy_to_device,
     _create_category_index_mapping,
     _split_pairs,
     parse_device_ids,
@@ -322,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)
@@ -387,12 +388,11 @@ 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 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