Add functions related to NVLink info

pkg/nvml/consts.go

@@ -20,6 +20,11 @@ import (
 	"github.com/NVIDIA/go-nvml/pkg/nvml"
 )
 
+// General untyped constants
+const (
+	NVLINK_MAX_LINKS = nvml.NVLINK_MAX_LINKS
+)
+
 // Return constants
 const (
 	SUCCESS                       = Return(nvml.SUCCESS)
@@ -131,3 +136,19 @@ const (
 	EventTypeSingleBitEccError = nvml.EventTypeSingleBitEccError
 	EventTypeDoubleBitEccError = nvml.EventTypeDoubleBitEccError
 )
+
+// GPU Topology enumeration
+const (
+	TOPOLOGY_INTERNAL   = GpuTopologyLevel(nvml.TOPOLOGY_INTERNAL)
+	TOPOLOGY_SINGLE     = GpuTopologyLevel(nvml.TOPOLOGY_SINGLE)
+	TOPOLOGY_MULTIPLE   = GpuTopologyLevel(nvml.TOPOLOGY_MULTIPLE)
+	TOPOLOGY_HOSTBRIDGE = GpuTopologyLevel(nvml.TOPOLOGY_HOSTBRIDGE)
+	TOPOLOGY_NODE       = GpuTopologyLevel(nvml.TOPOLOGY_NODE)
+	TOPOLOGY_SYSTEM     = GpuTopologyLevel(nvml.TOPOLOGY_SYSTEM)
+)
+
+// Generic enable/disable constants
+const (
+	FEATURE_DISABLED = EnableState(nvml.FEATURE_DISABLED)
+	FEATURE_ENABLED  = EnableState(nvml.FEATURE_ENABLED)
+)

pkg/nvml/device.go

@@ -178,3 +178,27 @@ func (d nvmlDevice) GetSupportedEventTypes() (uint64, Return) {
 	e, r := nvml.Device(d).GetSupportedEventTypes()
 	return e, Return(r)
 }
+
+// GetTopologyCommonAncestor retrieves the common ancestor for two devices.
+func (d nvmlDevice) GetTopologyCommonAncestor(o Device) (GpuTopologyLevel, Return) {
+	other, ok := o.(nvmlDevice)
+	if !ok {
+		return 0, ERROR_INVALID_ARGUMENT
+	}
+
+	l, r := nvml.Device(d).GetTopologyCommonAncestor(nvml.Device(other))
+	return GpuTopologyLevel(l), Return(r)
+}
+
+// GetNvLinkState retrieves the state of the device's NvLink for the link specified.
+func (d nvmlDevice) GetNvLinkState(link int) (EnableState, Return) {
+	s, r := nvml.Device(d).GetNvLinkState(link)
+	return EnableState(s), Return(r)
+}
+
+// GetNvLinkRemotePciInfo retrieves the PCI information for the remote node on a NvLink link.
+// Note: pciSubSystemId is not filled in this function and is indeterminate.
+func (d nvmlDevice) GetNvLinkRemotePciInfo(link int) (PciInfo, Return) {
+	p, r := nvml.Device(d).GetNvLinkRemotePciInfo(link)
+	return PciInfo(p), Return(r)
+}

pkg/nvml/types.go

@@ -59,8 +59,11 @@ type Device interface {
 	GetMigMode() (int, int, Return)
 	GetMinorNumber() (int, Return)
 	GetName() (string, Return)
+	GetNvLinkRemotePciInfo(int) (PciInfo, Return)
+	GetNvLinkState(int) (EnableState, Return)
 	GetPciInfo() (PciInfo, Return)
 	GetSupportedEventTypes() (uint64, Return)
+	GetTopologyCommonAncestor(Device) (GpuTopologyLevel, Return)
 	GetUUID() (string, Return)
 	IsMigDeviceHandle() (bool, Return)
 	RegisterEvents(uint64, EventSet) Return
@@ -145,3 +148,9 @@ type DeviceArchitecture nvml.DeviceArchitecture
 
 // BrandType represents the brand of a GPU device
 type BrandType nvml.BrandType
+
+// GpuTopologyLevel represents level relationships within a system between two GPUs
+type GpuTopologyLevel nvml.GpuTopologyLevel
+
+// EnableState represents a generic enable/disable enum
+type EnableState nvml.EnableState