Touch-ups

Author: Andy-Jost

cuda_core/cuda/core/experimental/_memory/_buffer.pyx

@@ -39,11 +39,7 @@ cdef class Buffer:
     Support for data interchange mechanisms are provided by DLPack.
     """
     def __cinit__(self):
-        self._ptr = 0
-        self._size = 0
-        self._mr = None
-        self._ptr_obj = None
-        self._alloc_stream = None
+        self._clear()
 
     def _clear(self):
         self._ptr = 0

cuda_core/cuda/core/experimental/_memory/_dmr.pxd

@@ -3,9 +3,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from cuda.bindings cimport cydriver
-
 from cuda.core.experimental._memory._buffer cimport MemoryResource
-from cuda.core.experimental._memory._ipc cimport IPCAllocationHandle, IPCData
+from cuda.core.experimental._memory._ipc cimport IPCData
 
 
 cdef class DeviceMemoryResource(MemoryResource):

cuda_core/cuda/core/experimental/_memory/_dmr.pyx

@@ -14,7 +14,6 @@ from cuda.core.experimental._memory cimport _ipc
 from cuda.core.experimental._memory._ipc cimport IPCAllocationHandle, IPCData
 from cuda.core.experimental._stream cimport default_stream, Stream
 from cuda.core.experimental._utils.cuda_utils cimport (
-    _check_driver_error as raise_if_driver_error,
     check_or_create_options,
     HANDLE_RETURN,
 )
@@ -75,9 +74,8 @@ class DeviceMemoryResourceAttributes:
                 mr = self._mr()
                 if mr is None:
                     raise RuntimeError("DeviceMemoryResource is expired")
-                # TODO: this implementation does not allow lowering to Cython + nogil
-                err, value = driver.cuMemPoolGetAttribute(mr.handle, attr_enum)
-                raise_if_driver_error(err)
+                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
@@ -117,6 +115,15 @@ class DeviceMemoryResourceAttributes:
     del mempool_property
 
 
+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 class DeviceMemoryResource(MemoryResource):
     """
     Create a device memory resource managing a stream-ordered memory pool.

cuda_core/cuda/core/experimental/_memory/_vmm.py

@@ -4,7 +4,7 @@
 
 import platform
 from dataclasses import dataclass, field
-from typing import Iterable, Literal, Optional, Union
+from typing import Iterable, Literal, Union
 
 from cuda.core.experimental._memory._buffer import Buffer, MemoryResource
 from cuda.core.experimental._stream import Stream
@@ -65,8 +65,8 @@ class VirtualMemoryResourceOptions:
     handle_type: VirtualMemoryHandleTypeT = "posix_fd"
     granularity: VirtualMemoryGranularityT = "recommended"
     gpu_direct_rdma: bool = False
-    addr_hint: Optional[int] = 0
-    addr_align: Optional[int] = None
+    addr_hint: int | None = 0
+    addr_align: int | None = None
     peers: Iterable[int] = field(default_factory=tuple)
     self_access: VirtualMemoryAccessTypeT = "rw"
     peer_access: VirtualMemoryAccessTypeT = "rw"