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

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

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._launch_config cimport LaunchConfig 

10from cuda.core._kernel_arg_handler cimport ParamHolder 

11from cuda.core._module cimport Kernel 

12from cuda.core._resource_handles cimport as_cu 

13from cuda.core._stream cimport Stream_accept, Stream 

14from cuda.core._utils.cuda_utils cimport ( 

15 check_or_create_options, 

16 HANDLE_RETURN, 

17) 

18from cuda.core._module import Kernel (empty)

19from cuda.core._stream import Stream (empty)

20from math import prod (empty)

21  

22  

23def launch(stream: Stream | GraphBuilder | IsStreamT, config: LaunchConfig, kernel: Kernel, *kernel_args): (empty)

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

25 object with launch-time configuration. 

26  

27 Parameters 

28 ---------- 

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

30 The stream establishing the stream ordering semantic of a 

31 launch. 

32 config : :obj:`LaunchConfig` 

33 Launch configurations inline with options provided by 

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

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

36 Kernel to launch. 

37 *kernel_args : Any 

38 Variable length argument list that is provided to the 

39 launching kernel. 

40  

41 """ 

42 cdef Stream s = Stream_accept(stream, allow_stream_protocol=True) 88 ctx1bcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123a456789!#$%'()*+,-./:;=?@[]^_`{|}

43 cdef LaunchConfig conf = check_or_create_options(LaunchConfig, config, "launch config") 88 ctx1bcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123a456789!#$%'()*+,-./:;=?@[]^_`{|}

44  

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

46 # TODO: merge with HelperKernelParams? 

47 cdef ParamHolder ker_args = ParamHolder(kernel_args) 89 ctx1~bcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123a456789!#$%'()*+,-./:;=?@[]^_`{|}

48 cdef void** args_ptr = <void**><uintptr_t>(ker_args.ptr) 88 ctx1bcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123a456789!#$%'()*+,-./:;=?@[]^_`{|}

49  

50 cdef Kernel ker = <Kernel>kernel 88 ctx1bcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123a456789!#$%'()*+,-./:;=?@[]^_`{|}

51 cdef cydriver.CUfunction func_handle = <cydriver.CUfunction>as_cu(ker._h_kernel) 88 ctx1bcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123a456789!#$%'()*+,-./:;=?@[]^_`{|}

52  

53 drv_cfg = conf._to_native_launch_config() 88 ctx1bcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123a456789!#$%'()*+,-./:;=?@[]^_`{|}

54 drv_cfg.hStream = as_cu(s._h_stream) 88 ctx1bcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123a456789!#$%'()*+,-./:;=?@[]^_`{|}

55 if conf.cooperative_launch: 88 ctx1bcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123a456789!#$%'()*+,-./:;=?@[]^_`{|}

56 _check_cooperative_launch(kernel, conf, s) 1 ctx1a

57 with nogil: 88 ctx1bcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123a456789!#$%'()*+,-./:;=?@[]^_`{|}

58 HANDLE_RETURN(cydriver.cuLaunchKernelEx(&drv_cfg, func_handle, args_ptr, NULL)) 88 ctx1bcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123a456789!#$%'()*+,-./:;=?@[]^_`{|}

59  

60  

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

62 dev = stream.device 1 ctx1a

63 num_sm = dev.properties.multiprocessor_count 1 ctx1a

64 max_grid_size = ( 

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

66 ) 

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

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

69 # We explicitly ask users to adjust. 

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

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