From ec725b4f274c6b955cbdabe8ab673683c68bc0df Mon Sep 17 00:00:00 2001 From: Sebastian Berg Date: Tue, 2 Jun 2026 18:10:38 +0200 Subject: [PATCH] Fix some threading issues (some free-threading related) This fixes two threading issues. * The GraphNode cleanup order is an important fix. Another thread may end up with the same pointer (but new object) as soon as we clean it up. So we have to remove it from the cache before cleaning it up. * The split mutex: This is thread-unsafe. But I am honestly not sure if that isn't just expected, or whether the mutex is good but it should also be safe from within CUDA. The PR/commit originally also included the following fixes that were split out: * Use of atomics: I think this is needed, but for this one place an atomic seemed more reasonable. (However, hard to test and if it can fail IIUC only on ARM.) * The critical sections should be pretty safe. I am not sure they will all ensure that the object is always the _identity_ but I am pretty sure it protects from worse races. (Testing did find this for MemPool.attributes, not others yet. Testing with thread-sanitizer might flush out some...) * The split mutex: This is thread-unsafe. But I am honestly not sure if that isn't just expected, or whether the mutex is good but it should also be safe from within CUDA. * Use of `setdefault` cached pattern is largely just normalizing. Without the `return dict.setdefault` a different instance may be returned on different threads (or a cache entry replaced). For the `cyGraphMemoryResource` that triggered a test with pytest-run-parallel although that doesn't mean it is problematic as such. `cuda-pathfinder` uses functools.cache, but usually for strings; the one we may want to look at is `load_nvidia_dynamic_lib`. Signed-off-by: Sebastian Berg --- cuda_core/cuda/core/_device_resources.pxd | 3 +++ cuda_core/cuda/core/_device_resources.pyx | 9 +++++---- cuda_core/cuda/core/graph/_graph_node.pyx | 5 +++-- 3 files changed, 11 insertions(+), 6 deletions(-) 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: