wip: update vfio-manage to choose best VFIO driver

Signed-off-by: Christopher Desiniotis <[email protected]>

Author: cdesiniotis

internal/nvpci/modalias.go

@@ -0,0 +1,264 @@
+package nvpci
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/sirupsen/logrus"
+	"golang.org/x/sys/unix"
+)
+
+// modAlias is a decomposed version of string like this
+//
+// vNNNNNNNNdNNNNNNNNsvNNNNNNNNsdNNNNNNNNbcNNscNNiNN
+//
+// (followed by a space or newline). The "NNNN" are always of the
+// length in the example unless replaced with a wildcard ("*"),
+// but we make no assumptions about length.
+type modAlias struct {
+	vendor     string // v
+	device     string // d
+	subvendor  string // sv
+	subdevice  string // sd
+	baseClass  string // bc
+	subClass   string // sc
+	interface_ string // i
+}
+
+// vfioAlias represents an entry from the modules.alias file for a vfio driver
+type vfioAlias struct {
+	modAlias *modAlias // The modalias pattern
+	driver   string    // The vfio driver name
+}
+
+// parseModAliasString parses a modalias string in the format:
+// vNNNNNNNNdNNNNNNNNsvNNNNNNNNsdNNNNNNNNbcNNscNNiNN
+// where N can be hex digits or wildcards (*)
+func parseModAliasString(input string) (*modAlias, error) {
+	if input == "" {
+		return nil, fmt.Errorf("modalias string is empty")
+	}
+
+	// Trim any whitespace or newlines
+	input = strings.TrimSpace(input)
+
+	// Trim the leading "pci:" prefix in the modalias file
+	split := strings.SplitN(input, ":", 2)
+	if len(split) != 2 {
+		return nil, fmt.Errorf("unexpected number of parts in modalias after trimming pci: prefix: %s", input)
+	}
+	input = split[1]
+
+	// Parse vendor ID (starts with 'v')
+	if !strings.HasPrefix(input, "v") {
+		return nil, fmt.Errorf("modalias must start with 'v', got: %s", input)
+	}
+
+	ma := &modAlias{}
+	remaining := input[1:] // skip 'v'
+
+	// Extract vendor
+	vendor, remaining, err := extractField(remaining, "d")
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse vendor: %w", err)
+	}
+	ma.vendor = vendor
+
+	// Extract device
+	device, remaining, err := extractField(remaining, "sv")
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse device: %w", err)
+	}
+	ma.device = device
+
+	// Extract subvendor
+	subvendor, remaining, err := extractField(remaining, "sd")
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse subvendor: %w", err)
+	}
+	ma.subvendor = subvendor
+
+	// Extract subdevice
+	subdevice, remaining, err := extractField(remaining, "bc")
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse subdevice: %w", err)
+	}
+	ma.subdevice = subdevice
+
+	// Extract base class
+	baseClass, remaining, err := extractField(remaining, "sc")
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse base class: %w", err)
+	}
+	ma.baseClass = baseClass
+
+	// Extract subclass
+	subClass, remaining, err := extractField(remaining, "i")
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse subclass: %w", err)
+	}
+	ma.subClass = subClass
+
+	// Extract interface (remaining content)
+	ma.interface_ = remaining
+
+	return ma, nil
+}
+
+// extractField extracts the value before the next delimiter from the input string.
+// Returns the extracted value, the remaining string (without the delimiter), and any error.
+func extractField(input, delimiter string) (string, string, error) {
+	idx := strings.Index(input, delimiter)
+	if idx == -1 {
+		// Delimiter not found - this could be the last field
+		return input, "", nil
+	}
+
+	value := input[:idx]
+	remaining := input[idx+len(delimiter):]
+
+	return value, remaining, nil
+}
+
+func getKernelVersion() (string, error) {
+	var uname unix.Utsname
+	if err := unix.Uname(&uname); err != nil {
+		return "", err
+	}
+
+	// Convert C-style byte array to Go string
+	release := make([]byte, 0, len(uname.Release))
+	for _, c := range uname.Release {
+		if c == 0 {
+			break
+		}
+		release = append(release, byte(c))
+	}
+
+	return string(release), nil
+}
+
+// parseVFIOAliases parses the modules.alias file and extracts all entries
+// that contain "vfio_pci" in the driver name
+func parseVFIOAliases(content string) []vfioAlias {
+	var aliases []vfioAlias
+
+	lines := strings.Split(content, "\n")
+	for _, line := range lines {
+		line = strings.TrimSpace(line)
+
+		if !strings.HasPrefix(line, "alias vfio_pci:") {
+			continue
+		}
+
+		split := strings.SplitN(line, " ", 3)
+		if len(split) != 3 {
+			logrus.Warnf("malformed modules.alias 'vfio_pci' line: %q", line)
+			continue
+		}
+		modAliasStr := split[1]
+		modAlias, err := parseModAliasString(modAliasStr)
+		if err != nil {
+			logrus.Warnf("failed to parse modalias string %q: %v", modAliasStr, err)
+			continue
+		}
+		logrus.Infof("modalias for %s: %+v", modAliasStr, modAlias)
+
+		driver := split[2]
+		aliases = append(aliases, vfioAlias{
+			modAlias: modAlias,
+			driver:   driver,
+		})
+	}
+
+	return aliases
+}
+
+// findBestMatch finds the best matching VFIO driver for the given modalias
+// by comparing against all available vfio alias patterns
+func findBestMatch(deviceModAlias *modAlias, aliases []vfioAlias) string {
+	var bestDriver string
+	bestSpecificity := -1
+
+	for _, alias := range aliases {
+		if matches, specificity := matchModalias(deviceModAlias, alias.modAlias); matches {
+			if specificity > bestSpecificity {
+				bestDriver = alias.driver
+				bestSpecificity = specificity
+			}
+		}
+	}
+
+	return bestDriver
+}
+
+// matchModalias checks if a device modalias matches a pattern from modules.alias
+// Returns true if it matches and a specificity score (higher is more specific)
+func matchModalias(deviceModAlias, patternModAlias *modAlias) (bool, int) {
+	specificity := 0
+
+	// Compare each field - wildcards in pattern match anything
+	// More specific matches (fewer wildcards) get higher specificity scores
+
+	if matches, score := matchField(deviceModAlias.vendor, patternModAlias.vendor); !matches {
+		return false, 0
+	} else {
+		specificity += score
+	}
+
+	if matches, score := matchField(deviceModAlias.device, patternModAlias.device); !matches {
+		return false, 0
+	} else {
+		specificity += score
+	}
+
+	if matches, score := matchField(deviceModAlias.subvendor, patternModAlias.subvendor); !matches {
+		return false, 0
+	} else {
+		specificity += score
+	}
+
+	if matches, score := matchField(deviceModAlias.subdevice, patternModAlias.subdevice); !matches {
+		return false, 0
+	} else {
+		specificity += score
+	}
+
+	if matches, score := matchField(deviceModAlias.baseClass, patternModAlias.baseClass); !matches {
+		return false, 0
+	} else {
+		specificity += score
+	}
+
+	if matches, score := matchField(deviceModAlias.subClass, patternModAlias.subClass); !matches {
+		return false, 0
+	} else {
+		specificity += score
+	}
+
+	if matches, score := matchField(deviceModAlias.interface_, patternModAlias.interface_); !matches {
+		return false, 0
+	} else {
+		specificity += score
+	}
+
+	return true, specificity
+}
+
+// matchField matches a single field from the modalias
+// Wildcards (*) in the pattern match any value
+// Returns true if it matches and a score (higher for exact matches)
+func matchField(deviceValue, patternValue string) (bool, int) {
+	// Wildcard in pattern matches anything
+	if patternValue == "*" || strings.Contains(patternValue, "*") {
+		return true, 0
+	}
+
+	// Exact match
+	if deviceValue == patternValue {
+		return true, len(patternValue) // Score based on field length
+	}
+
+	// No match
+	return false, 0
+}

