Initial weakref work

Author: vzhurba01

cuda_core/cuda/core/experimental/_graph.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+import weakref
 from typing import TYPE_CHECKING, Optional, Tuple
 
 if TYPE_CHECKING:
@@ -42,34 +43,55 @@ class GraphBuilder:
     :obj:`~_device.Device`, or a :obj:`~_stream.stream` object
     """
 
+    class _MembersNeededForFinalize:
+        __slots__ = ("stream", "graph")
+
+        def __init__(self, graph_builder_obj, stream_obj):
+            self.stream = stream_obj
+            self.graph = None
+            weakref.finalize(graph_builder_obj, self.close)
+
+        def close(self):
+            # FIXME: Are the stream and graph builder racing for the weakref callback?
+            #        If so, maybe we need to enforce that all capture is completed
+            status = handle_return(driver.cuStreamGetCaptureInfo(self.stream.handle))[0]
+            if status != driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_NONE:
+                # Callback routine needs to end capture for error free handling
+                handle_return(driver.cuStreamEndCapture(self.stream.handle))
+            if self.graph:
+                handle_return(driver.cuGraphDestroy(self.graph))
+            self.graph = None
+
+    __slots__ = ("__weakref__", "_mnff", "_is_primary", "_capturing")
+
     def __init__(self):
         raise NotImplementedError(
             "directly creating a Graph object can be ambiguous. Please either "
             "call Device.create_graph() or stream.creating_graph()"
         )
 
     @staticmethod
-    def _init(stream, is_root=True):
+    def _init(stream, _is_primary=True):
         self = GraphBuilder.__new__(GraphBuilder)
         # TODO: I need to know if we own this stream object.
         #       If from Device(), then we can destroy it on close
         #       If from Stream, then we can't
-        self._stream = stream
         self._capturing = False
-        self._is_root = is_root  # TODO: Is this info needed?
+        self._is_primary = _is_primary
+        self._mnff = GraphBuilder._MembersNeededForFinalize(self, stream)
         return self
 
     def _check_capture_stream_provided(self, *args, **kwargs):
-        if self._stream == None:
+        if self._mnff.stream == None:
             raise RuntimeError("Tried to use a stream capture operation on a graph builder without a stream")
 
     @property
     def legacy_stream_capture(self) -> Stream:
-        return self._stream
+        return self._mnff.stream
 
     @property
-    def is_root_builder(self) -> bool:
-        return self._is_root_builder
+    def is_primary(self) -> bool:
+        return self._is_primary
 
     @precondition(_check_capture_stream_provided)
     def begin_capture(self, mode="global"):
@@ -84,12 +106,12 @@ def begin_capture(self, mode="global"):
         else:
             raise ValueError(f"Only 'global', 'local' or 'relaxed' capture mode are supported, got {capture_mode}")
 
-        handle_return(driver.cuStreamBeginCapture(self._stream.handle, capture_mode))
+        handle_return(driver.cuStreamBeginCapture(self._mnff.stream.handle, capture_mode))
         self._capturing = True
 
     @precondition(_check_capture_stream_provided)
     def is_capture_active(self) -> bool:
-        result = handle_return(driver.cuStreamGetCaptureInfo(self._stream.handle))
+        result = handle_return(driver.cuStreamGetCaptureInfo(self._mnff.stream.handle))
 
         capture_status = result[0]
         if capture_status == driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_NONE:
@@ -107,34 +129,33 @@ def is_capture_active(self) -> bool:
     @precondition(_check_capture_stream_provided)
     def end_capture(self):
         if not self._capturing:
-            raise RuntimeError("Stream is not capturing")
-
-        self._graph = handle_return(driver.cuStreamEndCapture(self._stream.handle))
+            raise RuntimeError("Stream is not capturing. Did you forget to call begin_capture()?")
+        self._mnff.graph = handle_return(driver.cuStreamEndCapture(self.stream.handle))
         self._capturing = False
 
     def debug_dot_print(self, path, options: Optional[DebugPrintOptions] = None):
         # TODO: We should be able to print one while the capture is happening right? Just need to make sure driver version is new enough.
-        if self._graph == None:
+        if self._mnff.graph == None:
             raise RuntimeError("Graph needs to be built before generating a DOT debug file")
 
         # TODO: Apply each option to the value
         options_value = 0
 
-        handle_return(driver.cuGraphDebugDotPrint(self._graph, path, options_value))
+        handle_return(driver.cuGraphDebugDotPrint(self._mnff.graph, path, options_value))
 
     def fork(self, count) -> Tuple[Graph, ...]:
         if count <= 1:
             raise ValueError(f"Invalid fork count: expecting >= 2, got {count}")
 
         # 1. Record an event on our stream
-        event = self._stream.record()
+        event = self._mnff.stream.record()
 
         # TODO: Steps 2,3,4 can be combined under a single loop
 
         # 2. Create a streams for each of the new forks
         # TODO: Optimization where one of the fork stream is allowed to use
         # TODO: Should use the same stream options as initial stream??
-        fork_stream = [self._stream.device.create_stream() for i in range(count)]
+        fork_stream = [self._mnff.stream.device.create_stream() for i in range(count)]
 
         # 3. Have each new stream wait on our singular event
         for stream in fork_stream:
@@ -145,15 +166,21 @@ def fork(self, count) -> Tuple[Graph, ...]:
         event.close()
 
         # 5. Create new graph builders for each new stream fork
-        return [GraphBuilder._init(stream=stream, is_root=False) for stream in fork_stream]
+        return [GraphBuilder._init(stream=stream, is_primary=False) for stream in fork_stream]
 
     def join(self, *graph_builders):
         if len(graph_builders) < 1:
             raise ValueError("Must specify which graphs should join but none were given")
 
+        # Assert that none of the graph_builders are primary
+        for graph in graph_builders:
+            if graph.is_primary:
+                raise ValueError("The primary graph builder should not be joined. Others builders should instead be joined onto it.")
+
         for graph in graph_builders:
-            self._stream.wait(graph.legacy_stream_capture)
-            # TODO: Can we close each of those new streams? Do we need weakref?
+            self._mnff.stream.wait(graph.legacy_stream_capture)
+            # TODO: Do we close them now or let weakref handle it during garbage collection?
+            #       This is a perf question, is there a good default?
             graph.close()
 
     def create_conditional_handle(self, default_value=None):
@@ -169,10 +196,10 @@ def switch(self, handle, count):
         pass
 
     def close(self):
-        if self._capturing:
-            raise RuntimeError("Trying to close a graph builder who is still capturing")
-        # Can we call this directly? Or relying on weakref enough?
-        self._stream.close()
+        if self._mnff.capturing:
+            # Explicitly trying to close a graph builder who is still capturing is not allowed
+            raise RuntimeError("Trying to close a graph builder who is still capturing. Did you forget to call end_capture()?")
+        self._mnff.close()
 
 
 class Graph: