Coverage for cuda/core/system/_nvlink.pxi: 3.70%

27 statements  

« prev     ^ index     » next       coverage.py v7.15.0, created at 2026-07-03 01:38 +0000

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4  

5  

6_NVLINK_VERSION_MAPPING = { 

7 nvml.NvlinkVersion.VERSION_1_0: (1, 0), 

8 nvml.NvlinkVersion.VERSION_2_0: (2, 0), 

9 nvml.NvlinkVersion.VERSION_2_2: (2, 2), 

10 nvml.NvlinkVersion.VERSION_3_0: (3, 0), 

11 nvml.NvlinkVersion.VERSION_3_1: (3, 1), 

12 nvml.NvlinkVersion.VERSION_4_0: (4, 0), 

13 nvml.NvlinkVersion.VERSION_5_0: (5, 0), 

14} 

15  

16_NVLINK_VERSION_6_0 = getattr(nvml.NvlinkVersion, "VERSION_6_0", None) 

17if _NVLINK_VERSION_6_0 is not None: 

18 _NVLINK_VERSION_MAPPING[_NVLINK_VERSION_6_0] = (6, 0) 

19  

20  

21class _NvlinkInfoMeta(type): 

22 @property 

23 @deprecated( 

24 version="1.1.0", 

25 reason="Use Device.get_nvlink_count instead to get the actual number of Nvlinks available on a specific device." 

26 ) 

27 def max_links(cls): 

28 """ 

29 The statically-defined maximum number of Nvlinks available. Defined in 

30 upstream NVML as ``NVML_NVLINK_MAX_LINKS``. 

31  

32 To find the actual number of Nvlinks available on a device, use 

33 :py:attr:`Device.get_nvlink_count`. 

34 """ 

35 # This will always return 18, even on CTK 13.3 where it should be 36. 

36 return nvml.NVLINK_MAX_LINKS 1a

37  

38  

39cdef class _NvlinkInfo: 

40 """ 

41 Nvlink information for a device. 

42 """ 

43 cdef Device _device 

44 cdef int _link 

45  

46 def __init__(self, device: Device, link: int): 

47 self._device = device 

48 self._link = link 

49  

50 @property 

51 def version(self) -> tuple[int, int]: 

52 """ 

53 Retrieves the NvLink version for the device and link. 

54  

55 For all products with NvLink support. 

56  

57 Returns 

58 ------- 

59 tuple[int, int] 

60 The Nvlink version as a tuple of (major, minor). 

61 """ 

62 version = nvml.device_get_nvlink_version(self._device._handle, self._link) 

63 if version == nvml.NvlinkVersion.VERSION_INVALID: 

64 raise RuntimeError("Invalid NvLink version returned for device") 

65 try: 

66 return _NVLINK_VERSION_MAPPING[version] 

67 except KeyError: 

68 raise RuntimeError(f"Unknown NvLink version {version} returned for device") from None 

69  

70 @property 

71 def state(self) -> bool: 

72 """ 

73 Retrieves the state of the device's Nvlink for the device and link specified. 

74  

75 For Pascal™ or newer fully supported devices. 

76  

77 For all products with Nvlink support. 

78  

79 Returns 

80 ------- 

81 bool 

82 `True` if the Nvlink is active. 

83 """ 

84 return ( 

85 nvml.device_get_nvlink_state(self._device._handle, self._link) == nvml.EnableState.FEATURE_ENABLED 

86 ) 

87  

88  

89class NvlinkInfo(_NvlinkInfo, metaclass=_NvlinkInfoMeta): 

90 pass