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

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 Iterator, MutableSet, Set (empty)

19from typing import Any (empty)

20  

21  

22# ---- Python MutableSet wrapper ---------------------------------------------- 

23  

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

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

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

27  

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

29  

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

31 self._core = _AdjacencySetCore(node, is_fwd) 57 ctx1pq3gLMNcdjorsRtuvwSTUVxyzABCWXDEYFGZHI0JK1OPmik5aebf24Qhn

32  

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

34 @classmethod (empty)

35 def _from_iterable(cls, it) -> set[GraphNode]: 

36 return set(it) 1 ctx1a

37  

38 # --- abstract methods required by MutableSet --- 

39  

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

41 if not isinstance(x, GraphNode): 16 ctx1pqgcdjomikaebfhn

42 return False 

43 return (<_AdjacencySetCore>self._core).contains(<GraphNode>x) 17 ctx1lpqgcdjomikaebfhn

44  

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

46 return iter((<_AdjacencySetCore>self._core).query()) 45 ctx13gLMNcdrsRtuvwSTUVxyzABCWXDEYFGZHI0JK1OPaebQh

47  

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

49 return (<_AdjacencySetCore>self._core).count() 8 ctx13cdaebQh

50  

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

52 if not isinstance(value, GraphNode): 11 ctx1jmkaef24Qhn

53 raise TypeError( 

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

55 (<_AdjacencySetCore>self._core).check_mutation(value) 11 ctx1jmkaef24Qhn

56 if value in self: 8 ctx1jmkaefhn

57 return 1 ctx1a

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

59  

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

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

62 if value not in self: 4 ctx1kaeh

63 return 2 ctx1ah

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

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

66  

67 # --- override for bulk efficiency --- 

68  

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

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

71 (<_AdjacencySetCore>self._core).check_owner_mutable() 8 ctx1gcdiabf4

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

73 if members: 7 ctx1gcdiabf

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

75  

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

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

78 (<_AdjacencySetCore>self._core).check_owner_mutable() 2 ctx1la

79 if it is self: 1 ctx1a

80 self.clear() 

81 else: 

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

83 if to_remove: 1 ctx1a

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

85 return self 1 ctx1a

86  

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

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

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

90 nodes = [] 7 ctx1gcdjiab

91 for other in others: 7 ctx1gcdjiab

92 if isinstance(other, GraphNode): 7 ctx1gcdjiab

93 nodes.append(other) 

94 else: 

95 for n in other: 7 ctx1gcdjiab

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

97 raise TypeError( 

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

99 nodes.append(n) 7 ctx1gcdjiab

100 for n in nodes: 7 ctx1gcdjiab

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

102 if not nodes: 7 ctx1gcdjiab

103 return 1 ctx1b

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

105 if new: 7 ctx1gcdjiab

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

107  

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

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

110 self.update(it) 1 ctx1a

111 return self 1 ctx1a

112  

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

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

115  

116  

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

118  

119# Signature shared by driver_get_preds and driver_get_succs. 

120ctypedef cydriver.CUresult (*_adj_fn_t)( 

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

122  

123  

124cdef class _AdjacencySetCore: 

125 """Cythonized core implementing AdjacencySetProxy""" 

126 cdef: 

127 GraphNodeHandle _h_node 

128 GraphHandle _h_graph 

129 _adj_fn_t _query_fn 

130 bint _is_fwd 

131  

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

133 self._h_node = node._h_node 58 ctx1lpq3gLMNcdjorsRtuvwSTUVxyzABCWXDEYFGZHI0JK1OPmik5aebf24Qhn

134 self._h_graph = graph_node_get_graph(node._h_node) 57 ctx1pq3gLMNcdjorsRtuvwSTUVxyzABCWXDEYFGZHI0JK1OPmik5aebf24Qhn

135 self._is_fwd = is_fwd 57 ctx1pq3gLMNcdjorsRtuvwSTUVxyzABCWXDEYFGZHI0JK1OPmik5aebf24Qhn

136 self._query_fn = driver_get_succs if is_fwd else driver_get_preds 57 ctx1pq3gLMNcdjorsRtuvwSTUVxyzABCWXDEYFGZHI0JK1OPmik5aebf24Qhn

137  

138 cdef inline void _resolve_edge( (empty)

139 self, GraphNode other, 

140 cydriver.CUgraphNode* c_from, 

141 cydriver.CUgraphNode* c_to) noexcept: 

142 if self._is_fwd: 13 ctx1gcdjmikaebfhn

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

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

145 else: 

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

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

148  

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

150 if as_cu(self._h_graph) == NULL: 16 ctx1gcdjmikaebf24Qhn

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

152 if as_cu(self._h_node) == NULL: 16 ctx1gcdjmikaebf24Qhn

153 raise RuntimeError("GraphNode has been destroyed") 1 ctx14

154  

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

156 self.check_owner_mutable() 16 ctx1gcdjmikaebf24Qhn

157 GN_check_valid(other) 16 ctx1gcdjmikaebf24Qhn

158 if other._is_entry: 14 ctx1gcdjmikaebf2hn

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

160 if as_cu(graph_node_get_graph(other._h_node)) != as_cu(self._h_graph): 14 ctx1gcdjmikaebf2hn

161 raise ValueError("Graph nodes must belong to the same GraphDefinition") 1 ctx12

162  

163 cdef list query(self): (empty)

164 cdef cydriver.CUgraphNode c_node = as_cu(self._h_node) 47 ctx13gLMNcdrsRtuvwSTUVxyzABCWXDEYFGZHI0JK1OPiaebfQh

165 if c_node == NULL: 47 ctx13gLMNcdrsRtuvwSTUVxyzABCWXDEYFGZHI0JK1OPiaebfQh

166 return [] 2 ctx13Q

167 cdef cydriver.CUgraphNode stack_buf[16] 

168 cdef cydriver.CUgraphNode* nodes 

169 cdef size_t count = 0 45 ctx1gLMNcdrsRtuvwSTUVxyzABCWXDEYFGZHI0JK1OPiaebfh

170 cdef size_t i 

171 with nogil: 45 ctx1gLMNcdrsRtuvwSTUVxyzABCWXDEYFGZHI0JK1OPiaebfh

172 HANDLE_RETURN(self._query_fn(c_node, NULL, &count)) 45 ctx1gLMNcdrsRtuvwSTUVxyzABCWXDEYFGZHI0JK1OPiaebfh

173 if count == 0: 45 ctx1gLMNcdrsRtuvwSTUVxyzABCWXDEYFGZHI0JK1OPiaebfh

174 return [] 25 ctx1gRSTUVWXDEYFGZHI0JK1iaebf

175 cdef vector[cydriver.CUgraphNode] nodes_vec 

176 if count <= 16: 33 ctx1gLMNcdrstuvwxyzABCDEFGHIJKOPaebfh

177 nodes = stack_buf 31 ctx1gLMNrstuvwxyzABCDEFGHIJKOPaebfh

178 else: 

179 nodes_vec.resize(count) 2 ctx1cd

180 nodes = nodes_vec.data() 2 ctx1cd

181 with nogil: 33 ctx1gLMNcdrstuvwxyzABCDEFGHIJKOPaebfh

182 HANDLE_RETURN(self._query_fn(c_node, nodes, &count)) 33 ctx1gLMNcdrstuvwxyzABCDEFGHIJKOPaebfh

183 return [GraphNode._create(self._h_graph, nodes[i]) 33 ctx1gLMNcdrstuvwxyzABCDEFGHIJKOPaebfh

184 for i in range(count)] 33 ctx1gLMNcdrstuvwxyzABCDEFGHIJKOPaebfh

185  

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

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

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

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

190 return False 1 ctx1h

191 cdef cydriver.CUgraphNode stack_buf[16] 

192 cdef cydriver.CUgraphNode* nodes 

193 cdef size_t count = 0 16 ctx1pqgcdjomikaebfhn

194 cdef size_t i 

195 with nogil: 16 ctx1pqgcdjomikaebfhn

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

197 if count == 0: 16 ctx1pqgcdjomikaebfhn

198 return False 13 ctx1gcdjmikaebfhn

199 cdef vector[cydriver.CUgraphNode] nodes_vec 

200 if count <= 16: 11 ctx1pqcdokaebfh

201 nodes = stack_buf 9 ctx1pqokaebfh

202 else: 

203 nodes_vec.resize(count) 2 ctx1cd

204 nodes = nodes_vec.data() 2 ctx1cd

205 with nogil: 11 ctx1pqcdokaebfh

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

207 for i in range(count): 11 ctx1pqcdokaebfh

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

209 return True 8 ctx1pqcdokae

210 return False 6 ctx1oaebfh

211  

212 cdef Py_ssize_t count(self): (empty)

213 cdef cydriver.CUgraphNode c_node = as_cu(self._h_node) 8 ctx13cdaebQh

214 if c_node == NULL: 8 ctx13cdaebQh

215 return 0 2 ctx13Q

216 cdef size_t n = 0 6 ctx1cdaebh

217 with nogil: 6 ctx1cdaebh

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

219 return <Py_ssize_t>n 6 ctx1cdaebh

220  

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

222 cdef cydriver.CUgraphNode c_from, c_to 

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

224 with nogil: 8 ctx1jmkaefhn

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

226  

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

228 cdef size_t n = len(nodes) 7 ctx1gcdjiab

229 cdef vector[cydriver.CUgraphNode] from_vec 

230 cdef vector[cydriver.CUgraphNode] to_vec 

231 from_vec.resize(n) 7 ctx1gcdjiab

232 to_vec.resize(n) 7 ctx1gcdjiab

233 cdef size_t i 

234 for i in range(n): 7 ctx1gcdjiab

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

236 with nogil: 7 ctx1gcdjiab

237 HANDLE_RETURN(driver_add_edges( 7 ctx1gcdjiab

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

239  

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

241 cdef cydriver.CUgraphNode c_from, c_to 

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

243 with nogil: 3 ctx1kae

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

245  

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

247 cdef size_t n = len(nodes) 5 ctx1cdabf

248 cdef vector[cydriver.CUgraphNode] from_vec 

249 cdef vector[cydriver.CUgraphNode] to_vec 

250 from_vec.resize(n) 5 ctx1cdabf

251 to_vec.resize(n) 5 ctx1cdabf

252 cdef size_t i 

253 for i in range(n): 5 ctx1cdabf

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

255 with nogil: 5 ctx1cdabf

256 HANDLE_RETURN(driver_remove_edges( 5 ctx1cdabf

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

258  

259  

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

261  

262cdef inline cydriver.CUresult driver_get_preds( (empty)

263 cydriver.CUgraphNode node, cydriver.CUgraphNode* out, 

264 size_t* count) noexcept nogil: 

265 IF CUDA_CORE_BUILD_MAJOR >= 13: 

266 return cydriver.cuGraphNodeGetDependencies(node, out, NULL, count) 39 ctx1pqgLMNcorsRtuvwSTUVxyzABCWXDEYFGZOPiebf

267 ELSE: 

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

269  

270  

271cdef inline cydriver.CUresult driver_get_succs( (empty)

272 cydriver.CUgraphNode node, cydriver.CUgraphNode* out, 

273 size_t* count) noexcept nogil: 

274 IF CUDA_CORE_BUILD_MAJOR >= 13: 

275 return cydriver.cuGraphNodeGetDependentNodes(node, out, NULL, count) 29 ctx1gdjrstuvwxyzABCHI0JK1mkaebfhn

276 ELSE: 

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

278  

279  

280cdef inline cydriver.CUresult driver_add_edges( (empty)

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

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

283 IF CUDA_CORE_BUILD_MAJOR >= 13: 

284 return cydriver.cuGraphAddDependencies( 13 ctx1gcdjmikaebfhn

285 graph, from_arr, to_arr, NULL, count) 

286 ELSE: 

287 return cydriver.cuGraphAddDependencies( 

288 graph, from_arr, to_arr, count) 

289  

290  

291cdef inline cydriver.CUresult driver_remove_edges( (empty)

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

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

294 IF CUDA_CORE_BUILD_MAJOR >= 13: 

295 return cydriver.cuGraphRemoveDependencies( 7 ctx1cdkaebf

296 graph, from_arr, to_arr, NULL, count) 

297 ELSE: 

298 return cydriver.cuGraphRemoveDependencies( 

299 graph, from_arr, to_arr, count)