From 0e556f6287669d40f439500eda93678293e475d3 Mon Sep 17 00:00:00 2001 From: Yimo Jiang Date: Mon, 10 Aug 2026 07:09:22 +0000 Subject: [PATCH 1/3] fix(cuda.core): avoid truncating graph queries --- .../cuda/core/graph/_adjacency_set_proxy.pyx | 25 ++++-------- .../cuda/core/graph/_graph_definition.pyx | 39 ++++++++----------- .../tests/graph/test_graph_definition.py | 29 ++++++++++++++ 3 files changed, 53 insertions(+), 40 deletions(-) diff --git a/cuda_core/cuda/core/graph/_adjacency_set_proxy.pyx b/cuda_core/cuda/core/graph/_adjacency_set_proxy.pyx index e1762321ce0..0f31a5afe9c 100644 --- a/cuda_core/cuda/core/graph/_adjacency_set_proxy.pyx +++ b/cuda_core/cuda/core/graph/_adjacency_set_proxy.pyx @@ -144,14 +144,12 @@ cdef class _AdjacencySetCore: cdef cydriver.CUgraphNode c_node = as_cu(self._h_node) if c_node == NULL: return [] - cdef cydriver.CUgraphNode buf[16] - cdef size_t count = 16 + cdef size_t count = 0 cdef size_t i with nogil: - HANDLE_RETURN(self._query_fn(c_node, buf, &count)) - if count <= 16: - return [GraphNode._create(self._h_graph, buf[i]) - for i in range(count)] + HANDLE_RETURN(self._query_fn(c_node, NULL, &count)) + if count == 0: + return [] cdef vector[cydriver.CUgraphNode] nodes_vec nodes_vec.resize(count) with nogil: @@ -165,25 +163,16 @@ cdef class _AdjacencySetCore: cdef cydriver.CUgraphNode target = as_cu(other._h_node) if c_node == NULL or target == NULL: return False - cdef cydriver.CUgraphNode buf[16] - cdef size_t count = 16 + cdef size_t count = 0 cdef size_t i with nogil: - HANDLE_RETURN(self._query_fn(c_node, buf, &count)) - - # Fast path for small sets. - if count <= 16: - for i in range(count): - if buf[i] == target: - return True + HANDLE_RETURN(self._query_fn(c_node, NULL, &count)) + if count == 0: return False - - # Fallback for large sets. cdef vector[cydriver.CUgraphNode] nodes_vec nodes_vec.resize(count) with nogil: HANDLE_RETURN(self._query_fn(c_node, nodes_vec.data(), &count)) - assert count == nodes_vec.size() for i in range(count): if nodes_vec[i] == target: return True diff --git a/cuda_core/cuda/core/graph/_graph_definition.pyx b/cuda_core/cuda/core/graph/_graph_definition.pyx index 46896899ecd..e4bed7eef15 100644 --- a/cuda_core/cuda/core/graph/_graph_definition.pyx +++ b/cuda_core/cuda/core/graph/_graph_definition.pyx @@ -361,19 +361,17 @@ cdef class GraphDefinition: All nodes in the graph. """ cdef vector[cydriver.CUgraphNode] nodes_vec - nodes_vec.resize(128) - cdef size_t num_nodes = 128 + cdef size_t num_nodes = 0 with nogil: - HANDLE_RETURN(cydriver.cuGraphGetNodes(as_cu(self._h_graph), nodes_vec.data(), &num_nodes)) + HANDLE_RETURN(cydriver.cuGraphGetNodes(as_cu(self._h_graph), NULL, &num_nodes)) if num_nodes == 0: return set() - if num_nodes > 128: - nodes_vec.resize(num_nodes) - with nogil: - HANDLE_RETURN(cydriver.cuGraphGetNodes(as_cu(self._h_graph), nodes_vec.data(), &num_nodes)) + nodes_vec.resize(num_nodes) + with nogil: + HANDLE_RETURN(cydriver.cuGraphGetNodes(as_cu(self._h_graph), nodes_vec.data(), &num_nodes)) return {GraphNode._create(self._h_graph, nodes_vec[i]) for i in range(num_nodes)} @@ -388,31 +386,28 @@ cdef class GraphDefinition: """ cdef vector[cydriver.CUgraphNode] from_nodes cdef vector[cydriver.CUgraphNode] to_nodes - from_nodes.resize(128) - to_nodes.resize(128) - cdef size_t num_edges = 128 + cdef size_t num_edges = 0 with nogil: IF CUDA_CORE_BUILD_MAJOR >= 13: HANDLE_RETURN(cydriver.cuGraphGetEdges( - as_cu(self._h_graph), from_nodes.data(), to_nodes.data(), NULL, &num_edges)) + as_cu(self._h_graph), NULL, NULL, NULL, &num_edges)) ELSE: HANDLE_RETURN(cydriver.cuGraphGetEdges( - as_cu(self._h_graph), from_nodes.data(), to_nodes.data(), &num_edges)) + as_cu(self._h_graph), NULL, NULL, &num_edges)) if num_edges == 0: return set() - if num_edges > 128: - from_nodes.resize(num_edges) - to_nodes.resize(num_edges) - with nogil: - IF CUDA_CORE_BUILD_MAJOR >= 13: - HANDLE_RETURN(cydriver.cuGraphGetEdges( - as_cu(self._h_graph), from_nodes.data(), to_nodes.data(), NULL, &num_edges)) - ELSE: - HANDLE_RETURN(cydriver.cuGraphGetEdges( - as_cu(self._h_graph), from_nodes.data(), to_nodes.data(), &num_edges)) + from_nodes.resize(num_edges) + to_nodes.resize(num_edges) + with nogil: + IF CUDA_CORE_BUILD_MAJOR >= 13: + HANDLE_RETURN(cydriver.cuGraphGetEdges( + as_cu(self._h_graph), from_nodes.data(), to_nodes.data(), NULL, &num_edges)) + ELSE: + HANDLE_RETURN(cydriver.cuGraphGetEdges( + as_cu(self._h_graph), from_nodes.data(), to_nodes.data(), &num_edges)) return { (GraphNode._create(self._h_graph, from_nodes[i]), diff --git a/cuda_core/tests/graph/test_graph_definition.py b/cuda_core/tests/graph/test_graph_definition.py index 0aeb5a9d527..7ac4b243026 100644 --- a/cuda_core/tests/graph/test_graph_definition.py +++ b/cuda_core/tests/graph/test_graph_definition.py @@ -633,6 +633,35 @@ def test_succ(nonempty_graph_spec): assert actual == spec.expected_succ[name], f"succ mismatch for node {name}" +@pytest.mark.agent_authored(model="gpt-5.4") +def test_large_adjacency_set_is_not_truncated(init_cuda): + """Adjacency queries return and remove edges beyond the old 16-edge buffer.""" + g = GraphDefinition() + hub = g.empty() + successors = [g.empty() for _ in range(20)] + hub.succ.update(successors) + + assert len(hub.succ) == 20 + assert set(hub.succ) == set(successors) + assert successors[-1] in hub.succ + + hub.succ.clear() + assert len(hub.succ) == 0 + assert g.edges() == set() + + +@pytest.mark.agent_authored(model="gpt-5.4") +def test_large_graph_queries_are_not_truncated(init_cuda): + """Graph queries return nodes and edges beyond the old 128-item buffers.""" + g = GraphDefinition() + nodes = [g.empty() for _ in range(130)] + nodes[0].succ.update(nodes[1:]) + nodes[1].succ.add(nodes[2]) + + assert g.nodes() == set(nodes) + assert len(g.edges()) == 130 + + def test_node_graph_property(nonempty_graph_spec): """Every node's .graph property returns the parent GraphDefinition.""" spec = nonempty_graph_spec From 2617a4ee1c2368e98970e3481a4d991ff49aca7f Mon Sep 17 00:00:00 2001 From: Yimo Jiang Date: Mon, 10 Aug 2026 08:22:08 +0000 Subject: [PATCH 2/3] perf(cuda.core): retain adjacency stack buffer --- .../cuda/core/graph/_adjacency_set_proxy.pyx | 25 +++++++++++++------ 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/cuda_core/cuda/core/graph/_adjacency_set_proxy.pyx b/cuda_core/cuda/core/graph/_adjacency_set_proxy.pyx index 0f31a5afe9c..971e418a428 100644 --- a/cuda_core/cuda/core/graph/_adjacency_set_proxy.pyx +++ b/cuda_core/cuda/core/graph/_adjacency_set_proxy.pyx @@ -144,6 +144,8 @@ cdef class _AdjacencySetCore: cdef cydriver.CUgraphNode c_node = as_cu(self._h_node) if c_node == NULL: return [] + cdef cydriver.CUgraphNode stack_buf[16] + cdef cydriver.CUgraphNode* nodes cdef size_t count = 0 cdef size_t i with nogil: @@ -151,11 +153,14 @@ cdef class _AdjacencySetCore: if count == 0: return [] cdef vector[cydriver.CUgraphNode] nodes_vec - nodes_vec.resize(count) + if count <= 16: + nodes = stack_buf + else: + nodes_vec.resize(count) + nodes = nodes_vec.data() with nogil: - HANDLE_RETURN(self._query_fn( - c_node, nodes_vec.data(), &count)) - return [GraphNode._create(self._h_graph, nodes_vec[i]) + HANDLE_RETURN(self._query_fn(c_node, nodes, &count)) + return [GraphNode._create(self._h_graph, nodes[i]) for i in range(count)] cdef bint contains(self, GraphNode other): @@ -163,6 +168,8 @@ cdef class _AdjacencySetCore: cdef cydriver.CUgraphNode target = as_cu(other._h_node) if c_node == NULL or target == NULL: return False + cdef cydriver.CUgraphNode stack_buf[16] + cdef cydriver.CUgraphNode* nodes cdef size_t count = 0 cdef size_t i with nogil: @@ -170,11 +177,15 @@ cdef class _AdjacencySetCore: if count == 0: return False cdef vector[cydriver.CUgraphNode] nodes_vec - nodes_vec.resize(count) + if count <= 16: + nodes = stack_buf + else: + nodes_vec.resize(count) + nodes = nodes_vec.data() with nogil: - HANDLE_RETURN(self._query_fn(c_node, nodes_vec.data(), &count)) + HANDLE_RETURN(self._query_fn(c_node, nodes, &count)) for i in range(count): - if nodes_vec[i] == target: + if nodes[i] == target: return True return False From b9b6e7f67e300a14f3706b0a38c192297168261a Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Mon, 10 Aug 2026 10:24:09 -0700 Subject: [PATCH 3/3] test(cuda.core): cover large predecessor graph queries Verify exact edge identities so graph query regressions cannot pass through count-only checks. --- .../tests/graph/test_graph_definition.py | 30 ++++++++++++------- 1 file changed, 19 insertions(+), 11 deletions(-) diff --git a/cuda_core/tests/graph/test_graph_definition.py b/cuda_core/tests/graph/test_graph_definition.py index 7ac4b243026..a273e8b6a01 100644 --- a/cuda_core/tests/graph/test_graph_definition.py +++ b/cuda_core/tests/graph/test_graph_definition.py @@ -633,24 +633,30 @@ def test_succ(nonempty_graph_spec): assert actual == spec.expected_succ[name], f"succ mismatch for node {name}" -@pytest.mark.agent_authored(model="gpt-5.4") -def test_large_adjacency_set_is_not_truncated(init_cuda): +@pytest.mark.parametrize("adjacency_name", ("pred", "succ")) +@pytest.mark.agent_authored(model="gpt-5.6") +def test_large_adjacency_set_is_not_truncated(init_cuda, adjacency_name): """Adjacency queries return and remove edges beyond the old 16-edge buffer.""" g = GraphDefinition() hub = g.empty() - successors = [g.empty() for _ in range(20)] - hub.succ.update(successors) + neighbors = [g.empty() for _ in range(20)] + adjacency = getattr(hub, adjacency_name) + adjacency.update(neighbors) - assert len(hub.succ) == 20 - assert set(hub.succ) == set(successors) - assert successors[-1] in hub.succ + expected_edges = ( + {(node, hub) for node in neighbors} if adjacency_name == "pred" else {(hub, node) for node in neighbors} + ) + assert len(adjacency) == 20 + assert set(adjacency) == set(neighbors) + assert neighbors[-1] in adjacency + assert g.edges() == expected_edges - hub.succ.clear() - assert len(hub.succ) == 0 + adjacency.clear() + assert len(adjacency) == 0 assert g.edges() == set() -@pytest.mark.agent_authored(model="gpt-5.4") +@pytest.mark.agent_authored(model="gpt-5.6") def test_large_graph_queries_are_not_truncated(init_cuda): """Graph queries return nodes and edges beyond the old 128-item buffers.""" g = GraphDefinition() @@ -658,8 +664,10 @@ def test_large_graph_queries_are_not_truncated(init_cuda): nodes[0].succ.update(nodes[1:]) nodes[1].succ.add(nodes[2]) + expected_edges = {(nodes[0], node) for node in nodes[1:]} + expected_edges.add((nodes[1], nodes[2])) assert g.nodes() == set(nodes) - assert len(g.edges()) == 130 + assert g.edges() == expected_edges def test_node_graph_property(nonempty_graph_spec):