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

19 statements  

« prev     ^ index     » next       coverage.py v7.14.0, created at 2026-05-22 01:37 +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  

17cdef class NvlinkInfo: 

18 """ 

19 Nvlink information for a device. 

20 """ 

21 cdef Device _device 

22 cdef int _link 

23  

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

25 self._device = device 1a

26 self._link = link 1a

27  

28 @property 

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

30 """ 

31 Retrieves the NvLink version for the device and link. 

32  

33 For all products with NvLink support. 

34  

35 Returns 

36 ------- 

37 tuple[int, int] 

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

39 """ 

40 version = nvml.device_get_nvlink_version(self._device._handle, self._link) 1a

41 if version == nvml.NvlinkVersion.VERSION_INVALID: 

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

43 try: 

44 return _NVLINK_VERSION_MAPPING[version] 

45 except KeyError: 

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

47  

48 @property 

49 def state(self) -> bool: 

50 """ 

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

52  

53 For Pascal™ or newer fully supported devices. 

54  

55 For all products with Nvlink support. 

56  

57 Returns 

58 ------- 

59 bool 

60 `True` if the Nvlink is active. 

61 """ 

62 return ( 

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

64 ) 

65  

66 max_links = nvml.NVLINK_MAX_LINKS