Merge branch 'main' into bump_cuda_core

Author: leofang

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":

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

cuda_core/tests/memory_ipc/test_memory_ipc.py

@@ -38,13 +38,15 @@ def test_main(self, ipc_device, ipc_memory_resource):
 
         # Verify that the buffer was modified.
         helper.verify_buffer(flipped=True)
+        buffer.close()
 
     def child_main(self, device, mr, queue):
         device.set_current()
         buffer = queue.get(timeout=CHILD_TIMEOUT_SEC)
         helper = IPCBufferTestHelper(device, buffer)
         helper.verify_buffer(flipped=False)
         helper.fill_buffer(flipped=True)
+        buffer.close()
 
 
 class TestIPCMempoolMultiple:
@@ -78,6 +80,8 @@ def test_main(self, ipc_device, ipc_memory_resource):
         # Verify that the buffers were modified.
         IPCBufferTestHelper(device, buffer1).verify_buffer(flipped=False)
         IPCBufferTestHelper(device, buffer2).verify_buffer(flipped=True)
+        buffer1.close()
+        buffer2.close()
 
     def child_main(self, device, mr, idx, queue):
         # Note: passing the mr registers it so that buffers can be passed
@@ -89,6 +93,8 @@ def child_main(self, device, mr, idx, queue):
             IPCBufferTestHelper(device, buffer1).fill_buffer(flipped=False)
         elif idx == 2:
             IPCBufferTestHelper(device, buffer2).fill_buffer(flipped=True)
+        buffer1.close()
+        buffer2.close()
 
 
 class TestIPCSharedAllocationHandleAndBufferDescriptors:
@@ -110,10 +116,10 @@ def test_main(self, ipc_device, ipc_memory_resource):
         p2.start()
 
         # Allocate and share memory.
-        buf1 = mr.allocate(NBYTES)
-        buf2 = mr.allocate(NBYTES)
-        q1.put(buf1.get_ipc_descriptor())
-        q2.put(buf2.get_ipc_descriptor())
+        buffer1 = mr.allocate(NBYTES)
+        buffer2 = mr.allocate(NBYTES)
+        q1.put(buffer1.get_ipc_descriptor())
+        q2.put(buffer2.get_ipc_descriptor())
 
         # Wait for children.
         p1.join(timeout=CHILD_TIMEOUT_SEC)
@@ -122,8 +128,10 @@ def test_main(self, ipc_device, ipc_memory_resource):
         assert p2.exitcode == 0
 
         # Verify results.
-        IPCBufferTestHelper(device, buf1).verify_buffer(starting_from=1)
-        IPCBufferTestHelper(device, buf2).verify_buffer(starting_from=2)
+        IPCBufferTestHelper(device, buffer1).verify_buffer(starting_from=1)
+        IPCBufferTestHelper(device, buffer2).verify_buffer(starting_from=2)
+        buffer1.close()
+        buffer2.close()
 
     def child_main(self, device, alloc_handle, idx, queue):
         """Fills a shared memory buffer."""
@@ -134,6 +142,7 @@ def child_main(self, device, alloc_handle, idx, queue):
         buffer_descriptor = queue.get(timeout=CHILD_TIMEOUT_SEC)
         buffer = Buffer.from_ipc_descriptor(mr, buffer_descriptor)
         IPCBufferTestHelper(device, buffer).fill_buffer(starting_from=idx)
+        buffer.close()
 
 
 class TestIPCSharedAllocationHandleAndBufferObjects:
@@ -154,10 +163,10 @@ def test_main(self, ipc_device, ipc_memory_resource):
         p2.start()
 
         # Allocate and share memory.
-        buf1 = mr.allocate(NBYTES)
-        buf2 = mr.allocate(NBYTES)
-        q1.put(buf1)
-        q2.put(buf2)
+        buffer1 = mr.allocate(NBYTES)
+        buffer2 = mr.allocate(NBYTES)
+        q1.put(buffer1)
+        q2.put(buffer2)
 
         # Wait for children.
         p1.join(timeout=CHILD_TIMEOUT_SEC)
@@ -166,8 +175,10 @@ def test_main(self, ipc_device, ipc_memory_resource):
         assert p2.exitcode == 0
 
         # Verify results.
-        IPCBufferTestHelper(device, buf1).verify_buffer(starting_from=1)
-        IPCBufferTestHelper(device, buf2).verify_buffer(starting_from=2)
+        IPCBufferTestHelper(device, buffer1).verify_buffer(starting_from=1)
+        IPCBufferTestHelper(device, buffer2).verify_buffer(starting_from=2)
+        buffer1.close()
+        buffer2.close()
 
     def child_main(self, device, alloc_handle, idx, queue):
         """Fills a shared memory buffer."""
@@ -179,3 +190,4 @@ def child_main(self, device, alloc_handle, idx, queue):
         # Now get buffers.
         buffer = queue.get(timeout=CHILD_TIMEOUT_SEC)
         IPCBufferTestHelper(device, buffer).fill_buffer(starting_from=idx)
+        buffer.close()

cuda_core/tests/memory_ipc/test_send_buffers.py

@@ -47,6 +47,7 @@ def test_ipc_send_buffers(ipc_device, nmrs):
     for buffer in buffers:
         helper = IPCBufferTestHelper(device, buffer)
         helper.verify_buffer(flipped=True)
+        buffer.close()
 
 
 def child_main(device, buffers):
@@ -55,3 +56,4 @@ def child_main(device, buffers):
         helper = IPCBufferTestHelper(device, buffer)
         helper.verify_buffer(flipped=False)
         helper.fill_buffer(flipped=True)
+        buffer.close()

cuda_core/tests/memory_ipc/test_serialize.py

@@ -48,6 +48,8 @@ def test_main(self, ipc_device, ipc_memory_resource):
         # Confirm buffers were modified.
         IPCBufferTestHelper(device, buffer1).verify_buffer(flipped=True)
         IPCBufferTestHelper(device, buffer2).verify_buffer(flipped=True)
+        buffer1.close()
+        buffer2.close()
 
     def child_main(self, conn):
         # Set up the device.
@@ -67,6 +69,8 @@ def child_main(self, conn):
         # Modify the buffers.
         IPCBufferTestHelper(device, buffer1).fill_buffer(flipped=True)
         IPCBufferTestHelper(device, buffer2).fill_buffer(flipped=True)
+        buffer1.close()
+        buffer2.close()
 
 
 class TestObjectSerializationWithMR:
@@ -97,6 +101,7 @@ def test_main(self, ipc_device, ipc_memory_resource):
 
         # Confirm buffer was modified.
         IPCBufferTestHelper(device, buffer).verify_buffer(flipped=True)
+        buffer.close()
 
     def child_main(self, pipe, _):
         device = Device()
@@ -110,6 +115,7 @@ def child_main(self, pipe, _):
         buffer = pipe[0].get(timeout=CHILD_TIMEOUT_SEC)
         assert buffer.memory_resource.handle == mr.handle
         IPCBufferTestHelper(device, buffer).fill_buffer(flipped=True)
+        buffer.close()
 
 
 def test_object_passing(ipc_device, ipc_memory_resource):
@@ -138,6 +144,7 @@ def test_object_passing(ipc_device, ipc_memory_resource):
     assert process.exitcode == 0
 
     helper.verify_buffer(flipped=True)
+    buffer.close()
 
 
 def child_main(alloc_handle, mr1, buffer_desc, buffer1):
@@ -178,3 +185,6 @@ def child_main(alloc_handle, mr1, buffer_desc, buffer1):
     helper1.verify_buffer(flipped=True)
     helper2.verify_buffer(flipped=True)
     helper3.verify_buffer(flipped=True)
+
+    # Close any one buffer.
+    buffer1.close()

cuda_core/tests/memory_ipc/test_workerpool.py

@@ -38,11 +38,13 @@ def test_main(self, ipc_device, nmrs):
 
         for buffer in buffers:
             IPCBufferTestHelper(device, buffer).verify_buffer(flipped=True)
+            buffer.close()
 
     def process_buffer(self, buffer):
         device = Device(buffer.memory_resource.device_id)
         device.set_current()
         IPCBufferTestHelper(device, buffer).fill_buffer(flipped=True)
+        buffer.close()
 
 
 class TestIpcWorkerPoolUsingIPCDescriptors:
@@ -73,13 +75,15 @@ def test_main(self, ipc_device, nmrs):
 
         for buffer in buffers:
             IPCBufferTestHelper(device, buffer).verify_buffer(flipped=True)
+            buffer.close()
 
     def process_buffer(self, mr_idx, buffer_desc):
         mr = self.mrs[mr_idx]
         device = Device(mr.device_id)
         device.set_current()
         buffer = Buffer.from_ipc_descriptor(mr, buffer_desc)
         IPCBufferTestHelper(device, buffer).fill_buffer(flipped=True)
+        buffer.close()
 
 
 class TestIpcWorkerPoolUsingRegistry:
@@ -110,8 +114,10 @@ def test_main(self, ipc_device, nmrs):
 
         for buffer in buffers:
             IPCBufferTestHelper(device, buffer).verify_buffer(flipped=True)
+            buffer.close()
 
     def process_buffer(self, device, buffer_s):
         device.set_current()
         buffer = pickle.loads(buffer_s)  # noqa: S301
         IPCBufferTestHelper(device, buffer).fill_buffer(flipped=True)
+        buffer.close()

cuda_core/tests/test_stream.py

@@ -4,7 +4,7 @@
 import pytest
 from cuda.core.experimental import Device, Stream, StreamOptions
 from cuda.core.experimental._event import Event
-from cuda.core.experimental._stream import LEGACY_DEFAULT_STREAM, PER_THREAD_DEFAULT_STREAM, default_stream
+from cuda.core.experimental._stream import LEGACY_DEFAULT_STREAM, PER_THREAD_DEFAULT_STREAM
 from cuda.core.experimental._utils.cuda_utils import driver
 
 
@@ -107,11 +107,6 @@ def test_per_thread_default_stream():
     assert isinstance(PER_THREAD_DEFAULT_STREAM, Stream)
 
 
-def test_default_stream():
-    stream = default_stream()
-    assert isinstance(stream, Stream)
-
-
 def test_stream_subclassing(init_cuda):
     class MyStream(Stream):
         pass