Coverage for cuda / bindings / _internal / utils.pyx: 67%

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

2# 

3# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE 

4  

5cimport cpython 

6from libc.stdint cimport intptr_t 

7from libcpp.utility cimport move 

8from cython.operator cimport dereference as deref 

9  

10  

11cdef bint is_nested_sequence(data): 

12 if not cpython.PySequence_Check(data): 

13 return False 

14 else: 

15 for i in data: 

16 if not cpython.PySequence_Check(i): 

17 return False 

18 else: 

19 return True 

20  

21  

22cdef void* get_buffer_pointer(buf, Py_ssize_t size, readonly=True) except*: 

23 """The caller must ensure ``buf`` is alive when the returned pointer is in use.""" 

24 cdef void* bufPtr 

25 cdef int flags = cpython.PyBUF_ANY_CONTIGUOUS 

26 if not readonly: 

27 flags |= cpython.PyBUF_WRITABLE 

28 cdef int status = -1 

29 cdef cpython.Py_buffer view 

30  

31 if isinstance(buf, int): 

32 bufPtr = <void*><intptr_t>buf 

33 else: # try buffer protocol 

34 try: 

35 status = cpython.PyObject_GetBuffer(buf, &view, flags) 

36 # when the caller does not provide a size, it is set to -1 at generate-time by cybind 

37 if size != -1: 

38 assert view.len == size 

39 assert view.ndim == 1 

40 except Exception as e: 

41 adj = "writable " if not readonly else "" 

42 raise ValueError( 

43 "buf must be either a Python int representing the pointer " 

44 f"address to a valid buffer, or a 1D contiguous {adj}" 

45 "buffer, of size bytes") from e 

46 else: 

47 bufPtr = view.buf 

48 finally: 

49 if status == 0: 

50 cpython.PyBuffer_Release(&view) 

51  

52 return bufPtr 

53  

54  

55# Cython can't infer the ResT overload when it is wrapped in nullable_unique_ptr, 

56# so we need a dummy (__unused) input argument to help it 

57cdef int get_resource_ptr(nullable_unique_ptr[vector[ResT]] &in_out_ptr, object obj, ResT* __unused) except 1: 

58 if cpython.PySequence_Check(obj): 

59 vec = new vector[ResT](len(obj)) 

60 # set the ownership immediately to avoid leaking the `vec` memory in 

61 # case of exception in the following loop 

62 in_out_ptr.reset(vec, True) 

63 for i in range(len(obj)): 

64 deref(vec)[i] = obj[i] 

65 else: 

66 in_out_ptr.reset(<vector[ResT]*><intptr_t>obj, False) 

67 return 0 

68  

69  

70cdef int get_resource_ptrs(nullable_unique_ptr[ vector[PtrT*] ] &in_out_ptr, object obj, PtrT* __unused) except 1: 

71 if cpython.PySequence_Check(obj): 

72 vec = new vector[PtrT*](len(obj)) 

73 # set the ownership immediately to avoid leaking the `vec` memory in 

74 # case of exception in the following loop 

75 in_out_ptr.reset(vec, True) 

76 for i in range(len(obj)): 

77 deref(vec)[i] = <PtrT*><intptr_t>(obj[i]) 

78 else: 

79 in_out_ptr.reset(<vector[PtrT*]*><intptr_t>obj, False) 

80 return 0 

81  

82  

83cdef int get_nested_resource_ptr(nested_resource[ResT] &in_out_ptr, object obj, ResT* __unused) except 1: 

84 cdef nullable_unique_ptr[ vector[intptr_t] ] nested_ptr 

85 cdef nullable_unique_ptr[ vector[vector[ResT]] ] nested_res_ptr 

86 cdef vector[intptr_t]* nested_vec = NULL 

87 cdef vector[vector[ResT]]* nested_res_vec = NULL 

88 cdef size_t i = 0, length = 0 

89 cdef intptr_t addr 

90  

91 if is_nested_sequence(obj): 

92 length = len(obj) 

93 nested_res_vec = new vector[vector[ResT]](length) 

94 nested_vec = new vector[intptr_t](length) 

95 # set the ownership immediately to avoid leaking memory in case of 

96 # exception in the following loop 

97 nested_res_ptr.reset(nested_res_vec, True) 

98 nested_ptr.reset(nested_vec, True) 

99 for i, obj_i in enumerate(obj): 

100 if ResT is char: 

101 obj_i_bytes = (<str?>(obj_i)).encode() 

102 str_len = <size_t>(len(obj_i_bytes)) + 1 # including null termination 

103 deref(nested_res_vec)[i].resize(str_len) 

104 obj_i_ptr = <char*>(obj_i_bytes) 

105 # cast to size_t explicitly to work around a potentially Cython bug 

106 deref(nested_res_vec)[i].assign(obj_i_ptr, obj_i_ptr + <size_t>str_len) 

107 else: 

108 deref(nested_res_vec)[i] = obj_i 

109 deref(nested_vec)[i] = <intptr_t>(deref(nested_res_vec)[i].data()) 

110 elif cpython.PySequence_Check(obj): 

111 length = len(obj) 

112 nested_vec = new vector[intptr_t](length) 

113 nested_ptr.reset(nested_vec, True) 

114 for i, addr in enumerate(obj): 

115 deref(nested_vec)[i] = addr 

116 nested_res_ptr.reset(NULL, False) 

117 else: 

118 # obj is an int (ResT**) 

119 nested_res_ptr.reset(NULL, False) 

120 nested_ptr.reset(<vector[intptr_t]*><intptr_t>obj, False) 

121  

122 in_out_ptr.ptrs = move(nested_ptr) 

123 in_out_ptr.nested_resource_ptr = move(nested_res_ptr) 

124 return 0 

125  

126  

127class FunctionNotFoundError(RuntimeError): pass 

128  

129class NotSupportedError(RuntimeError): pass