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
« 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
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}
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)
21cdef class NvlinkInfo:
22 """
23 Nvlink information for a device.
24 """
25 cdef Device _device
26 cdef int _link
28 def __init__(self, device: Device, link: int):
29 self._device = device 1a
30 self._link = link 1a
32 @property
33 def version(self) -> tuple[int, int]:
34 """
35 Retrieves the NvLink version for the device and link.
37 For all products with NvLink support.
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
52 @property
53 def state(self) -> bool:
54 """
55 Retrieves the state of the device's Nvlink for the device and link specified.
57 For Pascal™ or newer fully supported devices.
59 For all products with Nvlink support.
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 )
70 max_links = nvml.NVLINK_MAX_LINKS