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

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 typing import TYPE_CHECKING (empty)

22  

23if TYPE_CHECKING: (empty)

24 from cuda.core._device import Device 

25 from cuda.core.graph import GraphBuilder 

26 from cuda.core.typing import DevicePointerType (empty)

27  

28__all__ = ['GraphMemoryResource'] (empty)

29  

30  

31cdef class GraphMemoryResourceAttributes: 

32 cdef: 

33 int _device_id 

34  

35 def __init__(self, *args, **kwargs) -> None: 

36 raise RuntimeError("GraphMemoryResourceAttributes cannot be instantiated directly. Please use MemoryResource APIs.") 1 ctx1y

37  

38 @classmethod (empty)

39 def _init(cls, device_id: int) -> GraphMemoryResourceAttributes: 

40 cdef GraphMemoryResourceAttributes self = GraphMemoryResourceAttributes.__new__(cls) 4 ctx1abcd

41 self._device_id = device_id 4 ctx1abcd

42 return self 4 ctx1abcd

43  

44 def __repr__(self) -> str: 

45 return f"{self.__class__.__name__}(%s)" % ", ".join( 1 ctx1d

46 f"{attr}={getattr(self, attr)}" for attr in dir(self) 1 ctx1d

47 if not attr.startswith("_") 1 ctx1d

48 ) 

49  

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

51 with nogil: 4 ctx1abcd

52 HANDLE_RETURN(cydriver.cuDeviceGetGraphMemAttribute(self._device_id, attr_enum, value)) 5 ctx1eabcd

53 return 0 4 ctx1abcd

54  

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

56 with nogil: 3 ctx1abc

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

58 return 0 3 ctx1abc

59  

60 @property 

61 def reserved_mem_current(self) -> int: 

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

63 cdef cydriver.cuuint64_t value 

64 self._getattribute(cydriver.CUgraphMem_attribute.CU_GRAPH_MEM_ATTR_RESERVED_MEM_CURRENT, &value) 4 ctx1abcd

65 return int(value) 4 ctx1abcd

66  

67 @property 

68 def reserved_mem_high(self) -> int: 

69 """ 

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

71 reset it to the current usage. 

72 """ 

73 cdef cydriver.cuuint64_t value 

74 self._getattribute(cydriver.CUgraphMem_attribute.CU_GRAPH_MEM_ATTR_RESERVED_MEM_HIGH, &value) 4 ctx1abcd

75 return int(value) 4 ctx1abcd

76  

77 @reserved_mem_high.setter 

78 def reserved_mem_high(self, value: int) -> None: 

79 if value != 0: 3 ctx1abc

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

81 cdef cydriver.cuuint64_t zero = 0 3 ctx1abc

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

83  

84 @property 

85 def used_mem_current(self) -> int: 

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

87 cdef cydriver.cuuint64_t value 

88 self._getattribute(cydriver.CUgraphMem_attribute.CU_GRAPH_MEM_ATTR_USED_MEM_CURRENT, &value) 4 ctx1abcd

89 return int(value) 4 ctx1abcd

90  

91 @property 

92 def used_mem_high(self) -> int: 

93 """ 

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

95 the current usage. 

96 """ 

97 cdef cydriver.cuuint64_t value 

98 self._getattribute(cydriver.CUgraphMem_attribute.CU_GRAPH_MEM_ATTR_USED_MEM_HIGH, &value) 4 ctx1abcd

99 return int(value) 4 ctx1abcd

100  

101 @used_mem_high.setter 

102 def used_mem_high(self, value: int) -> None: 

103 if value != 0: 3 ctx1abc

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

105 cdef cydriver.cuuint64_t zero = 0 3 ctx1abc

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

107  

108  

109cdef class cyGraphMemoryResource(MemoryResource): 

110 def __cinit__(self, int device_id) -> None: 

111 self._device_id = device_id 1 ctx1u

112  

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

114 """ 

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

116 """ 

117 cdef Stream s = Stream_accept(stream) 18 ctx1fghopqrstijklmnabc

118 return GMR_allocate(self, size, s) 18 ctx1fghopqrstijklmnabc

