Coverage for cuda/core/graph/_host_callback.pyx: 90.48%

42 statements  

« prev     ^ index     » next       coverage.py v7.16.0, created at 2026-09-03 02:41 +0000

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4  

5from libc.stdint cimport uintptr_t 

6from libc.stdlib cimport malloc 

7from libc.string cimport memcpy as c_memcpy 

8  

9from cuda.bindings cimport cydriver 

10  

11from cuda.core._resource_handles cimport ( 

12 OpaqueHandle, 

13 make_opaque_malloc, 

14 make_opaque_py, 

15) 

16  

17import sys 

18import ctypes as ct 

19  

20  

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

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

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

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

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

26_CUHOSTFN_HINT = ( 

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

28 if sys.platform != "win32" 

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

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

31) 

32  

33  

34def _cuhostfn_type_error(detail): 

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

36 return TypeError( 1klmnoj

37 f"host callback {detail}; CUDA requires a callback matching CUhostFn " 1klmnoj

38 f"(void (*)(void*)), declared as {_CUHOSTFN_HINT}. " 1klmnoj

39 "Alternatively, pass a Python callable." 

40 ) 

41  

42  

43def _validate_ctypes_host_callback(fn): 

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

45  

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

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

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

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

50 """ 

51 restype = fn.restype 1kcdeqrstlmnojuvwxyzAfBbghia

52 argtypes = fn.argtypes 1kcdeqrstlmnojuvwxyzAfBbghia

53 if restype is not None or argtypes is None or tuple(argtypes) != (ct.c_void_p,): 1kcdeqrstlmnojuvwxyzAfBbghia

54 raise _cuhostfn_type_error( 1klmnoj

55 f"has prototype restype={restype!r}, argtypes={argtypes!r}") 1klmnoj

56  

57  

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

59 (<object>data)() 

60  

61  

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

63 return fn == <cydriver.CUhostFn>_py_host_trampoline 1DEFGHIJKLbMNOPQ

64  

65  

66cdef void _resolve_host_callback( 

67 object fn, object user_data, 

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

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

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

71 the owners that keep it alive. 

72  

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

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

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

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

77  

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

79 address is passed to CUDA. 

80 """ 

81 if isinstance(fn, ct._CFuncPtr): 2T k c d U R e q r s t l m n o j D $bu V W v X w Y x Z y 0 z 1 A 2 E 3 4 f B F G 5 6 7 8 9 ! # $ % ' H ( I ) J K * + , - L . / b M : ; = g C ? N O h @ [ P ] ^ i _ ` Q { | a S } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b

82 _validate_ctypes_host_callback(fn) 1kcdeqrstlmnojuvwxyzAfBbghia

83 out_fn[0] = <cydriver.CUhostFn><uintptr_t>ct.cast(fn, ct.c_void_p).value 1cdeqrstjuvwxyzAfBbghia

84 if user_data is None: 1cdeqrstjuvwxyzAfBbghia

85 out_user_data[0] = NULL 1qrstjuvwxyzAB

86 elif isinstance(user_data, int): 1cdefbghia

87 out_user_data[0] = <void*><uintptr_t>user_data 

88 else: 

89 buf = bytes(user_data) 1cdefbghia

90 if len(buf): 1cdefbghia

91 out_user_data[0] = malloc(len(buf)) 1cdefbghia

92 if out_user_data[0] == NULL: 1cdefbghia

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

94 c_memcpy(out_user_data[0], <const char*>buf, len(buf)) 1cdefbghia

95 out_data_owner[0] = make_opaque_malloc(out_user_data[0]) 1cdefbghia

96 else: 

97 out_user_data[0] = NULL 

98 else: 

99 if not callable(fn): 2T U R D $bV W X Y Z 0 1 2 E 3 4 F G 5 6 7 8 9 ! # $ % ' H ( I ) J K * + , - L . / M : ; = C ? N O @ [ P ] ^ _ ` Q { | a S } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b

100 raise TypeError( 1C

101 f"callback must be callable, got {type(fn).__name__}") 1C

102 if user_data is not None: 2T U R D $bV W X Y Z 0 1 2 E 3 4 F G 5 6 7 8 9 ! # $ % ' H ( I ) J K * + , - L . / M : ; = C ? N O @ [ P ] ^ _ ` Q { | a S } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b

103 raise ValueError( 2R $ba S

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

105 out_fn[0] = <cydriver.CUhostFn>_py_host_trampoline 2T U R D V W X Y Z 0 1 2 E 3 4 F G 5 6 7 8 9 ! # $ % ' H ( I ) J K * + , - L . / M : ; = C ? N O @ [ P ] ^ _ ` Q { | S } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b

106 out_user_data[0] = <void*>fn 2T U R D V W X Y Z 0 1 2 E 3 4 F G 5 6 7 8 9 ! # $ % ' H ( I ) J K * + , - L . / M : ; = C ? N O @ [ P ] ^ _ ` Q { | S } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b

107  

108 out_fn_owner[0] = make_opaque_py(fn) 2T c d U R e q r s t j D u V W v X w Y x Z y 0 z 1 A 2 E 3 4 f B F G 5 6 7 8 9 ! # $ % ' H ( I ) J K * + , - L . / b M : ; = g C ? N O h @ [ P ] ^ i _ ` Q { | a S } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b