Skip to content

Commit b11ec0f

Browse files
committed
examples(cuda.core): add a chunked device-to-NumPy copy with HOST_COPY_FN
1 parent a67ad70 commit b11ec0f

2 files changed

Lines changed: 64 additions & 0 deletions

File tree

‎cuda_core/docs/source/examples.rst‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,9 @@ Kernels and launch configuration
3838
- :cuda-core-example:`host_launch.py <host_launch.py>`
3939
orders a host function between two kernel launches on a stream and reads a
4040
pinned buffer from it.
41+
- :cuda-core-example:`host_copy.py <host_copy.py>`
42+
copies a device buffer into a NumPy array in chunks through two pinned
43+
staging slots, with the host copies done by ``HOST_COPY_FN``.
4144

4245
Memory management and interop
4346
-----------------------------

‎cuda_core/examples/host_copy.py‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
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

Comments
 (0)