119  

120 def deallocate( (empty)

121 self, 

122 ptr: DevicePointerType, 

123 size_t size, 

124 *, 

125 stream: Stream | GraphBuilder 

126 ) -> None: 

127 """ 

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

129 """ 

130 cdef Stream s = Stream_accept(stream) 

131 return GMR_deallocate(ptr, size, s) 

132  

133 def close(self) -> None: (empty)

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

135 pass 

136  

137 def trim(self) -> None: (empty)

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

139 with nogil: 3 ctx1abc

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

141  

142 @property 

143 def attributes(self) -> GraphMemoryResourceAttributes: 

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

145 return GraphMemoryResourceAttributes._init(self._device_id) 4 ctx1abcd

146  

147 @property 

148 def device_id(self) -> int: 

149 """The associated device ordinal.""" 

150 return self._device_id 10 ctx1ijklmnabcv

151  

152 @property 

153 def is_device_accessible(self) -> bool: 

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

155 return True 1 ctx1v

156  

157 @property 

158 def is_host_accessible(self) -> bool: 

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

160 return False 1 ctx1v

161  

162  

163cdef dict _mem_resource_cache = {} (empty)

164  

165class GraphMemoryResource(cyGraphMemoryResource): (empty)

166 """ 

167 A memory resource for memory related to graphs. 

168  

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

170 set of status queries. 

171  

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

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

174  

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

176 capturing is enabled results in a runtime error. 

177  

178 Parameters 

179 ---------- 

180 device_id: int | Device 

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

182 """ 

183  

184 def __new__(cls, device_id: int | Device) -> GraphMemoryResource: (empty)

185 cdef int c_device_id = getattr(device_id, 'device_id', device_id) 23 ctx1fghowpqrxstijklmnabcvdu

186 return cls._create(c_device_id) 23 ctx1fghowpqrxstijklmnabcvdu

187  

188 @classmethod (empty)

189 def _create(cls, int device_id) -> GraphMemoryResource: 

190 # we use a dict currently, because functools.cache is currently less 

191 # thread-safe see also: https://github.com/python/cpython/issues/150708 

192 res = _mem_resource_cache.get(device_id) 23 ctx1fghowpqrxstijklmnabcvdu

193 if res is not None: 23 ctx1fghowpqrxstijklmnabcvdu

194 return res 23 ctx1fghowpqrxstijklmnabcvdu

195  

196 # create new instance, but in case of a race may return another: 

197 new = cyGraphMemoryResource.__new__(cls, device_id) 1 ctx1u

198 return _mem_resource_cache.setdefault(device_id, new) 1 ctx1u

199  

200  

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

202# A result of CU_STREAM_CAPTURE_STATUS_INVALIDATED is considered an error. 

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

204 cdef cydriver.CUstreamCaptureStatus capturing 

205 HANDLE_RETURN(cydriver.cuStreamIsCapturing(s, &capturing)) 18 ctx1fghopqrstijklmnabc

206 if capturing != cydriver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_ACTIVE: 18 ctx1fghopqrstijklmnabc

207 raise RuntimeError("GraphMemoryResource cannot perform memory operations on " 3 ctx1fgh

208 "a non-capturing stream.") 

209  

210  

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

212 cdef cydriver.CUstream s = as_cu(stream._h_stream) 18 ctx1fghopqrstijklmnabc

213 cdef DevicePtrHandle h_ptr 

214 with nogil: 18 ctx1fghopqrstijklmnabc

215 check_capturing(s) 18 ctx1fghopqrstijklmnabc

216 h_ptr = deviceptr_alloc_async(size, stream._h_stream) 18 ctx1fghopqrstijklmnabc

217 if not h_ptr: 18 ctx1fghopqrstijklmnabc

218 HANDLE_RETURN(get_last_error()) 

219 raise RuntimeError( 

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

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

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

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

224 ) 

225 return Buffer_from_deviceptr_handle(h_ptr, size, self, None) 18 ctx1fghopqrstijklmnabc

226  

227  

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

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

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

231 with nogil: 

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