Add IOMMUFD support

pkg/nvpci/mock.go

@@ -84,6 +84,13 @@ func (m *MockNvpci) AddMockA100(address string, numaNode int, sriov *SriovInfo)
 		return err
 	}
 
+	vfioDev := filepath.Join(deviceDir, "vfio-dev")
+	vfioFD := filepath.Join(vfioDev, "vfio8")
+	err = os.MkdirAll(vfioFD, 0755)
+	if err != nil {
+		return err
+	}
+
 	iommuGroup := 20
 	_, err = os.Create(filepath.Join(deviceDir, strconv.Itoa(iommuGroup)))
 	if err != nil {

pkg/nvpci/nvpci.go

@@ -112,6 +112,7 @@ type NvidiaPCIDevice struct {
 	Device     uint16
 	DeviceName string
 	Driver     string
+	IommuFD    int
 	IommuGroup int
 	NumaNode   int
 	Config     *ConfigSpace
@@ -290,6 +291,11 @@ func (p *nvpci) getGPUByPciBusID(address string, cache map[string]*NvidiaPCIDevi
 		return nil, fmt.Errorf("unable to detect IOMMU group for %s: %w", address, err)
 	}
 
+	iommuFD, err := getIOMMUFD(devicePath)
+	if err != nil {
+		return nil, fmt.Errorf("unable to get IOMMU FD for %s: %w", address, err)
+	}
+
 	numa, err := os.ReadFile(path.Join(devicePath, "numa_node"))
 	if err != nil {
 		return nil, fmt.Errorf("unable to read PCI NUMA node for %s: %v", address, err)
@@ -373,6 +379,7 @@ func (p *nvpci) getGPUByPciBusID(address string, cache map[string]*NvidiaPCIDevi
 		Class:      uint32(classID),
 		Device:     uint16(deviceID),
 		Driver:     driver,
+		IommuFD:    int(iommuFD),
 		IommuGroup: int(iommuGroup),
 		NumaNode:   int(numaNode),
 		Config:     config,
@@ -521,6 +528,31 @@ func getDriver(devicePath string) (string, error) {
 	return "", err
 }
 
+// /sys/bus/pci/devices/0000:df:00.0/vfio-dev/vfio0.
+func getIOMMUFD(devicePath string) (int, error) {
+	vfioDevDir := filepath.Join(devicePath, "vfio-dev")
+	// Read the directory; expect exactly one entry
+	entries, err := os.ReadDir(vfioDevDir)
+	switch {
+	case os.IsNotExist(err):
+		return -1, nil
+	case err == nil:
+		if len(entries) == 0 {
+			return -1, fmt.Errorf("no VFIO device file found in %s", vfioDevDir)
+		}
+		name := entries[0].Name()
+		// Strip the "vfio" prefix to get the numeric part
+		idxStr := strings.TrimPrefix(name, "vfio")
+		idx, err := strconv.Atoi(idxStr)
+		if err != nil {
+			return -1, fmt.Errorf("failed to parse VFIO index from %q: %w", name, err)
+		}
+
+		return idx, nil
+	}
+	return -1, err
+}
+
 func getIOMMUGroup(devicePath string) (int64, error) {
 	var iommuGroup int64
 	iommu, err := filepath.EvalSymlinks(path.Join(devicePath, "iommu_group"))

pkg/nvpci/nvpci_test.go

@@ -74,6 +74,33 @@ func TestNvpci(t *testing.T) {
 	_, err = nvpci.GetGPUByIndex(1)
 	require.Error(t, err, "No error returned when getting GPU at invalid index")
 }
+func TestNvpciIOMMUFD(t *testing.T) {
+	testCases := []struct {
+		Description string
+		IOMMUFD     int
+	}{
+		{
+			Description: "IOMMUFD 8",
+			IOMMUFD:     8,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.Description, func(t *testing.T) {
+			nvpci, err := NewMockNvpci()
+			require.Nil(t, err, "Error creating NewMockNvpci")
+			defer nvpci.Cleanup()
+
+			err = nvpci.AddMockA100("0000:80:05.1", 0, nil)
+			require.Nil(t, err, "Error adding Mock A100 device to MockNvpci")
+
+			devices, err := nvpci.GetGPUs()
+			require.Nil(t, err, "Error getting GPUs")
+			require.Equal(t, 1, len(devices), "Wrong number of GPU devices")
+			require.Equal(t, 8, devices[0].IommuFD, "Wrong IOMMUFD found for device")
+		})
+	}
+}
 
 func TestNvpciNUMANode(t *testing.T) {
 	testCases := []struct {