give our ABC class (MemoryResource) fast access in C/Cython

Author: leofang

cuda_core/cuda/core/experimental/_memory.pyx

@@ -39,9 +39,6 @@ if TYPE_CHECKING:
     from ._device import Device
     import uuid
 
-# TODO: define a memory property mixin class and make Buffer and
-# MemoryResource both inherit from it
-
 
 PyCapsule = TypeVar("PyCapsule")
 """Represent the capsule type."""
@@ -50,7 +47,56 @@ DevicePointerT = Union[driver.CUdeviceptr, int, None]
 """A type union of :obj:`~driver.CUdeviceptr`, `int` and `None` for hinting :attr:`Buffer.handle`."""
 
 
-cdef class Buffer:
+cdef class _cyBuffer:
+    """
+    Internal only. Responsible for offering fast C method access.
+    """
+    cdef:
+        intptr_t _ptr
+        size_t _size
+        _cyMemoryResource _mr
+        object _ptr_obj
+
+
+cdef class _cyMemoryResource:
+    """
+    Internal only. Responsible for offering fast C method access.
+    """
+    cdef Buffer _allocate(self, size_t size, cyStream stream):
+        raise NotImplementedError
+
+    cdef void _deallocate(self, intptr_t ptr, size_t size, cyStream stream) noexcept:
+        raise NotImplementedError
+
+
+class MemoryResourceAttributes(abc.ABC):
+
+    @property
+    @abc.abstractmethod
+    def is_device_accessible(self) -> bool:
+        """bool: True if buffers allocated by this resource can be accessed on the device."""
+        ...
+
+    @property
+    @abc.abstractmethod
+    def is_host_accessible(self) -> bool:
+        """bool: True if buffers allocated by this resource can be accessed on the host."""
+        ...
+
+    @property
+    @abc.abstractmethod
+    def device_id(self) -> int:
+        """int: The device ordinal for which this memory resource is responsible.
+
+        Raises
+        ------
+        RuntimeError
+            If the resource is not bound to a specific device.
+        """
+        ...
+
+
+cdef class Buffer(_cyBuffer, MemoryResourceAttributes):
     """Represent a handle to allocated memory.
 
     This generic object provides a unified representation for how
@@ -59,12 +105,7 @@ cdef class Buffer:
 
     Support for data interchange mechanisms are provided by DLPack.
     """
-
-    cdef:
-        intptr_t _ptr
-        size_t _size
-        object _mr
-        object _ptr_obj
+    cdef dict __dict__  # required if inheriting from both Cython/Python classes
 
     def __init__(self, *args, **kwargs):
         raise RuntimeError("Buffer objects cannot be instantiated directly. Please use MemoryResource APIs.")
@@ -96,17 +137,12 @@ cdef class Buffer:
             The stream object to use for asynchronous deallocation. If None,
             the behavior depends on the underlying memory resource.
         """
-        cdef _cyMemoryResource cy_mr
         if self._ptr and self._mr is not None:
-            if isinstance(self._mr, _cyMemoryResource):
-                # FIXME
-                if stream is None:
-                    stream = Stream.__new__(Stream)
-                    (<cyStream>(stream))._handle = <cydriver.CUstream>(0)
-                cy_mr = <_cyMemoryResource>(self._mr)
-                cy_mr._deallocate(self._ptr, self._size, <cyStream>stream)
-            else:
-                self._mr.deallocate(self._ptr, self._size, stream)
+            # To be fixed in NVIDIA/cuda-python#1032
+            if stream is None:
+                stream = Stream.__new__(Stream)
+                (<cyStream>(stream))._handle = <cydriver.CUstream>(0)
+            self._mr._deallocate(self._ptr, self._size, <cyStream>stream)
             self._ptr = 0
             self._mr = None
             self._ptr_obj = None
@@ -303,18 +339,7 @@ cdef class Buffer:
         return Buffer._init(ptr, size, mr=mr)
 
 
-cdef class _cyMemoryResource:
-    """
-    Internal only. Responsible for offering fast C method access.
-    """
-    cdef Buffer _allocate(self, size_t size, cyStream stream):
-        raise NotImplementedError
-
-    cdef int _deallocate(self, intptr_t ptr, size_t size, cyStream stream) except?-1:
-        raise NotImplementedError
-
-
-class MemoryResource(abc.ABC):
+cdef class MemoryResource(_cyMemoryResource, MemoryResourceAttributes, abc.ABC):
     """Abstract base class for memory resources that manage allocation and deallocation of buffers.
 
     Subclasses must implement methods for allocating and deallocation, as well as properties
