Adjust deallocation stream for legacy memory resources to avoid platform-dependent errors. Add dependence on mempool_device where needed for certain tests.

Author: Andy-Jost

cuda_core/cuda/core/experimental/_memory/_legacy.py

@@ -62,7 +62,8 @@ def deallocate(self, ptr: DevicePointerT, size, stream):
         stream : Stream
             The stream on which to perform the deallocation synchronously.
         """
-        stream.sync()
+        if stream is not None:
+            stream.sync()
         (err,) = driver.cuMemFreeHost(ptr)
         raise_if_driver_error(err)
 
@@ -97,10 +98,11 @@ def allocate(self, size, stream=None) -> Buffer:
             stream = default_stream()
         err, ptr = driver.cuMemAlloc(size)
         raise_if_driver_error(err)
-        return Buffer._init(ptr, size, self)
+        return Buffer._init(ptr, size, self, stream)
 
     def deallocate(self, ptr, size, stream):
-        stream.sync()
+        if stream is not None:
+            stream.sync()
         (err,) = driver.cuMemFree(ptr)
         raise_if_driver_error(err)
 

cuda_core/tests/conftest.py

@@ -110,4 +110,16 @@ def ipc_memory_resource(ipc_device):
     return mr
 
 
+@pytest.fixture
+def mempool_device():
+    """Obtains a device suitable for mempool tests, or skips."""
+    device = Device()
+    device.set_current()
+
+    if not device.properties.memory_pools_supported:
+        pytest.skip("Device does not support mempool operations")
+
+    return device
+
+
 skipif_need_cuda_headers = pytest.mark.skipif(helpers.CUDA_INCLUDE_PATH is None, reason="need CUDA header")

cuda_core/tests/test_graph_mem.py

@@ -76,10 +76,10 @@ def free(self, buffers):
 
 
 @pytest.mark.parametrize("mode", ["no_graph", "global", "thread_local", "relaxed"])
-def test_graph_alloc(init_cuda, mode):
+def test_graph_alloc(mempool_device, mode):
     """Test basic graph capture with memory allocated and deallocated by GraphMemoryResource."""
     NBYTES = 64
-    device = Device()
+    device = mempool_device
     stream = device.create_stream()
     dmr = DeviceMemoryResource(device)
     gmr = GraphMemoryResource(device)
@@ -118,10 +118,10 @@ def apply_kernels(mr, stream, out):
 
 @pytest.mark.skipif(IS_WINDOWS or IS_WSL, reason="auto_free_on_launch not supported on Windows")
 @pytest.mark.parametrize("mode", ["global", "thread_local", "relaxed"])
-def test_graph_alloc_with_output(init_cuda, mode):
+def test_graph_alloc_with_output(mempool_device, mode):
     """Test for memory allocated in a graph being used outside the graph."""
     NBYTES = 64
-    device = Device()
+    device = mempool_device
     stream = device.create_stream()
     gmr = GraphMemoryResource(device)
 
@@ -157,8 +157,8 @@ def test_graph_alloc_with_output(init_cuda, mode):
 
 
 @pytest.mark.parametrize("mode", ["global", "thread_local", "relaxed"])
-def test_graph_mem_set_attributes(init_cuda, mode):
-    device = Device()
+def test_graph_mem_set_attributes(mempool_device, mode):
+    device = mempool_device
     stream = device.create_stream()
     gmr = GraphMemoryResource(device)
     mman = GraphMemoryTestManager(gmr, stream, mode)

cuda_core/tests/test_memory.py

@@ -39,18 +39,6 @@
 POOL_SIZE = 2097152  # 2MB size
 
 
-@pytest.fixture(scope="function")
-def mempool_device():
-    """Obtains a device suitable for mempool tests, or skips."""
-    device = Device()
-    device.set_current()
-
-    if not device.properties.memory_pools_supported:
-        pytest.skip("Device does not support mempool operations")
-
-    return device
-
-
 class DummyDeviceMemoryResource(MemoryResource):
     def __init__(self, device):
         self.device = device