Merge branch 'main' into rparolin/skip_ipc_on_wsl

Author: leofang

.github/workflows/ci.yml

@@ -193,6 +193,7 @@ jobs:
 
   checks:
     name: Check job status
+    if: always()
     runs-on: ubuntu-latest
     needs:
       - test-linux-64

cuda_bindings/setup.py

@@ -329,7 +329,7 @@ def do_cythonize(extensions):
     return cythonize(
         extensions,
         nthreads=nthreads,
-        compiler_directives=dict(language_level=3, embedsignature=True, binding=True),
+        compiler_directives=dict(language_level=3, embedsignature=True, binding=True, freethreading_compatible=True),
         **extra_cythonize_kwargs,
     )
 

cuda_bindings/tests/cython/build_tests.bat

@@ -5,5 +5,5 @@ REM SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
 setlocal
 	set CL=%CL% /I"%CUDA_HOME%\include"
-	cythonize -3 -i %~dp0test_*.pyx
+	cythonize -3 -i -Xfreethreading_compatible=True %~dp0test_*.pyx
 endlocal

cuda_bindings/tests/cython/build_tests.sh

@@ -14,4 +14,4 @@ else
   exit 1
 fi
 
-cythonize -3 -i ${SCRIPTPATH}/test_*.pyx
+cythonize -3 -i -Xfreethreading_compatible=True ${SCRIPTPATH}/test_*.pyx

cuda_core/build_hooks.py

@@ -101,7 +101,7 @@ def get_cuda_paths():
         verbose=True,
         language_level=3,
         nthreads=nthreads,
-        compiler_directives={"embedsignature": True, "warn.deprecated.IF": False},
+        compiler_directives={"embedsignature": True, "warn.deprecated.IF": False, "freethreading_compatible": True},
         compile_time_env=compile_time_env,
     )
 

cuda_core/cuda/core/_version.py

@@ -2,4 +2,4 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-__version__ = "0.3.3a0"
+__version__ = "0.4.0"

cuda_core/cuda/core/experimental/_device.pyx

