Add method IsImexCapable to Device interface

Signed-off-by: Carlos Eduardo Arango Gutierrez <[email protected]>

Author: ArangoGutierrez

pkg/nvlib/device/device.go

@@ -32,6 +32,7 @@ type Device interface {
 	GetMigDevices() ([]MigDevice, error)
 	GetMigProfiles() ([]MigProfile, error)
 	GetPCIBusID() (string, error)
+	IsImexEnabled() (bool, error)
 	IsMigCapable() (bool, error)
 	IsMigEnabled() (bool, error)
 	VisitMigDevices(func(j int, m MigDevice) error) error
@@ -208,6 +209,33 @@ func (d *device) IsMigEnabled() (bool, error) {
 	return (mode == nvml.DEVICE_MIG_ENABLE), nil
 }
 
+// IsImexEnabled checks if a device has IMEX capabilities.
+func (d *device) IsImexEnabled() (bool, error) {
+	if d.lib.hasSymbol("nvmlDeviceGetGpuFabricInfo") {
+		_, ret := d.GetGpuFabricInfo()
+		if ret == nvml.ERROR_NOT_SUPPORTED {
+			return false, nil
+		}
+		if ret != nvml.SUCCESS {
+			return false, fmt.Errorf("error getting GPU Fabric Info: %v", ret)
+		}
+		return true, nil
+	}
+
+	if d.lib.hasSymbol("nvmlDeviceGetGpuFabricInfoV") {
+		_, ret := d.GetGpuFabricInfoV().V2()
+		if ret == nvml.ERROR_NOT_SUPPORTED {
+			return false, nil
+		}
+		if ret != nvml.SUCCESS {
+			return false, fmt.Errorf("error getting GPU Fabric Info: %v", ret)
+		}
+		return true, nil
+	}
+
+	return false, nil
+}
+
 // VisitMigDevices walks a top-level device and invokes a callback function for each MIG device configured on it.
 func (d *device) VisitMigDevices(visit func(int, MigDevice) error) error {
 	capable, err := d.IsMigCapable()