internal/nvpci/nvpci.go

@@ -54,15 +54,18 @@ func (w *nvpciWrapper) UnbindFromDriver(dev *nvpci.NvidiaPCIDevice) error {
 }
 
 func (d *nvidiaPCIDevice) bindToVFIODriver() error {
-	// TODO: Instead of always binding to vfio-pci, check if a vfio variant module
-	// should be used instead. This is required for GB200 where the nvgrace-gpu-vfio-pci
-	// module must be used instead of vfio-pci.
-	if d.Driver != vfioPCIDriverName {
+	vfioDriverName, err := d.findBestVFIOVariant()
+	if err != nil {
+		return fmt.Errorf("failed to find best vfio variant driver: %w", err)
+	}
+	logrus.Infof("best vfio variant driver for %s: %s", d.Address, vfioDriverName)
+
+	if d.Driver != vfioDriverName {
 		if err := unbind(d.Address); err != nil {
 			return fmt.Errorf("failed to unbind device %s: %w", d.Address, err)
 		}
-		if err := bind(d.Address, vfioPCIDriverName); err != nil {
-			return fmt.Errorf("failed to bind device %s to %s: %w", d.Address, vfioPCIDriverName, err)
+		if err := bind(d.Address, vfioDriverName); err != nil {
+			return fmt.Errorf("failed to bind device %s to %s: %w", d.Address, vfioDriverName, err)
 		}
 	}
 
@@ -78,15 +81,15 @@ func (d *nvidiaPCIDevice) bindToVFIODriver() error {
 	if err != nil {
 		return fmt.Errorf("failed to get driver for graphics auxiliary device %s: %w", auxDev, err)
 	}
-	if auxDevDriver == vfioPCIDriverName {
+	if auxDevDriver == vfioDriverName {
 		return nil
 	}
 
 	if err := unbind(auxDev); err != nil {
 		return fmt.Errorf("failed to unbind graphics auxiliary device %s: %w", auxDev, err)
 	}
-	if err := bind(auxDev, vfioPCIDriverName); err != nil {
-		return fmt.Errorf("failed to bind graphics auxiliary device %s to %s: %w", auxDev, vfioPCIDriverName, err)
+	if err := bind(auxDev, vfioDriverName); err != nil {
+		return fmt.Errorf("failed to bind graphics auxiliary device %s to %s: %w", auxDev, vfioDriverName, err)
 	}
 
 	return nil
@@ -190,3 +193,57 @@ func unbind(device string) error {
 
 	return nil
 }
+
+/* findBestVFIOVariant:
+ *
+ * Find the "best" match of all vfio_pci aliases for dev in the host
+ * modules.alias file. This uses the algorithm of finding every
+ * modules.alias line that begins with "vfio_pci:", then picking the
+ * one that matches the device's own modalias value (from the file of
+ * that name in the device's sysfs directory) with the fewest
+ * "wildcards" (* character, meaning "match any value for this
+ * attribute").
+ */
+func (d *nvidiaPCIDevice) findBestVFIOVariant() (string, error) {
+	modAliasPath := filepath.Join(d.Path, "modalias")
+	modAliasContent, err := os.ReadFile(modAliasPath)
+	if err != nil {
+		return "", fmt.Errorf("failed to read modalias file for %s: %w", d.Address, err)
+	}
+
+	modAliasStr := strings.TrimSpace(string(modAliasContent))
+	modAlias, err := parseModAliasString(modAliasStr)
+	if err != nil {
+		return "", fmt.Errorf("failed to parse modalias string %q for device %q: %w", modAliasStr, d.Address, err)
+	}
+	logrus.Debugf("modalias for device %q: %+v", d.Address, modAlias)
+
+	kernelVersion, err := getKernelVersion()
+	if err != nil {
+		return "", fmt.Errorf("failed to get kernel version: %w", err)
+	}
+	logrus.Debugf("kernel version: %s", kernelVersion)
+
+	modulesAliasFilePath := filepath.Join("/lib/modules", kernelVersion, "modules.alias")
+	modulesAliasContent, err := os.ReadFile(modulesAliasFilePath)
+	if err != nil {
+		return "", fmt.Errorf("failed to read file %s: %w", modulesAliasFilePath, err)
+	}
+
+	// Parse modules.alias and find all vfio_pci entries
+	vfioAliases := parseVFIOAliases(string(modulesAliasContent))
+	if len(vfioAliases) == 0 {
+		logrus.Warnf("No vfio_pci entries found in modules.alias file, falling back to default vfio-pci driver")
+		return vfioPCIDriverName, nil
+	}
+
+	// Find the best matching VFIO driver for this device
+	bestMatch := findBestMatch(modAlias, vfioAliases)
+	if bestMatch == "" {
+		logrus.Warnf("No matching vfio driver found for device %s, falling back to default vfio-pci", d.Address)
+		return vfioPCIDriverName, nil
+	}
+
+	logrus.Infof("Best VFIO driver for device %s: %s", d.Address, bestMatch)
+	return bestMatch, nil
+}