@@ -323,14 +348,10 @@ class MemoryResource(abc.ABC):
     hold a reference to self, the buffer properties are retrieved simply by looking up the underlying
     memory resource's respective property.)
     """
+    cdef dict __dict__  # required if inheriting from both Cython/Python classes
 
-    @abc.abstractmethod
-    def __init__(self, *args, **kwargs):
-        """Initialize the memory resource.
-
-        Subclasses may use additional arguments to configure the resource.
-        """
-        ...
+    cdef void _deallocate(self, intptr_t ptr, size_t size, cyStream stream) noexcept:
+        self.deallocate(ptr, size, stream)
 
     @abc.abstractmethod
     def allocate(self, size_t size, stream: Stream = None) -> Buffer:
@@ -370,30 +391,6 @@ class MemoryResource(abc.ABC):
         """
         ...
 
-    @property
-    @abc.abstractmethod
-    def is_device_accessible(self) -> bool:
-        """bool: True if buffers allocated by this resource can be accessed on the device."""
-        ...
-
-    @property
-    @abc.abstractmethod
-    def is_host_accessible(self) -> bool:
-        """bool: True if buffers allocated by this resource can be accessed on the host."""
-        ...
-
-    @property
-    @abc.abstractmethod
-    def device_id(self) -> int:
-        """int: The device ordinal for which this memory resource is responsible.
-
-        Raises
-        ------
-        RuntimeError
-            If the resource is not bound to a specific device.
-        """
-        ...
-
 
 # IPC is currently only supported on Linux. On other platforms, the IPC handle
 # type is set equal to the no-IPC handle type.
@@ -570,7 +567,7 @@ class DeviceMemoryResourceAttributes:
 _ipc_registry = {}
 
 
-cdef class DeviceMemoryResource(_cyMemoryResource, MemoryResource):
+cdef class DeviceMemoryResource(MemoryResource):
     """
     Create a device memory resource managing a stream-ordered memory pool.
 
@@ -657,7 +654,7 @@ cdef class DeviceMemoryResource(_cyMemoryResource, MemoryResource):
         bint _is_mapped
         object _uuid
         IPCAllocationHandle _alloc_handle
-        dict __dict__  # TODO: check if we still need this
+        dict __dict__  # required if inheriting from both Cython/Python classes
         object __weakref__
 
     def __cinit__(self):
@@ -917,14 +914,13 @@ cdef class DeviceMemoryResource(_cyMemoryResource, MemoryResource):
             stream = default_stream()
         return self._allocate(size, <cyStream>stream)
 
-    cdef int _deallocate(self, intptr_t ptr, size_t size, cyStream stream) except?-1:
+    cdef void _deallocate(self, intptr_t ptr, size_t size, cyStream stream) noexcept:
         cdef cydriver.CUstream s = stream._handle
         cdef cydriver.CUdeviceptr devptr = <cydriver.CUdeviceptr>ptr
         with nogil:
             HANDLE_RETURN(cydriver.cuMemFreeAsync(devptr, s))
-        return 0
 
-    def deallocate(self, ptr: DevicePointerT, size_t size, stream: Stream = None):
+    cpdef deallocate(self, ptr: DevicePointerT, size_t size, stream: Stream = None):
         """Deallocate a buffer previously allocated by this resource.
 
         Parameters