|
| 1 | +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 2 | +# |
| 3 | +# SPDX-License-Identifier: Apache-2.0 |
| 4 | + |
| 5 | +# ################################################################################ |
| 6 | +# |
| 7 | +# This example copies a device buffer into a NumPy array in chunks through two |
| 8 | +# pinned staging slots. Each chunk is a device-to-pinned copy followed by |
| 9 | +# HOST_COPY_FN, which copies the slot into the array without the GIL. Each slot |
| 10 | +# has its own stream, so the copy into one slot overlaps the host copy out of |
| 11 | +# the other, and stream order keeps a slot from being reused too early. |
| 12 | +# |
| 13 | +# ################################################################################ |
| 14 | + |
| 15 | +# /// script |
| 16 | +# dependencies = ["cuda_bindings", "cuda_core", "numpy"] |
| 17 | +# /// |
| 18 | + |
| 19 | +from contextlib import closing |
| 20 | + |
| 21 | +import numpy as np |
| 22 | + |
| 23 | +from cuda.core import HOST_COPY_FN, Buffer, Device, LegacyPinnedMemoryResource, host_launch, pack_host_copy_args |
| 24 | + |
| 25 | + |
| 26 | +def main(): |
| 27 | + dev = Device(0) |
| 28 | + dev.set_current() |
| 29 | + chunk = 4 << 20 |
| 30 | + total = 16 * chunk |
| 31 | + result = np.empty(total, dtype=np.uint8) |
| 32 | + pinned_mr = LegacyPinnedMemoryResource() |
| 33 | + |
| 34 | + with ( |
| 35 | + closing(dev.create_stream()) as stream0, |
| 36 | + closing(dev.create_stream()) as stream1, |
| 37 | + dev.allocate(total, stream=stream0) as device_buffer, |
| 38 | + pinned_mr.allocate(chunk) as slot0, |
| 39 | + pinned_mr.allocate(chunk) as slot1, |
| 40 | + ): |
| 41 | + device_buffer.fill(0x5A, stream=stream0) |
| 42 | + stream1.wait(stream0) |
| 43 | + |
| 44 | + streams = (stream0, stream1) |
| 45 | + slots = (slot0, slot1) |
| 46 | + for i, offset in enumerate(range(0, total, chunk)): |
| 47 | + stream, slot = streams[i % 2], slots[i % 2] |
| 48 | + src = Buffer.from_handle(int(device_buffer.handle) + offset, chunk, owner=device_buffer) |
| 49 | + src.copy_to(slot, stream=stream) |
| 50 | + dst = memoryview(result)[offset : offset + chunk] |
| 51 | + host_launch(stream, HOST_COPY_FN, user_data=pack_host_copy_args(dst, slot, chunk)) |
| 52 | + |
| 53 | + stream0.sync() |
| 54 | + stream1.sync() |
| 55 | + |
| 56 | + assert (result == 0x5A).all() |
| 57 | + print(f"Copied {total >> 20} MiB into a NumPy array through two pinned slots.") |
| 58 | + |
| 59 | + |
| 60 | +if __name__ == "__main__": |
| 61 | + main() |
0 commit comments