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

22 statements  

« prev     ^ index     » next       coverage.py v7.14.1, created at 2026-06-13 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  

21cdef class NvlinkInfo: 

22 """ 

23 Nvlink information for a device. 

24 """ 

25 cdef Device _device 

26 cdef int _link 

27  

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

29 self._device = device 1a

30 self._link = link 1a

31  

32 @property 

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

34 """ 

35 Retrieves the NvLink version for the device and link. 

36  

37 For all products with NvLink support. 

38  

39 Returns 

40 ------- 

41 tuple[int, int] 

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

43 """ 

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

45 if version == nvml.NvlinkVersion.VERSION_INVALID: 

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

47 try: 

48 return _NVLINK_VERSION_MAPPING[version] 

49 except KeyError: 

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

51  

52 @property 

53 def state(self) -> bool: 

54 """ 

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

56  

57 For Pascal™ or newer fully supported devices. 

58  

59 For all products with Nvlink support. 

60  

61 Returns 

62 ------- 

63 bool 

64 `True` if the Nvlink is active. 

65 """ 

66 return ( 1a

67 nvml.device_get_nvlink_state(self._device._handle, self._link) == nvml.EnableState.FEATURE_ENABLED 1a

68 ) 

69  

70 max_links = nvml.NVLINK_MAX_LINKS