Coverage for cuda/core/_launcher.pyx: 70.45%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4  

5from libc.stdint cimport uintptr_t 

6  

7from cuda.bindings cimport cydriver 

8  

9from cuda.core._host_callback cimport _resolve_host_callback 

10from cuda.core._launch_config cimport LaunchConfig 

11from cuda.core._kernel_arg_handler cimport ParamHolder 

12from cuda.core._module cimport Kernel 

13from cuda.core._rt cimport as_cu, launch_host_function, OpaqueHandle 

14from cuda.core._stream cimport Stream_accept, Stream 

15from cuda.core._utils.cuda_utils cimport ( 

16 check_or_create_options, 

17 HANDLE_RETURN, 

18) 

19from cuda.core._module import Kernel 

20from cuda.core._stream import Stream 

21from math import prod 

22from typing import TYPE_CHECKING 

23  

24if TYPE_CHECKING: 

25 from collections.abc import Callable 

26  

27 from cuda.core._utils.pycompat import BufferProtocol 

28 from cuda.core.graph import GraphBuilder 

29 from cuda.core.typing import IsStreamType 

30  

31__all__ = ['launch', 'host_launch'] 

32  

33  

34def launch( 

35 stream: Stream | GraphBuilder | IsStreamType, 

36 LaunchConfig config not None: LaunchConfig, 

37 Kernel kernel not None: Kernel, 

38 *kernel_args 

39) -> None: 

40 """Launches a :obj:`~_module.Kernel` 

41 object with launch-time configuration. 

42  

43 Parameters 

44 ---------- 

45 stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder` 

46 The stream establishing the stream ordering semantic of a 

47 launch. 

48 config : :obj:`LaunchConfig` 

49 Launch configurations inline with options provided by 

50 :obj:`~_launcher.LaunchConfig` dataclass. 

51 kernel : :obj:`~_module.Kernel` 

52 Kernel to launch. 

53 *kernel_args : Any 

54 Variable length argument list that is provided to the 

55 launching kernel. 

56  

57 """ 

58 cdef Stream s = Stream_accept(stream, allow_stream_protocol=True) 154 ctx2l m n o p q r s t u v w x y z A B C D E F G H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ aba bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b

59 cdef LaunchConfig conf = check_or_create_options(LaunchConfig, config, "launch config") 154 ctx2l m n o p q r s t u v w x y z A B C D E F G H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ aba bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b

60  

61 # TODO: can we ensure kernel_args is valid/safe to use here? 

62 # TODO: merge with HelperKernelParams? 

63 cdef ParamHolder ker_args = ParamHolder(kernel_args) 154 ctx2l m n o p q r s t u v w x y z A B C D E F G H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ aba bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b

64 cdef void** args_ptr = <void**><uintptr_t>(ker_args.ptr) 154 ctx2l m n o p q r s t u v w x y z A B C D E F G H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ aba bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b

65  

66 cdef cydriver.CUfunction func_handle = <cydriver.CUfunction>as_cu(kernel._h_kernel) 154 ctx2l m n o p q r s t u v w x y z A B C D E F G H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ aba bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b

67  

68 drv_cfg = conf._to_native_launch_config() 154 ctx2l m n o p q r s t u v w x y z A B C D E F G H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ aba bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b

69 drv_cfg.hStream = as_cu(s._h_stream) 154 ctx2l m n o p q r s t u v w x y z A B C D E F G H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ aba bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b

70 if conf.is_cooperative: 154 ctx2l m n o p q r s t u v w x y z A B C D E F G H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ aba bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b

71 _check_cooperative_launch(kernel, conf, s) 1 ctx1a

72 with nogil: 154 ctx2l m n o p q r s t u v w x y z A B C D E F G H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ aba bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b

73 HANDLE_RETURN(cydriver.cuLaunchKernelEx(&drv_cfg, func_handle, args_ptr, NULL)) 154 ctx2l m n o p q r s t u v w x y z A B C D E F G H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ aba bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b

74  

75  

76def host_launch(stream: Stream | GraphBuilder | IsStreamType, fn: Callable[..., object], *, user_data: int | BufferProtocol | None = None) -> None: 

