diff --git a/docs/content/docs/features/large-problems.mdx b/docs/content/docs/features/large-problems.mdx index 11a6804..dc6cca2 100644 --- a/docs/content/docs/features/large-problems.mdx +++ b/docs/content/docs/features/large-problems.mdx @@ -97,7 +97,11 @@ out = src( the first one gathers the sketches and runs the QR. Everything runs in the calling process, one thread per GPU, so nothing changes for the caller: the call returns the output cores as with one GPU, errors are raised as usual, and the inputs are -read once and shared. +read once and shared. Each thread runs on a non-blocking CUDA stream of its own GPU, +not on CuPy's legacy default stream, where a copy between two GPUs waits for the +copy in the other direction and the exchanges between GPUs would get half the +bandwidth. The stream first waits for the work already queued on the caller's +current stream, so inputs given as CuPy arrays may still be in flight. The sketch index is split: each GPU owns a contiguous block of the `chi_out` sketch columns and computes only those columns of every environment and of every diff --git a/src/src_method/_group.py b/src/src_method/_group.py index 20c6602..413a11f 100644 --- a/src/src_method/_group.py +++ b/src/src_method/_group.py @@ -28,6 +28,7 @@ enable_peer_access, get_xp, use_device, + use_stream, ) if TYPE_CHECKING: @@ -125,14 +126,14 @@ def run(self, fn: Callable[[int], T]) -> list[T]: Exception: The first error raised by a rank, after every rank stopped. """ if self.size == 1: - with use_device(self.xp, self.devices[0]): + with use_device(self.xp, self.devices[0]), use_stream(self.xp): return [fn(0)] first: list[BaseException] = [] lock = threading.Lock() def worker(rank: int) -> T: try: - with use_device(self.xp, self.devices[rank]): + with use_device(self.xp, self.devices[rank]), use_stream(self.xp): return fn(rank) except BaseException as err: secondary = isinstance( diff --git a/src/src_method/utils/__init__.py b/src/src_method/utils/__init__.py index 91b2111..bed085c 100644 --- a/src/src_method/utils/__init__.py +++ b/src/src_method/utils/__init__.py @@ -27,6 +27,7 @@ to_host_async, to_numpy, use_device, + use_stream, ) from .linalg import truncated_qr @@ -56,4 +57,5 @@ "to_numpy", "truncated_qr", "use_device", + "use_stream", ] diff --git a/src/src_method/utils/_backend.py b/src/src_method/utils/_backend.py index ff78a37..3124d57 100644 --- a/src/src_method/utils/_backend.py +++ b/src/src_method/utils/_backend.py @@ -213,6 +213,42 @@ def use_device(xp: ModuleType, device: int) -> Iterator[None]: yield +# One per device, reused across calls: CuPy's pool hands a freed block only to the +# stream that allocated it, so a new stream per call would strand the blocks the +# previous one cached. +_streams: dict[int, Any] = {} + + +@contextmanager +def use_stream(xp: ModuleType) -> Iterator[None]: + """Run the block on a non-blocking stream of the current device. + + On the legacy null stream, a peer copy into one device serializes with the copy + in the other direction, which halves the bandwidth of every collective over + PCIe. The stream first waits for the work already queued on the previous + current stream, so inputs written there are ready, and is synchronized on a + normal exit, so whatever the block returns is ready on any stream. On the host + this is a no-op. + + Args: + xp: The array namespace. + + Yields: + Nothing. + """ + if is_host(xp): + yield + return + device = xp.cuda.Device().id + stream = _streams.get(device) or _streams.setdefault( + device, xp.cuda.Stream(non_blocking=True) + ) + stream.wait_event(xp.cuda.get_current_stream().record()) + with stream: + yield + stream.synchronize() + + # cudaErrorPeerAccessAlreadyEnabled: another library, or an earlier call, did it. _PEER_ALREADY_ENABLED = 704 diff --git a/tests/test_backend.py b/tests/test_backend.py index 62a1ca1..4e18e60 100644 --- a/tests/test_backend.py +++ b/tests/test_backend.py @@ -21,6 +21,7 @@ to_device_async, to_host_async, use_device, + use_stream, ) @@ -90,6 +91,11 @@ def test_host_devices_are_simulated(): assert current_device(np) == 0 +def test_host_use_stream_is_a_no_op(): + with use_stream(np): + assert isinstance(current_stream(np), NullStream) + + def test_copy_into_on_host(): dst = np.empty((2, 3), dtype=complex) diff --git a/tests/test_gpu_backend.py b/tests/test_gpu_backend.py index 0e851ba..9788162 100644 --- a/tests/test_gpu_backend.py +++ b/tests/test_gpu_backend.py @@ -13,6 +13,7 @@ import src_method._sweep as sweep_module from src_method import Resources, apply, compress, src +from src_method._group import DeviceGroup from src_method._kernels import SiteKernels from src_method._plan import ( GPU_MARGIN_FRACTION, @@ -475,3 +476,35 @@ def test_too_many_devices_raise() -> None: with pytest.raises(ValueError, match="visible"): src(*stack, chi_out=4, device="gpu", resources=Resources(devices=count + 1)) + + +def test_group_runs_on_a_non_blocking_stream() -> None: + """Ranks leave the null stream, keep one stream per device, and see prior work.""" + group = DeviceGroup("gpu", [0]) + data = cupy.zeros(1 << 24, dtype=cupy.float32) + data += 1 # queued on the null stream before the run + + def fn(_rank: int) -> tuple[int, float]: + stream = cupy.cuda.get_current_stream() + assert stream != cupy.cuda.Stream.null + return stream.ptr, float((data * 2).sum()) + + ((first, total),) = group.run(fn) + ((second, _),) = group.run(fn) + + assert total == 2 * data.size + assert first == second + assert cupy.cuda.get_current_stream() == cupy.cuda.Stream.null + + +@pytest.mark.usefixtures("two_gpus") +def test_two_gpus_run_on_their_own_streams() -> None: + """Each rank of a group runs on a non-blocking stream of its own device.""" + group = DeviceGroup("gpu", 2) + + def fn(rank: int) -> tuple[int, int]: + stream = cupy.cuda.get_current_stream() + assert stream != cupy.cuda.Stream.null + return rank, stream.device_id + + assert group.run(fn) == [(0, 0), (1, 1)]