Coverage for cuda / pathfinder / _headers / supported_nvidia_headers.py: 100%

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

2# SPDX-License-Identifier: Apache-2.0 

3 

4from typing import Final 

5 

6from cuda.pathfinder._utils.platform_aware import IS_WINDOWS 

7 

8SUPPORTED_HEADERS_CTK_COMMON = { 

9 "cccl": "cuda/std/version", 

10 "cublas": "cublas.h", 

11 "cudart": "cuda_runtime.h", 

12 "cufft": "cufft.h", 

13 "curand": "curand.h", 

14 "cusolver": "cusolverDn.h", 

15 "cusparse": "cusparse.h", 

16 "npp": "npp.h", 

17 "nvcc": "fatbinary_section.h", 

18 "nvfatbin": "nvFatbin.h", 

19 "nvjitlink": "nvJitLink.h", 

20 "nvjpeg": "nvjpeg.h", 

21 "nvrtc": "nvrtc.h", 

22 "nvvm": "nvvm.h", 

23} 

24 

25SUPPORTED_HEADERS_CTK_LINUX_ONLY = { 

26 "cufile": "cufile.h", 

27} 

28SUPPORTED_HEADERS_CTK_LINUX = SUPPORTED_HEADERS_CTK_COMMON | SUPPORTED_HEADERS_CTK_LINUX_ONLY 

29 

30SUPPORTED_HEADERS_CTK_WINDOWS_ONLY: dict[str, str] = {} 

31SUPPORTED_HEADERS_CTK_WINDOWS = SUPPORTED_HEADERS_CTK_COMMON | SUPPORTED_HEADERS_CTK_WINDOWS_ONLY 

32 

33SUPPORTED_HEADERS_CTK_ALL = ( 

34 SUPPORTED_HEADERS_CTK_COMMON | SUPPORTED_HEADERS_CTK_LINUX_ONLY | SUPPORTED_HEADERS_CTK_WINDOWS_ONLY 

35) 

36SUPPORTED_HEADERS_CTK: Final[dict[str, str]] = ( 

37 SUPPORTED_HEADERS_CTK_WINDOWS if IS_WINDOWS else SUPPORTED_HEADERS_CTK_LINUX 

38) 

39 

40SUPPORTED_SITE_PACKAGE_HEADER_DIRS_CTK = { 

41 "cccl": ( 

42 "nvidia/cu13/include/cccl", # cuda-toolkit[cccl]==13.* 

43 "nvidia/cuda_cccl/include", # cuda-toolkit[cccl]==12.* 

44 ), 

45 "cublas": ("nvidia/cu13/include", "nvidia/cublas/include"), 

46 "cudart": ("nvidia/cu13/include", "nvidia/cuda_runtime/include"), 

47 "cufft": ("nvidia/cu13/include", "nvidia/cufft/include"), 

48 "cufile": ("nvidia/cu13/include", "nvidia/cufile/include"), 

49 "curand": ("nvidia/cu13/include", "nvidia/curand/include"), 

50 "cusolver": ("nvidia/cu13/include", "nvidia/cusolver/include"), 

51 "cusparse": ("nvidia/cu13/include", "nvidia/cusparse/include"), 

52 "npp": ("nvidia/cu13/include", "nvidia/npp/include"), 

53 "nvcc": ("nvidia/cu13/include", "nvidia/cuda_nvcc/include"), 

54 "nvfatbin": ("nvidia/cu13/include", "nvidia/nvfatbin/include"), 

55 "nvjitlink": ("nvidia/cu13/include", "nvidia/nvjitlink/include"), 

56 "nvjpeg": ("nvidia/cu13/include", "nvidia/nvjpeg/include"), 

57 "nvrtc": ("nvidia/cu13/include", "nvidia/cuda_nvrtc/include"), 

58 "nvvm": ("nvidia/cu13/include", "nvidia/cuda_nvcc/nvvm/include"), 

59} 

60 

61SUPPORTED_HEADERS_NON_CTK_COMMON = { 

62 "cusparseLt": "cusparseLt.h", 

63 "cutensor": "cutensor.h", 

64} 

65SUPPORTED_HEADERS_NON_CTK_LINUX_ONLY = { 

66 "nvshmem": "nvshmem.h", 

67} 

68SUPPORTED_HEADERS_NON_CTK_WINDOWS_ONLY: Final[dict[str, str]] = {} 

69SUPPORTED_HEADERS_NON_CTK_LINUX = SUPPORTED_HEADERS_NON_CTK_COMMON | SUPPORTED_HEADERS_NON_CTK_LINUX_ONLY 

70SUPPORTED_HEADERS_NON_CTK_WINDOWS = SUPPORTED_HEADERS_NON_CTK_COMMON | SUPPORTED_HEADERS_NON_CTK_WINDOWS_ONLY 

71SUPPORTED_HEADERS_NON_CTK_ALL = ( 

72 SUPPORTED_HEADERS_NON_CTK_COMMON | SUPPORTED_HEADERS_NON_CTK_LINUX_ONLY | SUPPORTED_HEADERS_NON_CTK_WINDOWS_ONLY 

73) 

74SUPPORTED_HEADERS_NON_CTK: Final[dict[str, str]] = ( 

75 SUPPORTED_HEADERS_NON_CTK_WINDOWS if IS_WINDOWS else SUPPORTED_HEADERS_NON_CTK_LINUX 

76) 

77 

78SUPPORTED_SITE_PACKAGE_HEADER_DIRS_NON_CTK = { 

79 "cusparseLt": ("nvidia/cusparselt/include",), 

80 "cutensor": ("cutensor/include",), 

81 "nvshmem": ("nvidia/nvshmem/include",), 

82} 

83 

84SUPPORTED_INSTALL_DIRS_NON_CTK = { 

85 "nvshmem": ("/usr/include/nvshmem_*",), 

86}