Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 24 additions & 1 deletion cuda_core/cuda/core/_memory/_pinned_memory_resource.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,30 @@ class PinnedMemoryResource(_MemPool):
"""
def __init__(self, options: PinnedMemoryResourceOptions | None=None) -> None: ...
def allocate(self, size: int, *, stream: Stream | GraphBuilder) -> Buffer:
"""Allocate a host-pinned buffer asynchronously on the supplied stream."""
"""Allocate a host-pinned buffer asynchronously on the supplied stream.

Parameters
----------
size : int
The size of the buffer to allocate, in bytes.
stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
Keyword-only. The stream on which to perform the allocation
asynchronously. Must be passed explicitly; pass
``device.default_stream`` to use the default stream.

Returns
-------
Buffer
The allocated buffer object, which is accessible from the host and
from the device that the stream belongs to.

Raises
------
RuntimeError
If the stream's device does not support the requested host memory
pool. Use :class:`LegacyPinnedMemoryResource` when stream-ordered
allocation is not required.
"""
def __reduce__(self) -> tuple[object, ...]: ...
@staticmethod
def from_registry(uuid: uuid.UUID) -> PinnedMemoryResource:
Expand Down
25 changes: 24 additions & 1 deletion cuda_core/cuda/core/_memory/_pinned_memory_resource.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -113,7 +113,30 @@ cdef class PinnedMemoryResource(_MemPool):
_PMR_init(self, options)

def allocate(self, size_t size, *, stream: Stream | GraphBuilder) -> Buffer:
"""Allocate a host-pinned buffer asynchronously on the supplied stream."""
"""Allocate a host-pinned buffer asynchronously on the supplied stream.

Parameters
----------
size : int
The size of the buffer to allocate, in bytes.
stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
Keyword-only. The stream on which to perform the allocation
asynchronously. Must be passed explicitly; pass
``device.default_stream`` to use the default stream.

Returns
-------
Buffer
The allocated buffer object, which is accessible from the host and
from the device that the stream belongs to.

Raises
------
RuntimeError
If the stream's device does not support the requested host memory
pool. Use :class:`LegacyPinnedMemoryResource` when stream-ordered
allocation is not required.
"""
MP_check_open(self)
if self.is_mapped:
raise TypeError("Cannot allocate from a mapped IPC-enabled memory resource")
Expand Down
Loading