Coverage for cuda / core / _memory / _legacy.py: 64.52%

62 statements  

« prev     ^ index     » next       coverage.py v7.14.0, created at 2026-05-22 01:37 +0000

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from __future__ import annotations 

6 

7from typing import TYPE_CHECKING 

8 

9if TYPE_CHECKING: 

10 from cuda.core._memory._buffer import DevicePointerType 

11 from cuda.core._stream import Stream 

12 

13from cuda.core._memory._buffer import Buffer, MemoryResource 

14from cuda.core._utils.cuda_utils import ( 

15 _check_driver_error as raise_if_driver_error, 

16) 

17from cuda.core._utils.cuda_utils import ( 

18 driver, 

19) 

20 

21__all__ = ["LegacyPinnedMemoryResource"] 

22 

23 

24class LegacyPinnedMemoryResource(MemoryResource): 

25 """Create a pinned memory resource that uses legacy cuMemAllocHost/cudaMallocHost 

26 APIs. 

27 """ 

28 

29 # TODO: support creating this MR with flags that are later passed to cuMemHostAlloc? 

30 

31 def allocate(self, size, *, stream: Stream | None = None) -> Buffer: 

32 """Allocate a buffer of the requested size. 

33 

34 ``cuMemAllocHost`` is synchronous, so this resource ignores any 

35 supplied stream. The argument is accepted (and validated when 

36 non-``None``) for interface conformance with stream-ordered 

37 memory resources. 

38 

39 Parameters 

40 ---------- 

41 size : int 

42 The size of the buffer to allocate, in bytes. 

43 stream : Stream, optional 

44 Keyword-only. Validated when provided but otherwise unused. 

45 

46 Returns 

47 ------- 

48 Buffer 

49 The allocated buffer object, which is accessible on both host and device. 

50 """ 

51 from cuda.core._stream import Stream_accept 1bcdefghijklmnopqrstuvwxyzABCDEFGHIJK:,-.LMN/OPQRSTUVWXYZ0123456789!#$%'()*+

52 

53 if stream is not None: 1bcdefghijklmnopqrstuvwxyzABCDEFGHIJK:,-.LMN/OPQRSTUVWXYZ0123456789!#$%'()*+

54 Stream_accept(stream) 

55 if size: 1bcdefghijklmnopqrstuvwxyzABCDEFGHIJK:,-.LMN/OPQRSTUVWXYZ0123456789!#$%'()*+

56 err, ptr = driver.cuMemAllocHost(size) 1bcdefghijklmnopqrstuvwxyzABCDEFGHIJK:,-.LMN/OPQRSTUVWXYZ0123456789!#$%'()*+

57 raise_if_driver_error(err) 1bcdefghijklmnopqrstuvwxyzABCDEFGHIJK:,-.LMN/OPQRSTUVWXYZ0123456789!#$%'()*+

58 else: 

59 ptr = 0 

60 return Buffer._init(ptr, size, self) 1bcdefghijklmnopqrstuvwxyzABCDEFGHIJK:,-.LMN/OPQRSTUVWXYZ0123456789!#$%'()*+

61 

62 def deallocate(self, ptr: DevicePointerType, size, *, stream: Stream | None = None): 

63 """Deallocate a buffer previously allocated by this resource. 

64 

65 Parameters 

66 ---------- 

67 ptr : :obj:`~_memory.DevicePointerType` 

68 The pointer or handle to the buffer to deallocate. 

69 size : int 

70 The size of the buffer to deallocate, in bytes. 

71 stream : Stream, optional 

72 Keyword-only. If provided, ``stream.sync()`` is called before the 

73 host allocation is freed. ``None`` skips the sync. 

74 """ 

75 from cuda.core._stream import Stream_accept 1bcdefghijklmnopqrstuvwxyzABCDEFGHIJK,-.LMN/OPQRSTUVWXYZ0123456789!#$%'()*+

76 

77 if stream is not None: 1bcdefghijklmnopqrstuvwxyzABCDEFGHIJK,-.LMN/OPQRSTUVWXYZ0123456789!#$%'()*+

78 Stream_accept(stream).sync() 1bcdefghijklmnopqrstuvwxyzABCDEFGHIJK,-.LMN/OPQRSTUVWXYZ0123456789!#$%'()*+

79 

80 if size: 1bcdefghijklmnopqrstuvwxyzABCDEFGHIJK,-.LMN/OPQRSTUVWXYZ0123456789!#$%'()*+

81 (err,) = driver.cuMemFreeHost(ptr) 1bcdefghijklmnopqrstuvwxyzABCDEFGHIJK,-.LMN/OPQRSTUVWXYZ0123456789!#$%'()*+

82 raise_if_driver_error(err) 1bcdefghijklmnopqrstuvwxyzABCDEFGHIJK,-.LMN/OPQRSTUVWXYZ0123456789!#$%'()*+

83 

84 @property 

85 def is_device_accessible(self) -> bool: 

86 """bool: this memory resource provides device-accessible buffers.""" 

87 return True 1bcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!#$%'()*+

88 

89 @property 

90 def is_host_accessible(self) -> bool: 

91 """bool: this memory resource provides host-accessible buffers.""" 

92 return True 1bcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!#$%'()*+

93 

94 @property 

95 def device_id(self) -> int: 

96 """This memory resource is not bound to any GPU.""" 

97 raise RuntimeError("a pinned memory resource is not bound to any GPU") 

98 

99 

100class _SynchronousMemoryResource(MemoryResource): 

101 __slots__ = ("_device_id",) 

102 

103 def __init__(self, device_id): 

104 from .._device import Device 

105 

106 self._device_id = Device(device_id).device_id 

107 

108 def allocate(self, size, *, stream: Stream | None = None) -> Buffer: 

109 # cuMemAlloc is synchronous; stream is accepted (and validated) 

110 # for interface conformance but not used. 

111 from cuda.core._stream import Stream_accept 

112 

113 if stream is not None: 

114 Stream_accept(stream) 

115 if size: 

116 err, ptr = driver.cuMemAlloc(size) 

117 raise_if_driver_error(err) 

118 else: 

119 ptr = 0 

120 return Buffer._init(ptr, size, self) 

121 

122 def deallocate(self, ptr, size, *, stream: Stream | None = None): 

123 from cuda.core._stream import Stream_accept 

124 

125 if stream is not None: 

126 Stream_accept(stream).sync() 

127 if size: 

128 (err,) = driver.cuMemFree(ptr) 

129 raise_if_driver_error(err) 

130 

131 @property 

132 def is_device_accessible(self) -> bool: 

133 return True 

134 

135 @property 

136 def is_host_accessible(self) -> bool: 

137 return False 

138 

139 @property 

140 def device_id(self) -> int: 

141 return self._device_id