Coverage for cuda/core/_host_callback.pyx: 72.04%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4  

5from cpython.buffer cimport ( 

6 PyBUF_SIMPLE, 

7 PyBUF_WRITABLE, 

8 PyBuffer_Release, 

9 PyObject_CheckBuffer, 

10 PyObject_GetBuffer, 

11 Py_buffer, 

12) 

13from libc.stdint cimport uintptr_t 

14from libc.stdlib cimport malloc 

15from libc.string cimport memcpy as c_memcpy 

16  

17from cuda.bindings cimport cydriver 

18  

19from cuda.core._memory._buffer cimport Buffer, Buffer_check_open 

20from cuda.core._rt cimport ( 

21 OpaqueHandle, 

22 make_opaque_malloc, 

23 make_opaque_py, 

24) 

25  

26import sys 

27import ctypes as ct 

28from _ctypes import CFuncPtr 

29  

30import numpy 

31  

32from cuda.core._utils.pycompat import BufferProtocol 

33  

34__all__ = ['HOST_COPY_FN', 'pack_host_copy_args'] 

35  

36  

37# CUhostFn is `void (CUDA_CB *)(void*)`. CUDA_CB is __stdcall on Windows and 

38# empty elsewhere, but ctypes only honors that distinction when it builds a 

39# callback on 32-bit x86 Windows, which cuda.core does not support: on win-64 

40# and ARM64 both CFUNCTYPE and WINFUNCTYPE produce a FFI_DEFAULT_ABI thunk. The 

41# declared result and argument types are all that remain worth checking. 

42_CUHOSTFN_HINT = ( 

43 "ctypes.CFUNCTYPE(None, ctypes.c_void_p)" 

44 if sys.platform != "win32" 

45 else "ctypes.CFUNCTYPE(None, ctypes.c_void_p) or " 

46 "ctypes.WINFUNCTYPE(None, ctypes.c_void_p)" 

47) 

48  

49  

50def _cuhostfn_type_error(detail): 

51 """Build the rejection message for a non-conforming ctypes callback.""" 

52 return TypeError( 8 ctx1TUVWXrYZ

53 f"host callback {detail}; CUDA requires a callback matching CUhostFn " 8 ctx1TUVWXrYZ

54 f"(void (*)(void*)), declared as {_CUHOSTFN_HINT}. " 8 ctx1TUVWXrYZ

55 "Alternatively, pass a Python callable." 

56 ) 

57  

58  

59def _validate_ctypes_host_callback(fn): 

60 """Reject ctypes callbacks whose declared prototype is not CUhostFn. 

61  

62 ``restype`` and ``argtypes`` are the prototype the caller declared, and are 

63 what CUDA calls through. A function pointer taken from a shared library 

64 keeps ctypes' defaults -- a ``c_int`` result and unspecified arguments -- 

65 until the caller declares otherwise, so it must be declared to be accepted. 

66 """ 

67 restype = fn.restype 41 ctx1TijkEFGHUVWXrIJKLMNOlPQhmnoexyabczRuvdAYZ

68 argtypes = fn.argtypes 41 ctx1TijkEFGHUVWXrIJKLMNOlPQhmnoexyabczRuvdAYZ

69 if restype is not None or argtypes is None or tuple(argtypes) != (ct.c_void_p,): 41 ctx1TijkEFGHUVWXrIJKLMNOlPQhmnoexyabczRuvdAYZ

70 raise _cuhostfn_type_error( 8 ctx1TUVWXrYZ

71 f"has prototype restype={restype!r}, argtypes={argtypes!r}") 8 ctx1TUVWXrYZ

72  

73  

74cdef void _py_host_trampoline(void* data) noexcept with gil: 

75 (<object>data)() 

76  

77  

78cdef bint _is_py_host_trampoline(cydriver.CUhostFn fn) noexcept nogil: 

79 return fn == <cydriver.CUhostFn>_py_host_trampoline 13 ctx123456789!h#$%

80  

81  

82ctypedef struct _HostCopyArgs: 

83 uintptr_t dst 

84 uintptr_t src 

85 size_t nbytes 

86  

87  

88cdef void _host_copy(void* data) noexcept nogil: 

89 cdef _HostCopyArgs* args = <_HostCopyArgs*>data 

90 if args.nbytes: 

