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

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

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._resource_handles 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 (empty)

19from typing import Any, TypeVar (empty)

20  

21_S = TypeVar("_S") (empty)

22  

23  

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

25  

26class AdjacencySetProxy(MutableSet[GraphNode]): (empty)

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

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

29  

30 __slots__ = ("_core",) (empty)

31  

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

33 self._core = _AdjacencySetCore(node, is_fwd) 61 ctx1pq7gLMNcdjorsVtuvwWXYZxyzABC01DE2FG3HI4JK5OPQRSTmik9aebf68Uhn

34  

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

36 @classmethod (empty)

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: (empty)

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

44 return False 

45 return (<_AdjacencySetCore>self._core).contains(<GraphNode>x) 16 ctx1pqgcdjomikaebfhn

46  

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

48 return iter((<_AdjacencySetCore>self._core).query()) 49 ctx17gLMNcdrsVtuvwWXYZxyzABC01DE2FG3HI4JK5OPQRSTaebUh

49  

50 def __len__(self) -> int: (empty)

51 return (<_AdjacencySetCore>self._core).count() 8 ctx17cdaebUh

52  

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

54 if not isinstance(value, GraphNode): 11 ctx1jmkaef68Uhn

55 raise TypeError( 

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

57 (<_AdjacencySetCore>self._core).check_mutation(value) 11 ctx1jmkaef68Uhn

58 if value in self: 8 ctx1jmkaefhn

59 return 1 ctx1a

60 (<_AdjacencySetCore>self._core).add_edge(<GraphNode>value) 8 ctx1jmkaefhn

61  

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

63 (<_AdjacencySetCore>self._core).check_owner_mutable() 4 ctx1kaeh

64 if value not in self: 4 ctx1kaeh

65 return 2 ctx1ah

66 (<_AdjacencySetCore>self._core).check_mutation(value) 3 ctx1kae

67 (<_AdjacencySetCore>self._core).remove_edge(<GraphNode>value) 3 ctx1kae

68  

69 # --- override for bulk efficiency --- 

70  

71 def clear(self) -> None: (empty)

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

73 (<_AdjacencySetCore>self._core).check_owner_mutable() 8 ctx1gcdiabf8

74 members = (<_AdjacencySetCore>self._core).query() 7 ctx1gcdiabf

75 if members: 7 ctx1gcdiabf

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

77  

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

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

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

81 if it is self: 1 ctx1a

82 self.clear() 

83 else: 

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

85 if to_remove: 1 ctx1a

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

87 return self 1 ctx1a

88  

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

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

91 (<_AdjacencySetCore>self._core).check_owner_mutable() 7 ctx1gcdjiab

92 nodes = [] 7 ctx1gcdjiab

93 for other in others: 7 ctx1gcdjiab

94 if isinstance(other, GraphNode): 7 ctx1gcdjiab

95 nodes.append(other) 

96 else: 

97 for n in other: 7 ctx1gcdjiab

98 if not isinstance(n, GraphNode): 7 ctx1gcdjiab

99 raise TypeError( 

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

101 nodes.append(n) 7 ctx1gcdjiab

102 for n in nodes: 7 ctx1gcdjiab

103 (<_AdjacencySetCore>self._core).check_mutation(n) 7 ctx1gcdjiab

104 if not nodes: 7 ctx1gcdjiab

105 return 1 ctx1b

106 new = [n for n in nodes if n not in self] 7 ctx1gcdjiab

107 if new: 7 ctx1gcdjiab

108 (<_AdjacencySetCore>self._core).add_edges(new) 7 ctx1gcdjiab

109  

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

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

112 self.update(it) 1 ctx1a

113 return self 1 ctx1a

114  

115 def __repr__(self) -> str: (empty)

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

117  

118  

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

120  

121# Signature shared by driver_get_preds and driver_get_succs. 

122ctypedef cydriver.CUresult (*_adj_fn_t)( 

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

124  

125  

126cdef class _AdjacencySetCore: 

127 """Cythonized core implementing AdjacencySetProxy""" 

128 cdef: 

129 GraphNodeHandle _h_node 

130 GraphHandle _h_graph 

131 _adj_fn_t _query_fn 

132 bint _is_fwd 

133  

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

135 self._h_node = node._h_node 61 ctx1pq7gLMNcdjorsVtuvwWXYZxyzABC01DE2FG3HI4JK5OPQRSTmik9aebf68Uhn

136 self._h_graph = graph_node_get_graph(node._h_node) 61 ctx1pq7gLMNcdjorsVtuvwWXYZxyzABC01DE2FG3HI4JK5OPQRSTmik9aebf68Uhn

137 self._is_fwd = is_fwd 61 ctx1pq7gLMNcdjorsVtuvwWXYZxyzABC01DE2FG3HI4JK5OPQRSTmik9aebf68Uhn

138 self._query_fn = driver_get_succs if is_fwd else driver_get_preds 61 ctx1pq7gLMNcdjorsVtuvwWXYZxyzABC01DE2FG3HI4JK5OPQRSTmik9aebf68Uhn

139  

140 cdef inline void _resolve_edge( (empty)

141 self, GraphNode other, 

142 cydriver.CUgraphNode* c_from, 

143 cydriver.CUgraphNode* c_to) noexcept: 

144 if self._is_fwd: 13 ctx1gcdjmikaebfhn

145 c_from[0] = as_cu(self._h_node) 10 ctx1gdjmkabfhn

146 c_to[0] = as_cu(other._h_node) 10 ctx1gdjmkabfhn

147 else: 

148 c_from[0] = as_cu(other._h_node) 5 ctx1ciebf

149 c_to[0] = as_cu(self._h_node) 5 ctx1ciebf

150  

151 cdef inline void check_owner_mutable(self) except *: (empty)

152 if as_cu(self._h_graph) == NULL: 16 ctx1gcdjmikaebf68Uhn

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

154 if as_cu(self._h_node) == NULL: 16 ctx1gcdjmikaebf68Uhn

155 raise RuntimeError("GraphNode has been destroyed") 1 ctx18

156  

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

158 self.check_owner_mutable() 16 ctx1gcdjmikaebf68Uhn

159 GN_check_valid(other) 16 ctx1gcdjmikaebf68Uhn

160 if other._is_entry: 14 ctx1gcdjmikaebf6hn

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

162 if as_cu(graph_node_get_graph(other._h_node)) != as_cu(self._h_graph): 14 ctx1gcdjmikaebf6hn

163 raise ValueError("Graph nodes must belong to the same GraphDefinition") 1 ctx16

164  

165 cdef list query(self): (empty)

166 cdef cydriver.CUgraphNode c_node = as_cu(self._h_node) 51 ctx17gLMNcdrsVtuvwWXYZxyzABC01DE2FG3HI4JK5OPQRSTiaebfUh

167 if c_node == NULL: 51 ctx17gLMNcdrsVtuvwWXYZxyzABC01DE2FG3HI4JK5OPQRSTiaebfUh

168 return [] 2 ctx17U

169 cdef cydriver.CUgraphNode stack_buf[16] 

170 cdef cydriver.CUgraphNode* nodes 

171 cdef size_t count = 0 49 ctx1gLMNcdrsVtuvwWXYZxyzABC01DE2FG3HI4JK5OPQRSTiaebfh

172 cdef size_t i 

173 with nogil: 49 ctx1gLMNcdrsVtuvwWXYZxyzABC01DE2FG3HI4JK5OPQRSTiaebfh

174 HANDLE_RETURN(self._query_fn(c_node, NULL, &count)) 49 ctx1gLMNcdrsVtuvwWXYZxyzABC01DE2FG3HI4JK5OPQRSTiaebfh

175 if count == 0: 49 ctx1gLMNcdrsVtuvwWXYZxyzABC01DE2FG3HI4JK5OPQRSTiaebfh

176 return [] 25 ctx1gVWXYZ01DE2FG3HI4JK5iaebf

177 cdef vector[cydriver.CUgraphNode] nodes_vec 

178 if count <= 16: 37 ctx1gLMNcdrstuvwxyzABCDEFGHIJKOPQRSTaebfh

179 nodes = stack_buf 35 ctx1gLMNrstuvwxyzABCDEFGHIJKOPQRSTaebfh

180 else: 

181 nodes_vec.resize(count) 2 ctx1cd

182 nodes = nodes_vec.data() 2 ctx1cd

183 with nogil: 37 ctx1gLMNcdrstuvwxyzABCDEFGHIJKOPQRSTaebfh

184 HANDLE_RETURN(self._query_fn(c_node, nodes, &count)) 37 ctx1gLMNcdrstuvwxyzABCDEFGHIJKOPQRSTaebfh

185 return [GraphNode._create(self._h_graph, nodes[i]) 37 ctx1gLMNcdrstuvwxyzABCDEFGHIJKOPQRSTaebfh

186 for i in range(count)] 37 ctx1gLMNcdrstuvwxyzABCDEFGHIJKOPQRSTaebfh

187  

188 cdef bint contains(self, GraphNode other): (empty)

189 cdef cydriver.CUgraphNode c_node = as_cu(self._h_node) 16 ctx1pqgcdjomikaebfhn

190 cdef cydriver.CUgraphNode target = as_cu(other._h_node) 16 ctx1pqgcdjomikaebfhn

191 if c_node == NULL or target == NULL: 16 ctx1pqgcdjomikaebfhn

192 return False 1 ctx1h

193 cdef cydriver.CUgraphNode stack_buf[16] 

194 cdef cydriver.CUgraphNode* nodes 

195 cdef size_t count = 0 16 ctx1pqgcdjomikaebfhn

196 cdef size_t i 

197 with nogil: 16 ctx1pqgcdjomikaebfhn

198 HANDLE_RETURN(self._query_fn(c_node, NULL, &count)) 16 ctx1pqgcdjomikaebfhn

199 if count == 0: 16 ctx1pqgcdjomikaebfhn

200 return False 13 ctx1gcdjmikaebfhn

201 cdef vector[cydriver.CUgraphNode] nodes_vec 

202 if count <= 16: 11 ctx1pqcdokaebfh

203 nodes = stack_buf 9 ctx1pqokaebfh

204 else: 

205 nodes_vec.resize(count) 2 ctx1cd

206 nodes = nodes_vec.data() 2 ctx1cd

207 with nogil: 11 ctx1pqcdokaebfh

208 HANDLE_RETURN(self._query_fn(c_node, nodes, &count)) 11 ctx1pqcdokaebfh

209 for i in range(count): 11 ctx1pqcdokaebfh

210 if nodes[i] == target: 11 ctx1pqcdokaebfh

211 return True 8 ctx1pqcdokae

212 return False 6 ctx1oaebfh

213  

214 cdef Py_ssize_t count(self): (empty)

215 cdef cydriver.CUgraphNode c_node = as_cu(self._h_node) 8 ctx17cdaebUh

216 if c_node == NULL: 8 ctx17cdaebUh

217 return 0 2 ctx17U

218 cdef size_t n = 0 6 ctx1cdaebh

219 with nogil: 6 ctx1cdaebh

220 HANDLE_RETURN(self._query_fn(c_node, NULL, &n)) 6 ctx1cdaebh

221 return <Py_ssize_t>n 6 ctx1cdaebh

222  

223 cdef void add_edge(self, GraphNode other): (empty)

224 cdef cydriver.CUgraphNode c_from, c_to 

225 self._resolve_edge(other, &c_from, &c_to) 8 ctx1jmkaefhn

226 with nogil: 8 ctx1jmkaefhn

227 HANDLE_RETURN(driver_add_edges(as_cu(self._h_graph), &c_from, &c_to, 1)) 8 ctx1jmkaefhn

228  

229 cdef void add_edges(self, list nodes): (empty)

230 cdef size_t n = len(nodes) 7 ctx1gcdjiab

231 cdef vector[cydriver.CUgraphNode] from_vec 

232 cdef vector[cydriver.CUgraphNode] to_vec 

233 from_vec.resize(n) 7 ctx1gcdjiab

234 to_vec.resize(n) 7 ctx1gcdjiab

235 cdef size_t i 

236 for i in range(n): 7 ctx1gcdjiab

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

238 with nogil: 7 ctx1gcdjiab

239 HANDLE_RETURN(driver_add_edges( 7 ctx1gcdjiab

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

241  

242 cdef void remove_edge(self, GraphNode other): (empty)

243 cdef cydriver.CUgraphNode c_from, c_to 

244 self._resolve_edge(other, &c_from, &c_to) 3 ctx1kae

245 with nogil: 3 ctx1kae

246 HANDLE_RETURN(driver_remove_edges(as_cu(self._h_graph), &c_from, &c_to, 1)) 3 ctx1kae

247  

248 cdef void remove_edges(self, list nodes): (empty)

249 cdef size_t n = len(nodes) 5 ctx1cdabf

250 cdef vector[cydriver.CUgraphNode] from_vec 

251 cdef vector[cydriver.CUgraphNode] to_vec 

252 from_vec.resize(n) 5 ctx1cdabf

253 to_vec.resize(n) 5 ctx1cdabf

254 cdef size_t i 

255 for i in range(n): 5 ctx1cdabf

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

257 with nogil: 5 ctx1cdabf

258 HANDLE_RETURN(driver_remove_edges( 5 ctx1cdabf

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

260  

261  

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

263  

264cdef inline cydriver.CUresult driver_get_preds( (empty)

265 cydriver.CUgraphNode node, cydriver.CUgraphNode* out, 

266 size_t* count) noexcept nogil: 

267 IF CUDA_CORE_BUILD_MAJOR >= 13: 

268 return cydriver.cuGraphNodeGetDependencies(node, out, NULL, count) 43 ctx1pqgLMNcorsVtuvwWXYZxyzABC01DE2FG3OPQRSTiebf

269 ELSE: 

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

271  

272  

273cdef inline cydriver.CUresult driver_get_succs( (empty)

274 cydriver.CUgraphNode node, cydriver.CUgraphNode* out, 

275 size_t* count) noexcept nogil: 

276 IF CUDA_CORE_BUILD_MAJOR >= 13: 

277 return cydriver.cuGraphNodeGetDependentNodes(node, out, NULL, count) 29 ctx1gdjrstuvwxyzABCHI4JK5mkaebfhn

278 ELSE: 

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

280  

281  

282cdef inline cydriver.CUresult driver_add_edges( (empty)

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

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

285 IF CUDA_CORE_BUILD_MAJOR >= 13: 

286 return cydriver.cuGraphAddDependencies( 13 ctx1gcdjmikaebfhn

287 graph, from_arr, to_arr, NULL, count) 

288 ELSE: 

289 return cydriver.cuGraphAddDependencies( 

290 graph, from_arr, to_arr, count) 

291  

292  

293cdef inline cydriver.CUresult driver_remove_edges( (empty)

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

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

296 IF CUDA_CORE_BUILD_MAJOR >= 13: 

297 return cydriver.cuGraphRemoveDependencies( 7 ctx1cdkaebf

298 graph, from_arr, to_arr, NULL, count) 

299 ELSE: 

300 return cydriver.cuGraphRemoveDependencies( 

301 graph, from_arr, to_arr, count)