Coverage for cuda / core / _memory / _managed_memory_resource.pyx: 88.64%

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 cuda.bindings cimport cydriver 

8  

9from cuda.core._memory._memory_pool cimport _MemPool, _MemPoolOptions 

10from cuda.core._utils.cuda_utils cimport ( 

11 HANDLE_RETURN, 

12 check_or_create_options, 

13) 

14  

15from dataclasses import dataclass (empty)

16import threading (empty)

17import warnings (empty)

18  

19__all__ = ['ManagedMemoryResource', 'ManagedMemoryResourceOptions'] (empty)

20  

21  

22@dataclass (empty)

23cdef class ManagedMemoryResourceOptions: 

24 """Customizable :obj:`~_memory.ManagedMemoryResource` options. 

25  

26 Attributes 

27 ---------- 

28 preferred_location : int | None, optional 

29 The preferred device location for the managed memory. 

30 Use a device ID (0, 1, 2, ...) for device preference, -1 for CPU/host, 

31 or None to let the driver decide. 

32 (Default to None) 

33 """ 

34 preferred_location: int | None = None (empty)

35  

36  

37cdef class ManagedMemoryResource(_MemPool): 

38 """ 

39 A managed memory resource managing a stream-ordered memory pool. 

40  

41 Managed memory is accessible from both the host and device, with automatic 

42 migration between them as needed. 

43  

44 Parameters 

45 ---------- 

46 options : ManagedMemoryResourceOptions 

47 Memory resource creation options. 

48  

49 If set to `None`, the memory resource uses the driver's current 

50 stream-ordered memory pool. If no memory pool is set as current, 

51 the driver's default memory pool is used. 

52  

53 If not set to `None`, a new memory pool is created, which is owned by 

54 the memory resource. 

55  

56 When using an existing (current or default) memory pool, the returned 

57 managed memory resource does not own the pool (`is_handle_owned` is 

58 `False`), and closing the resource has no effect. 

59  

60 Notes 

61 ----- 

62 IPC (Inter-Process Communication) is not currently supported for managed 

63 memory pools. 

64 """ 

65  

66 def __init__(self, options=None): 

67 cdef ManagedMemoryResourceOptions opts = check_or_create_options( 13 ctx1mbcdaefghijkl

68 ManagedMemoryResourceOptions, options, "ManagedMemoryResource options", 

69 keep_none=True 

70 ) 

71 cdef _MemPoolOptions opts_base = _MemPoolOptions() 13 ctx1mbcdaefghijkl

72  

73 cdef int device_id = -1 13 ctx1mbcdaefghijkl

74 cdef object preferred_location = None 13 ctx1mbcdaefghijkl

75 if opts: 13 ctx1mbcdaefghijkl

76 preferred_location = opts.preferred_location 12 ctx1bcdaefghijkl

77 if preferred_location is not None: 12 ctx1bcdaefghijkl

78 device_id = preferred_location 1 ctx1c

79 opts_base._use_current = False 12 ctx1bcdaefghijkl

80  

81 opts_base._ipc_enabled = False # IPC not supported for managed memory pools 13 ctx1mbcdaefghijkl

82  

83 IF CUDA_CORE_BUILD_MAJOR >= 13: 

84 # Set location based on preferred_location 

85 if preferred_location is None: 13 ctx1mbcdaefghijkl

86 # Let the driver decide 

87 opts_base._location = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_NONE 12 ctx1mbdaefghijkl

88 elif device_id == -1: 1 ctx1c

89 # CPU/host preference 

90 opts_base._location = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST 

91 else: 

92 # Device preference 

93 opts_base._location = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE 1 ctx1c

94  

95 opts_base._type = cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_MANAGED 13 ctx1mbcdaefghijkl

96  

97 super().__init__(device_id, opts_base) 13 ctx1mbcdaefghijkl

98 _check_concurrent_managed_access() 13 ctx1mbcdaefghijkl

99 ELSE: 

100 raise RuntimeError("ManagedMemoryResource requires CUDA 13.0 or later") 

101  

102 @property 

103 def is_device_accessible(self) -> bool: 

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

105 return True 2 ctx1mb

106  

107 @property 

108 def is_host_accessible(self) -> bool: 

109 """Return True. This memory resource provides host-accessible buffers.""" 

110 return True 2 ctx1mb

111  

112  

113cdef bint _concurrent_access_warned = False (empty)

114cdef object _concurrent_access_lock = threading.Lock() (empty)

115  

116  

117cdef inline _check_concurrent_managed_access(): (empty)

118 """Warn once if the platform lacks concurrent managed memory access.""" 

119 global _concurrent_access_warned 

120 if _concurrent_access_warned: 13 ctx1mbcdaefghijkl

121 return 12 ctx1mbcdefghijkl

122  

123 cdef int c_concurrent = 0 1 ctx1a

124 with _concurrent_access_lock: 1 ctx1a

125 if _concurrent_access_warned: 1 ctx1a

126 return 

127  

128 # concurrent_managed_access is a system-level attribute for sm_60 and 

129 # later, so any device will do. 

130 with nogil: 1 ctx1a

131 HANDLE_RETURN(cydriver.cuDeviceGetAttribute( 1 ctx1a

132 &c_concurrent, 

133 cydriver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_CONCURRENT_MANAGED_ACCESS, 

134 0)) 

135 if not c_concurrent: 1 ctx1a

136 warnings.warn( 

137 "This platform does not support concurrent managed memory access " 

138 "(Device.properties.concurrent_managed_access is False). Host access to any managed " 

139 "allocation is forbidden while any GPU kernel is in flight, even " 

140 "if the kernel does not touch that allocation. Failing to " 

141 "synchronize before host access will cause a segfault. " 

142 "See: https://docs.nvidia.com/cuda/cuda-c-programming-guide/" 

143 "index.html#gpu-exclusive-access-to-managed-memory", 

144 UserWarning, 

145 stacklevel=3 

146 ) 

147  

148 _concurrent_access_warned = True 1 ctx1a

149  

150  

151def reset_concurrent_access_warning(): (empty)

152 """Reset the concurrent access warning flag for testing purposes.""" 

153 global _concurrent_access_warned 

154 _concurrent_access_warned = False