Fix some threading issues (some free-threading related) (NVIDIA#2162)

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 <sebastianb@nvidia.com>

Author: seberg

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)

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:

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(<uintptr_t>self._h_node.get(), None)
         invalidate_graph_node(self._h_node)
+        with nogil:
+            HANDLE_RETURN(cydriver.cuGraphDestroyNode(node))
 
     @property
     def pred(self) -> AdjacencySetProxy: