Skip to content
Merged
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
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@ from cuda.core.experimental._utils.cuda_utils cimport (

from dataclasses import dataclass
from typing import Optional, TYPE_CHECKING
import cython
import platform # no-cython-lint
import uuid
import weakref
Expand Down Expand Up @@ -49,8 +48,8 @@ cdef class DeviceMemoryResourceOptions:
Maximum pool size. When set to 0, defaults to a system-dependent value.
(Default to 0)
"""
ipc_enabled : cython.bint = False
max_size : cython.size_t = 0
ipc_enabled : bool = False
max_size : int = 0


cdef class DeviceMemoryResourceAttributes:
Expand All @@ -66,63 +65,70 @@ cdef class DeviceMemoryResourceAttributes:
self._mr_weakref = mr
return self

@DMRA_mempool_attribute(bool)
cdef int _getattribute(self, cydriver.CUmemPool_attribute attr_enum, void* value) except?-1:
cdef DeviceMemoryResource mr = <DeviceMemoryResource>(self._mr_weakref())
if mr is None:
raise RuntimeError("DeviceMemoryResource is expired")
cdef cydriver.CUmemoryPool pool_handle = mr._handle
with nogil:
HANDLE_RETURN(cydriver.cuMemPoolGetAttribute(pool_handle, attr_enum, value))
return 0

@property
def reuse_follow_event_dependencies(self):
"""Allow memory to be reused when there are event dependencies between streams."""
cdef int value
self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_FOLLOW_EVENT_DEPENDENCIES, &value)
return bool(value)

@DMRA_mempool_attribute(bool)
@property
def reuse_allow_opportunistic(self):
"""Allow reuse of completed frees without dependencies."""
cdef int value
self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_ALLOW_OPPORTUNISTIC, &value)
return bool(value)

@DMRA_mempool_attribute(bool)
@property
def reuse_allow_internal_dependencies(self):
"""Allow insertion of new stream dependencies for memory reuse."""
cdef int value
self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_ALLOW_INTERNAL_DEPENDENCIES, &value)
return bool(value)

@DMRA_mempool_attribute(int)
@property
def release_threshold(self):
"""Amount of reserved memory to hold before OS release."""
cdef cydriver.cuuint64_t value
self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD, &value)
return int(value)

@DMRA_mempool_attribute(int)
@property
def reserved_mem_current(self):
"""Current amount of backing memory allocated."""
cdef cydriver.cuuint64_t value
self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT, &value)
return int(value)

@DMRA_mempool_attribute(int)
@property
def reserved_mem_high(self):
"""High watermark of backing memory allocated."""
cdef cydriver.cuuint64_t value
self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH, &value)
return int(value)

@DMRA_mempool_attribute(int)
@property
def used_mem_current(self):
"""Current amount of memory in use."""
cdef cydriver.cuuint64_t value
self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_USED_MEM_CURRENT, &value)
return int(value)

@DMRA_mempool_attribute(int)
@property
def used_mem_high(self):
"""High watermark of memory in use."""


cdef DMRA_mempool_attribute(property_type: type):
def decorator(stub):
attr_enum = getattr(
driver.CUmemPool_attribute, f"CU_MEMPOOL_ATTR_{stub.__name__.upper()}"
)

def fget(DeviceMemoryResourceAttributes self) -> property_type:
cdef DeviceMemoryResource mr = <DeviceMemoryResource> self._mr_weakref()
if mr is None:
raise RuntimeError("DeviceMemoryResource is expired")
value = DMRA_getattribute(<cydriver.CUmemoryPool><uintptr_t> mr.handle,
<cydriver.CUmemPool_attribute><uintptr_t> attr_enum)
return property_type(value)
return property(fget=fget, doc=stub.__doc__)
return decorator


cdef int DMRA_getattribute(
cydriver.CUmemoryPool pool_handle, cydriver.CUmemPool_attribute attr_enum
):
cdef int value
with nogil:
HANDLE_RETURN(cydriver.cuMemPoolGetAttribute(pool_handle, attr_enum, <void *> &value))
return value
cdef cydriver.cuuint64_t value
self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_USED_MEM_HIGH, &value)
return int(value)


cdef class DeviceMemoryResource(MemoryResource):
Expand Down
Loading