Ensure allocation stream is used for buffer deallocation if no explicit stream is provided (#1032)

* release gil

* release gil for record

* use __dealloc__ in event/stream

* nit: remove print

* fix linter error

* reduce further the number of Python objects held by Stream

* replace a few more __del__ by __dealloc__

* improve __cuda_stream__ performance

* cythonize Buffer & DMR (WIP - not working!)

* minor fixes - still failing

* fix casting and avoid repetive assignment

* ensure we have C access for DeviceMemoryResource

* restore the contracts for now (the deallocation stream should be fixed)

* fully cythonize Stream

* make linter happy

* ensure allocation stream can be used for deallocation; suppress casting warning

* int -> intptr_t -> void* is safer

Author: leofang

cuda_core/cuda/core/experimental/_memory.pyx

@@ -56,6 +56,7 @@ cdef class _cyBuffer:
         size_t _size
         _cyMemoryResource _mr
         object _ptr_obj
+        cyStream _alloc_stream
 
 
 cdef class _cyMemoryResource:
@@ -107,22 +108,31 @@ cdef class Buffer(_cyBuffer, MemoryResourceAttributes):
     """
     cdef dict __dict__  # required if inheriting from both Cython/Python classes
 
+    def __cinit__(self):
+        self._ptr = 0
+        self._size = 0
+        self._mr = None
+        self._ptr_obj = None
+        self._alloc_stream = None
+
     def __init__(self, *args, **kwargs):
         raise RuntimeError("Buffer objects cannot be instantiated directly. Please use MemoryResource APIs.")
 
     @classmethod
-    def _init(cls, ptr: DevicePointerT, size_t size, mr: MemoryResource | None = None):
+    def _init(cls, ptr: DevicePointerT, size_t size, mr: MemoryResource | None = None, stream: Stream | None = None):
         cdef Buffer self = Buffer.__new__(cls)
         self._ptr = <intptr_t>(int(ptr))
         self._ptr_obj = ptr
         self._size = size
         self._mr = mr
+        self._alloc_stream = <cyStream>(stream) if stream is not None else None
         return self
 
     def __dealloc__(self):
-        self.close()
+        self.close(self._alloc_stream)
 
     def __reduce__(self):
+        # Must not serialize the parent's stream!
         return Buffer.from_ipc_descriptor, (self.memory_resource, self.get_ipc_descriptor())
 
     cpdef close(self, stream: Stream = None):
@@ -137,15 +147,21 @@ cdef class Buffer(_cyBuffer, MemoryResourceAttributes):
             The stream object to use for asynchronous deallocation. If None,
             the behavior depends on the underlying memory resource.
         """
+        cdef cyStream s
         if self._ptr and self._mr is not None:
-            # 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)
+                if self._alloc_stream is not None:
+                    s = self._alloc_stream
+                else:
+                    # TODO: remove this branch when from_handle takes a stream
+                    s = <cyStream>(default_stream())
+            else:
+                s = <cyStream>stream
+            self._mr._deallocate(self._ptr, self._size, s)
             self._ptr = 0
             self._mr = None
             self._ptr_obj = None
+            self._alloc_stream = None
 
     @property
     def handle(self) -> DevicePointerT:
@@ -206,16 +222,19 @@ cdef class Buffer(_cyBuffer, MemoryResourceAttributes):
         return IPCBufferDescriptor._init(data_b, self.size)
 
     @classmethod
-    def from_ipc_descriptor(cls, mr: DeviceMemoryResource, ipc_buffer: IPCBufferDescriptor) -> Buffer:
+    def from_ipc_descriptor(cls, mr: DeviceMemoryResource, ipc_buffer: IPCBufferDescriptor, stream: Stream = None) -> Buffer:
         """Import a buffer that was exported from another process."""
         if not mr.is_ipc_enabled:
             raise RuntimeError("Memory resource is not IPC-enabled")
+        if stream is None:
+            # Note: match this behavior to DeviceMemoryResource.allocate()
+            stream = default_stream()
         cdef cydriver.CUmemPoolPtrExportData share_data
         memcpy(share_data.reserved, <const void*><const char*>(ipc_buffer._reserved), sizeof(share_data.reserved))
         cdef cydriver.CUdeviceptr ptr
         with nogil:
             HANDLE_RETURN(cydriver.cuMemPoolImportPointer(&ptr, mr._mempool_handle, &share_data))
-        return Buffer.from_handle(<intptr_t>ptr, ipc_buffer.size, mr)
+        return Buffer._init(<intptr_t>ptr, ipc_buffer.size, mr, stream)
 
     def copy_to(self, dst: Buffer = None, *, stream: Stream) -> Buffer:
         """Copy from this buffer to the dst buffer asynchronously on the given stream.
@@ -336,6 +355,7 @@ cdef class Buffer(_cyBuffer, MemoryResourceAttributes):
         mr : :obj:`~_memory.MemoryResource`, optional
             Memory resource associated with the buffer
         """
+        # TODO: It is better to take a stream for latter deallocation
         return Buffer._init(ptr, size, mr=mr)
 
 
@@ -839,7 +859,7 @@ cdef class DeviceMemoryResource(MemoryResource):
         cdef int handle = int(alloc_handle)
         with nogil:
             HANDLE_RETURN(cydriver.cuMemPoolImportFromShareableHandle(
-                &(self._mempool_handle), <void*>handle, _IPC_HANDLE_TYPE, 0)
+                &(self._mempool_handle), <void*><intptr_t>(handle), _IPC_HANDLE_TYPE, 0)
             )
         if uuid is not None:
             registered = self.register(uuid)
