Skip to content

Commit 2231df7

Browse files
Andy-Jostclaude
andcommitted
cuda.core: share VirtualMemoryResource buffers between processes
Buffers of a VirtualMemoryResource whose handle_type is "posix_fd" can now be exported and imported like pool-backed buffers. Buffer.ipc_descriptor exports every physical allocation that backs a VirtualMemoryBuffer as a file descriptor, Buffer.from_ipc_descriptor imports the descriptor with a VirtualMemoryResource of the receiving process, and a buffer or its descriptor can be sent through multiprocessing directly. A grown buffer imports as one contiguous range with the same byte layout. The import reuses MemAllocationHandle: import_mem_allocation_handle wraps cuMemImportFromShareableHandle in the same box and cuMemRelease deleter as create_mem_allocation_handle, with the importer's access descriptors. VirtualMemoryIPCBufferDescriptor carries one IPCAllocationHandle and one size per allocation and owns the file descriptors; multiprocessing duplicates them into the receiving process. VirtualMemoryResource pickles as (device, options), which lets a buffer pickle as (resource, descriptor). Fabric handles are not shared yet: no available test system reports fabric support or exposes an IMEX channel, so that path could not be verified. Part of #2980. Co-Authored-By: Claude Fable 5.1 <[email protected]>
1 parent ea58e1b commit 2231df7

16 files changed

Lines changed: 699 additions & 28 deletions

File tree

‎cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,35 @@ what makes that safe: the allocation is released exactly once, when its last map
199199
its own `cuMemCreate` reference. It is not called for buffers from `allocate()`, whose ranges
200200
free themselves. A subclass override of `deallocate()` therefore does not run for them.
201201

202+
## Sharing across processes
203+
204+
A `VirtualMemoryBuffer` is shared the way pool-backed buffers are, through
205+
`Buffer.ipc_descriptor` and `Buffer.from_ipc_descriptor`, when the resource's
206+
`handle_type` can travel between processes (`posix_fd` on Linux;
207+
`is_ipc_enabled` reports it). The exporter calls `cuMemExportToShareableHandle`
208+
once per mapping that covers the buffer's size, in address order, and the
209+
descriptor carries the handles with each allocation's size. Each file
210+
descriptor is owned by an `IPCAllocationHandle` and closed when the descriptor
211+
goes; `multiprocessing` duplicates it into the receiving process. The sizes
212+
travel because the OS handle does not expose them to the importer; a wrong size
213+
fails the importer's `cuMemMap`.
214+
215+
The importer is a `VirtualMemoryResource` of the receiving process with the
216+
same `handle_type`. `import_mem_allocation_handle` calls
217+
`cuMemImportFromShareableHandle` and wraps the result in the same box and
218+
deleter as `create_mem_allocation_handle`: the driver gives an imported handle
219+
one reference and frees the memory when all references are released and no
220+
mapping remains, so `cuMemRelease` is the right teardown for both. The access
221+
descriptors stored in the box are the importer's, built from its options for
222+
its device; the exporter's access does not travel. The import then reserves
223+
`sum(sizes)` with the importer's alignment, maps each allocation in order, and
224+
builds a range and a `VirtualMemoryBuffer` exactly as `allocate()` does, so a
225+
grown buffer imports as one contiguous range with the same byte layout, and the
226+
imported buffer frees itself through the same deleters. A descriptor imported
227+
in the exporting process yields an alias of the exporter's memory at a new
228+
address. A `VirtualMemoryResource` pickles as (device, options), which is what
229+
lets a `Buffer` pickle as (resource, descriptor).
230+
202231
## Why `modify_allocation` returns a new buffer
203232

204233
The input buffer stays open and aliases the result. The chunks the input already mapped are

‎cuda_core/cuda/core/_cpp/rt/api.hpp‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -262,6 +262,18 @@ MemAllocationHandle create_mem_allocation_handle(size_t size, const CUmemAllocat
262262
// Size of the allocation; the only size cuMemMap accepts for it.
263263
size_t mem_allocation_size(const MemAllocationHandle& h) noexcept;
264264

265+
// Import a physical allocation that another process exported with
266+
// cuMemExportToShareableHandle. `os_handle` is the shareable handle in the
267+
// form cuMemImportFromShareableHandle expects for `handle_type` (a POSIX file
268+
// descriptor is its value cast to void*). `size` is the allocation's size,
269+
// which the OS handle does not expose to the importer; the caller takes it
270+
// from the exporter, and a wrong size fails the later cuMemMap. The result is
271+
// a MemAllocationHandle like one from create_mem_allocation_handle: the
272+
// access descriptors are applied to every mapping, and the last reference
273+
// calls cuMemRelease. Returns empty handle on error (caller must check).
274+
MemAllocationHandle import_mem_allocation_handle(void* os_handle, CUmemAllocationHandleType handle_type,
275+
size_t size, const CUmemAccessDesc* descs, size_t count);
276+
265277
// Reserve an address range via cuMemAddressReserve. Pass alignment 0 for the
266278
// driver default. When the last reference is released, cuMemAddressFree is
267279
// called with the exact reserved pair. Returns empty handle on error.

‎cuda_core/cuda/core/_cpp/rt/driver_api.hpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,7 @@ namespace cuda_core::rt {
135135
X(cuMemMap, 10020) \
136136
X(cuMemUnmap, 10020) \
137137
X(cuMemSetAccess, 10020) \
138+
X(cuMemImportFromShareableHandle, 10020) \
138139
/* cuda-bindings requests 7000 (PTDS) or 2000 (legacy) */ \
139140
X(cuStreamSynchronize, 7000) \
140141
X(cuStreamIsCapturing, 10000) \

‎cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp‎

Lines changed: 30 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,22 @@ const Box* box_of(const Handle& h) noexcept {
5252
reinterpret_cast<const char*>(h.get()) - offsetof(Box, resource));
5353
}
5454

55+
// Wrap an allocation handle this process holds one reference to, whether
56+
// cuMemCreate or cuMemImportFromShareableHandle produced it. The last
57+
// reference releases it; the memory is freed once no mapping remains.
58+
MemAllocationHandle wrap_mem_allocation(CUmemGenericAllocationHandle handle, size_t size,
59+
std::vector<CUmemAccessDesc>&& access) {
60+
auto box = std::shared_ptr<const MemAllocationBox>(
61+
new MemAllocationBox{{handle}, size, std::move(access)},
62+
[](const MemAllocationBox* b) {
63+
GILReleaseGuard gil;
64+
pw_cuMemRelease(b->resource.raw);
65+
delete b;
66+
}
67+
);
68+
return MemAllocationHandle(box, &box->resource);
69+
}
70+
5571
} // namespace
5672

5773
// ============================================================================
@@ -68,15 +84,20 @@ MemAllocationHandle create_mem_allocation_handle(size_t size, const CUmemAllocat
6884
if (CUDA_SUCCESS != (err = DRIVER_CALL(cuMemCreate, &handle, size, &prop, 0))) {
6985
return {};
7086
}
71-
auto box = std::shared_ptr<const MemAllocationBox>(
72-
new MemAllocationBox{{handle}, size, std::move(access)},
73-
[](const MemAllocationBox* b) {
74-
GILReleaseGuard gil;
75-
pw_cuMemRelease(b->resource.raw);
76-
delete b;
77-
}
78-
);
79-
return MemAllocationHandle(box, &box->resource);
87+
return wrap_mem_allocation(handle, size, std::move(access));
88+
}
89+
90+
MemAllocationHandle import_mem_allocation_handle(void* os_handle, CUmemAllocationHandleType handle_type,
91+
size_t size, const CUmemAccessDesc* descs, size_t count) {
92+
// Copy the descriptors before the driver call so a failed copy leaves
93+
// nothing to undo.
94+
std::vector<CUmemAccessDesc> access(descs, descs + count);
95+
GILReleaseGuard gil;
96+
CUmemGenericAllocationHandle handle = 0;
97+
if (CUDA_SUCCESS != (err = DRIVER_CALL(cuMemImportFromShareableHandle, &handle, os_handle, handle_type))) {
98+
return {};
99+
}
100+
return wrap_mem_allocation(handle, size, std::move(access));
80101
}
81102

82103
size_t mem_allocation_size(const MemAllocationHandle& h) noexcept {

‎cuda_core/cuda/core/_memory/_buffer.pyi‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,9 @@
22

33
from typing import TypedDict
44

5+
from _typeshed import Incomplete
56
from cuda.core._memory._copy_enums import CopyOptions
6-
from cuda.core._memory._device_memory_resource import DeviceMemoryResource
77
from cuda.core._memory._ipc import IPCBufferDescriptor
8-
from cuda.core._memory._pinned_memory_resource import PinnedMemoryResource
98
from cuda.core._stream import Stream
109
from cuda.core._utils.pycompat import BufferProtocol
1110
from cuda.core.graph import GraphBuilder
@@ -82,13 +81,18 @@ class Buffer:
8281
:class:`Buffer` is closed or garbage collected.
8382
"""
8483
@classmethod
85-
def from_ipc_descriptor(cls, mr: DeviceMemoryResource | PinnedMemoryResource, ipc_descriptor: IPCBufferDescriptor, *, stream: Stream) -> Buffer:
84+
def from_ipc_descriptor(cls, mr: Incomplete, ipc_descriptor: IPCBufferDescriptor, *, stream: Stream) -> Buffer:
8685
"""Import a buffer that was exported from another process.
8786
8887
Parameters
8988
----------
90-
mr : :obj:`~_memory.DeviceMemoryResource` | :obj:`~_memory.PinnedMemoryResource`
91-
The IPC-enabled memory resource matching the exporting process.
89+
mr : :obj:`~_memory.DeviceMemoryResource` | :obj:`~_memory.PinnedMemoryResource` | :obj:`~_memory.VirtualMemoryResource`
90+
The IPC-enabled memory resource matching the exporting process. A
91+
descriptor exported from a :class:`VirtualMemoryBuffer` is imported
92+
with a :class:`VirtualMemoryResource` of this process whose
93+
``handle_type`` matches the exporter's; the import maps the shared
94+
physical memory for that resource's device with that resource's
95+
access options.
9296
ipc_descriptor : :obj:`~_memory.IPCBufferDescriptor`
9397
The descriptor exported from another process.
9498
stream : :obj:`~_stream.Stream`

‎cuda_core/cuda/core/_memory/_buffer.pyx‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -388,15 +388,21 @@ cdef class Buffer:
388388

389389
@classmethod
390390
def from_ipc_descriptor(
391-
cls, mr: DeviceMemoryResource | PinnedMemoryResource, ipc_descriptor: IPCBufferDescriptor,
391+
cls, mr: DeviceMemoryResource | PinnedMemoryResource | VirtualMemoryResource,
392+
ipc_descriptor: IPCBufferDescriptor,
392393
*, stream: Stream
393394
) -> Buffer:
394395
"""Import a buffer that was exported from another process.
395396

396397
Parameters
397398
----------
398-
mr : :obj:`~_memory.DeviceMemoryResource` | :obj:`~_memory.PinnedMemoryResource`
399-
The IPC-enabled memory resource matching the exporting process.
399+
mr : :obj:`~_memory.DeviceMemoryResource` | :obj:`~_memory.PinnedMemoryResource` | :obj:`~_memory.VirtualMemoryResource`
400+
The IPC-enabled memory resource matching the exporting process. A
401+
descriptor exported from a :class:`VirtualMemoryBuffer` is imported
402+
with a :class:`VirtualMemoryResource` of this process whose
403+
``handle_type`` matches the exporter's; the import maps the shared
404+
physical memory for that resource's device with that resource's
405+
access options.
400406
ipc_descriptor : :obj:`~_memory.IPCBufferDescriptor`
401407
The descriptor exported from another process.
402408
stream : :obj:`~_stream.Stream`

‎cuda_core/cuda/core/_memory/_ipc.pxd‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,10 +53,18 @@ cdef class IPCAllocationHandle:
5353
cpdef close(self)
5454

5555

56+
# Descriptor of a VirtualMemoryResource buffer: one exported handle per
57+
# physical allocation that backs the buffer, in address order.
58+
cdef class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor):
59+
cdef readonly int _handle_type
60+
cdef readonly tuple _chunk_sizes
61+
cdef readonly tuple _handles
62+
63+
5664
# Buffer IPC Implementation
5765
# -------------------------
5866
cdef IPCBufferDescriptor Buffer_get_ipc_descriptor(Buffer)
59-
cdef Buffer Buffer_from_ipc_descriptor(cls, _MemPool, IPCBufferDescriptor, stream)
67+
cdef Buffer Buffer_from_ipc_descriptor(cls, object, IPCBufferDescriptor, stream)
6068

6169

6270
# _MemPool IPC Implementation

‎cuda_core/cuda/core/_memory/_ipc.pyi‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# This file was generated by stubgen-pyx v0.2.22 from cuda_core/cuda/core/_memory/_ipc.pyx
22

33
import uuid
4+
from typing import Any
45

56
__all__ = []
67

@@ -39,7 +40,8 @@ class IPCBufferDescriptor:
3940
def size(self) -> int: ...
4041

4142
class IPCAllocationHandle:
42-
"""Shareable handle to an IPC-enabled device memory pool."""
43+
"""Shareable OS handle to an IPC-enabled memory pool or to a physical allocation
44+
of a :class:`VirtualMemoryResource` buffer."""
4345
def __init__(self, *arg, **kwargs) -> None: ...
4446
@classmethod
4547
def _init(cls, handle: int, uuid: uuid.UUID | None) -> IPCAllocationHandle: ...
@@ -54,5 +56,29 @@ class IPCAllocationHandle:
5456
@property
5557
def uuid(self) -> uuid.UUID: ...
5658

59+
class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor):
60+
"""Serializable object describing a :class:`VirtualMemoryBuffer` that can be shared between processes.
61+
62+
The descriptor holds one exported handle per physical allocation that
63+
backs the buffer, in address order, with each allocation's size. With
64+
POSIX file descriptors the handles are :class:`IPCAllocationHandle`
65+
objects: the descriptor owns them and closes them when it is released, and
66+
``multiprocessing`` duplicates them into the receiving process. Plain
67+
``pickle`` cannot carry file descriptors.
68+
69+
Note
70+
----
71+
The sizes are controlled by the exporting peer. Receivers must treat them
72+
as untrusted and import only through :meth:`Buffer.from_ipc_descriptor`,
73+
which fails instead of mapping a size the driver rejects.
74+
"""
75+
_handle_type: int
76+
_chunk_sizes: tuple[Any, ...]
77+
_handles: tuple[Any, ...]
78+
79+
@staticmethod
80+
def _from_exports(handle_type: int, chunk_sizes: tuple[int, ...], handles: tuple[Any, ...], size: int) -> VirtualMemoryIPCBufferDescriptor: ...
81+
def __reduce__(self) -> tuple[object, ...]: ...
82+
5783
def _reduce_allocation_handle(alloc_handle: IPCAllocationHandle) -> tuple[object, ...]: ...
5884
def _reconstruct_allocation_handle(cls: type, df: object, uuid: uuid.UUID | None) -> IPCAllocationHandle: ...

‎cuda_core/cuda/core/_memory/_ipc.pyx‎

Lines changed: 60 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,8 @@ cdef inline int IPCAllocationHandle_check_open(IPCAllocationHandle self) except
116116

117117

118118
cdef class IPCAllocationHandle:
119-
"""Shareable handle to an IPC-enabled device memory pool."""
119+
"""Shareable OS handle to an IPC-enabled memory pool or to a physical allocation
120+
of a :class:`VirtualMemoryResource` buffer."""
120121

121122
def __init__(self, *arg, **kwargs) -> None:
122123
raise RuntimeError("IPCAllocationHandle objects cannot be instantiated directly. Please use MemoryResource APIs.")
@@ -169,12 +170,52 @@ def _reconstruct_allocation_handle(cls: type, df: object, uuid: uuid.UUID | None
169170
multiprocessing.reduction.register(IPCAllocationHandle, _reduce_allocation_handle)
170171

171172

173+
cdef class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor):
174+
"""Serializable object describing a :class:`VirtualMemoryBuffer` that can be shared between processes.
175+
176+
The descriptor holds one exported handle per physical allocation that
177+
backs the buffer, in address order, with each allocation's size. With
178+
POSIX file descriptors the handles are :class:`IPCAllocationHandle`
179+
objects: the descriptor owns them and closes them when it is released, and
180+
``multiprocessing`` duplicates them into the receiving process. Plain
181+
``pickle`` cannot carry file descriptors.
182+
183+
Note
184+
----
185+
The sizes are controlled by the exporting peer. Receivers must treat them
186+
as untrusted and import only through :meth:`Buffer.from_ipc_descriptor`,
187+
which fails instead of mapping a size the driver rejects.
188+
"""
189+
190+
@staticmethod
191+
def _from_exports(
192+
handle_type: int, chunk_sizes: tuple[int, ...], handles: tuple, size: int
193+
) -> VirtualMemoryIPCBufferDescriptor:
194+
cdef VirtualMemoryIPCBufferDescriptor self = VirtualMemoryIPCBufferDescriptor.__new__(
195+
VirtualMemoryIPCBufferDescriptor)
196+
self._payload = b""
197+
self._size = size
198+
self._handle_type = handle_type
199+
self._chunk_sizes = tuple(chunk_sizes)
200+
self._handles = tuple(handles)
201+
return self
202+
203+
def __reduce__(self) -> tuple[object, ...]:
204+
return VirtualMemoryIPCBufferDescriptor._from_exports, (
205+
self._handle_type, self._chunk_sizes, self._handles, self._size)
206+
207+
172208
# Buffer IPC Implementation
173209
# -------------------------
174210
cdef IPCBufferDescriptor Buffer_get_ipc_descriptor(Buffer self):
175211
Buffer_check_open(self)
176212
if not self.memory_resource.is_ipc_enabled:
177213
raise RuntimeError("Memory resource is not IPC-enabled")
214+
if not isinstance(self.memory_resource, _MemPool):
215+
# VirtualMemoryBuffer overrides ipc_descriptor; a plain Buffer wrapped
216+
# with Buffer.from_handle(mr=<VirtualMemoryResource>) lands here.
217+
raise TypeError(
218+
f"a Buffer that {type(self.memory_resource).__name__} did not allocate cannot be shared")
178219
cdef cydriver.CUmemPoolPtrExportData data
179220
with nogil:
180221
HANDLE_RETURN(
@@ -186,10 +227,24 @@ cdef IPCBufferDescriptor Buffer_get_ipc_descriptor(Buffer self):
186227
return IPCBufferDescriptor._init(data_b, self.size)
187228

188229
cdef Buffer Buffer_from_ipc_descriptor(
189-
cls, _MemPool mr, IPCBufferDescriptor ipc_descriptor, stream
230+
cls, object mr, IPCBufferDescriptor ipc_descriptor, stream
190231
):
191232
"""Import a buffer that was exported from another process."""
192-
MP_check_open(mr)
233+
if isinstance(ipc_descriptor, VirtualMemoryIPCBufferDescriptor):
234+
# Imported here rather than cimported: _virtual_memory_resource
235+
# cimports this module.
236+
from cuda.core._memory._virtual_memory_resource import VirtualMemoryResource
237+
if not isinstance(mr, VirtualMemoryResource):
238+
raise TypeError(
239+
"this descriptor was exported from a VirtualMemoryResource buffer and must be "
240+
f"imported with a VirtualMemoryResource, not {type(mr).__name__}")
241+
return mr._import_ipc_buffer(ipc_descriptor, stream)
242+
if not isinstance(mr, _MemPool):
243+
raise TypeError(
244+
"this descriptor was exported from a memory pool buffer and must be imported with a "
245+
f"DeviceMemoryResource or PinnedMemoryResource, not {type(mr).__name__}")
246+
cdef _MemPool pool = <_MemPool>mr
247+
MP_check_open(pool)
193248
if not mr.is_ipc_enabled:
194249
raise RuntimeError("Memory resource is not IPC-enabled")
195250
cdef size_t payload_size = len(ipc_descriptor._payload)
@@ -201,7 +256,7 @@ cdef Buffer Buffer_from_ipc_descriptor(
201256
)
202257
cdef Stream s = Stream_accept(stream)
203258
cdef DevicePtrHandle h_ptr = deviceptr_import_ipc(
204-
mr._h_pool,
259+
pool._h_pool,
205260
ipc_descriptor.payload_ptr(),
206261
s._h_stream
207262
)
@@ -220,7 +275,7 @@ cdef Buffer Buffer_from_ipc_descriptor(
220275
f"IPC buffer descriptor size ({claimed_size}) exceeds "
221276
f"mapped allocation extent ({mapped_size} bytes)"
222277
)
223-
return Buffer_from_deviceptr_handle(h_ptr, claimed_size, mr, ipc_descriptor)
278+
return Buffer_from_deviceptr_handle(h_ptr, claimed_size, pool, ipc_descriptor)
224279

225280

226281
# _MemPool IPC Implementation

0 commit comments

Comments
 (0)