91 c_memcpy(<void*>args.dst, <const void*>args.src, args.nbytes) 

92  

93  

94cdef object _host_copy_fn_pointer(): 

95 return ct.CFUNCTYPE(None, ct.c_void_p)(<uintptr_t><cydriver.CUhostFn>_host_copy) 

96  

97  

98HOST_COPY_FN: CFuncPtr = _host_copy_fn_pointer() 

99  

100  

101cdef uintptr_t _host_copy_address(object obj, size_t nbytes, bint writable, str name) except? 0: 

102 cdef Buffer buf 

103 cdef Py_buffer view 

104 # These export the buffer protocol over their own storage, so a value 

105 # meant as an address would otherwise pack the address of that storage. 

106 if isinstance(obj, (bool, numpy.integer, ct._SimpleCData)): 18 ctx2a b c f p S w t s RbSbq ' B C D 0 g

107 raise TypeError(f"{name} must not be a {type(obj).__name__}; pass int(addr) for an address") 3 ctx1BCD

108 if isinstance(obj, int): 18 ctx2a b c f p S w t s RbSbq ' B C D 0 g

109 return <uintptr_t>obj 5 ctx2c S s RbSb

110 if isinstance(obj, Buffer): 14 ctx1abfpwtsq'BCD0g

111 buf = <Buffer>obj 6 ctx1afpwtg

112 Buffer_check_open(buf) 6 ctx1afpwtg

113 # A Buffer without a memory resource could only be checked with a 

114 # driver query, which is not allowed on a host callback thread. 

115 if buf._memory_resource is not None and not buf._memory_resource.is_host_accessible: 5 ctx1afptg

116 raise ValueError(f"{name} is a Buffer that is not host-accessible") 1 ctx1t

117 if nbytes > buf._size: 4 ctx1afpg

118 raise ValueError(f"nbytes ({nbytes}) exceeds the size of {name} ({buf._size})") 1 ctx1p

119 return <uintptr_t>buf.handle 3 ctx1afg

120 if not PyObject_CheckBuffer(obj): 14 ctx1abfpwtsq'BCD0g

121 raise TypeError( 1 ctx10

122 f"{name} must be an int, a Buffer, or a buffer-protocol object, got {type(obj).__name__}") 1 ctx10

123 PyObject_GetBuffer(obj, &view, PyBUF_SIMPLE | PyBUF_WRITABLE if writable else PyBUF_SIMPLE) 13 ctx1abfpwtsq'BCDg

124 try: 12 ctx1abfpwtsqBCDg

125 if nbytes > <size_t>view.len: 12 ctx1abfpwtsqBCDg

126 raise ValueError(f"nbytes ({nbytes}) exceeds the size of {name} ({view.len})") 1 ctx1p

127 return <uintptr_t>view.buf 12 ctx1abfpwtsqBCDg

128 finally: 

129 PyBuffer_Release(&view) 

130  

131  

132def pack_host_copy_args( 

133 dst: int | Buffer | BufferProtocol, src: int | Buffer | BufferProtocol, nbytes: int) -> bytes: 

