Merge branch 'graph-node-repr' into graph-public-api

Author: Andy-Jost

cuda_core/cuda/core/_cpp/REGISTRY_DESIGN.md

@@ -0,0 +1,51 @@
+# Handle and Object Registries
+
+When Python-managed objects round-trip through the CUDA driver (e.g.,
+querying a graph's nodes and getting back raw `CUgraphNode` pointers),
+we need to recover the original Python object rather than creating a
+duplicate.
+
+This document describes the approach used to achieve this. The pattern
+is driven mainly by needs arising in the context of CUDA graphs, but
+it is general and can be extended to other object types as needs arise.
+
+This solves the same problem as pybind11's `registered_instances` map
+and is sometimes called the Identity Map pattern. Two registries work
+together to map a raw driver handle all the way back to the original
+Python object. Both use weak references so they
+do not prevent cleanup. Entries are removed either explicitly (via
+`destroy()` or a Box destructor) or implicitly when the weak reference
+expires.
+
+## Level 1: Driver Handle -> Resource Handle (C++)
+
+`HandleRegistry` in `resource_handles.cpp` maps a raw CUDA handle
+(e.g., `CUevent`, `CUkernel`, `CUgraphNode`) to the `weak_ptr` that
+owns it. When a `_ref` constructor receives a raw handle, it
+checks the registry first. If found, it returns the existing
+`shared_ptr`, preserving the Box and its metadata (e.g., `EventBox`
+carries timing/IPC flags, `KernelBox` carries the library dependency).
+
+Without this level, a round-tripped handle would produce a new Box
+with default metadata, losing information that was set at creation.
+
+Instances: `event_registry`, `kernel_registry`, `graph_node_registry`.
+
+## Level 2: Resource Handle -> Python Object (Cython)
+
+`_node_registry` in `_graph_node.pyx` is a `WeakValueDictionary`
+mapping a resource address (`shared_ptr::get()`) to a Python
+`GraphNode` object. When `GraphNode._create` receives a handle from
+Level 1, it checks this registry. If found, it returns the existing
+Python object.
+
+Without this level, each driver round-trip would produce a distinct
+Python object for the same logical node, resulting in surprising
+behavior:
+
+```python
+a = g.empty()
+a.succ = {b}
+b2, = a.succ     # queries driver, gets back CUgraphNode for b
+assert b2 is b   # fails without Level 2 registry
+```

cuda_core/cuda/core/_cpp/resource_handles.cpp

@@ -174,13 +174,8 @@ class HandleRegistry {
     }
 
     void unregister_handle(const Key& key) noexcept {
-        try {
-            std::lock_guard<std::mutex> lock(mutex_);
-            auto it = map_.find(key);
-            if (it != map_.end() && it->second.expired()) {
-                map_.erase(it);
-            }
-        } catch (...) {}
+        std::lock_guard<std::mutex> lock(mutex_);
+        map_.erase(key);
     }
 
     Handle lookup(const Key& key) {
@@ -393,6 +388,7 @@ ContextHandle get_event_context(const EventHandle& h) noexcept {
     return h ? get_box(h)->h_context : ContextHandle{};
 }
 
+// See REGISTRY_DESIGN.md (Level 1: Driver Handle -> Resource Handle)
 static HandleRegistry<CUevent, EventHandle> event_registry;
 
 EventHandle create_event_handle(const ContextHandle& h_ctx, unsigned int flags,
@@ -899,6 +895,7 @@ static const KernelBox* get_box(const KernelHandle& h) {
     );
 }
 
+// See REGISTRY_DESIGN.md (Level 1: Driver Handle -> Resource Handle)
 static HandleRegistry<CUkernel, KernelHandle> kernel_registry;
 
 KernelHandle create_kernel_handle(const LibraryHandle& h_library, const char* name) {
@@ -969,6 +966,7 @@ static const GraphNodeBox* get_box(const GraphNodeHandle& h) {
     );
 }
 
+// See REGISTRY_DESIGN.md (Level 1: Driver Handle -> Resource Handle)
 static HandleRegistry<CUgraphNode, GraphNodeHandle> graph_node_registry;
 
 GraphNodeHandle create_graph_node_handle(CUgraphNode node, const GraphHandle& h_graph) {
@@ -989,7 +987,7 @@ GraphHandle graph_node_get_graph(const GraphNodeHandle& h) noexcept {
     return h ? get_box(h)->h_graph : GraphHandle{};
 }
 
-void invalidate_graph_node_handle(const GraphNodeHandle& h) noexcept {
+void invalidate_graph_node(const GraphNodeHandle& h) noexcept {
     if (h) {
         CUgraphNode node = get_box(h)->resource;
         if (node) {

cuda_core/cuda/core/_cpp/resource_handles.hpp

@@ -416,7 +416,7 @@ GraphNodeHandle create_graph_node_handle(CUgraphNode node, const GraphHandle& h_
 GraphHandle graph_node_get_graph(const GraphNodeHandle& h) noexcept;
 
 // Zero the CUgraphNode resource inside the handle, marking it invalid.
-void invalidate_graph_node_handle(const GraphNodeHandle& h) noexcept;
+void invalidate_graph_node(const GraphNodeHandle& h) noexcept;
 
 // ============================================================================
 // Graphics resource handle functions

cuda_core/cuda/core/_resource_handles.pxd

@@ -186,7 +186,7 @@ cdef GraphHandle create_graph_handle_ref(cydriver.CUgraph graph, const GraphHand
 # Graph node handles
 cdef GraphNodeHandle create_graph_node_handle(cydriver.CUgraphNode node, const GraphHandle& h_graph) except+ nogil
 cdef GraphHandle graph_node_get_graph(const GraphNodeHandle& h) noexcept nogil
-cdef void invalidate_graph_node_handle(const GraphNodeHandle& h) noexcept nogil
+cdef void invalidate_graph_node(const GraphNodeHandle& h) noexcept nogil
 
 # Graphics resource handles
 cdef GraphicsResourceHandle create_graphics_resource_handle(

cuda_core/cuda/core/_resource_handles.pyx

@@ -159,7 +159,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
         cydriver.CUgraphNode node, const GraphHandle& h_graph) except+ nogil
     GraphHandle graph_node_get_graph "cuda_core::graph_node_get_graph" (
         const GraphNodeHandle& h) noexcept nogil
-    void invalidate_graph_node_handle "cuda_core::invalidate_graph_node_handle" (
+    void invalidate_graph_node "cuda_core::invalidate_graph_node" (
         const GraphNodeHandle& h) noexcept nogil
 
     # Graphics resource handles

cuda_core/cuda/core/graph/_adjacency_set_proxy.pyx

@@ -39,7 +39,7 @@ class AdjacencySetProxy(MutableSet):
     def __contains__(self, x):
         if not isinstance(x, GraphNode):
             return False
-        return x in (<_AdjacencySetCore>self._core).query()
+        return (<_AdjacencySetCore>self._core).contains(<GraphNode>x)
 
     def __iter__(self):
         return iter((<_AdjacencySetCore>self._core).query())
@@ -87,13 +87,13 @@ class AdjacencySetProxy(MutableSet):
             if isinstance(other, GraphNode):
                 nodes.append(other)
             else:
-                nodes.extend(other)
+                for n in other:
+                    if not isinstance(n, GraphNode):
+                        raise TypeError(
+                            f"expected GraphNode, got {type(n).__name__}")
+                    nodes.append(n)
         if not nodes:
             return
-        for n in nodes:
-            if not isinstance(n, GraphNode):
-                raise TypeError(
-                    f"expected GraphNode, got {type(n).__name__}")
         new = [n for n in nodes if n not in self]
         if new:
             (<_AdjacencySetCore>self._core).add_edges(new)
@@ -143,11 +143,14 @@ cdef class _AdjacencySetCore:
         cdef cydriver.CUgraphNode c_node = as_cu(self._h_node)
         if c_node == NULL:
             return []
-        cdef size_t count = 0
+        cdef cydriver.CUgraphNode buf[16]
+        cdef size_t count = 16
+        cdef size_t i
         with nogil:
-            HANDLE_RETURN(self._query_fn(c_node, NULL, &count))
-        if count == 0:
-            return []
+            HANDLE_RETURN(self._query_fn(c_node, buf, &count))
+        if count <= 16:
+            return [GraphNode._create(self._h_graph, buf[i])
+                    for i in range(count)]
         cdef vector[cydriver.CUgraphNode] nodes_vec
         nodes_vec.resize(count)
         with nogil:
@@ -156,6 +159,35 @@ cdef class _AdjacencySetCore:
         return [GraphNode._create(self._h_graph, nodes_vec[i])
                 for i in range(count)]
 
+    cdef bint contains(self, GraphNode other):
+        cdef cydriver.CUgraphNode c_node = as_cu(self._h_node)
+        cdef cydriver.CUgraphNode target = as_cu(other._h_node)
+        if c_node == NULL or target == NULL:
+            return False
+        cdef cydriver.CUgraphNode buf[16]
+        cdef size_t count = 16
+        cdef size_t i
+        with nogil:
+            HANDLE_RETURN(self._query_fn(c_node, buf, &count))
+
+        # Fast path for small sets.
+        if count <= 16:
+            for i in range(count):
+                if buf[i] == target:
+                    return True
+            return False
+
+        # Fallback for large sets.
+        cdef vector[cydriver.CUgraphNode] nodes_vec
+        nodes_vec.resize(count)
+        with nogil:
+            HANDLE_RETURN(self._query_fn(c_node, nodes_vec.data(), &count))
+        assert count == nodes_vec.size()
+        for i in range(count):
+            if nodes_vec[i] == target:
+                return True
+        return False
+
     cdef Py_ssize_t count(self):
         cdef cydriver.CUgraphNode c_node = as_cu(self._h_node)
         if c_node == NULL:

cuda_core/cuda/core/graph/_graph_def.pyx

@@ -159,6 +159,16 @@ cdef class GraphDef:
         """
         return self._entry.launch(config, kernel, *args)
 
+    def empty(self) -> "EmptyNode":
+        """Add an entry-point empty node (no dependencies).
+
+        Returns
+        -------
+        EmptyNode
+            A new EmptyNode with no dependencies.
+        """
+        return self._entry.join()
+
     def join(self, *nodes) -> "EmptyNode":
         """Create an empty node that depends on all given nodes.
 
@@ -322,18 +332,20 @@ cdef class GraphDef:
         set of GraphNode
             All nodes in the graph.
         """
-        cdef size_t num_nodes = 0
+        cdef vector[cydriver.CUgraphNode] nodes_vec
+        nodes_vec.resize(128)
+        cdef size_t num_nodes = 128
 
         with nogil:
-            HANDLE_RETURN(cydriver.cuGraphGetNodes(as_cu(self._h_graph), NULL, &num_nodes))
+            HANDLE_RETURN(cydriver.cuGraphGetNodes(as_cu(self._h_graph), nodes_vec.data(), &num_nodes))
 
         if num_nodes == 0:
             return set()
 
-        cdef vector[cydriver.CUgraphNode] nodes_vec
-        nodes_vec.resize(num_nodes)
-        with nogil:
-            HANDLE_RETURN(cydriver.cuGraphGetNodes(as_cu(self._h_graph), nodes_vec.data(), &num_nodes))
+        if num_nodes > 128:
+            nodes_vec.resize(num_nodes)
+            with nogil:
+                HANDLE_RETURN(cydriver.cuGraphGetNodes(as_cu(self._h_graph), nodes_vec.data(), &num_nodes))
 
         return set(GraphNode._create(self._h_graph, nodes_vec[i]) for i in range(num_nodes))
 
@@ -346,21 +358,12 @@ cdef class GraphDef:
             Each element is a (from_node, to_node) pair representing
             a dependency edge in the graph.
         """
-        cdef size_t num_edges = 0
-
-        with nogil:
-            IF CUDA_CORE_BUILD_MAJOR >= 13:
-                HANDLE_RETURN(cydriver.cuGraphGetEdges(as_cu(self._h_graph), NULL, NULL, NULL, &num_edges))
-            ELSE:
-                HANDLE_RETURN(cydriver.cuGraphGetEdges(as_cu(self._h_graph), NULL, NULL, &num_edges))
-
-        if num_edges == 0:
-            return set()
-
         cdef vector[cydriver.CUgraphNode] from_nodes
         cdef vector[cydriver.CUgraphNode] to_nodes
-        from_nodes.resize(num_edges)
-        to_nodes.resize(num_edges)
+        from_nodes.resize(128)
+        to_nodes.resize(128)
+        cdef size_t num_edges = 128
+
         with nogil:
             IF CUDA_CORE_BUILD_MAJOR >= 13:
                 HANDLE_RETURN(cydriver.cuGraphGetEdges(
@@ -369,6 +372,20 @@ cdef class GraphDef:
                 HANDLE_RETURN(cydriver.cuGraphGetEdges(
                     as_cu(self._h_graph), from_nodes.data(), to_nodes.data(), &num_edges))
 
+        if num_edges == 0:
+            return set()
+
+        if num_edges > 128:
+            from_nodes.resize(num_edges)
+            to_nodes.resize(num_edges)
+            with nogil:
+                IF CUDA_CORE_BUILD_MAJOR >= 13:
+                    HANDLE_RETURN(cydriver.cuGraphGetEdges(
+                        as_cu(self._h_graph), from_nodes.data(), to_nodes.data(), NULL, &num_edges))
+                ELSE:
+                    HANDLE_RETURN(cydriver.cuGraphGetEdges(
+                        as_cu(self._h_graph), from_nodes.data(), to_nodes.data(), &num_edges))
+
         return set(
             (GraphNode._create(self._h_graph, from_nodes[i]),
              GraphNode._create(self._h_graph, to_nodes[i]))