Skip to content

Commit f3c0af0

Browse files
saiyambhararasaiyambharara
committed
feat(cuda.core): grant every supporting device access to PinnedMemoryResource pools
Signed-off-by: Saiyam Bharara <[email protected]>
1 parent 8fa2bf7 commit f3c0af0

4 files changed

Lines changed: 50 additions & 33 deletions

File tree

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -74,10 +74,10 @@ class PinnedMemoryResource(_MemPool):
7474
:class:`LegacyPinnedMemoryResource`.
7575
7676
The driver makes a host memory pool accessible only from the host by
77-
default, so the resource grants the device that is current when it is
78-
created read and write access to the pool. Kernels on other devices cannot
79-
access its buffers. Without options, the grant applies to the driver's
80-
current pool, which the whole process shares.
77+
default, so the resource grants read and write access to the pool to every
78+
device that supports host memory pools (host NUMA memory pools if the pool
79+
has a NUMA node), as ``cudaMallocHost`` does. Without options, the grant
80+
applies to the driver's current pool, which the whole process shares.
8181
8282
To create an IPC-Enabled memory resource (MR) that is capable of sharing
8383
allocations between processes, specify ``ipc_enabled=True`` in the initializer
@@ -173,7 +173,7 @@ class PinnedMemoryResource(_MemPool):
173173
"""The host NUMA node ID used for pool placement, or -1 for OS-managed placement."""
174174
@property
175175
def is_device_accessible(self) -> bool:
176-
"""Return True. The device that is current when the resource is created can read and write its buffers."""
176+
"""Return True. Every device that supports this kind of host memory pool can read and write its buffers."""
177177
@property
178178
def is_host_accessible(self) -> bool:
179179
"""Return True. This memory resource provides host-accessible buffers."""

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

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ from cuda.core._utils.cuda_utils cimport (
2222
check_or_create_options,
2323
HANDLE_RETURN,
2424
)
25+
from libcpp.vector cimport vector
2526

2627
from dataclasses import dataclass
2728
import multiprocessing
@@ -102,10 +103,10 @@ cdef class PinnedMemoryResource(_MemPool):
102103
:class:`LegacyPinnedMemoryResource`.
103104
104105
The driver makes a host memory pool accessible only from the host by
105-
default, so the resource grants the device that is current when it is
106-
created read and write access to the pool. Kernels on other devices cannot
107-
access its buffers. Without options, the grant applies to the driver's
108-
current pool, which the whole process shares.
106+
default, so the resource grants read and write access to the pool to every
107+
device that supports host memory pools (host NUMA memory pools if the pool
108+
has a NUMA node), as ``cudaMallocHost`` does. Without options, the grant
109+
applies to the driver's current pool, which the whole process shares.
109110
110111
To create an IPC-Enabled memory resource (MR) that is capable of sharing
111112
allocations between processes, specify ``ipc_enabled=True`` in the initializer
@@ -246,7 +247,7 @@ cdef class PinnedMemoryResource(_MemPool):
246247

247248
@property
248249
def is_device_accessible(self) -> bool:
249-
"""Return True. The device that is current when the resource is created can read and write its buffers."""
250+
"""Return True. Every device that supports this kind of host memory pool can read and write its buffers."""
250251
return True
251252

252253
@property
@@ -274,7 +275,8 @@ cdef inline _PMR_init(PinnedMemoryResource self, options):
274275
cdef size_t max_size = 0
275276
cdef cydriver.CUmemLocationType loc_type
276277
cdef int numa_id = -1
277-
cdef cydriver.CUmemAccessDesc access
278+
cdef cydriver.CUmemAccessDesc desc
279+
cdef vector[cydriver.CUmemAccessDesc] access
278280

279281
if opts is not None:
280282
ipc_enabled = opts.ipc_enabled
@@ -321,15 +323,18 @@ cdef inline _PMR_init(PinnedMemoryResource self, options):
321323
)
322324

323325
# A pool on a host location is accessible only from the host by default
324-
# (see cuMemPoolCreate), so grant the current device read/write access. A
325-
# device without host pool support cannot allocate from it anyway.
326-
dev = Device()
327-
if _host_pool_supported(dev, numa_id):
328-
access.location = cumemlocation_from_id(
329-
cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, dev.device_id)
330-
access.flags = cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
326+
# (see cuMemPoolCreate), so grant read/write access to every device that
327+
# supports this kind of host pool, as cudaMallocHost memory is. Other
328+
# devices cannot use the pool.
329+
desc.flags = cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
330+
for dev in Device.get_all_devices():
331+
if _host_pool_supported(dev, numa_id):
332+
desc.location = cumemlocation_from_id(
333+
cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, dev.device_id)
334+
access.push_back(desc)
335+
if not access.empty():
331336
with nogil:
332-
HANDLE_RETURN(cydriver.cuMemPoolSetAccess(as_cu(self._h_pool), &access, 1))
337+
HANDLE_RETURN(cydriver.cuMemPoolSetAccess(as_cu(self._h_pool), access.data(), access.size()))
333338