77 """Enqueue a host function on a stream in stream order. 

78  

79 The function runs on a driver-owned thread once all work enqueued on the 

80 stream before this call has completed. Work enqueued after this call waits 

81 until the function returns. This call returns when the function is enqueued; 

82 observe completion with :meth:`Stream.sync` or an :class:`Event` recorded 

83 after it. 

84  

85 Passing a :class:`~graph.GraphBuilder` forwards to 

86 :meth:`~graph.GraphBuilder.callback`. 

87  

88 Parameters 

89 ---------- 

90 stream : Stream or graph.GraphBuilder or IsStreamType 

91 The stream establishing execution order, or a graph builder capturing 

92 the callback. 

93 fn : callable or ctypes function pointer 

94 A Python callable taking no arguments, which runs with the GIL, or a 

95 ctypes function pointer declared as ``CUhostFn``, for example 

96 ``ctypes.CFUNCTYPE(None, ctypes.c_void_p)``, which ``cuda.core`` 

97 passes to the driver as is. A pointer to compiled code runs without 

98 the GIL; a ctypes callback that wraps a Python function acquires the 

99 GIL through ctypes and is subject to the same deadlock rule as a 

100 Python callable. :data:`HOST_COPY_FN` is such a pointer, for a host 

101 copy that runs without the GIL. ``cuda.core`` keeps ``fn`` alive until 

102 the function has run and releases it afterwards on the main thread. 

103 user_data : int or bytes-like, optional 

104 Only for ctypes function pointers. An ``int`` is a raw pointer that 

105 the caller must keep valid until the function has run. Bytes-like data 

106 is copied; ``cuda.core`` owns the copy until the function has run and 

107 releases it afterwards on the main thread. 

108  

109 Raises 

110 ------ 

111 TypeError 

112 If ``stream`` is ``None`` or not a stream, or if ``fn`` is not 

113 callable or its ctypes prototype is not ``CUhostFn``. 

114 ValueError 

115 If ``user_data`` is given for a Python callable, or if ``fn`` is 

116 :data:`HOST_COPY_FN` and ``user_data`` is not a packed record. 

117 RuntimeError 

118 If ``stream`` is a :class:`Stream` that is capturing a graph, or a 

119 :class:`~graph.GraphBuilder` that is closed or not building. 

120  

121 Notes 

122 ----- 

123 An exception raised by a Python callable cannot propagate to the caller. 

124 It is reported through :func:`sys.unraisablehook` and the stream continues. 

125  

126 .. warning:: 

127  

128 The function must not call CUDA, directly or through a library. The 

129 driver does not permit such calls. Such a call may fail with an error, 

130 or it may deadlock. 

131  

132 A Python callable, or a ctypes callback that wraps a Python function, 

133 acquires the GIL on the driver's thread. A thread 

134 that holds the GIL while blocking inside a CUDA call waiting for this 

135 stream therefore deadlocks with the callback. ``cuda.core`` and 

136 ``cuda.bindings`` release the GIL around CUDA calls. A library that 

137 calls CUDA with the GIL held must not be used to wait on a stream with 

138 a pending Python callback. 

139 """ 

140 from cuda.core.graph._graph_builder import GraphBuilder 22 ctx2b c `bd :b;b=b?be f g {bh i _b@b|b[b]b^bj k

141  

142 if isinstance(stream, GraphBuilder): 22 ctx2b c `bd :b;b=b?be f g {bh i _b@b|b[b]b^bj k

143 return stream.callback(fn, user_data=user_data) 2 ctx2`b{b

144  

145 cdef Stream s = Stream_accept(stream, allow_stream_protocol=True) 20 ctx2b c d :b;b=b?be f g h i _b@b|b[b]b^bj k

146 cdef cydriver.CUstream c_stream = as_cu(s._h_stream) 19 ctx2b c d :b;b=b?be f g h i _b@b[b]b^bj k

147 cdef cydriver.CUstreamCaptureStatus status 

148 with nogil: 19 ctx2b c d :b;b=b?be f g h i _b@b[b]b^bj k

149 HANDLE_RETURN(cydriver.cuStreamIsCapturing(c_stream, &status)) 19 ctx2b c d :b;b=b?be f g h i _b@b[b]b^bj k

150 if status == cydriver.CU_STREAM_CAPTURE_STATUS_ACTIVE: 19 ctx2b c d :b;b=b?be f g h i _b@b[b]b^bj k

151 raise RuntimeError( 1 ctx2_b

152 "host_launch cannot enqueue on a stream that is capturing a graph; " 

153 "pass the GraphBuilder, or call GraphBuilder.callback") 

154  

155 cdef cydriver.CUhostFn c_fn 

156 cdef void* c_user_data = NULL 18 ctx2b c d :b;b=b?be f g h i @b[b]b^bj k

157 cdef OpaqueHandle fn_owner, data_owner 

158 _resolve_host_callback(fn, user_data, &c_fn, &c_user_data, &fn_owner, &data_owner) 18 ctx2b c d :b;b=b?be f g h i @b[b]b^bj k

159 with nogil: 10 ctx1bcdefghijk

160 HANDLE_RETURN(launch_host_function(c_stream, c_fn, c_user_data, fn_owner, data_owner)) 10 ctx1bcdefghijk

161 return None 10 ctx1bcdefghijk

162  

163  

164cdef _check_cooperative_launch(kernel: Kernel, config: LaunchConfig, stream: Stream): 

165 dev = stream.device 1 ctx1a

166 num_sm = dev.properties.multiprocessor_count 1 ctx1a

167 max_grid_size = ( 

168 kernel.occupancy.max_active_blocks_per_multiprocessor(prod(config.block), config.shmem_size) * num_sm 1 ctx1a

169 ) 

170 if prod(config.grid) > max_grid_size: 1 ctx1a

171 # For now let's try not to be smart and adjust the grid size behind users' back. 

172 # We explicitly ask users to adjust. 

173 x, y, z = config.grid 1 ctx1a

174 raise ValueError(f"The specified grid size ({x} * {y} * {z}) exceeds the limit ({max_grid_size})") 1 ctx1a