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

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 (empty)

6 

7from typing import TYPE_CHECKING (empty)

8 

9if TYPE_CHECKING: (empty)

10 from cuda.core._stream import Stream 

11 from cuda.core.graph import GraphBuilder 

12 from cuda.core.typing import DevicePointerType 

13 

14from cuda.core._memory._buffer import Buffer, MemoryResource (empty)

15from cuda.core._utils.cuda_utils import ( (empty)

16 _check_driver_error as raise_if_driver_error, 

17) 

18from cuda.core._utils.cuda_utils import ( (empty)

19 driver, 

20) 

21 

22__all__ = ["LegacyPinnedMemoryResource"] (empty)

23 

24 

25class LegacyPinnedMemoryResource(MemoryResource): (empty)

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

27 APIs. 

28 """ 

29 

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

31 

32 def allocate(self, size: int, *, stream: Stream | GraphBuilder | None = None) -> Buffer: (empty)

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

34 

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

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

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

38 memory resources. 

39 

40 Parameters 

41 ---------- 

42 size : int 

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

44 stream : Stream, optional 

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

46 

47 Returns 

48 ------- 

49 Buffer 

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

51 """ 

52 from cuda.core._stream import Stream_accept 76 ctx1cdefghijklmnopqrstuvwxyzABCDEFGHIJKL=?./MNO:PQRSTUVWXYZ0123456789!#$%'()*+;,

53 

54 if stream is not None: 76 ctx1cdefghijklmnopqrstuvwxyzABCDEFGHIJKL=?./MNO:PQRSTUVWXYZ0123456789!#$%'()*+;,

55 Stream_accept(stream) 

56 if size: 76 ctx1cdefghijklmnopqrstuvwxyzABCDEFGHIJKL=?./MNO:PQRSTUVWXYZ0123456789!#$%'()*+;,

57 err, ptr = driver.cuMemAllocHost(size) 75 ctx1cdefghijklmnopqrstuvwxyzABCDEFGHIJKL=?./MNO:PQRSTUVWXYZ0123456789!#$%'()*+,

58 raise_if_driver_error(err) 75 ctx1cdefghijklmnopqrstuvwxyzABCDEFGHIJKL=?./MNO:PQRSTUVWXYZ0123456789!#$%'()*+,

59 else: 

60 ptr = 0 1 ctx1;

61 return Buffer._init(ptr, size, self) 76 ctx1cdefghijklmnopqrstuvwxyzABCDEFGHIJKL=?./MNO:PQRSTUVWXYZ0123456789!#$%'()*+;,

62 

63 def deallocate(self, ptr: DevicePointerType, size: int, *, stream: Stream | GraphBuilder | None = None) -> None: (empty)

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

65 

66 Parameters 

67 ---------- 

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

69 The pointer or handle to the buffer to deallocate. 

70 size : int 

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

72 stream : Stream, optional 

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

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

75 """ 

76 from cuda.core._stream import Stream_accept 75 ctx1cdefghijklmnopqrstuvwxyzABCDEFGHIJKL./MNO:PQRSTUVWXYZ0123456789!#$%'()*+;@,

77 

78 if stream is not None: 75 ctx1cdefghijklmnopqrstuvwxyzABCDEFGHIJKL./MNO:PQRSTUVWXYZ0123456789!#$%'()*+;@,

79 Stream_accept(stream).sync() 75 ctx1cdefghijklmnopqrstuvwxyzABCDEFGHIJKL./MNO:PQRSTUVWXYZ0123456789!#$%'()*+;@,

80 

81 if size: 75 ctx1cdefghijklmnopqrstuvwxyzABCDEFGHIJKL./MNO:PQRSTUVWXYZ0123456789!#$%'()*+;@,

82 (err,) = driver.cuMemFreeHost(ptr) 74 ctx1cdefghijklmnopqrstuvwxyzABCDEFGHIJKL./MNO:PQRSTUVWXYZ0123456789!#$%'()*+@,

83 raise_if_driver_error(err) 74 ctx1cdefghijklmnopqrstuvwxyzABCDEFGHIJKL./MNO:PQRSTUVWXYZ0123456789!#$%'()*+@,

84 

85 @property (empty)

86 def is_device_accessible(self) -> bool: (empty)

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

88 return True 70 ctx1cdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!#$%'()*+,

89 

90 @property (empty)

91 def is_host_accessible(self) -> bool: (empty)

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

93 return True 70 ctx1cdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!#$%'()*+,

94 

95 @property (empty)

96 def device_id(self) -> int: (empty)

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

98 raise RuntimeError("a pinned memory resource is not bound to any GPU") 1 ctx1[

99 

100 

101class _SynchronousMemoryResource(MemoryResource): (empty)

102 __slots__ = ("_device_id",) (empty)

103 

104 def __init__(self, device_id: int) -> None: (empty)

105 from .._device import Device 2 ctx1b-

106 

107 self._device_id = Device(device_id).device_id 2 ctx1b-

108 

109 def allocate(self, size: int, *, stream: Stream | GraphBuilder | None = None) -> Buffer: (empty)

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

111 # for interface conformance but not used. 

112 from cuda.core._stream import Stream_accept 2 ctx1b-

113 

114 if stream is not None: 2 ctx1b-

115 Stream_accept(stream) 1 ctx1b

116 if size: 2 ctx1b-

117 err, ptr = driver.cuMemAlloc(size) 2 ctx1b-

118 raise_if_driver_error(err) 2 ctx1b-

119 else: 

120 ptr = 0 1 ctx1b

121 return Buffer._init(ptr, size, self) 2 ctx1b-

122 

123 def deallocate(self, ptr: DevicePointerType, size: int, *, stream: Stream | GraphBuilder | None = None) -> None: (empty)

124 from cuda.core._stream import Stream_accept 2 ctx1b-

125 

126 if stream is not None: 2 ctx1b-

127 Stream_accept(stream).sync() 2 ctx1b-

128 if size: 2 ctx1b-

129 (err,) = driver.cuMemFree(ptr) 2 ctx1b-

130 raise_if_driver_error(err) 2 ctx1b-

131 

132 @property (empty)

133 def is_device_accessible(self) -> bool: (empty)

134 return True 1 ctx1b

135 

136 @property (empty)

137 def is_host_accessible(self) -> bool: (empty)

138 return False 1 ctx1b

139 

140 @property (empty)

141 def device_id(self) -> int: (empty)

142 return self._device_id 1 ctx1b