@@ -16,7 +16,7 @@ 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, default_stream
+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 (
     ComputeCapability,
@@ -25,6 +25,7 @@ from cuda.core.experimental._utils.cuda_utils import (
     handle_return,
     runtime,
 )
+from cuda.core.experimental._stream cimport default_stream
 
 
 # TODO: I prefer to type these as "cdef object" and avoid accessing them from within Python,

cuda_core/cuda/core/experimental/_memory.pyx

@@ -12,6 +12,7 @@ from libc.string cimport memset, memcpy
 from cuda.bindings cimport cydriver
 
 from cuda.core.experimental._stream cimport Stream as cyStream
+from cuda.core.experimental._stream cimport default_stream
 from cuda.core.experimental._utils.cuda_utils cimport (
     _check_driver_error as raise_if_driver_error,
     check_or_create_options,
@@ -30,7 +31,7 @@ import platform
 import weakref
 
 from cuda.core.experimental._dlpack import DLDeviceType, make_py_capsule
-from cuda.core.experimental._stream import Stream, default_stream
+from cuda.core.experimental._stream import Stream
 from cuda.core.experimental._utils.cuda_utils import ( driver, Transaction, get_binding_version )
 
 if platform.system() == "Linux":
@@ -73,6 +74,8 @@ cdef class _cyMemoryResource:
 
 class MemoryResourceAttributes(abc.ABC):
 
+    __slots__ = ()
+
     @property
     @abc.abstractmethod
     def is_device_accessible(self) -> bool:
@@ -107,8 +110,6 @@ cdef class Buffer(_cyBuffer, MemoryResourceAttributes):
 
     Support for data interchange mechanisms are provided by DLPack.
     """
-    cdef dict __dict__  # required if inheriting from both Cython/Python classes
-
     def __cinit__(self):
         self._ptr = 0
         self._size = 0
@@ -369,8 +370,6 @@ cdef class MemoryResource(_cyMemoryResource, MemoryResourceAttributes, 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
-
     cdef void _deallocate(self, intptr_t ptr, size_t size, cyStream stream) noexcept:
         self.deallocate(ptr, size, stream)
 
@@ -585,7 +584,7 @@ class DeviceMemoryResourceAttributes:
 # This enables buffer serialization, as buffers can reduce to a pair
 # of comprising the memory resource UUID (the key into this registry)
 # and the serialized buffer descriptor.
-_ipc_registry = {}
+cdef object _ipc_registry = weakref.WeakValueDictionary()
 
 
 cdef class DeviceMemoryResource(MemoryResource):
@@ -675,7 +674,6 @@ cdef class DeviceMemoryResource(MemoryResource):
         bint _is_mapped
         object _uuid
         IPCAllocationHandle _alloc_handle
-        dict __dict__  # required if inheriting from both Cython/Python classes
         object __weakref__
 
     def __cinit__(self):
@@ -759,8 +757,6 @@ cdef class DeviceMemoryResource(MemoryResource):
                 with nogil:
                     HANDLE_RETURN(cydriver.cuMemPoolDestroy(self._mempool_handle))
         finally:
-            if self.is_mapped:
-                self.unregister()
             self._dev_id = cydriver.CU_DEVICE_INVALID
             self._mempool_handle = NULL
             self._attributes = None
@@ -806,13 +802,6 @@ cdef class DeviceMemoryResource(MemoryResource):
         self._uuid = uuid
         return self
 
-    def unregister(self):
-        """Unregister this mapped memory resource."""
-        assert self.is_mapped
-        if _ipc_registry is not None:  # can occur during shutdown catastrophe
-            with contextlib.suppress(KeyError):
-                del _ipc_registry[self.uuid]
-
     @property
     def uuid(self) -> Optional[uuid.UUID]:
         """
@@ -1019,9 +1008,7 @@ class LegacyPinnedMemoryResource(MemoryResource):
     APIs.
     """
 
-    def __init__(self):
-        # TODO: support flags from cuMemHostAlloc?
-        self._handle = None
+    # TODO: support creating this MR with flags that are later passed to cuMemHostAlloc?
 
     def allocate(self, size_t size, stream: Stream = None) -> Buffer:
         """Allocate a buffer of the requested size.
@@ -1080,7 +1067,6 @@ class _SynchronousMemoryResource(MemoryResource):
     __slots__ = ("_dev_id",)
 
     def __init__(self, device_id : int | Device):
-        self._handle = None
         self._dev_id = getattr(device_id, 'device_id', device_id)
 
     def allocate(self, size, stream=None) -> Buffer:

cuda_core/cuda/core/experimental/_stream.pxd

@@ -22,3 +22,6 @@ cdef class Stream:
     cpdef close(self)
     cdef int _get_context(self) except?-1 nogil
     cdef int _get_device_and_context(self) except?-1
+
+
+cdef Stream default_stream()

cuda_core/cuda/core/experimental/_stream.pyx

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 from libc.stdint cimport uintptr_t, INT32_MIN
+from libc.stdlib cimport strtol, getenv
 
 from cuda.bindings cimport cydriver
 
@@ -388,11 +389,16 @@ cdef class Stream:
         return GraphBuilder._init(stream=self, is_stream_owner=False)
 
 
-LEGACY_DEFAULT_STREAM = Stream._legacy_default()
-PER_THREAD_DEFAULT_STREAM = Stream._per_thread_default()
+# c-only python objects, not public
+cdef Stream C_LEGACY_DEFAULT_STREAM = Stream._legacy_default()
+cdef Stream C_PER_THREAD_DEFAULT_STREAM = Stream._per_thread_default()
 
+# standard python objects, public
+LEGACY_DEFAULT_STREAM = C_LEGACY_DEFAULT_STREAM
+PER_THREAD_DEFAULT_STREAM = C_PER_THREAD_DEFAULT_STREAM
 
-def default_stream():
+
+cdef Stream default_stream():
     """Return the default CUDA :obj:`~_stream.Stream`.
 
     The type of default stream returned depends on if the environment
@@ -403,8 +409,14 @@ def default_stream():
 
     """
     # TODO: flip the default
-    use_ptds = int(os.environ.get("CUDA_PYTHON_CUDA_PER_THREAD_DEFAULT_STREAM", 0))
+    cdef const char* use_ptds_raw = getenv("CUDA_PYTHON_CUDA_PER_THREAD_DEFAULT_STREAM")
+
+    cdef int use_ptds = 0
+    if use_ptds_raw != NULL:
+        use_ptds = strtol(use_ptds_raw, NULL, 10)
+
+    # value is non-zero, including for weird stuff like 123foo
     if use_ptds:
-        return PER_THREAD_DEFAULT_STREAM
+        return C_PER_THREAD_DEFAULT_STREAM
     else:
-        return LEGACY_DEFAULT_STREAM
+        return C_LEGACY_DEFAULT_STREAM