Create IOMMU symlink(s) for TPU devices.

The symlinks mirrors the relationship between TPU devices and its respective
IOMMU group and on the host.

Before the change, iommu_group has been treated as a normal file, which doesn't
create such symlink for the devices.

PiperOrigin-RevId: 597939650

pkg/sentry/fsimpl/sys/pci.go

@@ -43,7 +43,7 @@ var (
 	// Files allowlisted for host passthrough. These files are read-only.
 	sysDevicesFiles = map[string]any{
 		"vendor": nil, "device": nil, "subsystem_vendor": nil, "subsystem_device": nil,
-		"revision": nil, "class": nil, "numa_node": nil, "iommu_group": nil,
+		"revision": nil, "class": nil, "numa_node": nil,
 		"resource": nil, "pci_address": nil, "dev": nil, "driver_version": nil,
 		"reset_count": nil, "write_open_count": nil, "status": nil,
 		"is_device_owned": nil, "device_owner": nil, "framework_version": nil,
@@ -102,7 +102,7 @@ func (fs *filesystem) newPCIDevicesDir(ctx context.Context, creds *auth.Credenti
 
 // Recursively build out sysfs directories according to the allowlisted files,
 // directories, and symlinks defined in this package.
-func (fs *filesystem) mirrorPCIBusDeviceDir(ctx context.Context, creds *auth.Credentials, dir string) (map[string]kernfs.Inode, error) {
+func (fs *filesystem) mirrorPCIBusDeviceDir(ctx context.Context, creds *auth.Credentials, dir string, iommuGroups map[string]string) (map[string]kernfs.Inode, error) {
 	subs := map[string]kernfs.Inode{}
 	dents, err := hostDirEntries(dir)
 	if err != nil {
@@ -119,7 +119,7 @@ func (fs *filesystem) mirrorPCIBusDeviceDir(ctx context.Context, creds *auth.Cre
 			if match := sysDevicesDirRegex.MatchString(dent); !match {
 				continue
 			}
-			contents, err := fs.mirrorPCIBusDeviceDir(ctx, creds, dentPath)
+			contents, err := fs.mirrorPCIBusDeviceDir(ctx, creds, dentPath, iommuGroups)
 			if err != nil {
 				return nil, err
 			}
@@ -129,22 +129,44 @@ func (fs *filesystem) mirrorPCIBusDeviceDir(ctx context.Context, creds *auth.Cre
 				subs[dent] = fs.newHostFile(ctx, creds, defaultSysMode, dentPath)
 			}
 		case unix.S_IFLNK:
-			// Both the device and PCI address entries are links to the original PCI
-			// device directory that's at the same place earlier in the dir tree.
-			if match := pciDeviceRegex.MatchString(dent); !(match || dent == "device") {
+			linkContent := ""
+			switch {
+			case pciDeviceRegex.MatchString(dent) || dent == "device":
+				pciDeviceName, err := pciDeviceName(dir)
+				if err != nil {
+					return nil, err
+				}
+				// Both the device and PCI address entries are links to the original PCI
+				// device directory that's at the same place earlier in the dir tree.
+				linkContent = fmt.Sprintf("../../../%s", pciDeviceName)
+			case dent == "iommu_group":
+				pciDeviceName, err := pciDeviceName(dir)
+				if err != nil {
+					return nil, err
+				}
+				iommuGroupNum, exist := iommuGroups[pciDeviceName]
+				if !exist {
+					return nil, fmt.Errorf("no IOMMU group is found for device %v", pciDeviceName)
+				}
+				linkContent = fmt.Sprintf("../../../kernel/iommu_groups/%s", iommuGroupNum)
+			default:
 				continue
 			}
-			pciDeviceName := pciDeviceRegex.FindString(dir)
-			if pciDeviceName == "" {
-				return nil, fmt.Errorf("could not populate sysfs pci symlink %s", dir)
-			}
-			linkContent := fmt.Sprintf("../../../%s", pciDeviceName)
 			subs[dent] = kernfs.NewStaticSymlink(ctx, creds, linux.UNNAMED_MAJOR, fs.devMinor, fs.NextIno(), linkContent)
 		}
 	}
 	return subs, nil
 }
 
+// Infer a PCI device's name from its path.
+func pciDeviceName(pciDevicePath string) (string, error) {
+	pciDeviceName := pciDeviceRegex.FindString(pciDevicePath)
+	if pciDeviceName == "" {
+		return "", fmt.Errorf("no valid device name for the device path at %v", pciDevicePath)
+	}
+	return pciDeviceName, nil
+}
+
 func hostFileMode(path string) (uint32, error) {
 	fd, err := unix.Openat(-1, path, unix.O_RDONLY|unix.O_NOFOLLOW|unix.O_PATH, 0)
 	if err != nil {

pkg/sentry/fsimpl/sys/sys.go

@@ -19,6 +19,7 @@ import (
 	"bytes"
 	"fmt"
 	"os"
+	"path"
 	"strconv"
 
 	"golang.org/x/sys/unix"
@@ -39,6 +40,7 @@ const (
 	defaultSysMode           = linux.FileMode(0444)
 	defaultSysDirMode        = linux.FileMode(0755)
 	defaultMaxCachedDentries = uint64(1000)
+	iommuGroupSysPath        = "/sys/kernel/iommu_groups/"
 )
 
 // FilesystemType implements vfs.FilesystemType.
@@ -127,7 +129,11 @@ func (fsType FilesystemType) GetFilesystem(ctx context.Context, vfsObj *vfs.Virt
 		idata := opts.InternalData.(*InternalData)
 		productName = idata.ProductName
 		if idata.EnableTPUProxyPaths {
-			pciMainBusSub, err := fs.mirrorPCIBusDeviceDir(ctx, creds, pciMainBusDevicePath)
+			deviceToIommuGroup, err := pciDeviceIOMMUGroups(iommuGroupSysPath)
+			if err != nil {
+				return nil, nil, err
+			}
+			pciMainBusSub, err := fs.mirrorPCIBusDeviceDir(ctx, creds, pciMainBusDevicePath, deviceToIommuGroup)
 			if err != nil {
 				return nil, nil, err
 			}
@@ -197,6 +203,34 @@ func cpuDir(ctx context.Context, fs *filesystem, creds *auth.Credentials) kernfs
 	return fs.newDir(ctx, creds, defaultSysDirMode, children)
 }
 
+// Returns a map from a PCI device name to its IOMMU group if available.
+func pciDeviceIOMMUGroups(iommuGroupsPath string) (map[string]string, error) {
+	// IOMMU groups are organizd as iommu_group_path/$GROUP, where $GROUP is
+	// the IOMMU group number of which the device is a memeber.
+	iommuGroupNums, err := hostDirEntries(iommuGroupsPath)
+	if err != nil {
+		// When IOMMU is not enabled, skip the rest of the process.
+		if err == unix.ENOENT {
+			return nil, nil
+		}
+		return nil, err
+	}
+	// The returned map from PCI device name to its IOMMU group.
+	iommuGroups := map[string]string{}
+	for _, iommuGroupNum := range iommuGroupNums {
+		groupDevicesPath := path.Join(iommuGroupsPath, iommuGroupNum, "devices")
+		pciDeviceNames, err := hostDirEntries(groupDevicesPath)
+		if err != nil {
+			return nil, err
+		}
+		// An IOMMU group may include multiple devices.
+		for _, pciDeviceName := range pciDeviceNames {
+			iommuGroups[pciDeviceName] = iommuGroupNum
+		}
+	}
+	return iommuGroups, nil
+}
+
 func kernelDir(ctx context.Context, fs *filesystem, creds *auth.Credentials) kernfs.Inode {
 	// Set up /sys/kernel/debug/kcov. Technically, debugfs should be
 	// mounted at debug/, but for our purposes, it is sufficient to keep it