@@ -889,6 +909,7 @@ cdef class DeviceMemoryResource(MemoryResource):
         buf._ptr_obj = None
         buf._size = size
         buf._mr = self
+        buf._alloc_stream = stream
         return buf
 
     def allocate(self, size_t size, stream: Stream = None) -> Buffer:
@@ -931,10 +952,9 @@ cdef class DeviceMemoryResource(MemoryResource):
             The size of the buffer to deallocate, in bytes.
         stream : Stream, optional
             The stream on which to perform the deallocation asynchronously.
-            If None, an internal stream is used.
+            If the buffer is deallocated without an explicit stream, the allocation stream
+            is used.
         """
-        if stream is None:
-            stream = default_stream()
         self._deallocate(<intptr_t>ptr, size, <cyStream>stream)
 
     @property
@@ -1017,11 +1037,13 @@ class LegacyPinnedMemoryResource(MemoryResource):
         Buffer
             The allocated buffer object, which is accessible on both host and device.
         """
+        if stream is None:
+            stream = default_stream()
         err, ptr = driver.cuMemAllocHost(size)
         raise_if_driver_error(err)
-        return Buffer._init(ptr, size, self)
+        return Buffer._init(ptr, size, self, stream)
 
-    def deallocate(self, ptr: DevicePointerT, size_t size, stream: Stream = None):
+    def deallocate(self, ptr: DevicePointerT, size_t size, stream: Stream):
         """Deallocate a buffer previously allocated by this resource.
 
         Parameters
@@ -1030,12 +1052,10 @@ class LegacyPinnedMemoryResource(MemoryResource):
             The pointer or handle to the buffer to deallocate.
         size : int
             The size of the buffer to deallocate, in bytes.
-        stream : Stream, optional
-            The stream on which to perform the deallocation asynchronously.
-            If None, no synchronization would happen.
+        stream : Stream
+            The stream on which to perform the deallocation synchronously.
         """
-        if stream:
-            stream.sync()
+        stream.sync()
         err, = driver.cuMemFreeHost(ptr)
         raise_if_driver_error(err)
 
@@ -1063,13 +1083,13 @@ class _SynchronousMemoryResource(MemoryResource):
         self._dev_id = getattr(device_id, 'device_id', device_id)
 
     def allocate(self, size, stream=None) -> Buffer:
+        if stream is None:
+            stream = default_stream()
         err, ptr = driver.cuMemAlloc(size)
         raise_if_driver_error(err)
         return Buffer._init(ptr, size, self)
 
-    def deallocate(self, ptr, size, stream=None):
-        if stream is None:
-            stream = default_stream()
+    def deallocate(self, ptr, size, stream):
         stream.sync()
         err, = driver.cuMemFree(ptr)
         raise_if_driver_error(err)