diff --git a/cuda_core/cuda/core/graph/_adjacency_set_proxy.pyx b/cuda_core/cuda/core/graph/_adjacency_set_proxy.pyx index e1762321ce0..971e418a428 100644 --- a/cuda_core/cuda/core/graph/_adjacency_set_proxy.pyx +++ b/cuda_core/cuda/core/graph/_adjacency_set_proxy.pyx @@ -144,20 +144,23 @@ 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 cydriver.CUgraphNode stack_buf[16] + cdef cydriver.CUgraphNode* nodes + 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) + 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): @@ -165,27 +168,24 @@ 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 cydriver.CUgraphNode stack_buf[16] + cdef cydriver.CUgraphNode* nodes + 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) + 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)) - assert count == nodes_vec.size() + 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 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..a273e8b6a01 100644 --- a/cuda_core/tests/graph/test_graph_definition.py +++ b/cuda_core/tests/graph/test_graph_definition.py @@ -633,6 +633,43 @@ def test_succ(nonempty_graph_spec): assert actual == spec.expected_succ[name], f"succ mismatch for node {name}" +@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() + neighbors = [g.empty() for _ in range(20)] + adjacency = getattr(hub, adjacency_name) + adjacency.update(neighbors) + + 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 + + adjacency.clear() + assert len(adjacency) == 0 + assert g.edges() == set() + + +@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() + nodes = [g.empty() for _ in range(130)] + 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 g.edges() == expected_edges + + def test_node_graph_property(nonempty_graph_spec): """Every node's .graph property returns the parent GraphDefinition.""" spec = nonempty_graph_spec