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

Author: Andy-Jost

cuda_core/cuda/core/experimental/_memory/_graph_memory_resource.pyx

@@ -148,16 +148,13 @@ cdef class cyGraphMemoryResource(MemoryResource):
 
 class GraphMemoryResource(cyGraphMemoryResource):
     """
-    A memory resource managing the graph-specific memory pool.
+    A memory resource for memory related to graphs.
 
-    Graph-captured memory operations use a special internal memory pool, which
-    is a per-device singleton. This class serves as the interface to that pool.
     The only supported operations are allocation, deallocation, and a limited
     set of status queries.
 
-    This memory resource should be used to allocate memory when graph capturing
-    is enabled. Using this when graphs are not being captured will result in a
-    runtime error.
+    This memory resource should be used when building graphs.  Using this when
+    graphs capture is not enabled will result in a runtime error.
 
     Conversely, allocating memory from a `DeviceMemoryResource` when graph
     capturing is enabled results in a runtime error.

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/cuda/core/experimental/_memory/_virtual_memory_resource.py

@@ -20,7 +20,6 @@
 )
 
 if TYPE_CHECKING:
-    from cuda.core.experimental._graph import GraphBuilder
     from cuda.core.experimental._stream import Stream
 
 __all__ = ["VirtualMemoryResourceOptions", "VirtualMemoryResource"]
@@ -465,7 +464,7 @@ def _build_access_descriptors(self, prop: driver.CUmemAllocationProp) -> list:
 
         return descs
 
-    def allocate(self, size: int, stream: Stream | GraphBuilder | None = None) -> Buffer:
+    def allocate(self, size: int, stream: Stream | None = None) -> Buffer:
         """
         Allocate a buffer of the given size using CUDA virtual memory.
 
@@ -549,7 +548,7 @@ def allocate(self, size: int, stream: Stream | GraphBuilder | None = None) -> Bu
         buf = Buffer.from_handle(ptr=ptr, size=aligned_size, mr=self)
         return buf
 
-    def deallocate(self, ptr: int, size: int, stream: Stream | GraphBuilder | None = None) -> None:
+    def deallocate(self, ptr: int, size: int, stream: Stream | None = None) -> None:
         """
         Deallocate memory on the device using CUDA VMM APIs.
         """

cuda_core/cuda/core/experimental/_stream.pyx

@@ -453,9 +453,9 @@ cdef cydriver.CUstream _handle_from_stream_protocol(obj) except*:
 # needed checks and returns the relevant stream.
 cdef Stream Stream_accept(arg, bint allow_stream_protocol=False):
     if isinstance(arg, Stream):
-        return <Stream> arg
+        return <Stream>(arg)
     elif isinstance(arg, GraphBuilder):
-        return <Stream> arg.stream
+        return <Stream>(arg.stream)
     elif allow_stream_protocol:
         try:
             stream = Stream._init(arg)
@@ -469,5 +469,5 @@ cdef Stream Stream_accept(arg, bint allow_stream_protocol=False):
                 stacklevel=2,
                 category=DeprecationWarning,
             )
-            return <Stream> stream
+            return <Stream>(stream)
     raise TypeError(f"Stream or GraphBuilder expected, got {type(arg).__name__}")

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

@@ -34,8 +34,7 @@ def _common_kernels_alloc():
         }
     }
     """
-    arch = "".join(f"{i}" for i in Device().compute_capability)
-    program_options = ProgramOptions(std="c++17", arch=f"sm_{arch}")
+    program_options = ProgramOptions(std="c++17", arch=f"sm_{Device().arch}")
     prog = Program(code, code_type="c++", options=program_options)
     mod = prog.compile("cubin", name_expressions=("set_zero", "add_one"))
     return mod
@@ -76,10 +75,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 +117,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 +156,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)
@@ -209,12 +208,12 @@ def test_graph_mem_set_attributes(init_cuda, mode):
 
 
 @pytest.mark.parametrize("mode", ["global", "thread_local", "relaxed"])
-def test_gmr_check_capture_state(init_cuda, mode):
+def test_gmr_check_capture_state(mempool_device, mode):
     """
     Test expected errors (and non-errors) using GraphMemoryResource with graph
     capture.
     """
-    device = Device()
+    device = mempool_device
     stream = device.create_stream()
     gmr = GraphMemoryResource(device)
 
@@ -233,12 +232,12 @@ def test_gmr_check_capture_state(init_cuda, mode):
 
 
 @pytest.mark.parametrize("mode", ["global", "thread_local", "relaxed"])
-def test_dmr_check_capture_state(init_cuda, mode):
+def test_dmr_check_capture_state(mempool_device, mode):
     """
     Test expected errors (and non-errors) using DeviceMemoryResource with graph
     capture.
     """
-    device = Device()
+    device = mempool_device
     stream = device.create_stream()
     dmr = DeviceMemoryResource(device)
 

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