Coverage for cuda/bindings/_example_helpers/common.py: 0.00%

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

2# SPDX-License-Identifier: Apache-2.0 

3 

4 

5import os 

6import sys 

7 

8import numpy as np 

9 

10from cuda import pathfinder 

11from cuda.bindings import driver as cuda 

12from cuda.bindings import nvrtc 

13from cuda.bindings import runtime as cudart 

14 

15from .helper_cuda import check_cuda_errors 

16 

17 

18def requirement_not_met(message): 

19 print(message, file=sys.stderr) # noqa: T201 

20 exitcode = os.environ.get("CUDA_BINDINGS_SKIP_EXAMPLE", "1") 

21 return sys.exit(int(exitcode)) 

22 

23 

24def check_compute_capability_too_low(dev_id, required_cc_major_minor): 

25 cc_major = check_cuda_errors( 

26 cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, dev_id) 

27 ) 

28 cc_minor = check_cuda_errors( 

29 cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, dev_id) 

30 ) 

31 have_cc_major_minor = (cc_major, cc_minor) 

32 if have_cc_major_minor < required_cc_major_minor: 

33 requirement_not_met( 

34 f"CUDA device compute capability too low: {have_cc_major_minor=!r}, {required_cc_major_minor=!r}" 

35 ) 

36 

37 

38class KernelHelper: 

39 def __init__(self, code, dev_id): 

40 include_dirs = [] 

41 for libname in ("cudart", "cccl"): 

42 hdr_dir = pathfinder.find_nvidia_header_directory(libname) 

43 if hdr_dir is None: 

44 requirement_not_met(f'pathfinder.find_nvidia_header_directory("{libname}") returned None') 

45 include_dirs.append(hdr_dir) 

46 

47 prog = check_cuda_errors(nvrtc.nvrtcCreateProgram(str.encode(code), b"sourceCode.cu", 0, None, None)) 

48 

49 # Initialize CUDA 

50 check_cuda_errors(cudart.cudaFree(0)) 

51 

52 major = check_cuda_errors( 

53 cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, dev_id) 

54 ) 

55 minor = check_cuda_errors( 

56 cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, dev_id) 

57 ) 

58 _, nvrtc_minor = check_cuda_errors(nvrtc.nvrtcVersion()) 

59 use_cubin = nvrtc_minor >= 1 

60 prefix = "sm" if use_cubin else "compute" 

61 arch_arg = bytes(f"--gpu-architecture={prefix}_{major}{minor}", "ascii") 

62 

63 opts = [ 

64 b"--fmad=true", 

65 arch_arg, 

66 b"--std=c++17", 

67 b"-default-device", 

68 ] 

69 for inc_dir in include_dirs: 

70 opts.append(f"--include-path={inc_dir}".encode()) 

71 

72 try: 

73 check_cuda_errors(nvrtc.nvrtcCompileProgram(prog, len(opts), opts)) 

74 except RuntimeError as err: 

75 log_size = check_cuda_errors(nvrtc.nvrtcGetProgramLogSize(prog)) 

76 log = b" " * log_size 

77 check_cuda_errors(nvrtc.nvrtcGetProgramLog(prog, log)) 

78 import sys 

79 

80 print(log.decode(), file=sys.stderr) # noqa: T201 

81 print(err, file=sys.stderr) # noqa: T201 

82 sys.exit(1) 

83 

84 if use_cubin: 

85 data_size = check_cuda_errors(nvrtc.nvrtcGetCUBINSize(prog)) 

86 data = b" " * data_size 

87 check_cuda_errors(nvrtc.nvrtcGetCUBIN(prog, data)) 

88 else: 

89 data_size = check_cuda_errors(nvrtc.nvrtcGetPTXSize(prog)) 

90 data = b" " * data_size 

91 check_cuda_errors(nvrtc.nvrtcGetPTX(prog, data)) 

92 

93 self.module = check_cuda_errors(cuda.cuModuleLoadData(np.char.array(data))) 

94 

95 def get_function(self, name): 

96 return check_cuda_errors(cuda.cuModuleGetFunction(self.module, name))