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
« 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
5from libc.stdint cimport uintptr_t
6from libc.stdlib cimport malloc
7from libc.string cimport memcpy as c_memcpy
9from cuda.bindings cimport cydriver
11from cuda.core._resource_handles cimport (
12 OpaqueHandle,
13 make_opaque_malloc,
14 make_opaque_py,
15)
17import sys
18import ctypes as ct
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)
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 )
43def _validate_ctypes_host_callback(fn):
44 """Reject ctypes callbacks whose declared prototype is not CUhostFn.
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
58cdef void _py_host_trampoline(void* data) noexcept with gil:
59 (<object>data)()
62cdef bint _is_py_host_trampoline(cydriver.CUhostFn fn) noexcept nogil:
63 return fn == <cydriver.CUhostFn>_py_host_trampoline 1DEFGHIJKLbMNOPQ
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.
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.
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
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