Coverage for cuda / core / experimental / _dlpack.pyx: 92%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4  

5from enum import IntEnum 

6  

7  

8cdef void pycapsule_deleter(object capsule) noexcept: 

9 cdef DLManagedTensor* dlm_tensor 

10 cdef DLManagedTensorVersioned* dlm_tensor_ver 

11 # Do not invoke the deleter on a used capsule. 

12 if cpython.PyCapsule_IsValid( 

13 capsule, DLPACK_TENSOR_UNUSED_NAME): 

14 dlm_tensor = <DLManagedTensor*>( 

15 cpython.PyCapsule_GetPointer( 

16 capsule, DLPACK_TENSOR_UNUSED_NAME)) 

17 if dlm_tensor.deleter: 

18 dlm_tensor.deleter(dlm_tensor) 

19 elif cpython.PyCapsule_IsValid( 

20 capsule, DLPACK_VERSIONED_TENSOR_UNUSED_NAME): 

21 dlm_tensor_ver = <DLManagedTensorVersioned*>( 

22 cpython.PyCapsule_GetPointer( 

23 capsule, DLPACK_VERSIONED_TENSOR_UNUSED_NAME)) 

24 if dlm_tensor_ver.deleter: 

25 dlm_tensor_ver.deleter(dlm_tensor_ver) 

26  

27  

28cdef void deleter(DLManagedTensor* tensor) noexcept with gil: 

29 if tensor: 

30 if tensor.dl_tensor.shape: 

31 stdlib.free(tensor.dl_tensor.shape) 

32 if tensor.manager_ctx: 

33 cpython.Py_DECREF(<object>tensor.manager_ctx) 

34 tensor.manager_ctx = NULL 

35 stdlib.free(tensor) 

36  

37  

38cdef void versioned_deleter(DLManagedTensorVersioned* tensor) noexcept with gil: 

39 if tensor: 

40 if tensor.dl_tensor.shape: 

41 stdlib.free(tensor.dl_tensor.shape) 

42 if tensor.manager_ctx: 

43 cpython.Py_DECREF(<object>tensor.manager_ctx) 

44 tensor.manager_ctx = NULL 

45 stdlib.free(tensor) 

46  

47  

48cdef inline DLManagedTensorVersioned* allocate_dlm_tensor_versioned() except? NULL: 

49 cdef DLManagedTensorVersioned* dlm_tensor_ver = NULL 

50 try: 

51 dlm_tensor_ver = <DLManagedTensorVersioned*>( 

52 stdlib.malloc(sizeof(DLManagedTensorVersioned))) 

53 dlm_tensor_ver.dl_tensor.shape = NULL 

54 dlm_tensor_ver.manager_ctx = NULL 

55 return dlm_tensor_ver 

56 except: 

57 if dlm_tensor_ver: 

58 stdlib.free(dlm_tensor_ver) 

59 raise 

60  

61  

62cdef inline DLManagedTensor* allocate_dlm_tensor() except? NULL: 

63 cdef DLManagedTensor* dlm_tensor = NULL 

64 try: 

65 dlm_tensor = <DLManagedTensor*>( 

66 stdlib.malloc(sizeof(DLManagedTensor))) 

67 dlm_tensor.dl_tensor.shape = NULL 

68 dlm_tensor.manager_ctx = NULL 

69 return dlm_tensor 

70 except: 

71 if dlm_tensor: 

72 stdlib.free(dlm_tensor) 

73 raise 

74  

75  

76cdef inline int setup_dl_tensor_layout(DLTensor* dl_tensor, object buf) except -1: 

77 dl_tensor.ndim = 1 

78 cdef int64_t* shape_strides = \ 

79 <int64_t*>stdlib.malloc(sizeof(int64_t) * 2) 

80 shape_strides[0] = <int64_t>buf.size 

81 shape_strides[1] = 1 # redundant 

82 dl_tensor.shape = shape_strides 

83 dl_tensor.strides = NULL 

84 dl_tensor.byte_offset = 0 

85 return 0 

86  

87  

88cdef inline int setup_dl_tensor_device(DLTensor* dl_tensor, object buf) except -1: 

89 cdef DLDevice* device = &dl_tensor.device 

90 # buf should be a Buffer instance 

91 if buf.is_device_accessible and not buf.is_host_accessible: 

92 device.device_type = _kDLCUDA 

93 device.device_id = buf.device_id 

94 elif buf.is_device_accessible and buf.is_host_accessible: 

95 device.device_type = _kDLCUDAHost 

96 device.device_id = 0 

97 elif not buf.is_device_accessible and buf.is_host_accessible: 

98 device.device_type = _kDLCPU 

99 device.device_id = 0 

100 else: # not buf.is_device_accessible and not buf.is_host_accessible 

101 raise BufferError("invalid buffer") 

102 return 0 

103  

104  

105cdef inline int setup_dl_tensor_dtype(DLTensor* dl_tensor) except -1 nogil: 

106 cdef DLDataType* dtype = &dl_tensor.dtype 

107 dtype.code = <uint8_t>kDLInt 

108 dtype.lanes = <uint16_t>1 

109 dtype.bits = <uint8_t>8 

110 return 0 

111  

112  

113cpdef object make_py_capsule(object buf, bint versioned): 

114 cdef DLManagedTensor* dlm_tensor = NULL 

115 cdef DLManagedTensorVersioned* dlm_tensor_ver = NULL 

116 cdef DLTensor* dl_tensor 

117 cdef void* tensor_ptr 

118 cdef const char* capsule_name 

119 cdef object ret = None 

120  

121 try: 

122 if versioned: 

123 dlm_tensor_ver = allocate_dlm_tensor_versioned() 

124 # Transfer the reference to manager_ctx 

125 cpython.Py_INCREF(buf) 

126 dlm_tensor_ver.manager_ctx = <void*>buf 

127 dlm_tensor_ver.deleter = versioned_deleter 

128 dlm_tensor_ver.version.major = DLPACK_MAJOR_VERSION 

129 dlm_tensor_ver.version.minor = DLPACK_MINOR_VERSION 

130 dlm_tensor_ver.flags = 0 

131 dl_tensor = &dlm_tensor_ver.dl_tensor 

132 tensor_ptr = dlm_tensor_ver 

133 capsule_name = DLPACK_VERSIONED_TENSOR_UNUSED_NAME 

134 else: 

135 dlm_tensor = allocate_dlm_tensor() 

136 # Transfer the reference to manager_ctx 

137 cpython.Py_INCREF(buf) 

138 dlm_tensor.manager_ctx = <void*>buf 

139 dlm_tensor.deleter = deleter 

140 dl_tensor = &dlm_tensor.dl_tensor 

141 tensor_ptr = dlm_tensor 

142 capsule_name = DLPACK_TENSOR_UNUSED_NAME 

143  

144 dl_tensor.data = <void*><intptr_t>(int(buf.handle)) 

145 setup_dl_tensor_layout(dl_tensor, buf) 

146 setup_dl_tensor_device(dl_tensor, buf) 

147 setup_dl_tensor_dtype(dl_tensor) 

148 ret = cpython.PyCapsule_New(tensor_ptr, capsule_name, pycapsule_deleter) 

149 except: 

150 if ret is None: 

151 deleter(dlm_tensor) 

152 versioned_deleter(dlm_tensor_ver) 

153 raise 

154 return ret 

155  

156  

157class DLDeviceType(IntEnum): 

158 kDLCPU = _kDLCPU 

159 kDLCUDA = _kDLCUDA 

160 kDLCUDAHost = _kDLCUDAHost 

161 kDLCUDAManaged = _kDLCUDAManaged