From 036e4a07c8ec8d9fb32630e2d581f886f2df61a5 Mon Sep 17 00:00:00 2001 From: Saiyam Bharara Date: Fri, 6 Nov 2026 23:07:54 +0000 Subject: [PATCH] fix(cuda.core): grant the current device access to PinnedMemoryResource pools Signed-off-by: Saiyam Bharara --- .../core/_memory/_pinned_memory_resource.pyi | 8 +++- .../core/_memory/_pinned_memory_resource.pyx | 38 ++++++++++++---- cuda_core/docs/source/release/1.3.0-notes.rst | 7 +++ cuda_core/tests/test_memory.py | 45 +++++++++++++++++++ 4 files changed, 89 insertions(+), 9 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyi b/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyi index 244553778b9..af909bb095a 100644 --- a/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyi +++ b/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyi @@ -73,6 +73,12 @@ class PinnedMemoryResource(_MemPool): is unsupported and stream-ordered allocation is not needed, use :class:`LegacyPinnedMemoryResource`. + The driver makes a host memory pool accessible only from the host by + default, so the resource grants the device that is current when it is + created read and write access to the pool. Kernels on other devices cannot + access its buffers. Without options, the grant applies to the driver's + current pool, which the whole process shares. + 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 @@ -167,7 +173,7 @@ class PinnedMemoryResource(_MemPool): """The host NUMA node ID used for pool placement, or -1 for OS-managed placement.""" @property def is_device_accessible(self) -> bool: - """Return True. This memory resource provides device-accessible buffers.""" + """Return True. The device that is current when the resource is created can read and write its buffers.""" @property def is_host_accessible(self) -> bool: """Return True. This memory resource provides host-accessible buffers.""" diff --git a/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyx b/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyx index 62ff20bf8a8..cbfaa4f7638 100644 --- a/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyx @@ -15,6 +15,8 @@ from cuda.core._memory._memory_pool cimport ( ) from cuda.core._memory cimport _ipc from cuda.core._memory._ipc cimport IPCAllocationHandle +from cuda.core._memory._location cimport cumemlocation_from_id +from cuda.core._rt cimport as_cu from cuda.core._stream cimport Stream, Stream_accept from cuda.core._utils.cuda_utils cimport ( check_or_create_options, @@ -99,6 +101,12 @@ cdef class PinnedMemoryResource(_MemPool): is unsupported and stream-ordered allocation is not needed, use :class:`LegacyPinnedMemoryResource`. + The driver makes a host memory pool accessible only from the host by + default, so the resource grants the device that is current when it is + created read and write access to the pool. Kernels on other devices cannot + access its buffers. Without options, the grant applies to the driver's + current pool, which the whole process shares. + 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 @@ -141,13 +149,7 @@ cdef class PinnedMemoryResource(_MemPool): 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: + if not _host_pool_supported(device, self._numa_id): raise RuntimeError( f"CUDA device {device.device_id} does not support the requested " "host memory pool for PinnedMemoryResource. Use " @@ -244,7 +246,7 @@ cdef class PinnedMemoryResource(_MemPool): @property def is_device_accessible(self) -> bool: - """Return True. This memory resource provides device-accessible buffers.""" + """Return True. The device that is current when the resource is created can read and write its buffers.""" return True @property @@ -253,6 +255,14 @@ cdef class PinnedMemoryResource(_MemPool): return True +cdef inline bint _host_pool_supported(device, int numa_id): + return ( + device.properties.host_numa_memory_pools_supported + if numa_id >= 0 + else device.properties.host_memory_pools_supported + ) + + cdef inline _PMR_init(PinnedMemoryResource self, options): from .._device import Device @@ -264,6 +274,7 @@ cdef inline _PMR_init(PinnedMemoryResource self, options): cdef size_t max_size = 0 cdef cydriver.CUmemLocationType loc_type cdef int numa_id = -1 + cdef cydriver.CUmemAccessDesc access if opts is not None: ipc_enabled = opts.ipc_enabled @@ -309,6 +320,17 @@ cdef inline _PMR_init(PinnedMemoryResource self, options): max_size, ) + # A pool on a host location is accessible only from the host by default + # (see cuMemPoolCreate), so grant the current device read/write access. A + # device without host pool support cannot allocate from it anyway. + dev = Device() + if _host_pool_supported(dev, numa_id): + access.location = cumemlocation_from_id( + cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, dev.device_id) + access.flags = cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + with nogil: + HANDLE_RETURN(cydriver.cuMemPoolSetAccess(as_cu(self._h_pool), &access, 1)) + def _deep_reduce_pinned_memory_resource(mr: object) -> tuple[object, ...]: check_multiprocessing_start_method() diff --git a/cuda_core/docs/source/release/1.3.0-notes.rst b/cuda_core/docs/source/release/1.3.0-notes.rst index 167d4eafb73..ad8a59865a0 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -142,6 +142,13 @@ New features Fixes and enhancements ---------------------- +- Kernels can read and write :class:`PinnedMemoryResource` buffers. The driver + makes a host memory pool accessible only from the host by default, so a + kernel that touched such a buffer failed with ``CUDA_ERROR_ILLEGAL_ADDRESS``. + The resource now grants the device that is current when it is created read + and write access to the pool. + (`#3205 `__) + - ``register()`` on a :class:`~_memory.DeviceMemoryResource` or a :class:`PinnedMemoryResource` now raises ``ValueError`` when the resource already has a different UUID. The check used to be an ``assert``, so under diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 953ded9f04a..fbafb1b2207 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -51,6 +51,7 @@ DeviceMemoryResource, DeviceMemoryResourceOptions, GraphMemoryResource, + LaunchConfig, LegacyPinnedMemoryResource, ManagedBuffer, ManagedMemoryResource, @@ -58,10 +59,13 @@ MemoryResource, PinnedMemoryResource, PinnedMemoryResourceOptions, + Program, + ProgramOptions, StreamOptions, VirtualMemoryBuffer, VirtualMemoryResource, VirtualMemoryResourceOptions, + launch, ) from cuda.core._dlpack import DLDeviceType from cuda.core._memory._ipc import IPCBufferDescriptor @@ -1242,6 +1246,47 @@ def test_pinned_memory_resource_initialization(init_cuda): buffer.close() +@pytest.mark.human_reviewed +@pytest.mark.parametrize("pool", ["current", "new", "new_numa"]) +def test_pinned_memory_resource_kernel_access(init_cuda, pool): + """A kernel on the current device reads and writes PinnedMemoryResource buffers (#3205).""" + device = init_cuda + skip_if_pinned_memory_unsupported(device) + if pool == "current": + options = None + elif pool == "new": + options = PinnedMemoryResourceOptions(max_size=POOL_SIZE) + else: + numa_id = device.properties.host_numa_id + if numa_id < 0 or not device.properties.host_numa_memory_pools_supported: + pytest.skip("Device does not support host NUMA memory pools") + options = PinnedMemoryResourceOptions(max_size=POOL_SIZE, numa_id=numa_id) + + code = r""" + extern "C" __global__ void double_plus_one(int* data, size_t n) { + size_t i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) data[i] = data[i] * 2 + 1; + } + """ + kernel = ( + Program(code, "c++", ProgramOptions(arch=f"sm_{device.arch}")).compile("cubin").get_kernel("double_plus_one") + ) + + n = 1024 + mr = create_pinned_memory_resource_or_xfail(options, xfail_device=device) + buffer = _allocate_pinned_buffer_or_xfail(mr, n * ctypes.sizeof(ctypes.c_int), device=device) + stream = device.default_stream + stream.sync() + data = (ctypes.c_int * n).from_address(int(buffer.handle)) + data[:] = range(n) + + launch(stream, LaunchConfig(grid=(n + 255) // 256, block=256), kernel, buffer, ctypes.c_size_t(n)) + stream.sync() + assert data[:] == [2 * i + 1 for i in range(n)] + buffer.close() + mr.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)."""