134 """Pack the arguments of :data:`HOST_COPY_FN` into ``user_data`` bytes. 

135  

136 The result is a record of three native-width fields in this order: 

137 ``uintptr_t dst``, ``uintptr_t src``, ``size_t nbytes``, with no padding. 

138 ``struct.pack("@PPN", dst, src, nbytes)`` builds the same record, and 

139 ``struct.unpack("@PPN", data)`` reads one back. A caller that keeps such a 

140 record in memory of its own can pass its address as an ``int`` 

141 ``user_data`` instead. 

142  

143 ``dst`` and ``src`` are not retained. The caller must keep both alive, 

144 and must not resize them (a ``bytearray``, for example), until the copy 

145 has run. 

146  

147 Parameters 

148 ---------- 

149 dst : int | :class:`Buffer` | :obj:`collections.abc.Buffer` 

150 The destination. Only an ``int`` is a host address. A buffer-protocol 

151 object must be contiguous and writable. 

152 src : int | :class:`Buffer` | :obj:`collections.abc.Buffer` 

153 The source. Only an ``int`` is a host address. A buffer-protocol 

154 object must be contiguous. 

155 nbytes : int 

156 The number of bytes to copy. Zero packs a copy that does nothing. 

157  

158 Returns 

159 ------- 

160 bytes 

161 The packed record, to pass as ``user_data``. 

162  

163 Raises 

164 ------ 

165 TypeError 

166 If ``dst`` or ``src`` is not an ``int``, a :class:`Buffer`, or a 

167 buffer-protocol object, or is a ``bool``, a NumPy integer scalar, or a 

168 ctypes scalar such as ``ctypes.c_void_p``. A buffer-protocol object 

169 that cannot export a contiguous view, or a writable one for ``dst``, 

170 raises its own error. 

171 ValueError 

172 If ``nbytes`` exceeds the size of ``dst`` or ``src``, if a 

173 :class:`Buffer` comes from a memory resource that is not 

174 host-accessible, if an address is null, or if the ranges overlap. 

175 OverflowError 

176 If ``nbytes`` or an ``int`` address is negative or does not fit in a 

177 native-width field. 

178 RuntimeError 

179 If ``dst`` or ``src`` is a closed :class:`Buffer`. 

180  

181 Notes 

182 ----- 

183 This function makes no CUDA calls, so a Python host callback can call it, 

184 for example to fill a record that a later :data:`HOST_COPY_FN` reads 

185 through an ``int`` ``user_data``. It checks host accessibility only for a 

186 :class:`Buffer` with a memory resource. For an ``int`` or a :class:`Buffer` 

187 without one, the caller must ensure that the memory is host-accessible. 

188 """ 

189 cdef size_t n = nbytes 20 ctx2a b c f p S w t s WbXbRbSbq ' B C D 0 g

190 cdef _HostCopyArgs args 

191 args.dst = _host_copy_address(dst, n, True, "dst") 18 ctx2a b c f p S w t s RbSbq ' B C D 0 g

192 args.src = _host_copy_address(src, n, False, "src") 14 ctx1abcfpSwtsqBCDg

193 args.nbytes = n 8 ctx1abcfSsqg

194 if n: 8 ctx1abcfSsqg

195 if args.dst == 0 or args.src == 0: 7 ctx1abfSsqg

196 raise ValueError("dst and src must not be null") 1 ctx1s

197 if (args.dst - args.src if args.dst > args.src else args.src - args.dst) < n: 6 ctx1abfSqg

198 raise ValueError("dst and src overlap") 1 ctx1q

199 return (<char*>&args)[:sizeof(args)] 7 ctx1abcfSqg

200  

201  

202cdef void _resolve_host_callback( 

203 object fn, object user_data, 

204 cydriver.CUhostFn* out_fn, void** out_user_data, 

205 OpaqueHandle* out_fn_owner, OpaqueHandle* out_data_owner) except *: 

206 """Resolve a Python callable or ctypes CFuncPtr into a C callback pair and 

207 the owners that keep it alive. 

208  

209 On return ``*out_fn`` / ``*out_user_data`` are ready to pass to 

210 ``cuGraphAddHostNode`` or ``cuLaunchHostFunc``. ``*out_fn_owner`` owns the 

211 callback object; ``*out_data_owner`` owns a copied ``user_data`` buffer and 

212 is left null otherwise. The caller attaches both owners to the graph node. 

213  

214 ctypes callbacks are validated against the ``CUhostFn`` ABI before their 

215 address is passed to CUDA. 

216 """ 

217 cdef bint is_host_copy 

