Memory refactor (#1205)

* Resolve a Cython build warning.

* Make memory module into a package.

* Rename cyStream to _cyStream for consistency.

* Move defs to memory.pxd header

* Separate VMM.

* Weaken dependencies from device to memory module.

* Move LegacyPinnedMemoryResource to a submodule.

* Move _SynchronousMemoryResource into a submodule.

* Partly separates the IPC implementation.

* Move IPC registry to ipc module.

* Collect and reorder DeviceMemoryResource properties.

* Move more IPC implementation out of DeviceMemoryResource.

* Minor refactoring.

* Move Buffer IPC implementation.

* Simplify the class hierarchy (remove _cyBuffer and _cyMemoryResource).

* Refactor to shrink Cython interface.

* Simplify Buffer close.

* Refactor DeviceMemoryResource.__init__.

* Move Buffer into a separate module.

* Refactors DeviceMemoryResource IPC implementation.

* Removes superfluous _uuid member of DeviceMemoryResource.

* Adds __all__ lists.

* Prepend underscore to submodules, add a test for package contents.

* Refactor IPC data of DMR into IPCData class.

* General clean up.

* Touch-ups

* Cythonize DeviceMemoryResourceAttributes.

* Restore previous behavior for DMR.from_allocation_handle when passed a file descriptor (caller closes the fd).

* Rename _mr to _memory_resource. Change pointer types from intptr_t to uintptr_t.

* Rename files _dmr.* and _vmm.py to avoid abbreviations.

Author: Andy-Jost

cuda_core/cuda/core/experimental/_device.pyx

@@ -6,16 +6,14 @@ cimport cpython
 from libc.stdint cimport uintptr_t
 
 from cuda.bindings cimport cydriver
-
 from cuda.core.experimental._utils.cuda_utils cimport HANDLE_RETURN
 
 import threading
-from typing import Union
+from typing import Union, TYPE_CHECKING
 
 from cuda.core.experimental._context import Context, ContextOptions
 from cuda.core.experimental._event import Event, EventOptions
 from cuda.core.experimental._graph import GraphBuilder
-from cuda.core.experimental._memory import Buffer, DeviceMemoryResource, MemoryResource, _SynchronousMemoryResource
 from cuda.core.experimental._stream import IsStreamT, Stream, StreamOptions
 from cuda.core.experimental._utils.clear_error_support import assert_type
 from cuda.core.experimental._utils.cuda_utils import (
@@ -27,7 +25,8 @@ from cuda.core.experimental._utils.cuda_utils import (
 )
 from cuda.core.experimental._stream cimport default_stream
 
-
+if TYPE_CHECKING:
+    from cuda.core.experimental._memory import Buffer, MemoryResource
 
 # TODO: I prefer to type these as "cdef object" and avoid accessing them from within Python,
 # but it seems it is very convenient to expose them for testing purposes...
@@ -996,8 +995,10 @@ class Device:
                         )
                     )
                 if attr == 1:
+                    from cuda.core.experimental._memory import DeviceMemoryResource
                     device._mr = DeviceMemoryResource(dev_id)
                 else:
+                    from cuda.core.experimental._memory import _SynchronousMemoryResource
                     device._mr = _SynchronousMemoryResource(dev_id)
 
                 device._has_inited = False
@@ -1131,6 +1132,7 @@ class Device:
 
     @memory_resource.setter
     def memory_resource(self, mr):
+        from cuda.core.experimental._memory import MemoryResource
         assert_type(mr, MemoryResource)
         self._mr = mr
 

cuda_core/cuda/core/experimental/_event.pyx

@@ -7,14 +7,13 @@ from __future__ import annotations
 cimport cpython
 from libc.stdint cimport uintptr_t
 from libc.string cimport memcpy
-
 from cuda.bindings cimport cydriver
-
 from cuda.core.experimental._utils.cuda_utils cimport (
     check_or_create_options,
     HANDLE_RETURN
 )
 
+import cython
 from dataclasses import dataclass
 import multiprocessing
 from typing import TYPE_CHECKING, Optional
@@ -287,7 +286,7 @@ cdef class IPCEventDescriptor:
         raise RuntimeError("IPCEventDescriptor objects cannot be instantiated directly. Please use Event APIs.")
 
     @classmethod
-    def _init(cls, reserved: bytes, busy_waited: bint):
+    def _init(cls, reserved: bytes, busy_waited: cython.bint):
         cdef IPCEventDescriptor self = IPCEventDescriptor.__new__(cls)
         self._reserved = reserved
         self._busy_waited = busy_waited