Coverage for cuda / core / _memory / _graph_memory_resource.pyx: 80.26%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4  

5from __future__ import annotations 

6  

7from libc.stdint cimport intptr_t 

8  

9from cuda.bindings cimport cydriver 

10from cuda.core._memory._buffer cimport Buffer, Buffer_from_deviceptr_handle, MemoryResource 

11from cuda.core._resource_handles cimport ( 

12 DevicePtrHandle, 

13 deviceptr_alloc_async, 

14 get_last_error, 

15 as_cu, 

16) 

17  

18from cuda.core._stream cimport Stream_accept, Stream 

19from cuda.core._utils.cuda_utils cimport HANDLE_RETURN 

20  

21from functools import cache (empty)

22  

23__all__ = ['GraphMemoryResource'] (empty)

24  

25  

26cdef class GraphMemoryResourceAttributes: 

27 cdef: 

28 int _device_id 

29  

30 def __init__(self, *args, **kwargs): 

31 raise RuntimeError("GraphMemoryResourceAttributes cannot be instantiated directly. Please use MemoryResource APIs.") 

32  

33 @classmethod (empty)

34 def _init(cls, device_id: int): 

35 cdef GraphMemoryResourceAttributes self = GraphMemoryResourceAttributes.__new__(cls) 3 ctx1abc

36 self._device_id = device_id 3 ctx1abc

37 return self 3 ctx1abc

38  

39 def __repr__(self): 

40 return f"{self.__class__.__name__}(%s)" % ", ".join( 

41 f"{attr}={getattr(self, attr)}" for attr in dir(self) 

42 if not attr.startswith("_") 

43 ) 

44  

45 cdef int _getattribute(self, cydriver.CUgraphMem_attribute attr_enum, void* value) except?-1: (empty)

46 with nogil: 3 ctx1abc

47 HANDLE_RETURN(cydriver.cuDeviceGetGraphMemAttribute(self._device_id, attr_enum, value)) 4 ctx1dabc

48 return 0 3 ctx1abc

49  

50 cdef int _setattribute(self, cydriver.CUgraphMem_attribute attr_enum, void* value) except?-1: (empty)

51 with nogil: 3 ctx1abc

52 HANDLE_RETURN(cydriver.cuDeviceSetGraphMemAttribute(self._device_id, attr_enum, value)) 3 ctx1abc

53 return 0 3 ctx1abc

54  

55 @property 

56 def reserved_mem_current(self): 

57 """Current amount of backing memory allocated.""" 

58 cdef cydriver.cuuint64_t value 

59 self._getattribute(cydriver.CUgraphMem_attribute.CU_GRAPH_MEM_ATTR_RESERVED_MEM_CURRENT, &value) 3 ctx1abc

60 return int(value) 3 ctx1abc

61  

62 @property 

63 def reserved_mem_high(self): 

64 """ 

65 High watermark of backing memory allocated. It can be set to zero to 

66 reset it to the current usage. 

67 """ 

68 cdef cydriver.cuuint64_t value 

69 self._getattribute(cydriver.CUgraphMem_attribute.CU_GRAPH_MEM_ATTR_RESERVED_MEM_HIGH, &value) 3 ctx1abc

70 return int(value) 3 ctx1abc

71  

72 @reserved_mem_high.setter 

73 def reserved_mem_high(self, value: int): 

74 if value != 0: 3 ctx1abc

75 raise AttributeError(f"Attribute 'reserved_mem_high' may only be set to zero (got {value}).") 3 ctx1abc

76 cdef cydriver.cuuint64_t zero = 0 3 ctx1abc

77 self._setattribute(cydriver.CUgraphMem_attribute.CU_GRAPH_MEM_ATTR_RESERVED_MEM_HIGH, &zero) 3 ctx1abc

78  

79 @property 

80 def used_mem_current(self): 

81 """Current amount of memory in use.""" 

82 cdef cydriver.cuuint64_t value 

83 self._getattribute(cydriver.CUgraphMem_attribute.CU_GRAPH_MEM_ATTR_USED_MEM_CURRENT, &value) 3 ctx1abc

84 return int(value) 3 ctx1abc

85  

86 @property 

87 def used_mem_high(self): 

88 """ 

89 High watermark of memory in use. It can be set to zero to reset it to 

90 the current usage. 

91 """ 

92 cdef cydriver.cuuint64_t value 

93 self._getattribute(cydriver.CUgraphMem_attribute.CU_GRAPH_MEM_ATTR_USED_MEM_HIGH, &value) 3 ctx1abc

94 return int(value) 3 ctx1abc

95  

96 @used_mem_high.setter 

97 def used_mem_high(self, value: int): 

98 if value != 0: 3 ctx1abc

99 raise AttributeError(f"Attribute 'used_mem_high' may only be set to zero (got {value}).") 3 ctx1abc

100 cdef cydriver.cuuint64_t zero = 0 3 ctx1abc

101 self._setattribute(cydriver.CUgraphMem_attribute.CU_GRAPH_MEM_ATTR_USED_MEM_HIGH, &zero) 3 ctx1abc

102  

103  

104cdef class cyGraphMemoryResource(MemoryResource): 

105 def __cinit__(self, int device_id): 

106 self._device_id = device_id 1 ctx1t

107  

108 def allocate(self, size_t size, *, stream: Stream | GraphBuilder) -> Buffer: (empty)

109 """ 

110 Allocate a buffer of the requested size. See documentation for :obj:`~_memory.MemoryResource`. 

111 """ 

112 cdef Stream s = Stream_accept(stream) 18 ctx1efgnopqrshijklmabc

113 return GMR_allocate(self, size, s) 18 ctx1efgnopqrshijklmabc

114  

115 def deallocate(self, ptr: "DevicePointerType", size_t size, *, stream: Stream | GraphBuilder): (empty)

116 """ 

117 Deallocate a buffer of the requested size. See documentation for :obj:`~_memory.MemoryResource`. 

118 """ 

119 cdef Stream s = Stream_accept(stream) 

120 return GMR_deallocate(ptr, size, s) 

121  

122 def close(self): (empty)

123 """No operation (provided for compatibility).""" 

124 pass 

125  

126 def trim(self): (empty)

127 """Free unused memory that was cached on the specified device for use with graphs back to the OS.""" 

128 with nogil: 3 ctx1abc

129 HANDLE_RETURN(cydriver.cuDeviceGraphMemTrim(self._device_id)) 3 ctx1abc

130  

131 @property 

132 def attributes(self) -> GraphMemoryResourceAttributes: 

133 """Asynchronous allocation attributes related to graphs.""" 

134 return GraphMemoryResourceAttributes._init(self._device_id) 3 ctx1abc

135  

136 @property 

137 def device_id(self) -> int: 

138 """The associated device ordinal.""" 

139 return self._device_id 9 ctx1hijklmabc

140  

141 @property 

142 def is_device_accessible(self) -> bool: 

143 """Return True. This memory resource provides device-accessible buffers.""" 

144 return True 

145  

146 @property 

147 def is_host_accessible(self) -> bool: 

148 """Return False. This memory resource does not provide host-accessible buffers.""" 

149 return False 

150  

151  

152class GraphMemoryResource(cyGraphMemoryResource): (empty)

153 """ 

154 A memory resource for memory related to graphs. 

155  

156 The only supported operations are allocation, deallocation, and a limited 

157 set of status queries. 

158  

159 This memory resource should be used when building graphs. Using this when 

160 graphs capture is not enabled will result in a runtime error. 

161  

162 Conversely, allocating memory from a `DeviceMemoryResource` when graph 

163 capturing is enabled results in a runtime error. 

164  

165 Parameters 

166 ---------- 

167 device_id: int | Device 

168 Device or Device ordinal for which a graph memory resource is obtained. 

169 """ 

170  

171 def __new__(cls, device_id: int | Device): (empty)

172 cdef int c_device_id = getattr(device_id, 'device_id', device_id) 21 ctx1efgnuopqvrshijklmabct

173 return cls._create(c_device_id) 21 ctx1efgnuopqvrshijklmabct

174  

175 @classmethod (empty)

176 @cache (empty)

177 def _create(cls, int device_id): 

178 return cyGraphMemoryResource.__new__(cls, device_id) 1 ctx1t

179  

180  

181# Raise an exception if the given stream is capturing. 

182# A result of CU_STREAM_CAPTURE_STATUS_INVALIDATED is considered an error. 

183cdef inline int check_capturing(cydriver.CUstream s) except?-1 nogil: (empty)

184 cdef cydriver.CUstreamCaptureStatus capturing 

185 HANDLE_RETURN(cydriver.cuStreamIsCapturing(s, &capturing)) 18 ctx1efgnopqrshijklmabc

186 if capturing != cydriver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_ACTIVE: 18 ctx1efgnopqrshijklmabc

187 raise RuntimeError("GraphMemoryResource cannot perform memory operations on " 3 ctx1efg

188 "a non-capturing stream.") 

189  

190  

191cdef inline Buffer GMR_allocate(cyGraphMemoryResource self, size_t size, Stream stream): (empty)

192 cdef cydriver.CUstream s = as_cu(stream._h_stream) 18 ctx1efgnopqrshijklmabc

193 cdef DevicePtrHandle h_ptr 

194 with nogil: 18 ctx1efgnopqrshijklmabc

195 check_capturing(s) 18 ctx1efgnopqrshijklmabc

196 h_ptr = deviceptr_alloc_async(size, stream._h_stream) 18 ctx1efgnopqrshijklmabc

197 if not h_ptr: 18 ctx1efgnopqrshijklmabc

198 HANDLE_RETURN(get_last_error()) 

199 raise RuntimeError( 

200 f"Failed to allocate {size} bytes from GraphMemoryResource: " 

201 "cuda-core returned an empty allocation handle without recording a CUDA error. " 

202 "This is an internal cuda-core error; please report it with your CUDA driver, " 

203 "CUDA Toolkit, and cuda-python versions." 

204 ) 

205 return Buffer_from_deviceptr_handle(h_ptr, size, self, None) 18 ctx1efgnopqrshijklmabc

206  

207  

208cdef inline void GMR_deallocate(intptr_t ptr, size_t size, Stream stream) noexcept: (empty)

209 cdef cydriver.CUstream s = as_cu(stream._h_stream) 

210 cdef cydriver.CUdeviceptr devptr = <cydriver.CUdeviceptr>ptr 

211 with nogil: 

212 HANDLE_RETURN(cydriver.cuMemFreeAsync(devptr, s))