218 if isinstance(fn, ct._CFuncPtr): 122 ctx2* T i j + ( k E F G H U V W X r 2 TbI , - J . K / L : M ; N = O ? @ [ 3 ] ^ l P 4 5 _ ` { | } ~ abbbcbdbeb6 fb7 gb8 9 hbibjbkbQ ! h # lbm 1 mbn nbob$ o pbqb% e ) rbsbtbubvbwbxbx yby a b c z R u v zbd A AbBbCbUbVbY Z DbEbFbGbHbIbJbKbLbMbNbObPbQb

219 _validate_ctypes_host_callback(fn) 41 ctx1TijkEFGHUVWXrIJKLMNOlPQhmnoexyabczRuvdAYZ

220 out_fn[0] = <cydriver.CUhostFn><uintptr_t>ct.cast(fn, ct.c_void_p).value 34 ctx1ijkEFGHrIJKLMNOlPQhmnoexyabczRuvdA

221 is_host_copy = out_fn[0] == <cydriver.CUhostFn>_host_copy 34 ctx1ijkEFGHrIJKLMNOlPQhmnoexyabczRuvdA

222 if user_data is None: 34 ctx1ijkEFGHrIJKLMNOlPQhmnoexyabczRuvdA

223 out_user_data[0] = NULL 15 ctx1EFGHrIJKLMNOPQR

224 elif isinstance(user_data, int): 19 ctx1ijklhmnoexyabczuvdA

225 out_user_data[0] = <void*><uintptr_t>user_data 4 ctx1xyzA

226 else: 

227 buf = bytes(user_data) 15 ctx1ijklhmnoeabcuvd

228 if is_host_copy and len(buf) != sizeof(_HostCopyArgs): 15 ctx1ijklhmnoeabcuvd

229 raise ValueError( 2 ctx1uv

230 f"HOST_COPY_FN user_data must be {sizeof(_HostCopyArgs)} bytes " 2 ctx1uv

231 f"from pack_host_copy_args, got {len(buf)}") 2 ctx1uv

232 if len(buf): 13 ctx1ijklhmnoeabcd

233 out_user_data[0] = malloc(len(buf)) 13 ctx1ijklhmnoeabcd

234 if out_user_data[0] == NULL: 13 ctx1ijklhmnoeabcd

235 raise MemoryError("failed to allocate user_data buffer") 

236 c_memcpy(out_user_data[0], <const char*>buf, len(buf)) 13 ctx1ijklhmnoeabcd

237 out_data_owner[0] = make_opaque_malloc(out_user_data[0]) 13 ctx1ijklhmnoeabcd

238 else: 

239 out_user_data[0] = NULL 

240 if is_host_copy and out_user_data[0] == NULL: 32 ctx1ijkEFGHrIJKLMNOlPQhmnoexyabczRdA

241 raise ValueError("HOST_COPY_FN requires user_data from pack_host_copy_args or a nonzero address") 2 ctx1zR

242 else: 

243 if not callable(fn): 83 ctx2* + ( 2 Tb, - . / : ; = ? @ [ 3 ] ^ 4 5 _ ` { | } ~ abbbcbdbeb6 fb7 gb8 9 hbibjbkb! # lb1 mbnbob$ pbqb% e ) rbsbtbubvbwbxbybzbd AbBbCbUbVbDbEbFbGbHbIbJbKbLbMbNbObPbQb

244 raise TypeError( 2 ctx21 Ub

245 f"callback must be callable, got {type(fn).__name__}") 2 ctx21 Ub

246 if user_data is not None: 82 ctx2* + ( 2 Tb, - . / : ; = ? @ [ 3 ] ^ 4 5 _ ` { | } ~ abbbcbdbeb6 fb7 gb8 9 hbibjbkb! # lb1 mbnbob$ pbqb% e ) rbsbtbubvbwbxbybzbd AbBbCbVbDbEbFbGbHbIbJbKbLbMbNbObPbQb

247 raise ValueError( 5 ctx2( Tbe ) Vb

248 "user_data is only supported with ctypes function pointers") 

249 out_fn[0] = <cydriver.CUhostFn>_py_host_trampoline 79 ctx2* + ( 2 , - . / : ; = ? @ [ 3 ] ^ 4 5 _ ` { | } ~ abbbcbdbeb6 fb7 gb8 9 hbibjbkb! # lb1 mbnbob$ pbqb% ) rbsbtbubvbwbxbybzbd AbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQb

250 out_user_data[0] = <void*>fn 79 ctx2* + ( 2 , - . / : ; = ? @ [ 3 ] ^ 4 5 _ ` { | } ~ abbbcbdbeb6 fb7 gb8 9 hbibjbkb! # lb1 mbnbob$ pbqb% ) rbsbtbubvbwbxbybzbd AbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQb

251  

252 out_fn_owner[0] = make_opaque_py(fn) 108 ctx2* i j + ( k E F G H r 2 I , - J . K / L : M ; N = O ? @ [ 3 ] ^ l P 4 5 _ ` { | } ~ abbbcbdbeb6 fb7 gb8 9 hbibjbkbQ ! h # lbm 1 mbn nbob$ o pbqb% e ) rbsbtbubvbwbxbx yby a b c zbd A AbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQb