fix(cuda.core): preserve ordering for asynchronous memory operations in tests

cuda_core/examples/memory_pool_resources.py

@@ -89,13 +89,13 @@ def main():
 
         managed_buffer = managed_mr.allocate(nbytes, stream=stream)
         pinned_buffer = pinned_mr.allocate(nbytes, stream=stream)
+        stream.sync()
 
         managed_array = np.from_dlpack(managed_buffer).view(np.float32)
         pinned_array = np.from_dlpack(pinned_buffer).view(np.float32)
 
         managed_array[:] = np.arange(size, dtype=dtype)
         managed_original = managed_array.copy()
-        stream.sync()
 
         managed_buffer.copy_to(pinned_buffer, stream=stream)
         stream.sync()

cuda_core/examples/strided_memory_view_constructors.py

@@ -18,7 +18,7 @@
 import cupy as cp
 import numpy as np
 
-from cuda.core import Device
+from cuda.core import Device, Stream
 from cuda.core.utils import StridedMemoryView
 
 
@@ -39,7 +39,8 @@ def main():
 
     device = Device()
     device.set_current()
-    stream = device.create_stream()
+    cupy_stream = cp.cuda.get_current_stream()
+    stream = Stream.from_handle(cupy_stream.ptr)
     buffer = None
 
     try:
@@ -63,7 +64,6 @@ def main():
         buffer = device.memory_resource.allocate(gpu_array.nbytes, stream=stream)
         buffer_array = cp.from_dlpack(buffer).view(dtype=cp.float32).reshape(gpu_array.shape)
         buffer_array[...] = gpu_array
-        device.sync()
 
         buffer_view = StridedMemoryView.from_buffer(
             buffer,
@@ -77,7 +77,6 @@ def main():
     finally:
         if buffer is not None:
             buffer.close(stream)
-        stream.close()
 
 
 if __name__ == "__main__":

cuda_core/tests/helpers/buffers.py

@@ -146,8 +146,8 @@ class PatternGen:
     Provides methods to fill a target buffer with  known test patterns and
     verify the expected values.
 
-    If a stream is provided, operations are synchronized with respect to that
-    stream.  Otherwise, they are synchronized over the device.
+    Operations are submitted to the supplied stream. Verification synchronizes
+    that stream before comparing results on the host.
 
     The test pattern is either a fixed value or a cyclic pattern generated from
     an 8-bit seed.  Only one of `value` or `seed` should be supplied.
@@ -158,11 +158,10 @@ class PatternGen:
     buffer and then perform a comparison.
     """
 
-    def __init__(self, device, size, stream=None):
+    def __init__(self, device, size, *, stream):
         self.device = device
         self.size = size
-        self.stream = stream if stream is not None else device.create_stream()
-        self.sync_target = stream if stream is not None else device
+        self.stream = Stream_accept(stream)
         self.pattern_buffers = {}
 
     def fill_buffer(self, buffer, seed=None, value=None):
@@ -179,7 +178,7 @@ def verify_buffer(self, buffer, seed=None, value=None):
         pattern_buffer = self._get_pattern_buffer(seed, value)
         ptr_expected = self._ptr(pattern_buffer)
         scratch_buffer.copy_from(buffer, stream=self.stream)
-        self.sync_target.sync()
+        self.stream.sync()
         assert libc.memcmp(ptr_test, ptr_expected, self.size) == 0
 
     @staticmethod

cuda_core/tests/memory/test_managed_ops.py

@@ -105,6 +105,7 @@ def managed_buffer(request, location_ops_device, location_ops_mr):
     size = _MANAGED_TEST_ALLOCATION_SIZE
     if request.param == "pool":
         buf = location_ops_mr.allocate(size, stream=location_ops_device.default_stream)
+        location_ops_device.default_stream.sync()
         yield buf
         buf.close()
     else:
@@ -216,8 +217,8 @@ def test_same_location(self, location_ops_device, location_ops_mr):
         from cuda.core.utils import prefetch_batch
 
         device = location_ops_device
-        bufs = [location_ops_mr.allocate(_MANAGED_TEST_ALLOCATION_SIZE, stream=device.default_stream) for _ in range(3)]
         stream = device.create_stream()
+        bufs = [location_ops_mr.allocate(_MANAGED_TEST_ALLOCATION_SIZE, stream=stream) for _ in range(3)]
 
         prefetch_batch(stream, bufs, device)
         stream.sync()
@@ -231,12 +232,12 @@ def test_per_buffer_location(self, location_ops_device, location_ops_mr):
         from cuda.core.utils import prefetch_batch
 
         device = location_ops_device
-        bufs = [location_ops_mr.allocate(_MANAGED_TEST_ALLOCATION_SIZE, stream=device.default_stream) for _ in range(2)]
+        stream = device.create_stream()
+        bufs = [location_ops_mr.allocate(_MANAGED_TEST_ALLOCATION_SIZE, stream=stream) for _ in range(2)]
         # Per-buffer prefetch locations are only observable when the buffers sit
         # on distinct physical pages; assert that here so a pool-packing change
         # fails loudly instead of silently migrating one shared page.
         assert _page_base(bufs[0]) != _page_base(bufs[1])
-        stream = device.create_stream()
 
         prefetch_batch(stream, bufs, [Host(), device])
         stream.sync()
@@ -258,8 +259,8 @@ def test_basic(self, location_ops_device, location_ops_mr):
         if not hasattr(driver, "cuMemDiscardBatchAsync"):
             pytest.skip("cuMemDiscardBatchAsync unavailable")
         device = location_ops_device
-        bufs = [location_ops_mr.allocate(_MANAGED_TEST_ALLOCATION_SIZE, stream=device.default_stream) for _ in range(3)]
         stream = device.create_stream()
+        bufs = [location_ops_mr.allocate(_MANAGED_TEST_ALLOCATION_SIZE, stream=stream) for _ in range(3)]
         prefetch_batch(stream, bufs, device)
         stream.sync()
         discard_batch(stream, bufs)
@@ -277,8 +278,8 @@ def test_same_location(self, location_ops_device, location_ops_mr):
         if not hasattr(driver, "cuMemDiscardAndPrefetchBatchAsync"):
             pytest.skip("cuMemDiscardAndPrefetchBatchAsync unavailable")
         device = location_ops_device
-        bufs = [location_ops_mr.allocate(_MANAGED_TEST_ALLOCATION_SIZE, stream=device.default_stream) for _ in range(2)]
         stream = device.create_stream()
+        bufs = [location_ops_mr.allocate(_MANAGED_TEST_ALLOCATION_SIZE, stream=stream) for _ in range(2)]
         prefetch_batch(stream, bufs, Host())
         stream.sync()
         discard_prefetch_batch(stream, bufs, device)
@@ -487,9 +488,9 @@ def test_instance_discard(self, location_ops_device, managed_buffer):
     def test_instance_discard_prefetch(self, discard_prefetch_device):
         device = discard_prefetch_device
         mr = create_managed_memory_resource_or_skip()
-        buf = mr.allocate(_MANAGED_TEST_ALLOCATION_SIZE, stream=device.default_stream)
+        stream = device.create_stream()
+        buf = mr.allocate(_MANAGED_TEST_ALLOCATION_SIZE, stream=stream)
         try:
-            stream = device.create_stream()
             buf.prefetch(Host(), stream=stream)
             stream.sync()
             buf.discard_prefetch(device, stream=stream)

cuda_core/tests/memory_ipc/test_errors.py

@@ -120,9 +120,11 @@ class TestImportOversizedBufferDescriptorSize(ChildErrorHarness):
     """Reject peer-supplied sizes larger than the mapped allocation extent."""
 
     def PARENT_ACTION(self, queue):
-        self.buffer = self.mr.allocate(NBYTES, stream=self.device.default_stream)
+        stream = self.device.default_stream
+        self.buffer = self.mr.allocate(NBYTES, stream=stream)
         payload, _ = self.buffer.ipc_descriptor.__reduce__()[1]
         oversized = IPCBufferDescriptor._init(payload, NBYTES * 100)
+        stream.sync()
         queue.put(oversized)
 
     def CHILD_ACTION(self, queue):
@@ -157,8 +159,10 @@ def PARENT_ACTION(self, queue):
         options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True)
         mr2 = DeviceMemoryResource(self.device, options=options)
         self._extra_mrs.append(mr2)
-        buffer = mr2.allocate(NBYTES, stream=self.device.default_stream)
-        queue.put([self.mr, buffer.ipc_descriptor])  # Note: mr does not own this buffer
+        stream = self.device.default_stream
+        self.buffer = mr2.allocate(NBYTES, stream=stream)
+        stream.sync()
+        queue.put([self.mr, self.buffer.ipc_descriptor])  # Note: mr does not own this buffer
 
     def CHILD_ACTION(self, queue):
         mr, buffer_desc = queue.get(timeout=CHILD_TIMEOUT_SEC)
@@ -175,7 +179,9 @@ class TestImportBuffer(ChildErrorHarness):
     def PARENT_ACTION(self, queue):
         # Note: if the buffer is not attached to something to prolong its life,
         # CUDA_ERROR_INVALID_CONTEXT is raised from Buffer.__del__
-        self.buffer = self.mr.allocate(NBYTES, stream=self.device.default_stream)
+        stream = self.device.default_stream
+        self.buffer = self.mr.allocate(NBYTES, stream=stream)
+        stream.sync()
         queue.put(self.buffer)
 
     def CHILD_ACTION(self, queue):
@@ -197,8 +203,10 @@ def PARENT_ACTION(self, queue):
         options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True)
         mr2 = DeviceMemoryResource(self.device, options=options)
         self._extra_mrs.append(mr2)
-        self.buffer = mr2.allocate(NBYTES, stream=self.device.default_stream)
+        stream = self.device.default_stream
+        self.buffer = mr2.allocate(NBYTES, stream=stream)
         buffer_s = pickle.dumps(self.buffer)
+        stream.sync()
         queue.put(buffer_s)  # Note: mr2 not sent
 
     def CHILD_ACTION(self, queue):

cuda_core/tests/memory_ipc/test_ipc_duplicate_import.py

@@ -71,7 +71,9 @@ def test_main(self, ipc_device, ipc_memory_resource):
         mr = ipc_memory_resource
 
         log("allocating buffer")
-        buffer = mr.allocate(NBYTES, stream=ipc_device.default_stream)
+        stream = ipc_device.default_stream
+        buffer = mr.allocate(NBYTES, stream=stream)
+        stream.sync()
 
         # Start the child process.
         log("starting child")

cuda_core/tests/memory_ipc/test_leaks.py

@@ -102,8 +102,10 @@ def __reduce__(self):
 def test_pass_object(ipc_device, ipc_memory_resource, launcher, getobject):
     """Check for fd leaks when an object is sent as a subprocess argument."""
     mr = ipc_memory_resource
+    stream = ipc_device.default_stream
     with CheckFDLeaks():
-        obj = getobject(mr, ipc_device.default_stream)
+        obj = getobject(mr, stream)
+        stream.sync()
         try:
             launcher(obj, number=2)
         finally: