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 runtime as cudart 

13from cuda.bindings._v2 import nvrtc 

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 = nvrtc.create_program(str.encode(code), b"sourceCode.cu") 

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 = nvrtc.version() 

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 nvrtc.compile_program(prog, opts) 

74 except nvrtc.NvrtcError as err: 

75 log = nvrtc.get_program_log(prog) 

76 import sys 

77 

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

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

80 sys.exit(1) 

81 

82 if use_cubin: 

83 data = nvrtc.get_cubin(prog) 

84 else: 

85 data = nvrtc.get_ptx(prog) 

86 

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

88 

89 def get_function(self, name): 

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