diff --git a/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyi b/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyi index a83cd8ea581..9cad97a1d0d 100644 --- a/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyi +++ b/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyi @@ -5,8 +5,11 @@ from __future__ import annotations import uuid from dataclasses import dataclass +from cuda.core._memory._buffer import Buffer from cuda.core._memory._ipc import IPCAllocationHandle from cuda.core._memory._memory_pool import _MemPool +from cuda.core._stream import Stream +from cuda.core.graph import GraphBuilder @dataclass @@ -63,6 +66,14 @@ class PinnedMemoryResource(_MemPool): Notes ----- + The device associated with ``stream`` must support host memory pools. If + ``numa_id`` is set or derived for IPC, it must support host NUMA memory pools. + You can query these capabilities through + ``Device.properties.host_memory_pools_supported`` and + ``Device.properties.host_numa_memory_pools_supported``. If the required pool + is unsupported and stream-ordered allocation is not needed, use + :class:`LegacyPinnedMemoryResource`. + To create an IPC-Enabled memory resource (MR) that is capable of sharing allocations between processes, specify ``ipc_enabled=True`` in the initializer option. When IPC is enabled and ``numa_id`` is not specified, the NUMA node @@ -76,6 +87,9 @@ class PinnedMemoryResource(_MemPool): def __init__(self, options: PinnedMemoryResourceOptions | dict[str, object] | None=None) -> None: ... + def allocate(self, size: int, *, stream: Stream | GraphBuilder) -> Buffer: + """Allocate a host-pinned buffer asynchronously on the supplied stream.""" + def __reduce__(self) -> tuple[object, ...]: ... diff --git a/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyx b/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyx index 4335fbb41c2..e5f89606330 100644 --- a/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyx @@ -5,9 +5,11 @@ from __future__ import annotations from cuda.bindings cimport cydriver -from cuda.core._memory._memory_pool cimport _MemPool, MP_init_create_pool, MP_init_current_pool +from cuda.core._memory._buffer cimport Buffer +from cuda.core._memory._memory_pool cimport _MemPool, _MP_allocate, MP_init_create_pool, MP_init_current_pool from cuda.core._memory cimport _ipc from cuda.core._memory._ipc cimport IPCAllocationHandle +from cuda.core._stream cimport Stream, Stream_accept from cuda.core._utils.cuda_utils cimport ( check_or_create_options, HANDLE_RETURN, @@ -20,6 +22,11 @@ import uuid from cuda.core._utils.cuda_utils import check_multiprocessing_start_method +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from cuda.core.graph import GraphBuilder + __all__ = ['PinnedMemoryResource', 'PinnedMemoryResourceOptions'] @@ -78,6 +85,14 @@ cdef class PinnedMemoryResource(_MemPool): Notes ----- + The device associated with ``stream`` must support host memory pools. If + ``numa_id`` is set or derived for IPC, it must support host NUMA memory pools. + You can query these capabilities through + ``Device.properties.host_memory_pools_supported`` and + ``Device.properties.host_numa_memory_pools_supported``. If the required pool + is unsupported and stream-ordered allocation is not needed, use + :class:`LegacyPinnedMemoryResource`. + To create an IPC-Enabled memory resource (MR) that is capable of sharing allocations between processes, specify ``ipc_enabled=True`` in the initializer option. When IPC is enabled and ``numa_id`` is not specified, the NUMA node @@ -91,6 +106,26 @@ cdef class PinnedMemoryResource(_MemPool): def __init__(self, options: PinnedMemoryResourceOptions | dict[str, object] | None = None) -> None: _PMR_init(self, options) + def allocate(self, size_t size, *, stream: Stream | GraphBuilder) -> Buffer: + """Allocate a host-pinned buffer asynchronously on the supplied stream.""" + if self.is_mapped: + raise TypeError("Cannot allocate from a mapped IPC-enabled memory resource") + cdef Stream s = Stream_accept(stream) + device = s.device + cdef bint supported = ( + device.properties.host_numa_memory_pools_supported + if self._numa_id >= 0 + else device.properties.host_memory_pools_supported + ) + + if not supported: + raise RuntimeError( + f"CUDA device {device.device_id} does not support the requested " + "host memory pool for PinnedMemoryResource. Use " + "LegacyPinnedMemoryResource if memory-pool features are not required." + ) + return _MP_allocate(self, size, s) + def __reduce__(self) -> tuple[object, ...]: return PinnedMemoryResource.from_registry, (self.uuid,) diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 6af4d025b03..5c98ec668ae 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -751,6 +751,21 @@ def test_pinned_memory_resource_initialization(init_cuda): buffer.close() +@pytest.mark.agent_authored(model="cursor-grok-4.5") +def test_pinned_memory_resource_rejects_unsupported_host_pool(init_cuda): + """allocate() must fail on devices without host memory pool support (see #2486).""" + device = init_cuda + if device.properties.host_memory_pools_supported: + pytest.skip("Device supports host memory pools") + + mr = PinnedMemoryResource(PinnedMemoryResourceOptions(max_size=POOL_SIZE)) + try: + with pytest.raises(RuntimeError, match="does not support.*LegacyPinnedMemoryResource"): + mr.allocate(1024, stream=device.default_stream) + finally: + mr.close() + + def test_managed_memory_resource_initialization(init_cuda): device = Device() skip_if_managed_memory_unsupported(device)