Coverage for cuda/core/graph/_host_callback.pyx: 83.87%
31 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-19 01:12 +0000
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-19 01:12 +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 GraphHandle,
13 OpaqueHandle,
14 graph_set_slot,
15 make_opaque_malloc,
16 make_opaque_py,
17)
18from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
20import ctypes as ct
23cdef void _py_host_trampoline(void* data) noexcept with gil:
24 (<object>data)()
27cdef bint _is_py_host_trampoline(cydriver.CUhostFn fn) noexcept nogil:
28 return fn == <cydriver.CUhostFn>_py_host_trampoline
31cdef void _resolve_host_callback(
32 object fn, object user_data,
33 cydriver.CUhostFn* out_fn, void** out_user_data,
34 OpaqueHandle* out_fn_owner, OpaqueHandle* out_data_owner) except *:
35 """Resolve a Python callable or ctypes CFuncPtr into a C callback pair and
36 the owners that keep it alive.
38 On return ``*out_fn`` / ``*out_user_data`` are ready to pass to
39 ``cuGraphAddHostNode`` or ``cuLaunchHostFunc``. ``*out_fn_owner`` owns the
40 callback object; ``*out_data_owner`` owns a copied ``user_data`` buffer and
41 is left null otherwise. The caller attaches the owners to the node's graph
42 slots.
43 """
44 if isinstance(fn, ct._CFuncPtr): 1abfec`ghijklmnopqrstduvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!#$%'()*+,-./:;=?@[]^
45 out_fn[0] = <cydriver.CUhostFn><uintptr_t>ct.cast(fn, ct.c_void_p).value 1abcgjlnprtdu
46 if user_data is None: 1abcgjlnprtdu
47 out_user_data[0] = NULL 1gjlnprtu
48 elif isinstance(user_data, int): 1abcd
49 out_user_data[0] = <void*><uintptr_t>user_data
50 else:
51 buf = bytes(user_data) 1abcd
52 if len(buf): 1_abcd
53 out_user_data[0] = malloc(len(buf)) 1abcd
54 if out_user_data[0] == NULL: 1abcd
55 raise MemoryError("failed to allocate user_data buffer")
56 c_memcpy(out_user_data[0], <const char*>buf, len(buf)) 1abcd
57 out_data_owner[0] = make_opaque_malloc(out_user_data[0]) 1abcd
58 else:
59 out_user_data[0] = NULL
60 else:
61 if user_data is not None: 1fe`hikmoqsvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!#$%'()*+,-./:;=?@[]^
62 raise ValueError( 1e`
63 "user_data is only supported with ctypes function pointers")
64 out_fn[0] = <cydriver.CUhostFn>_py_host_trampoline 1fehikmoqsvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!#$%'()*+,-./:;=?@[]^
65 out_user_data[0] = <void*>fn 1fehikmoqsvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!#$%'()*+,-./:;=?@[]^
67 out_fn_owner[0] = make_opaque_py(fn) 1abfecghijklmnopqrstduvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!#$%'()*+,-./:;=?@[]^
70cdef int _attach_host_callback_owners(
71 const GraphHandle& h_graph, cydriver.CUgraphNode node,
72 OpaqueHandle fn_owner, OpaqueHandle data_owner) except -1:
73 """Attach a resolved host callback's owners to its node's graph slots: the
74 callback in slot 0 and any copied ``user_data`` buffer in slot 1.
75 """
76 HANDLE_RETURN(graph_set_slot(h_graph, node, 0, fn_owner)) 1abfecghijklmnopqrstduvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!#$%'()*+,-./:;=?@[]^
77 if data_owner: 1abfecghijklmnopqrstduvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!#$%'()*+,-./:;=?@[]^
78 HANDLE_RETURN(graph_set_slot(h_graph, node, 1, data_owner)) 1abcd
79 return 0 1abfecghijklmnopqrstduvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!#$%'()*+,-./:;=?@[]^