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

28 statements  

« prev     ^ index     » next       coverage.py v7.15.2, created at 2026-07-19 01:12 +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 def max_links(cls): 

24 """ 

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

26 upstream NVML as ``NVML_NVLINK_MAX_LINKS``. 

27  

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

29 :py:attr:`Device.get_nvlink_count`. 

30  

31 .. version-deprecated:: 1.1.0 

32 This property is deprecated and will be removed in a future release. 

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

34 """ 

35 warnings.warn( 1a

36 "The `max_links` property is deprecated and will be removed in a future release. " 

37 "Use `Device.get_nvlink_count` instead.", 

38 DeprecationWarning, 1a

39 ) 

40  

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

42 return nvml.NVLINK_MAX_LINKS 1a

43  

44  

45cdef class _NvlinkInfo: 

46 """ 

47 Nvlink information for a device. 

48 """ 

49 cdef Device _device 

50 cdef int _link 

51  

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

53 self._device = device 

54 self._link = link 

55  

56 @property 

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

58 """ 

59 Retrieves the NvLink version for the device and link. 

60  

61 For all products with NvLink support. 

62  

63 Returns 

64 ------- 

65 tuple[int, int] 

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

67 """ 

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

69 if version == nvml.NvlinkVersion.VERSION_INVALID: 

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

71 try: 

72 return _NVLINK_VERSION_MAPPING[version] 

73 except KeyError: 

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

75  

76 @property 

77 def state(self) -> bool: 

78 """ 

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

80  

81 For Pascal™ or newer fully supported devices. 

82  

83 For all products with Nvlink support. 

84  

85 Returns 

86 ------- 

87 bool 

88 `True` if the Nvlink is active. 

89 """ 

90 return ( 

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

92 ) 

93  

94  

95class NvlinkInfo(_NvlinkInfo, metaclass=_NvlinkInfoMeta): 

96 pass