334339

335340
def _deep_reduce_pinned_memory_resource(mr: object) -> tuple[object, ...]:

‎cuda_core/docs/source/release/1.3.0-notes.rst‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -168,9 +168,10 @@ Fixes and enhancements
168168
- Kernels can read and write :class:`PinnedMemoryResource` buffers. The driver
169169
makes a host memory pool accessible only from the host by default, so a
170170
kernel that touched such a buffer failed with ``CUDA_ERROR_ILLEGAL_ADDRESS``.
171-
The resource now grants the device that is current when it is created read
172-
and write access to the pool.
173-
(`#3205 <https://github.com/NVIDIA/cuda-python/issues/3205>`__)
171+
The resource now grants read and write access to the pool to every device
172+
that supports host memory pools, as ``cudaMallocHost`` does.
173+
(`#3205 <https://github.com/NVIDIA/cuda-python/issues/3205>`__,
174+
`#3415 <https://github.com/NVIDIA/cuda-python/issues/3415>`__)
174175

175176
- ``register()`` on a :class:`~_memory.DeviceMemoryResource` or a
176177
:class:`PinnedMemoryResource` now raises ``ValueError`` when the resource

‎cuda_core/tests/test_memory.py‎

Lines changed: 22 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1247,18 +1247,25 @@ def test_pinned_memory_resource_initialization(init_cuda):
12471247

12481248

12491249
@pytest.mark.human_reviewed
1250+
@pytest.mark.parametrize("kernel_on", ["creating_device", "other_device"])
12501251
@pytest.mark.parametrize("pool", ["current", "new", "new_numa"])
1251-
def test_pinned_memory_resource_kernel_access(init_cuda, pool):
1252-
"""A kernel on the current device reads and writes PinnedMemoryResource buffers (#3205)."""
1253-
device = init_cuda
1252+
def test_pinned_memory_resource_kernel_access(init_cuda, pool, kernel_on):
1253+
"""A kernel on any device reads and writes PinnedMemoryResource buffers (#3205, #3415)."""
1254+
device = init_cuda # Device(0)
1255+
kernel_device = device
1256+
if kernel_on == "other_device":
1257+
if len(Device.get_all_devices()) < 2:
1258+
pytest.skip("Test requires at least 2 GPUs")
1259+
kernel_device = Device(1)
12541260
skip_if_pinned_memory_unsupported(device)
1261+
skip_if_pinned_memory_unsupported(kernel_device)
12551262
if pool == "current":
12561263
options = None
12571264
elif pool == "new":
12581265
options = PinnedMemoryResourceOptions(max_size=POOL_SIZE)
12591266
else:
12601267
numa_id = device.properties.host_numa_id
1261-
if numa_id < 0 or not device.properties.host_numa_memory_pools_supported:
1268+
if numa_id < 0 or not all(d.properties.host_numa_memory_pools_supported for d in (device, kernel_device)):
12621269
pytest.skip("Device does not support host NUMA memory pools")
12631270
options = PinnedMemoryResourceOptions(max_size=POOL_SIZE, numa_id=numa_id)
12641271

@@ -1268,20 +1275,24 @@ def test_pinned_memory_resource_kernel_access(init_cuda, pool):
12681275
if (i < n) data[i] = data[i] * 2 + 1;
12691276
}
12701277
"""
1271-
kernel = (
1272-
Program(code, "c++", ProgramOptions(arch=f"sm_{device.arch}")).compile("cubin").get_kernel("double_plus_one")
1273-
)
12741278

12751279
n = 1024
12761280
mr = create_pinned_memory_resource_or_xfail(options, xfail_device=device)
12771281
buffer = _allocate_pinned_buffer_or_xfail(mr, n * ctypes.sizeof(ctypes.c_int), device=device)
1278-
stream = device.default_stream
1279-
stream.sync()
1282+
device.default_stream.sync()
12801283
data = (ctypes.c_int * n).from_address(int(buffer.handle))
12811284
data[:] = range(n)
12821285

1283-
launch(stream, LaunchConfig(grid=(n + 255) // 256, block=256), kernel, buffer, ctypes.c_size_t(n))
1284-
stream.sync()
1286+
# The resource was created with device current. The kernel runs on kernel_device.
1287+
kernel_device.set_current()
1288+
try:
1289+
program = Program(code, "c++", ProgramOptions(arch=f"sm_{kernel_device.arch}"))
1290+
kernel = program.compile("cubin").get_kernel("double_plus_one")
1291+
stream = kernel_device.default_stream
1292+
launch(stream, LaunchConfig(grid=(n + 255) // 256, block=256), kernel, buffer, ctypes.c_size_t(n))
1293+
stream.sync()
1294+
finally:
1295+
device.set_current()
12851296
assert data[:] == [2 * i + 1 for i in range(n)]
12861297
buffer.close()
12871298
mr.close()

0 commit comments

Comments
 (0)