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

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 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 

19  

20import ctypes as ct 

21  

22  

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

24 (<object>data)() 

25  

26  

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

28 return fn == <cydriver.CUhostFn>_py_host_trampoline 

29  

30  

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. 

37  

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!#$%'()*+,-./:;=?@[]^

66  

67 out_fn_owner[0] = make_opaque_py(fn) 1abfecghijklmnopqrstduvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!#$%'()*+,-./:;=?@[]^

68  

69  

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!#$%'()*+,-./:;=?@[]^