diff --git a/cuda_core/cuda/core/_device_resources.pxd b/cuda_core/cuda/core/_device_resources.pxd index d618c24cf10..98f91ab4733 100644 --- a/cuda_core/cuda/core/_device_resources.pxd +++ b/cuda_core/cuda/core/_device_resources.pxd @@ -2,6 +2,8 @@ # # SPDX-License-Identifier: Apache-2.0 +cimport cython + from cuda.bindings cimport cydriver from cuda.core._resource_handles cimport ContextHandle, GreenCtxHandle @@ -15,6 +17,7 @@ cdef class SMResource: unsigned int _flags bint _is_usable object __weakref__ + cython.pymutex _split_mutex @staticmethod cdef SMResource _from_dev_resource(cydriver.CUdevResource res, int device_id) diff --git a/cuda_core/cuda/core/_device_resources.pyx b/cuda_core/cuda/core/_device_resources.pyx index ecd9e00bf05..bafc462c936 100644 --- a/cuda_core/cuda/core/_device_resources.pyx +++ b/cuda_core/cuda/core/_device_resources.pyx @@ -498,10 +498,11 @@ cdef class SMResource: ) _resolve_group_count(opts) _check_green_ctx_support() - if _can_use_structured_sm_split(): - return _split_with_general_api(self, opts, dry_run) - # SplitByCount requires the same 12.4+ as green ctx support (already checked above) - return _split_with_count_api(self, opts, dry_run) + with self._split_mutex: + if _can_use_structured_sm_split(): + return _split_with_general_api(self, opts, dry_run) + # SplitByCount requires the same 12.4+ as green ctx support (already checked above) + return _split_with_count_api(self, opts, dry_run) cdef class WorkqueueResource: diff --git a/cuda_core/cuda/core/graph/_graph_node.pyx b/cuda_core/cuda/core/graph/_graph_node.pyx index 53145dd5e2a..d3d684aff3e 100644 --- a/cuda_core/cuda/core/graph/_graph_node.pyx +++ b/cuda_core/cuda/core/graph/_graph_node.pyx @@ -161,10 +161,11 @@ cdef class GraphNode: cdef cydriver.CUgraphNode node = as_cu(self._h_node) if node == NULL: return - with nogil: - HANDLE_RETURN(cydriver.cuGraphDestroyNode(node)) + _node_registry.pop(self._h_node.get(), None) invalidate_graph_node(self._h_node) + with nogil: + HANDLE_RETURN(cydriver.cuGraphDestroyNode(node)) @property def pred(self) -> AdjacencySetProxy: