Coverage for cuda/core/graph/_adjacency_set_proxy.pyx: 76.05%

1# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4  

5"""Mutable-set proxy for graph node predecessors and successors.""" 

6  

7from libc.stddef cimport size_t 

8from libcpp.vector cimport vector 

9from cuda.bindings cimport cydriver 

10from cuda.core.graph._graph_node cimport GraphNode, GN_check_valid 

11from cuda.core._rt cimport ( 

12 GraphHandle, 

13 GraphNodeHandle, 

14 as_cu, 

15 graph_node_get_graph, 

16) 

17from cuda.core._utils.cuda_utils cimport HANDLE_RETURN 

18from collections.abc import Iterable, Iterator, MutableSet, Set 

19from typing import Any, TypeVar 

20  

21_S = TypeVar("_S") 

22  

23  

24# ---- Python MutableSet wrapper ---------------------------------------------- 

25  

26class AdjacencySetProxy(MutableSet[GraphNode]): 

27 """Mutable set proxy for a node's predecessors or successors. Mutations 

28 write through to the underlying CUDA graph.""" 

29  

30 __slots__ = ("_core",) 

31  

32 def __init__(self, GraphNode node not None: GraphNode, bint is_fwd) -> None: 

33 self._core = _AdjacencySetCore(node, is_fwd) 65 ctx1pq!gLMNcdkorsYtuvwZ012xyzABC34DE5FG6HI7JK8OPQRSTUVWmjl$aebf9#Xhin

34  

35 # Used by operators such as &|^ to create non-proxy views when needed. 

36 @classmethod 

37 def _from_iterable(cls, it: Iterable[_S]) -> set[_S]: 

38 return set(it) 1 ctx1a

39  

40 # --- abstract methods required by MutableSet --- 

41  

42 def __contains__(self, x: object) -> bool: 

43 if not isinstance(x, GraphNode): 17 ctx1pqgcdkomjlaebfhin

44 return False 1 ctx1i

45 return (<_AdjacencySetCore>self._core).contains(<GraphNode>x) 17 ctx1pqgcdkomjlaebfhin

46  

47 def __iter__(self) -> Iterator[GraphNode]: 

48 return iter((<_AdjacencySetCore>self._core).query()) 53 ctx1!gLMNcdrsYtuvwZ012xyzABC34DE5FG6HI7JK8OPQRSTUVWaebXhi

49  

50 def __len__(self) -> int: 

51 return (<_AdjacencySetCore>self._core).count() 9 ctx1!cdaebXhi

52  

53 def add(self, GraphNode value not None: GraphNode) -> None: 

54 if not isinstance(value, GraphNode): 12 ctx1kmlaef9#Xhin

55 raise TypeError( 

56 f"expected GraphNode, got {type(value).__name__}") 

57 (<_AdjacencySetCore>self._core).check_mutation(value) 12 ctx1kmlaef9#Xhin

58 if value in self: 9 ctx1kmlaefhin

59 return 1 ctx1a

60 (<_AdjacencySetCore>self._core).add_edge(<GraphNode>value) 9 ctx1kmlaefhin

61  

62 def discard(self, value: object) -> None: 

63 # Non-GraphNode values are a no-op, matching MutableSet semantics. 

64 (<_AdjacencySetCore>self._core).check_owner_mutable() 5 ctx1laehi

65 if value not in self: 5 ctx1laehi

66 return 3 ctx1ahi

67 (<_AdjacencySetCore>self._core).check_mutation(value) 3 ctx1lae

68 (<_AdjacencySetCore>self._core).remove_edge(<GraphNode>value) 3 ctx1lae

69  

70 # --- override for bulk efficiency --- 

71  

72 def clear(self) -> None: 

73 """Remove all edges in a single driver call.""" 

74 (<_AdjacencySetCore>self._core).check_owner_mutable() 8 ctx1gcdjabf#

75 members = (<_AdjacencySetCore>self._core).query() 7 ctx1gcdjabf

76 if members: 7 ctx1gcdjabf

77 (<_AdjacencySetCore>self._core).remove_edges(members) 5 ctx1cdabf

78  

79 def __isub__(self, it: Set[Any]) -> "AdjacencySetProxy": 

80 """Remove edges to all nodes in *it* in a single driver call.""" 

81 (<_AdjacencySetCore>self._core).check_owner_mutable() 1 ctx1a

82 if it is self: 1 ctx1a

83 self.clear() 

84 else: 

85 to_remove = [v for v in it if isinstance(v, GraphNode) and v in self] 1 ctx1a

86 if to_remove: 1 ctx1a

87 (<_AdjacencySetCore>self._core).remove_edges(to_remove) 1 ctx1a

88 return self 1 ctx1a

89  

90 def update(self, *others) -> None: 

91 """Add edges to multiple nodes at once.""" 

92 (<_AdjacencySetCore>self._core).check_owner_mutable() 7 ctx1gcdkjab

93 nodes = [] 7 ctx1gcdkjab

94 for other in others: 7 ctx1gcdkjab

95 if isinstance(other, GraphNode): 7 ctx1gcdkjab

96 nodes.append(other) 

97 else: 

98 for n in other: 7 ctx1gcdkjab

99 if not isinstance(n, GraphNode): 7 ctx1gcdkjab

100 raise TypeError( 

101 f"expected GraphNode, got {type(n).__name__}") 

102 nodes.append(n) 7 ctx1gcdkjab

103 for n in nodes: 7 ctx1gcdkjab

104 (<_AdjacencySetCore>self._core).check_mutation(n) 7 ctx1gcdkjab

105 if not nodes: 7 ctx1gcdkjab

106 return 1 ctx1b

107 new = [n for n in nodes if n not in self] 7 ctx1gcdkjab

108 if new: 7 ctx1gcdkjab

109 (<_AdjacencySetCore>self._core).add_edges(new) 7 ctx1gcdkjab

110  

111 def __ior__(self, it: Set[Any]) -> "AdjacencySetProxy": # type: ignore[misc] 

112 """Add edges to all nodes in *it* in a single driver call.""" 

113 self.update(it) 1 ctx1a

114 return self 1 ctx1a

115  

116 def __repr__(self) -> str: 

117 return "{" + ", ".join(repr(n) for n in self) + "}" 1 ctx1a

118  

119  

120# ---- cdef core holding a function pointer ------------------------------------ 

121  

122# Signature shared by driver_get_preds and driver_get_succs. 

123ctypedef cydriver.CUresult (*_adj_fn_t)( 

124 cydriver.CUgraphNode, cydriver.CUgraphNode*, size_t*) noexcept nogil 

125  

126  

127cdef class _AdjacencySetCore: 

128 """Cythonized core implementing AdjacencySetProxy""" 

129 cdef: 

130 GraphNodeHandle _h_node 

131 GraphHandle _h_graph 

132 _adj_fn_t _query_fn 

133 bint _is_fwd 

134  

135 def __init__(self, GraphNode node, bint is_fwd): 

136 self._h_node = node._h_node 65 ctx1pq!gLMNcdkorsYtuvwZ012xyzABC34DE5FG6HI7JK8OPQRSTUVWmjl$aebf9#Xhin

137 self._h_graph = graph_node_get_graph(node._h_node) 65 ctx1pq!gLMNcdkorsYtuvwZ012xyzABC34DE5FG6HI7JK8OPQRSTUVWmjl$aebf9#Xhin

138 self._is_fwd = is_fwd 65 ctx1pq!gLMNcdkorsYtuvwZ012xyzABC34DE5FG6HI7JK8OPQRSTUVWmjl$aebf9#Xhin

139 self._query_fn = driver_get_succs if is_fwd else driver_get_preds 65 ctx1pq!gLMNcdkorsYtuvwZ012xyzABC34DE5FG6HI7JK8OPQRSTUVWmjl$aebf9#Xhin

140  

141 cdef inline void _resolve_edge( 

142 self, GraphNode other, 

143 cydriver.CUgraphNode* c_from, 

144 cydriver.CUgraphNode* c_to) noexcept: 

145 if self._is_fwd: 14 ctx1gcdkmjlaebfhin

146 c_from[0] = as_cu(self._h_node) 11 ctx1gdkmlabfhin

147 c_to[0] = as_cu(other._h_node) 11 ctx1gdkmlabfhin

148 else: 

149 c_from[0] = as_cu(other._h_node) 5 ctx1cjebf

150 c_to[0] = as_cu(self._h_node) 5 ctx1cjebf

151  

152 cdef inline void check_owner_mutable(self) except *: 

153 if as_cu(self._h_graph) == NULL: 17 ctx1gcdkmjlaebf9#Xhin

154 raise RuntimeError("GraphDefinition is no longer valid") 

155 if as_cu(self._h_node) == NULL: 17 ctx1gcdkmjlaebf9#Xhin

156 raise RuntimeError("GraphNode has been destroyed") 1 ctx1#

157  

158 cdef inline void check_mutation(self, GraphNode other) except *: 

159 self.check_owner_mutable() 17 ctx1gcdkmjlaebf9#Xhin

160 GN_check_valid(other) 17 ctx1gcdkmjlaebf9#Xhin

161 if other._is_entry: 15 ctx1gcdkmjlaebf9hin

162 raise ValueError("The virtual graph entry node cannot be used in an edge") 

163 if as_cu(graph_node_get_graph(other._h_node)) != as_cu(self._h_graph): 15 ctx1gcdkmjlaebf9hin

164 raise ValueError("Graph nodes must belong to the same GraphDefinition") 1 ctx19

165  

166 cdef list query(self): 

167 cdef cydriver.CUgraphNode c_node = as_cu(self._h_node) 55 ctx1!gLMNcdrsYtuvwZ012xyzABC34DE5FG6HI7JK8OPQRSTUVWjaebfXhi

168 if c_node == NULL: 55 ctx1!gLMNcdrsYtuvwZ012xyzABC34DE5FG6HI7JK8OPQRSTUVWjaebfXhi

169 return [] 2 ctx1!X

170 cdef cydriver.CUgraphNode stack_buf[16] 

171 cdef cydriver.CUgraphNode* nodes 

172 cdef size_t count = 0 53 ctx1gLMNcdrsYtuvwZ012xyzABC34DE5FG6HI7JK8OPQRSTUVWjaebfhi

173 cdef size_t i 

174 with nogil: 53 ctx1gLMNcdrsYtuvwZ012xyzABC34DE5FG6HI7JK8OPQRSTUVWjaebfhi

175 HANDLE_RETURN(self._query_fn(c_node, NULL, &count)) 53 ctx1gLMNcdrsYtuvwZ012xyzABC34DE5FG6HI7JK8OPQRSTUVWjaebfhi

176 if count == 0: 53 ctx1gLMNcdrsYtuvwZ012xyzABC34DE5FG6HI7JK8OPQRSTUVWjaebfhi

177 return [] 25 ctx1gYZ01234DE5FG6HI7JK8jaebf

178 cdef vector[cydriver.CUgraphNode] nodes_vec 

179 if count <= 16: 41 ctx1gLMNcdrstuvwxyzABCDEFGHIJKOPQRSTUVWaebfhi

180 nodes = stack_buf 39 ctx1gLMNrstuvwxyzABCDEFGHIJKOPQRSTUVWaebfhi

181 else: 

182 nodes_vec.resize(count) 2 ctx1cd

183 nodes = nodes_vec.data() 2 ctx1cd

184 with nogil: 41 ctx1gLMNcdrstuvwxyzABCDEFGHIJKOPQRSTUVWaebfhi

185 HANDLE_RETURN(self._query_fn(c_node, nodes, &count)) 41 ctx1gLMNcdrstuvwxyzABCDEFGHIJKOPQRSTUVWaebfhi

186 return [GraphNode._create(self._h_graph, nodes[i]) 41 ctx1gLMNcdrstuvwxyzABCDEFGHIJKOPQRSTUVWaebfhi

187 for i in range(count)] 41 ctx1gLMNcdrstuvwxyzABCDEFGHIJKOPQRSTUVWaebfhi

188  

189 cdef bint contains(self, GraphNode other): 

190 cdef cydriver.CUgraphNode c_node = as_cu(self._h_node) 17 ctx1pqgcdkomjlaebfhin

191 cdef cydriver.CUgraphNode target = as_cu(other._h_node) 17 ctx1pqgcdkomjlaebfhin

192 if c_node == NULL or target == NULL: 17 ctx1pqgcdkomjlaebfhin

193 return False 1 ctx1h

194 cdef cydriver.CUgraphNode stack_buf[16] 

195 cdef cydriver.CUgraphNode* nodes 

196 cdef size_t count = 0 17 ctx1pqgcdkomjlaebfhin

197 cdef size_t i 

198 with nogil: 17 ctx1pqgcdkomjlaebfhin

199 HANDLE_RETURN(self._query_fn(c_node, NULL, &count)) 17 ctx1pqgcdkomjlaebfhin

200 if count == 0: 17 ctx1pqgcdkomjlaebfhin

201 return False 14 ctx1gcdkmjlaebfhin

202 cdef vector[cydriver.CUgraphNode] nodes_vec 

203 if count <= 16: 11 ctx1pqcdolaebfh

204 nodes = stack_buf 9 ctx1pqolaebfh

205 else: 

206 nodes_vec.resize(count) 2 ctx1cd

207 nodes = nodes_vec.data() 2 ctx1cd

208 with nogil: 11 ctx1pqcdolaebfh

209 HANDLE_RETURN(self._query_fn(c_node, nodes, &count)) 11 ctx1pqcdolaebfh

210 for i in range(count): 11 ctx1pqcdolaebfh

211 if nodes[i] == target: 11 ctx1pqcdolaebfh

212 return True 8 ctx1pqcdolae

213 return False 6 ctx1oaebfh

214  

215 cdef Py_ssize_t count(self): 

216 cdef cydriver.CUgraphNode c_node = as_cu(self._h_node) 9 ctx1!cdaebXhi

217 if c_node == NULL: 9 ctx1!cdaebXhi

218 return 0 2 ctx1!X

219 cdef size_t n = 0 7 ctx1cdaebhi

220 with nogil: 7 ctx1cdaebhi

221 HANDLE_RETURN(self._query_fn(c_node, NULL, &n)) 7 ctx1cdaebhi

222 return <Py_ssize_t>n 7 ctx1cdaebhi

223  

224 cdef void add_edge(self, GraphNode other): 

225 cdef cydriver.CUgraphNode c_from, c_to 

226 self._resolve_edge(other, &c_from, &c_to) 9 ctx1kmlaefhin

227 with nogil: 9 ctx1kmlaefhin

228 HANDLE_RETURN(driver_add_edges(as_cu(self._h_graph), &c_from, &c_to, 1)) 9 ctx1kmlaefhin

229  

230 cdef void add_edges(self, list nodes): 

231 cdef size_t n = len(nodes) 7 ctx1gcdkjab

232 cdef vector[cydriver.CUgraphNode] from_vec 

233 cdef vector[cydriver.CUgraphNode] to_vec 

234 from_vec.resize(n) 7 ctx1gcdkjab

235 to_vec.resize(n) 7 ctx1gcdkjab

236 cdef size_t i 

237 for i in range(n): 7 ctx1gcdkjab

238 self._resolve_edge(<GraphNode>nodes[i], &from_vec[i], &to_vec[i]) 7 ctx1gcdkjab

239 with nogil: 7 ctx1gcdkjab

240 HANDLE_RETURN(driver_add_edges( 7 ctx1gcdkjab

241 as_cu(self._h_graph), from_vec.data(), to_vec.data(), n)) 

242  

243 cdef void remove_edge(self, GraphNode other): 

244 cdef cydriver.CUgraphNode c_from, c_to 

245 self._resolve_edge(other, &c_from, &c_to) 3 ctx1lae

246 with nogil: 3 ctx1lae

247 HANDLE_RETURN(driver_remove_edges(as_cu(self._h_graph), &c_from, &c_to, 1)) 3 ctx1lae

248  

249 cdef void remove_edges(self, list nodes): 

250 cdef size_t n = len(nodes) 5 ctx1cdabf

251 cdef vector[cydriver.CUgraphNode] from_vec 

252 cdef vector[cydriver.CUgraphNode] to_vec 

253 from_vec.resize(n) 5 ctx1cdabf

254 to_vec.resize(n) 5 ctx1cdabf

255 cdef size_t i 

256 for i in range(n): 5 ctx1cdabf

257 self._resolve_edge(<GraphNode>nodes[i], &from_vec[i], &to_vec[i]) 5 ctx1cdabf

258 with nogil: 5 ctx1cdabf

259 HANDLE_RETURN(driver_remove_edges( 5 ctx1cdabf

260 as_cu(self._h_graph), from_vec.data(), to_vec.data(), n)) 

261  

262  

263# ---- driver wrappers: absorb CUDA version differences ---- 

264  

265cdef inline cydriver.CUresult driver_get_preds( 

266 cydriver.CUgraphNode node, cydriver.CUgraphNode* out, 

267 size_t* count) noexcept nogil: 

268 IF CUDA_CORE_BUILD_MAJOR >= 13: 

269 return cydriver.cuGraphNodeGetDependencies(node, out, NULL, count) 46 ctx1pqgLMNcorsYtuvwZ012xyzABC34DE5FG6OPQRSTUVWjebf

270 ELSE: 

271 return cydriver.cuGraphNodeGetDependencies(node, out, count) 

272  

273  

274cdef inline cydriver.CUresult driver_get_succs( 

275 cydriver.CUgraphNode node, cydriver.CUgraphNode* out, 

276 size_t* count) noexcept nogil: 

277 IF CUDA_CORE_BUILD_MAJOR >= 13: 

278 return cydriver.cuGraphNodeGetDependentNodes(node, out, NULL, count) 30 ctx1gdkrstuvwxyzABCHI7JK8mlaebfhin

279 ELSE: 

280 return cydriver.cuGraphNodeGetDependentNodes(node, out, count) 

281  

282  

283cdef inline cydriver.CUresult driver_add_edges( 

284 cydriver.CUgraph graph, cydriver.CUgraphNode* from_arr, 

285 cydriver.CUgraphNode* to_arr, size_t count) noexcept nogil: 

286 IF CUDA_CORE_BUILD_MAJOR >= 13: 

287 return cydriver.cuGraphAddDependencies( 14 ctx1gcdkmjlaebfhin

288 graph, from_arr, to_arr, NULL, count) 

289 ELSE: 

290 return cydriver.cuGraphAddDependencies( 

291 graph, from_arr, to_arr, count) 

292  

293  

294cdef inline cydriver.CUresult driver_remove_edges( 

295 cydriver.CUgraph graph, cydriver.CUgraphNode* from_arr, 

296 cydriver.CUgraphNode* to_arr, size_t count) noexcept nogil: 

297 IF CUDA_CORE_BUILD_MAJOR >= 13: 

298 return cydriver.cuGraphRemoveDependencies( 7 ctx1cdlaebf

299 graph, from_arr, to_arr, NULL, count) 

300 ELSE: 

301 return cydriver.cuGraphRemoveDependencies( 

302 graph, from_arr, to_arr, count)