Coverage for cuda / core / experimental / _memory / _legacy.py: 70%

1# SPDX-FileCopyrightText: Copyright (c) 2024-2025 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 

9from cuda.core.experimental._memory._buffer import Buffer, MemoryResource 

10from cuda.core.experimental._utils.cuda_utils import ( 

11 _check_driver_error as raise_if_driver_error, 

12) 

13from cuda.core.experimental._utils.cuda_utils import ( 

14 driver, 

15) 

16 

17if TYPE_CHECKING: 

18 from cuda.core.experimental._memory.buffer import DevicePointerT 

19 

20__all__ = ["LegacyPinnedMemoryResource", "_SynchronousMemoryResource"] 

21 

22 

23class LegacyPinnedMemoryResource(MemoryResource): 

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

25 APIs. 

26 """ 

27 

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

29 

30 def allocate(self, size, stream=None) -> Buffer: 

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

32 

33 Parameters 

34 ---------- 

35 size : int 

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

37 stream : Stream, optional 

38 Currently ignored 

39 

40 Returns 

41 ------- 

42 Buffer 

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

44 """ 

45 if stream is None: 

46 from cuda.core.experimental._stream import default_stream 

47 

48 stream = default_stream() 

49 err, ptr = driver.cuMemAllocHost(size) 

50 raise_if_driver_error(err) 

51 return Buffer._init(ptr, size, self, stream) 

52 

53 def deallocate(self, ptr: DevicePointerT, size, stream): 

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

55 

56 Parameters 

57 ---------- 

58 ptr : :obj:`~_memory.DevicePointerT` 

59 The pointer or handle to the buffer to deallocate. 

60 size : int 

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

62 stream : Stream 

63 The stream on which to perform the deallocation synchronously. 

64 """ 

65 if stream is not None: 

66 stream.sync() 

67 (err,) = driver.cuMemFreeHost(ptr) 

68 raise_if_driver_error(err) 

69 

70 @property 

71 def is_device_accessible(self) -> bool: 

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

73 return True 

74 

75 @property 

76 def is_host_accessible(self) -> bool: 

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

78 return True 

79 

80 @property 

81 def device_id(self) -> int: 

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

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

84 

85 

86class _SynchronousMemoryResource(MemoryResource): 

87 __slots__ = ("_dev_id",) 

88 

89 def __init__(self, device_id): 

90 from .._device import Device 

91 

92 self._dev_id = Device(device_id).device_id 

93 

94 def allocate(self, size, stream=None) -> Buffer: 

95 if stream is None: 

96 from cuda.core.experimental._stream import default_stream 

97 

98 stream = default_stream() 

99 err, ptr = driver.cuMemAlloc(size) 

100 raise_if_driver_error(err) 

101 return Buffer._init(ptr, size, self, stream) 

102 

103 def deallocate(self, ptr, size, stream): 

104 if stream is not None: 

105 stream.sync() 

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

107 raise_if_driver_error(err) 

108 

109 @property 

110 def is_device_accessible(self) -> bool: 

111 return True 

112 

113 @property 

114 def is_host_accessible(self) -> bool: 

115 return False 

116 

117 @property 

118 def device_id(self) -> int: 

119 return self._dev_id