Merge branch 'NVIDIA:main' into main

Author: rluo8

.pre-commit-config.yaml

@@ -15,7 +15,7 @@ ci:
 # pre-commit autoupdate --freeze
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 5ba58aca0bd5bc7c0e1c0fc45af2e88d6a2bde83  # frozen: v0.14.10
+    rev: c60c980e561ed3e73101667fe8365c609d19a438  # frozen: v0.15.9
     hooks:
       - id: ruff-check
         args: [--fix, --show-fixes]
@@ -76,22 +76,22 @@ repos:
     - id: rst-inline-touching-normal
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: a66e98df7b4aeeb3724184b332785976d062b92e  # frozen: v1.19.1
+    rev: 8e5c80792e2ec0c87804d8ef915bf35e2caea6da  # frozen: v1.20.0
     hooks:
       - id: mypy
         name: mypy-pathfinder
         files: ^cuda_pathfinder/cuda/.*\.py$  # Exclude tests directory
         args: [--config-file=cuda_pathfinder/pyproject.toml]
 
   - repo: https://github.com/rhysd/actionlint
-    rev: "0933c147c9d6587653d45fdcb4c497c57a65f9af"  # frozen: v1.7.10
+    rev: "914e7df21a07ef503a81201c76d2b11c789d3fca"  # frozen: v1.7.12
     hooks:
       - id: actionlint
         args: ["-shellcheck="]
         exclude: ^\.github/workflows/coverage.yml$
 
   - repo: https://github.com/MarcoGorelli/cython-lint
-    rev: "d9ff7ce99ef4f2ae8fba93079ca9d76c4651d4ac"  # frozen: v0.18.0
+    rev: "7c6152f6c8f9087684ff2e09a9227941e233bafb"  # frozen: v0.19.0
     hooks:
       - id: cython-lint
         args: [--no-pycodestyle]

context7.json

@@ -1,5 +1,4 @@
 {
-  "description": "Python access to NVIDIA's CUDA platform, including cuda.bindings (low-level CUDA C API bindings), cuda.core (Pythonic CUDA runtime), and cuda.pathfinder (component discovery)",
   "url": "https://context7.com/nvidia/cuda-python",
   "public_key": "pk_gupaHhsdvsuT1j3BZpb7i"
 }

cuda_core/cuda/core/__init__.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
 
@@ -31,12 +31,6 @@
 from cuda.core import system, utils
 from cuda.core._device import Device
 from cuda.core._event import Event, EventOptions
-from cuda.core._graph import (
-    Graph,
-    GraphBuilder,
-    GraphCompleteOptions,
-    GraphDebugPrintOptions,
-)
 from cuda.core._graphics import GraphicsResource
 from cuda.core._launch_config import LaunchConfig
 from cuda.core._launcher import launch
@@ -69,3 +63,12 @@
     StreamOptions,
 )
 from cuda.core._tensor_map import TensorMapDescriptor, TensorMapDescriptorOptions
+from cuda.core.graph import (
+    Condition,
+    Graph,
+    GraphAllocOptions,
+    GraphBuilder,
+    GraphCompleteOptions,
+    GraphDebugPrintOptions,
+    GraphDef,
+)

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) {
@@ -957,7 +954,7 @@ GraphHandle create_graph_handle_ref(CUgraph graph, const GraphHandle& h_parent)
 
 namespace {
 struct GraphNodeBox {
-    CUgraphNode resource;
+    mutable CUgraphNode resource;
     GraphHandle h_graph;
 };
 }  // namespace
@@ -969,15 +966,37 @@ 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) {
+    if (node) {
+        if (auto h = graph_node_registry.lookup(node)) {
+            return h;
+        }
+    }
     auto box = std::make_shared<const GraphNodeBox>(GraphNodeBox{node, h_graph});
-    return GraphNodeHandle(box, &box->resource);
+    GraphNodeHandle h(box, &box->resource);
+    if (node) {
+        graph_node_registry.register_handle(node, h);
+    }
+    return h;
 }
 
 GraphHandle graph_node_get_graph(const GraphNodeHandle& h) noexcept {
     return h ? get_box(h)->h_graph : GraphHandle{};
 }
 
+void invalidate_graph_node(const GraphNodeHandle& h) noexcept {
+    if (h) {
+        CUgraphNode node = get_box(h)->resource;
+        if (node) {
+            graph_node_registry.unregister_handle(node);
+        }
+        get_box(h)->resource = nullptr;
+    }
+}
+
 // ============================================================================
 // Graphics Resource Handles
 // ============================================================================

cuda_core/cuda/core/_cpp/resource_handles.hpp

@@ -415,6 +415,9 @@ GraphNodeHandle create_graph_node_handle(CUgraphNode node, const GraphHandle& h_
 // Extract the owning graph handle from a node handle.
 GraphHandle graph_node_get_graph(const GraphNodeHandle& h) noexcept;
 
+// Zero the CUgraphNode resource inside the handle, marking it invalid.
+void invalidate_graph_node(const GraphNodeHandle& h) noexcept;
+
 // ============================================================================
 // Graphics resource handle functions
 // ============================================================================

cuda_core/cuda/core/_device.pyx

@@ -1,11 +1,10 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations
 
 cimport cpython
-from libc.stdint cimport uintptr_t
 
 from cuda.bindings cimport cydriver
 from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
@@ -24,7 +23,6 @@ from cuda.core._resource_handles cimport (
     as_cu,
 )
 
-from cuda.core._graph import GraphBuilder
 from cuda.core._stream import IsStreamT, Stream, StreamOptions
 from cuda.core._utils.clear_error_support import assert_type
 from cuda.core._utils.cuda_utils import (
@@ -1363,15 +1361,17 @@ class Device:
         self._check_context_initialized()
         handle_return(runtime.cudaDeviceSynchronize())
 
-    def create_graph_builder(self) -> GraphBuilder:
-        """Create a new :obj:`~_graph.GraphBuilder` object.
+    def create_graph_builder(self) -> "GraphBuilder":
+        """Create a new :obj:`~graph.GraphBuilder` object.
 
         Returns
         -------
-        :obj:`~_graph.GraphBuilder`
+        :obj:`~graph.GraphBuilder`
             Newly created graph builder object.
 
         """
+        from cuda.core.graph._graph_builder import GraphBuilder
+
         self._check_context_initialized()
         return GraphBuilder._init(stream=self.create_stream(), is_stream_owner=True)
 

cuda_core/cuda/core/_graph/__init__.py

@@ -26,7 +26,7 @@ def launch(stream: Stream | GraphBuilder | IsStreamT, config: LaunchConfig, kern
 
     Parameters
     ----------
-    stream : :obj:`~_stream.Stream` | :obj:`~_graph.GraphBuilder`
+    stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
         The stream establishing the stream ordering semantic of a
         launch.
     config : :obj:`LaunchConfig`