Skip to content

Commit 1f41281

Browse files
Andy-Jostclaude
andauthored
cuda.core: share VirtualMemoryResource buffers between processes (#3028)
* 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 <noreply@anthropic.com> * cuda.core: compare the descriptor's chunk count with a signed length The CI wheel build compiles with -Werror and -Wsign-compare, and the VMM import path compared a size_t count with the Py_ssize_t that len() returns. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * cuda.core: mark two VMM IPC tests thread-unsafe Under pytest-run-parallel the same test body runs in several threads. The alias test checks mapping state by address, and a freed address can be reused by another thread's allocation; the descriptor test checks a closed file descriptor number, which another thread can reuse. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * cuda.core: validate and harden VirtualMemoryResource IPC imports Address the review of the VirtualMemoryResource IPC path. Import: check every imported allocation's location against the resource's before mapping (the same device, or the same host location type) and raise a ValueError that names both; reject a descriptor with a size but no allocations; attach notes to driver failures that say which chunk failed and why a peer-controlled size fails at the map. The imported buffer no longer keeps the descriptor: the driver copies the shared allocation into a handle of its own and keeps no reference to the file descriptor, so the descriptors live only as long as the caller holds the descriptor object. Export: an imported buffer, and a buffer grown from one, cannot be exported again because the driver exports only allocations created with the requested handle type; say so instead of surfacing INVALID_VALUE. modify_allocation() rejects a config whose handle_type differs from the resource's, so every chunk of a buffer is exportable the same way. __reduce__ preserves subclasses. Tests: parametrize the memory_ipc suite over VirtualMemoryResource (VirtualMR), with the pool-only scenarios guarded or skipped; PatternGen works on the first `size` bytes of a larger buffer; the fd-leak harness asserts on the success path; new VMM-only tests for many-chunk ranges, a second device with and without peers, descriptor validation and rollback, re-export, empty buffers, and a DLPack consumer. Docs: file descriptor and memory lifetime, multiprocessing-only transport, trust boundary, device rule, re-export limit; VirtualMemoryIPCBufferDescriptor in api_private; the from_ipc_descriptor stub has its mr type again. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * cuda.core: no descriptor cache, no re-export, raw-fd transport for VMM IPC Apply the decisions on the VirtualMemoryResource IPC design. No exporter-side cache. VirtualMemoryBuffer.ipc_descriptor exports on every access and returns a new descriptor that owns its file descriptors; they close when the descriptor is released. Pickling a buffer for multiprocessing builds a transient descriptor, so a sender holds no file descriptors between sends. A spawned Process is the one exception: it receives the file descriptors when it is created, after pickling, so the transient descriptor stays on the spawning Popen until the Process object is released. A descriptor pins the physical memory while it exists, and nothing else does. No re-export of imported memory. Every physical allocation records whether cuMemImportFromShareableHandle produced it (mem_allocation_is_imported), and modify_allocation reuses the allocation handles of its input, so the mark reaches every alias, in-place grow, and move. Exporting such a buffer, directly or through pickling, raises a RuntimeError before any driver call that says the buffer contains memory imported from another process, that it cannot be exported again, and that the caller should forward the descriptor it imported from or copy into a buffer it owns. A Queue feeder thread reports the same error and drops the item. Raw-fd transport. VirtualMemoryIPCBufferDescriptor exposes handle_type, sizes, handles, and fds, and from_fds() rebuilds a descriptor in the receiving process from file descriptors received another way (the integers are duplicated; IPCAllocationHandle objects are shared). The descriptor's own pickling stays multiprocessing-only. No ipc_enabled option: handle_type is the switch and is_ipc_enabled the query, and the docs say so. Tests: file descriptor counts return to baseline on both sides, export under a lowered RLIMIT_NOFILE, the re-export errors for imported and derived buffers (direct, Pipe, Process arguments, Queue feeder thread), and raw-fd round trips through os.dup and a Unix socket with SCM_RIGHTS. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * docs(cuda.core): record the rejected retention design for VMM IPC forwarding Write down in VMM_DESIGN.md the opt-in design that would let an imported VirtualMemoryResource buffer be forwarded (keep the received handles with the imported buffer and duplicate them on export), why the option would belong on the resource, and its cost, so that it does not have to be re-derived. Not implemented; nothing needs it yet. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * test(cuda.core): import resource inside the fd-limit test The module-level import of the POSIX-only resource module made pytest fail to collect tests/memory_ipc/test_vmm_ipc.py on Windows, so the Windows rows ran no cuda.core tests at all. Every test in the module is skipped on Windows by its fixture; only the import has to move. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * cuda.core: VMM IPC review fixes; no sharing of host-located virtual memory Take ownership of each exported or duplicated file descriptor before appending its handle, so a failed append cannot close the descriptor twice. Validate imported chunk sizes against the minimum granularity that cuMemMap requires rather than the resource's granularity option. Name both causes in the note on a failed import: a handle that is not a CUDA allocation of the expected type, or an owning device that this process cannot use. Report is_ipc_enabled False for host-located resources: host NUMA virtual memory has no verified allocation path, so it is not shared yet; the docs and the design note say so. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * cuda.core: VMM IPC review fixes: posix_fd-only from_fds, scoped docs, no test skips - from_fds() takes no handle_type; a file descriptor can only carry posix_fd - The descriptor, resource, and release-note text scope the pin, pickle, and same-device rules to POSIX file descriptors, and the fds docstrings state the descriptor's lifetime rule - The 1.3.0 notes record that a per-call modify_allocation config must keep the resource's handle_type - VMM_DESIGN.md records the fabric direction: a descriptor class of its own - TestManyChunkBuffer and TestReexport no longer skip when the driver refuses the decoy's address hint; the grow relocates in that case too Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * test(cuda.core): collect garbage before each file-descriptor baseline in the VMM IPC tests init_cuda no longer runs gc.collect() at every teardown (#3053), so objects left by earlier tests (child pipes, queues, descriptors) can be collected in the middle of a test that counts process-wide file descriptors, and the count then drops below its baseline. Collect before taking the baseline. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * test(cuda.core): mark the VMM re-export test thread-unsafe TestReexport._derived asserts that a grow moved the buffer after a decoy took the adjacent range. Under pytest-run-parallel another thread can free that range between the decoy attempt and the grow, so the grow extends in place and the assertion fails. Mark the test thread_unsafe like the other address-probing tests in the file. Also record in VMM_DESIGN.md the numbers behind the no-cache export rule and the age of the multiprocessing machinery the Process path relies on. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
1 parent 1cd41ea commit 1f41281

26 files changed

Lines changed: 2284 additions & 140 deletions

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

Lines changed: 112 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -144,8 +144,9 @@ what makes that safe: the allocation is released exactly once, when its last map
144144
`location_type="host"` with a handle type other than `None`, which the driver rejects, and a
145145
request for GPUDirect RDMA on a device without support. `__init__` runs it after the VMM-support
146146
check; `modify_allocation` runs it on a per-call configuration, which must also name the
147-
resource's location. The resource reports `is_ipc_enabled = False`, which
148-
`Buffer.ipc_descriptor` reads.
147+
resource's location and `handle_type`. The resource reports `is_ipc_enabled` (True for
148+
`posix_fd` on Linux with device-located memory), which `Buffer.ipc_descriptor` and
149+
`Buffer.from_ipc_descriptor` check.
149150
- `cdef class VirtualMemoryBuffer(Buffer)` carries no extra state. It is created with
150151
`Buffer_from_deviceptr_handle(h_ptr, size, self, cls=VirtualMemoryBuffer)` and documented in
151152
`api.rst` like `ManagedBuffer`. It overrides `close(stream=None)` to reject a capturing stream
@@ -199,6 +200,115 @@ what makes that safe: the allocation is released exactly once, when its last map
199200
its own `cuMemCreate` reference. It is not called for buffers from `allocate()`, whose ranges
200201
free themselves. A subclass override of `deallocate()` therefore does not run for them.
201202

203+
## Sharing across processes
204+
205+
A `VirtualMemoryBuffer` is shared the way pool-backed buffers are, through
206+
`Buffer.ipc_descriptor` and `Buffer.from_ipc_descriptor`, when the resource's
207+
`handle_type` can travel between processes (`posix_fd` on Linux) and the memory
208+
lives on a device (`is_ipc_enabled` reports both; host-located memory is not
209+
shared yet). The exporter calls `cuMemExportToShareableHandle`
210+
once per mapping that covers the buffer's size, in address order, and the
211+
descriptor carries the handles with each allocation's size. Each file
212+
descriptor is owned by an `IPCAllocationHandle` and closed when the descriptor
213+
goes; `multiprocessing` duplicates it into the receiving process. The sizes
214+
travel because the OS handle does not expose them to the importer; a wrong size
215+
fails the importer's `cuMemMap`.
216+
217+
The importer is a `VirtualMemoryResource` of the receiving process with the
218+
same `handle_type`, for the device that owns the memory. `import_mem_allocation_handle`
219+
calls `cuMemImportFromShareableHandle` and wraps the result in the same box and
220+
deleter as `create_mem_allocation_handle`: the driver gives an imported handle
221+
one reference and frees the memory when all references are released and no
222+
mapping remains, so `cuMemRelease` is the right teardown for both. After each
223+
import, `cuMemGetAllocationPropertiesFromHandle` must report the resource's
224+
location (the same device, or the same host location type); a mismatch is a
225+
`ValueError` before anything is mapped, because the access descriptors would
226+
describe the wrong location. The access descriptors stored in the box are the
227+
importer's, built from its options for its device and peers; the exporter's
228+
access does not travel. The import then reserves `sum(sizes)` with the
229+
importer's alignment, maps each allocation in order, and builds a range and a
230+
`VirtualMemoryBuffer` exactly as `allocate()` does, so a grown buffer imports
231+
as one contiguous range with the same byte layout, and the imported buffer
232+
frees itself through the same deleters. A descriptor imported in the exporting
233+
process yields an alias of the exporter's memory at a new address. A
234+
`VirtualMemoryResource` pickles as (device, options), which is what lets a
235+
`Buffer` pickle as (resource, descriptor).
236+
237+
File descriptors and memory. A descriptor pins the physical memory while it
238+
exists, in every process that holds one, and nothing else does. The exporting
239+
buffer keeps no descriptor: `ipc_descriptor` exports again on every access and
240+
returns a new descriptor that owns one file descriptor per allocation, and
241+
`VirtualMemoryBuffer.__reduce__` builds a transient one while `multiprocessing`
242+
pickles the buffer, so a sender holds no file descriptors between sends. The
243+
numbers set that rule: PyTorch's expandable segments grow in 2 MB to 20 MB
244+
chunks, so a buffer of several GB holds several hundred allocations, against a
245+
default process limit of 1024 file descriptors. One
246+
transport needs more: a `Queue`, `Pipe`, or `Pool` duplicates the file
247+
descriptors while pickling (the resource sharer), but a spawned `Process`
248+
records their numbers during pickling and passes them to the child only when
249+
it is created (`spawnv_passfds`), so `__reduce__` keeps the transient
250+
descriptor on the spawning `Popen` object, which dies with the `Process`
251+
object. Without that, the numbers are reused by the spawn's own pipes
252+
("bad value(s) in fds_to_keep") or the child imports a stranger's fd. That
253+
machinery (`spawnv_passfds`, `get_spawning_popen()`) has not changed since
254+
Python 3.4 (2014). The
255+
driver copies
256+
each shared allocation into a handle of its own during the import and keeps no
257+
reference to the file descriptor, so the imported buffer does not keep the
258+
descriptor either; it records only that it was imported (`is_mapped`). A
259+
descriptor parked in a `Queue` or sent to a `Pool` pins the memory in the
260+
sender until the receiver unpickles it. Besides `multiprocessing`, a process
261+
with its own file descriptor passing sends the descriptor's `fds`, `sizes`,
262+
and `size` and rebuilds it with `VirtualMemoryIPCBufferDescriptor.from_fds()`,
263+
which duplicates the integers it is given and always builds a `posix_fd`
264+
descriptor.
265+
266+
Re-export. The driver exports only allocations this process created with the
267+
requested handle type; an imported allocation has none, and the driver reports
268+
INVALID_VALUE. cuda.core does not rely on that: `MemAllocationBox.imported`
269+
records provenance per allocation (`mem_allocation_is_imported`), and since
270+
`modify_allocation` reuses the allocation handles of the input, the mark
271+
travels into every alias, in-place grow, and move. `_export_ipc_descriptor`
272+
checks the allocations that cover the buffer before any driver call and raises
273+
a RuntimeError that says the buffer contains memory imported from another
274+
process, that such memory cannot be exported again, and that the caller should
275+
forward the descriptor it imported from or copy into a buffer it owns. The same
276+
exception surfaces from `pickle`/`multiprocessing` (a `Queue` feeder thread
277+
reports it through `Queue._on_queue_feeder_error` and drops the item; the queue
278+
stays usable). PyTorch applies the same contract to received CUDA tensors
279+
("Attempted to send CUDA tensor received from another process"). There is no
280+
retention option.
281+
282+
Forwarding an imported buffer would need the received file descriptors, because
283+
they are the only exportable form of imported memory. The design, not
284+
implemented because nothing needs it yet: a `VirtualMemoryResourceOptions`
285+
field that makes imports through that resource keep their received
286+
`IPCAllocationHandle`s with the imported buffer, and an export loop that uses
287+
the provenance mark to duplicate the kept handle for imported allocations while
288+
exporting created ones through the driver. The option belongs on the resource
289+
rather than on `from_ipc_descriptor`, because a pickled buffer imports through
290+
the resource it carries and the shared `Buffer.from_ipc_descriptor` signature
291+
stays unchanged; the consequence is that the sender's options decide whether a
292+
receiver can forward a pickled buffer. The cost is one open file descriptor per
293+
imported allocation for the imported buffer's lifetime, the same pressure on
294+
the process limit that the no-cache export rule avoids, which is why it would
295+
stay opt-in.
296+
297+
A `modify_allocation` config must keep the resource's `handle_type`: every
298+
chunk of a buffer must be exportable the same way.
299+
300+
Fabric handles. `CU_MEM_HANDLE_TYPE_FABRIC` exports a `CUmemFabricHandle`, 64
301+
bytes of plain data that travel through `pickle`, MPI, or a socket, need no
302+
close, and cannot hold a reference to the allocation while in flight. The
303+
importer needs access to the exporter's IMEX channel rather than a matching
304+
device index, and the memory may live on another node. The file-descriptor
305+
rules above (ownership, pinning, `pickle` raising, the same-device check) do
306+
not hold for it, so fabric sharing gets a descriptor class of its own instead
307+
of a widened `VirtualMemoryIPCBufferDescriptor`, whose `handles` stay
308+
`IPCAllocationHandle`s. `from_fds()` is POSIX-only by construction for the
309+
same reason and takes no handle type; a fabric descriptor is plain data and
310+
needs no out-of-band constructor. Tracked on the VMM epic (#2980).
311+
202312
## Why `modify_allocation` returns a new buffer
203313

204314
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: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -262,6 +262,25 @@ 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+
// Whether cuMemImportFromShareableHandle produced the allocation. The driver
266+
// exports only allocations this process created with a handle type, so an
267+
// imported allocation cannot be exported again; callers check this before
268+
// cuMemExportToShareableHandle instead of relying on the driver's error. The
269+
// mark travels with the allocation into every range that maps it.
270+
bool mem_allocation_is_imported(const MemAllocationHandle& h) noexcept;
271+
272+
// Import a physical allocation that another process exported with
273+
// cuMemExportToShareableHandle. `os_handle` is the shareable handle in the
274+
// form cuMemImportFromShareableHandle expects for `handle_type` (a POSIX file
275+
// descriptor is its value cast to void*). `size` is the allocation's size,
276+
// which the OS handle does not expose to the importer; the caller takes it
277+
// from the exporter, and a wrong size fails the later cuMemMap. The result is
278+
// a MemAllocationHandle like one from create_mem_allocation_handle: the
279+
// access descriptors are applied to every mapping, and the last reference
280+
// calls cuMemRelease. Returns empty handle on error (caller must check).
281+
MemAllocationHandle import_mem_allocation_handle(void* os_handle, CUmemAllocationHandleType handle_type,
282+
size_t size, const CUmemAccessDesc* descs, size_t count);
283+
265284
// Reserve an address range via cuMemAddressReserve. Pass alignment 0 for the
266285
// driver default. When the last reference is released, cuMemAddressFree is
267286
// 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: 36 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,10 @@ using namespace detail;
2929
namespace {
3030

3131
struct MemAllocationBox {
32-
MemAllocationValue resource; // from cuMemCreate
32+
MemAllocationValue resource; // from cuMemCreate or cuMemImportFromShareableHandle
3333
size_t size; // the only size cuMemMap accepts for it
3434
std::vector<CUmemAccessDesc> access; // applied to every mapping of this allocation
35+
bool imported; // another process created it; it cannot be exported again
3536
};
3637

3738
struct VaReservationBox {
@@ -61,6 +62,22 @@ const Box* box_of(const Handle& h) noexcept {
6162
reinterpret_cast<const char*>(h.get()) - offsetof(Box, resource));
6263
}
6364

65+
// Wrap an allocation handle this process holds one reference to, whether
66+
// cuMemCreate or cuMemImportFromShareableHandle produced it. The last
67+
// reference releases it; the memory is freed once no mapping remains.
68+
MemAllocationHandle wrap_mem_allocation(CUmemGenericAllocationHandle handle, size_t size,
69+
std::vector<CUmemAccessDesc>&& access, bool imported) {
70+
auto box = std::shared_ptr<const MemAllocationBox>(
71+
new MemAllocationBox{{handle}, size, std::move(access), imported},
72+
[](const MemAllocationBox* b) {
73+
GILReleaseGuard gil;
74+
pw_cuMemRelease(b->resource.raw);
75+
delete b;
76+
}
77+
);
78+
return MemAllocationHandle(box, &box->resource);
79+
}
80+
6481
} // namespace
6582

6683
// ============================================================================
@@ -77,21 +94,30 @@ MemAllocationHandle create_mem_allocation_handle(size_t size, const CUmemAllocat
7794
if (CUDA_SUCCESS != (err = DRIVER_CALL(cuMemCreate, &handle, size, &prop, 0))) {
7895
return {};
7996
}
80-
auto box = std::shared_ptr<const MemAllocationBox>(
81-
new MemAllocationBox{{handle}, size, std::move(access)},
82-
[](const MemAllocationBox* b) {
83-
GILReleaseGuard gil;
84-
pw_cuMemRelease(b->resource.raw);
85-
delete b;
86-
}
87-
);
88-
return MemAllocationHandle(box, &box->resource);
97+
return wrap_mem_allocation(handle, size, std::move(access), false);
98+
}
99+
100+
MemAllocationHandle import_mem_allocation_handle(void* os_handle, CUmemAllocationHandleType handle_type,
101+
size_t size, const CUmemAccessDesc* descs, size_t count) {
102+
// Copy the descriptors before the driver call so a failed copy leaves
103+
// nothing to undo.
104+
std::vector<CUmemAccessDesc> access(descs, descs + count);
105+
GILReleaseGuard gil;
106+
CUmemGenericAllocationHandle handle = 0;
107+
if (CUDA_SUCCESS != (err = DRIVER_CALL(cuMemImportFromShareableHandle, &handle, os_handle, handle_type))) {
108+
return {};
109+
}
110+
return wrap_mem_allocation(handle, size, std::move(access), true);
89111
}
90112

91113
size_t mem_allocation_size(const MemAllocationHandle& h) noexcept {
92114
return h ? box_of<MemAllocationBox>(h)->size : 0;
93115
}
94116

117+
bool mem_allocation_is_imported(const MemAllocationHandle& h) noexcept {
118+
return h ? box_of<MemAllocationBox>(h)->imported : false;
119+
}
120+
95121
// ============================================================================
96122
// Address reservations
97123
// ============================================================================

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

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ from cuda.core._memory._copy_enums import CopyOptions
66
from cuda.core._memory._device_memory_resource import DeviceMemoryResource
77
from cuda.core._memory._ipc import IPCBufferDescriptor
88
from cuda.core._memory._pinned_memory_resource import PinnedMemoryResource
9+
from cuda.core._memory._virtual_memory_resource import VirtualMemoryResource
910
from cuda.core._stream import Stream
1011
from cuda.core._utils.pycompat import BufferProtocol
1112
from cuda.core.graph import GraphBuilder
@@ -82,13 +83,18 @@ class Buffer:
8283
:class:`Buffer` is closed or garbage collected.
8384
"""
8485
@classmethod
85-
def from_ipc_descriptor(cls, mr: DeviceMemoryResource | PinnedMemoryResource, ipc_descriptor: IPCBufferDescriptor, *, stream: Stream) -> Buffer:
86+
def from_ipc_descriptor(cls, mr: DeviceMemoryResource | PinnedMemoryResource | VirtualMemoryResource, ipc_descriptor: IPCBufferDescriptor, *, stream: Stream) -> Buffer:
8687
"""Import a buffer that was exported from another process.
8788
8889
Parameters
8990
----------
90-
mr : :obj:`~_memory.DeviceMemoryResource` | :obj:`~_memory.PinnedMemoryResource`
91-
The IPC-enabled memory resource matching the exporting process.
91+
mr : :obj:`~_memory.DeviceMemoryResource` | :obj:`~_memory.PinnedMemoryResource` | :obj:`~_memory.VirtualMemoryResource`
92+
The IPC-enabled memory resource matching the exporting process. A
93+
descriptor exported from a :class:`VirtualMemoryBuffer` is imported
94+
with a :class:`VirtualMemoryResource` of this process whose
95+
``handle_type`` matches the exporter's; the import maps the shared
96+
physical memory for that resource's device with that resource's
97+
access options.
9298
ipc_descriptor : :obj:`~_memory.IPCBufferDescriptor`
9399
The descriptor exported from another process.
94100
stream : :obj:`~_stream.Stream`
@@ -103,7 +109,11 @@ class Buffer:
103109
"""
104110
@property
105111
def ipc_descriptor(self) -> IPCBufferDescriptor:
106-
"""Descriptor for sharing this buffer with other processes."""
112+
"""Descriptor for sharing this buffer with other processes.
113+
114+
A pool-backed buffer caches its descriptor; a
115+
:class:`VirtualMemoryBuffer` exports a new one on every access.
116+
"""
107117
def close(self, stream: Stream | GraphBuilder | None=None) -> None:
108118
"""Deallocate this buffer asynchronously on the given stream.
109119

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

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ from cuda.core._dlpack import classify_dl_device, make_py_capsule
5858
from cuda.core._device import Device
5959

6060
if TYPE_CHECKING:
61+
from cuda.core._memory._virtual_memory_resource import VirtualMemoryResource
6162
from cuda.core.graph import GraphBuilder
6263

6364

@@ -402,15 +403,21 @@ cdef class Buffer:
402403

403404
@classmethod
404405
def from_ipc_descriptor(
405-
cls, mr: DeviceMemoryResource | PinnedMemoryResource, ipc_descriptor: IPCBufferDescriptor,
406+
cls, mr: DeviceMemoryResource | PinnedMemoryResource | VirtualMemoryResource,
407+
ipc_descriptor: IPCBufferDescriptor,
406408
*, stream: Stream
407409
) -> Buffer:
408410
"""Import a buffer that was exported from another process.
409411

410412
Parameters
411413
----------
412-
mr : :obj:`~_memory.DeviceMemoryResource` | :obj:`~_memory.PinnedMemoryResource`
413-
The IPC-enabled memory resource matching the exporting process.
414+
mr : :obj:`~_memory.DeviceMemoryResource` | :obj:`~_memory.PinnedMemoryResource` | :obj:`~_memory.VirtualMemoryResource`
415+
The IPC-enabled memory resource matching the exporting process. A
416+
descriptor exported from a :class:`VirtualMemoryBuffer` is imported
417+
with a :class:`VirtualMemoryResource` of this process whose
418+
``handle_type`` matches the exporter's; the import maps the shared
419+
physical memory for that resource's device with that resource's
420+
access options.
414421
ipc_descriptor : :obj:`~_memory.IPCBufferDescriptor`
415422
The descriptor exported from another process.
416423
stream : :obj:`~_stream.Stream`
@@ -428,7 +435,11 @@ cdef class Buffer:
428435
@property
429436
@cython.critical_section
430437
def ipc_descriptor(self) -> IPCBufferDescriptor:
431-
"""Descriptor for sharing this buffer with other processes."""
438+
"""Descriptor for sharing this buffer with other processes.
439+
440+
A pool-backed buffer caches its descriptor; a
441+
:class:`VirtualMemoryBuffer` exports a new one on every access.
442+
"""
432443
Buffer_check_open(self)
433444
cdef object ipc_data
434445
if self._ipc_data is None:

‎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

0 commit comments

Comments
 (0)