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

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)

21from typing import TYPE_CHECKING (empty)

22  

23if TYPE_CHECKING: (empty)

24 from cuda.core.graph import GraphBuilder 

25 from cuda.core.typing import IsStreamType 

26  

27__all__ = ['launch'] (empty)

28  

29  

30def launch( (empty)

31 stream: Stream | GraphBuilder | IsStreamType, 

32 config: LaunchConfig, 

33 kernel: Kernel, 

34 *kernel_args 

35) -> None: 

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

37 object with launch-time configuration. 

38  

39 Parameters 

40 ---------- 

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

42 The stream establishing the stream ordering semantic of a 

43 launch. 

44 config : :obj:`LaunchConfig` 

45 Launch configurations inline with options provided by 

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

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

48 Kernel to launch. 

49 *kernel_args : Any 

50 Variable length argument list that is provided to the 

51 launching kernel. 

52  

53 """ 

54 cdef Stream s = Stream_accept(stream, allow_stream_protocol=True) 122 ctx2b c d e f g h i j k l 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 ! # $ % ' ( ) * + , - . a / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGb

55 cdef LaunchConfig conf = check_or_create_options(LaunchConfig, config, "launch config") 122 ctx2b c d e f g h i j k l 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 ! # $ % ' ( ) * + , - . a / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGb

56  

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

58 # TODO: merge with HelperKernelParams? 

59 cdef ParamHolder ker_args = ParamHolder(kernel_args) 122 ctx2b c d e f g h i j k l 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 ! # $ % ' ( ) * + , - . a / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGb

60 cdef void** args_ptr = <void**><uintptr_t>(ker_args.ptr) 122 ctx2b c d e f g h i j k l 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 ! # $ % ' ( ) * + , - . a / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGb

61  

62 cdef Kernel ker = <Kernel>kernel 122 ctx2b c d e f g h i j k l 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 ! # $ % ' ( ) * + , - . a / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGb

63 cdef cydriver.CUfunction func_handle = <cydriver.CUfunction>as_cu(ker._h_kernel) 122 ctx2b c d e f g h i j k l 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 ! # $ % ' ( ) * + , - . a / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGb

64  

65 drv_cfg = conf._to_native_launch_config() 122 ctx2b c d e f g h i j k l 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 ! # $ % ' ( ) * + , - . a / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGb

66 drv_cfg.hStream = as_cu(s._h_stream) 122 ctx2b c d e f g h i j k l 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 ! # $ % ' ( ) * + , - . a / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGb

67 if conf.is_cooperative: 122 ctx2b c d e f g h i j k l 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 ! # $ % ' ( ) * + , - . a / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGb

68 _check_cooperative_launch(kernel, conf, s) 1 ctx1a

69 with nogil: 122 ctx2b c d e f g h i j k l 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 ! # $ % ' ( ) * + , - . a / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGb

70 HANDLE_RETURN(cydriver.cuLaunchKernelEx(&drv_cfg, func_handle, args_ptr, NULL)) 122 ctx2b c d e f g h i j k l 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 ! # $ % ' ( ) * + , - . a / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGb

71  

72  

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

74 dev = stream.device 1 ctx1a

75 num_sm = dev.properties.multiprocessor_count 1 ctx1a

76 max_grid_size = ( 

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

78 ) 

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

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

81 # We explicitly ask users to adjust. 

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

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