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

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)

22from typing import TYPE_CHECKING (empty)

23  

24if TYPE_CHECKING: (empty)

25 from cuda.core._device import Device 

26 from cuda.core.graph import GraphBuilder (empty)

27 from cuda.core.typing import DevicePointerType 

28  

29__all__ = ['GraphMemoryResource'] (empty)

30  

31  

32cdef class GraphMemoryResourceAttributes: 

33 cdef: 

34 int _device_id 

35  

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

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

38  

39 @classmethod (empty)

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

41 cdef GraphMemoryResourceAttributes self = GraphMemoryResourceAttributes.__new__(cls) 4 ctx1abce

42 self._device_id = device_id 4 ctx1abce

43 return self 4 ctx1abce

44  

45 def __repr__(self) -> str: 

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

47 f"{attr}={getattr(self, attr)}" for attr in dir(self) 2 ctx1de

48 if not attr.startswith("_") 1 ctx1e

49 ) 

50  

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

52 with nogil: 4 ctx1abce

53 HANDLE_RETURN(cydriver.cuDeviceGetGraphMemAttribute(self._device_id, attr_enum, value)) 4 ctx1abce

54 return 0 4 ctx1abce

55  

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

57 with nogil: 3 ctx1abc

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

59 return 0 3 ctx1abc

60  

61 @property 

62 def reserved_mem_current(self) -> int: 

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

64 cdef cydriver.cuuint64_t value 

65 self._getattribute(cydriver.CUgraphMem_attribute.CU_GRAPH_MEM_ATTR_RESERVED_MEM_CURRENT, &value) 4 ctx1abce

66 return int(value) 4 ctx1abce

67  

68 @property 

69 def reserved_mem_high(self) -> int: 

70 """ 

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

72 reset it to the current usage. 

73 """ 

74 cdef cydriver.cuuint64_t value 

75 self._getattribute(cydriver.CUgraphMem_attribute.CU_GRAPH_MEM_ATTR_RESERVED_MEM_HIGH, &value) 4 ctx1abce

76 return int(value) 4 ctx1abce

77  

78 @reserved_mem_high.setter 

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

80 if value != 0: 3 ctx1abc

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

82 cdef cydriver.cuuint64_t zero = 0 3 ctx1abc

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

84  

85 @property 

86 def used_mem_current(self) -> int: 

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

88 cdef cydriver.cuuint64_t value 

89 self._getattribute(cydriver.CUgraphMem_attribute.CU_GRAPH_MEM_ATTR_USED_MEM_CURRENT, &value) 4 ctx1abce

90 return int(value) 4 ctx1abce

91  

92 @property 

93 def used_mem_high(self) -> int: 

94 """ 

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

96 the current usage. 

97 """ 

98 cdef cydriver.cuuint64_t value 

99 self._getattribute(cydriver.CUgraphMem_attribute.CU_GRAPH_MEM_ATTR_USED_MEM_HIGH, &value) 4 ctx1abce

100 return int(value) 4 ctx1abce

101  

102 @used_mem_high.setter 

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

104 if value != 0: 3 ctx1abc

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

106 cdef cydriver.cuuint64_t zero = 0 3 ctx1abc

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

108  

109  

110cdef class cyGraphMemoryResource(MemoryResource): 

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

112 self._device_id = device_id 1 ctx1v

113  

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

115 """ 

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

117 """ 

118 cdef Stream s = Stream_accept(stream) 18 ctx1fghopqrstijklmnabc

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

120  

121 def deallocate( (empty)

122 self, 

123 ptr: DevicePointerType, 

124 size_t size, 

125 *, 

126 stream: Stream | GraphBuilder 

127 ) -> None: 

128 """ 

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

130 """ 

131 cdef Stream s = Stream_accept(stream) 

132 return GMR_deallocate(ptr, size, s) 

133  

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

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

136 pass 

137  

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

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

140 with nogil: 3 ctx1abc

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

142  

143 @property 

144 def attributes(self) -> GraphMemoryResourceAttributes: 

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

146 return GraphMemoryResourceAttributes._init(self._device_id) 4 ctx1abce

147  

148 @property 

149 def device_id(self) -> int: 

150 """The associated device ordinal.""" 

151 return self._device_id 10 ctx1ijklmnabcu

152  

153 @property 

154 def is_device_accessible(self) -> bool: 

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

156 return True 1 ctx1u

157  

158 @property 

159 def is_host_accessible(self) -> bool: 

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

161 return False 1 ctx1u

162  

163  

164class GraphMemoryResource(cyGraphMemoryResource): (empty)

165 """ 

166 A memory resource for memory related to graphs. 

167  

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

169 set of status queries. 

170  

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

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

173  

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

175 capturing is enabled results in a runtime error. 

176  

177 Parameters 

178 ---------- 

179 device_id: int | Device 

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

181 """ 

182  

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

184 cdef int c_device_id = getattr(device_id, 'device_id', device_id) 23 ctx1fghowpqrxstijklmnabcuev

185 return cls._create(c_device_id) 23 ctx1fghowpqrxstijklmnabcuev

186  

187 @classmethod (empty)

188 @cache (empty)

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

190 return cyGraphMemoryResource.__new__(cls, device_id) 1 ctx1v

191  

192  

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

194# A result of CU_STREAM_CAPTURE_STATUS_INVALIDATED is considered an error. 

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

196 cdef cydriver.CUstreamCaptureStatus capturing 

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

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

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

200 "a non-capturing stream.") 

201  

202  

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

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

205 cdef DevicePtrHandle h_ptr 

206 with nogil: 18 ctx1fghopqrstijklmnabc

207 check_capturing(s) 18 ctx1fghopqrstijklmnabc

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

209 if not h_ptr: 18 ctx1fghopqrstijklmnabc

210 HANDLE_RETURN(get_last_error()) 

211 raise RuntimeError( 

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

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

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

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

216 ) 

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

218  

219  

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

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

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

223 with nogil: 

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