Merge branch 'graph-node-identity' into graph-node-repr

Author: Andy-Jost

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) {
@@ -989,7 +984,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/_graph/_graph_def/_graph_node.pyx

@@ -48,7 +48,7 @@ from cuda.core._resource_handles cimport (
     create_graph_handle_ref,
     create_graph_node_handle,
     graph_node_get_graph,
-    invalidate_graph_node_handle,
+    invalidate_graph_node,
 )
 from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, _parse_fill_value
 
@@ -63,11 +63,11 @@ from cuda.core import Device
 from cuda.core._graph._graph_def._adjacency_set_proxy import AdjacencySetProxy
 from cuda.core._utils.cuda_utils import driver, handle_return
 
-_node_cache = weakref.WeakValueDictionary()
+_node_registry = weakref.WeakValueDictionary()
 
 
-cdef inline GraphNode _cached(GraphNode n):
-    _node_cache[<uintptr_t>n._h_node.get()] = n
+cdef inline GraphNode _registered(GraphNode n):
+    _node_registry[<uintptr_t>n._h_node.get()] = n
     return n
 
 
@@ -153,8 +153,8 @@ cdef class GraphNode:
             return
         with nogil:
             HANDLE_RETURN(cydriver.cuGraphDestroyNode(node))
-        _node_cache.pop(<uintptr_t>self._h_node.get(), None)
-        invalidate_graph_node_handle(self._h_node)
+        _node_registry.pop(<uintptr_t>self._h_node.get(), None)
+        invalidate_graph_node(self._h_node)
 
     @property
     def pred(self):
@@ -532,8 +532,7 @@ cdef inline ConditionalNode _make_conditional_node(
     n._cond_type = cond_type
     n._branches = branches
 
-    return _cached(n)
-
+    return _registered(n)
 
 cdef inline GraphNode GN_create(GraphHandle h_graph, cydriver.CUgraphNode node):
     cdef GraphNodeHandle h_node = create_graph_node_handle(node, h_graph)
@@ -544,12 +543,12 @@ cdef inline GraphNode GN_create(GraphHandle h_graph, cydriver.CUgraphNode node):
         (<GraphNode>n)._h_node = h_node
         return n
 
-    # Return a cached object or create and cache a new one.
-    cached = _node_cache.get(<uintptr_t>h_node.get())
-    if cached is not None:
-        return <GraphNode>cached
+    # Return a registered object or create and register a new one.
+    registered = _node_registry.get(<uintptr_t>h_node.get())
+    if registered is not None:
+        return <GraphNode>registered
     else:
-        return _cached(GN_create_impl(h_node))
+        return _registered(GN_create_impl(h_node))
 
 
 cdef inline GraphNode GN_create_impl(GraphNodeHandle h_node):
@@ -617,7 +616,7 @@ cdef inline KernelNode GN_launch(GraphNode self, LaunchConfig conf, Kernel ker,
     _attach_user_object(as_cu(h_graph), <void*>new KernelHandle(ker._h_kernel),
                         <cydriver.CUhostFn>_destroy_kernel_handle_copy)
 
-    return _cached(KernelNode._create_with_params(
+    return _registered(KernelNode._create_with_params(
         create_graph_node_handle(new_node, h_graph),
         conf.grid, conf.block, conf.shmem_size,
         ker._h_kernel))
@@ -646,7 +645,7 @@ cdef inline EmptyNode GN_join(GraphNode self, tuple nodes):
         HANDLE_RETURN(cydriver.cuGraphAddEmptyNode(
             &new_node, as_cu(h_graph), deps_ptr, num_deps))
 
-    return _cached(EmptyNode._create_impl(create_graph_node_handle(new_node, h_graph)))
+    return _registered(EmptyNode._create_impl(create_graph_node_handle(new_node, h_graph)))
 
 
 cdef inline AllocNode GN_alloc(GraphNode self, size_t size, object options):
@@ -722,7 +721,7 @@ cdef inline AllocNode GN_alloc(GraphNode self, size_t size, object options):
         HANDLE_RETURN(cydriver.cuGraphAddMemAllocNode(
             &new_node, as_cu(h_graph), deps, num_deps, &alloc_params))
 
-    return _cached(AllocNode._create_with_params(
+    return _registered(AllocNode._create_with_params(
         create_graph_node_handle(new_node, h_graph), alloc_params.dptr, size,
         device_id, memory_type, tuple(peer_ids)))
 
@@ -742,7 +741,7 @@ cdef inline FreeNode GN_free(GraphNode self, cydriver.CUdeviceptr c_dptr):
         HANDLE_RETURN(cydriver.cuGraphAddMemFreeNode(
             &new_node, as_cu(h_graph), deps, num_deps, c_dptr))
 
-    return _cached(FreeNode._create_with_params(create_graph_node_handle(new_node, h_graph), c_dptr))
+    return _registered(FreeNode._create_with_params(create_graph_node_handle(new_node, h_graph), c_dptr))
 
 
 cdef inline MemsetNode GN_memset(
@@ -777,7 +776,7 @@ cdef inline MemsetNode GN_memset(
             &new_node, as_cu(h_graph), deps, num_deps,
             &memset_params, ctx))
 
-    return _cached(MemsetNode._create_with_params(
+    return _registered(MemsetNode._create_with_params(
         create_graph_node_handle(new_node, h_graph), c_dst,
         val, elem_size, width, height, pitch))
 
@@ -838,7 +837,7 @@ cdef inline MemcpyNode GN_memcpy(
         HANDLE_RETURN(cydriver.cuGraphAddMemcpyNode(
             &new_node, as_cu(h_graph), deps, num_deps, &params, ctx))
 
-    return _cached(MemcpyNode._create_with_params(
+    return _registered(MemcpyNode._create_with_params(
         create_graph_node_handle(new_node, h_graph), c_dst, c_src, size,
         c_dst_type, c_src_type))
 
@@ -865,7 +864,7 @@ cdef inline ChildGraphNode GN_embed(GraphNode self, GraphDef child_def):
 
     cdef GraphHandle h_embedded = create_graph_handle_ref(embedded_graph, h_graph)
 
-    return _cached(ChildGraphNode._create_with_params(
+    return _registered(ChildGraphNode._create_with_params(
         create_graph_node_handle(new_node, h_graph), h_embedded))
 
 
@@ -887,7 +886,7 @@ cdef inline EventRecordNode GN_record_event(GraphNode self, Event ev):
     _attach_user_object(as_cu(h_graph), <void*>new EventHandle(ev._h_event),
                         <cydriver.CUhostFn>_destroy_event_handle_copy)
 
-    return _cached(EventRecordNode._create_with_params(
+    return _registered(EventRecordNode._create_with_params(
         create_graph_node_handle(new_node, h_graph), ev._h_event))
 
 
@@ -909,7 +908,7 @@ cdef inline EventWaitNode GN_wait_event(GraphNode self, Event ev):
     _attach_user_object(as_cu(h_graph), <void*>new EventHandle(ev._h_event),
                         <cydriver.CUhostFn>_destroy_event_handle_copy)
 
-    return _cached(EventWaitNode._create_with_params(
+    return _registered(EventWaitNode._create_with_params(
         create_graph_node_handle(new_node, h_graph), ev._h_event))
 
 
@@ -936,6 +935,6 @@ cdef inline HostCallbackNode GN_callback(GraphNode self, object fn, object user_
             &new_node, as_cu(h_graph), deps, num_deps, &node_params))
 
     cdef object callable_obj = fn if not isinstance(fn, ct._CFuncPtr) else None
-    return _cached(HostCallbackNode._create_with_params(
+    return _registered(HostCallbackNode._create_with_params(
         create_graph_node_handle(new_node, h_graph), callable_obj,
         node_params.fn, node_params.userData))

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