Add support for v5pod and fix TPU v5 bugs.

V5 support was broken in a few different ways:
- tpu device files (/dev/vfio/X) were created with incorrect minor device nums.
- PCI paths on bus' other than 0000:00 were not supported.
- VFIO unmap was broken and not properly added to the seccomp allowlist.
- The VFIO main device file (/dev/vfio/vfio) did not account for overlapping
  device address ranges that correspond to different VFIO container groups.

Previously TPU support was tested on machines with single TPUs, which masked
most of these issues.

All these issues should be fixed by this change. Tested manually on GKE.

PiperOrigin-RevId: 656037853

images/tpu/jax/Dockerfile.x86_64

@@ -1,4 +1,4 @@
-FROM python:3.8
+FROM python:3.10
 
 RUN pip install --upgrade jax
 RUN pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

pkg/abi/tpu/tpu.go

@@ -41,6 +41,9 @@ const (
 
 	// TPUV5eDeviceID is the PCI device ID of TPU V5e hardware.
 	TPUV5eDeviceID = 0x0063
+
+	// TPUV5pDeviceID is the PCI device ID of TPU V5p hardware.
+	TPUV5pDeviceID = 0x0062
 )
 
 // TPUV4InterruptsMap maps BAR indices to valid register offsets.

pkg/sentry/devices/tpuproxy/device.go

@@ -15,8 +15,8 @@
 package tpuproxy
 
 import (
-	"fmt"
 	"path/filepath"
+	"strconv"
 
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/abi/linux"
@@ -32,6 +32,7 @@ import (
 const (
 	// VFIO_MINOR is the VFIO minor number from include/linux/miscdevice.h.
 	VFIO_MINOR = 196
+
 	// VFIOPath is the path to a VFIO device, it is usually used to
 	// construct a VFIO container.
 	VFIOPath = "/dev/vfio/vfio"
@@ -46,7 +47,10 @@ const (
 type tpuDevice struct {
 	mu sync.Mutex
 
+	// minor is the device minor number.
 	minor uint32
+	// num is the number of the device in the dev filesystem (e.g /dev/vfio/0).
+	num uint32
 }
 
 // Open implements vfs.Device.Open.
@@ -58,10 +62,10 @@ func (dev *tpuDevice) Open(ctx context.Context, mnt *vfs.Mount, d *vfs.Dentry, o
 	}
 	dev.mu.Lock()
 	defer dev.mu.Unlock()
-	devName := fmt.Sprintf("vfio/%d", dev.minor)
+	devName := filepath.Join("vfio", strconv.Itoa(int(dev.num)))
 	hostFD, err := devClient.OpenAt(ctx, devName, opts.Flags)
 	if err != nil {
-		ctx.Warningf("vfioDevice: failed to open host %s: %v", devName, err)
+		ctx.Warningf("tpuDevice: failed to open host %s: %v", devName, err)
 		return nil, err
 	}
 	fd := &tpuFD{
@@ -83,12 +87,7 @@ func (dev *tpuDevice) Open(ctx context.Context, mnt *vfs.Mount, d *vfs.Dentry, o
 }
 
 // device implements vfs.Device for /dev/vfio/vfio.
-type vfioDevice struct {
-	mu sync.Mutex
-
-	// +checklocks:mu
-	devAddrSet DevAddrSet
-}
+type vfioDevice struct{}
 
 // Open implements vfs.Device.Open.
 func (dev *vfioDevice) Open(ctx context.Context, mnt *vfs.Mount, d *vfs.Dentry, opts vfs.OpenOptions) (*vfs.FileDescription, error) {
@@ -98,9 +97,7 @@ func (dev *vfioDevice) Open(ctx context.Context, mnt *vfs.Mount, d *vfs.Dentry,
 		return nil, linuxerr.ENOENT
 	}
 
-	dev.mu.Lock()
-	defer dev.mu.Unlock()
-	name := fmt.Sprintf("vfio/%s", filepath.Base(VFIOPath))
+	name := filepath.Join("vfio", filepath.Base(VFIOPath))
 	hostFD, err := client.OpenAt(ctx, name, opts.Flags)
 	if err != nil {
 		ctx.Warningf("failed to open host file %s: %v", name, err)
@@ -125,9 +122,10 @@ func (dev *vfioDevice) Open(ctx context.Context, mnt *vfs.Mount, d *vfs.Dentry,
 }
 
 // RegisterTPUDevice registers devices implemented by this package in vfsObj.
-func RegisterTPUDevice(vfsObj *vfs.VirtualFilesystem, minor uint32) error {
+func RegisterTPUDevice(vfsObj *vfs.VirtualFilesystem, minor, deviceNum uint32) error {
 	return vfsObj.RegisterDevice(vfs.CharDevice, linux.VFIO_MAJOR, minor, &tpuDevice{
 		minor: minor,
+		num:   deviceNum,
 	}, &vfs.RegisterDeviceOptions{
 		GroupName: tpuDeviceGroupName,
 	})

pkg/sentry/devices/tpuproxy/seccomp_filter.go

@@ -95,6 +95,10 @@ func Filters() seccomp.SyscallRules {
 				seccomp.NonNegativeFD{},
 				seccomp.EqualTo(linux.VFIO_IOMMU_MAP_DMA),
 			},
+			seccomp.PerArg{
+				seccomp.NonNegativeFD{},
+				seccomp.EqualTo(linux.VFIO_IOMMU_UNMAP_DMA),
+			},
 			seccomp.PerArg{
 				seccomp.NonNegativeFD{},
 				seccomp.EqualTo(linux.VFIO_SET_IOMMU),

pkg/sentry/devices/tpuproxy/vfio.go

@@ -16,6 +16,7 @@ package tpuproxy
 
 import (
 	"fmt"
+	"sync"
 
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/abi/linux"
@@ -45,10 +46,15 @@ type vfioFD struct {
 	device     *vfioDevice
 	queue      waiter.Queue
 	memmapFile vfioFDMemmapFile
+
+	mu sync.Mutex
+	// +checklocks:mu
+	devAddrSet DevAddrSet
 }
 
 // Release implements vfs.FileDescriptionImpl.Release.
 func (fd *vfioFD) Release(context.Context) {
+	fd.unpinRange(DevAddrRange{0, ^uint64(0)})
 	fdnotifier.RemoveFD(fd.hostFD)
 	fd.queue.Notify(waiter.EventHUp)
 	unix.Close(int(fd.hostFD))
@@ -198,15 +204,16 @@ func (fd *vfioFD) iommuMapDma(ctx context.Context, t *kernel.Task, arg hostarch.
 	// Unmap the reserved range, which is no longer required.
 	unix.RawSyscall(unix.SYS_MUNMAP, m, uintptr(ar.Length()), 0)
 
-	fd.device.mu.Lock()
-	defer fd.device.mu.Unlock()
+	fd.mu.Lock()
+	defer fd.mu.Unlock()
+	dar := devAddr
 	for _, pr := range prs {
-		rlen := uint64(pr.Source.Length())
-		fd.device.devAddrSet.InsertRange(DevAddrRange{
-			devAddr,
-			devAddr + rlen,
+		r := uint64(pr.Source.Length())
+		fd.devAddrSet.InsertRange(DevAddrRange{
+			dar,
+			dar + r,
 		}, pr)
-		devAddr += rlen
+		dar += r
 	}
 	return n, nil
 }
@@ -221,22 +228,29 @@ func (fd *vfioFD) iommuUnmapDma(ctx context.Context, t *kernel.Task, arg hostarc
 		// gVisor working with TPU.
 		return 0, linuxerr.ENOSYS
 	}
-	n, err := IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_IOMMU_MAP_DMA, &dmaUnmap)
+	n, err := IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_IOMMU_UNMAP_DMA, &dmaUnmap)
 	if err != nil {
 		return 0, nil
 	}
-	fd.device.mu.Lock()
-	defer fd.device.mu.Unlock()
-	s := &fd.device.devAddrSet
+	if _, err := dmaUnmap.CopyOut(t, arg); err != nil {
+		return 0, err
+	}
+
 	r := DevAddrRange{Start: dmaUnmap.IOVa, End: dmaUnmap.IOVa + dmaUnmap.Size}
-	seg := s.LowerBoundSegment(r.Start)
+	fd.unpinRange(r)
+	return n, nil
+}
+
+func (fd *vfioFD) unpinRange(r DevAddrRange) {
+	fd.mu.Lock()
+	defer fd.mu.Unlock()
+	seg := fd.devAddrSet.LowerBoundSegment(r.Start)
 	for seg.Ok() && seg.Start() < r.End {
-		seg = s.Isolate(seg, r)
+		seg = fd.devAddrSet.Isolate(seg, r)
 		mm.Unpin([]mm.PinnedRange{seg.Value()})
-		gap := s.Remove(seg)
+		gap := fd.devAddrSet.Remove(seg)
 		seg = gap.NextSegment()
 	}
-	return n, nil
 }
 
 // VFIO extension.

pkg/sentry/fsimpl/sys/pci.go

@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"path"
 	regex "regexp"
+	"strings"
 
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/abi/linux"
@@ -29,17 +30,19 @@ import (
 )
 
 const (
-	pciMainBusDevicePath = "/sys/devices/pci0000:00"
-	accelDevice          = "accel"
-	vfioDevice           = "vfio-dev"
+	accelDevice        = "accel"
+	vfioDevice         = "vfio-dev"
+	sysDevicesMainPath = "/sys/devices"
 )
 
 var (
-	// Matches PCI device addresses in the main domain.
-	pciDeviceRegex = regex.MustCompile(`0000:([a-fA-F0-9]{2}|[a-fA-F0-9]{4}):[a-fA-F0-9]{2}\.[a-fA-F0-9]{1,2}`)
+	// pciBusRegex matches PCI bus addresses.
+	pciBusRegex = regex.MustCompile(`pci0000:[[:xdigit:]]{2}`)
+	// Matches PCI device addresses.
+	pciDeviceRegex = regex.MustCompile(`0000:([[:xdigit:]]{2}|[[:xdigit:]]{4}):[[:xdigit:]]{2}\.[[:xdigit:]]{1,2}`)
 	// Matches the directories for the main bus (i.e. pci000:00),
 	// individual devices (e.g. 00:00:04.0), accel (TPU v4), and vfio (TPU v5)
-	sysDevicesDirRegex = regex.MustCompile(`pci0000:00|accel|vfio|(0000:([a-fA-F0-9]{2}|[a-fA-F0-9]{4}):[a-fA-F0-9]{2}\.[a-fA-F0-9]{1,2})`)
+	sysDevicesDirRegex = regex.MustCompile(`pci0000:[[:xdigit:]]{2}|accel|vfio|vfio-dev|(0000:([[:xdigit:]]{2}|[[:xdigit:]]{4}):[[:xdigit:]]{2}\.[[:xdigit:]]{1,2})`)
 	// Files allowlisted for host passthrough. These files are read-only.
 	sysDevicesFiles = map[string]any{
 		"vendor": nil, "device": nil, "subsystem_vendor": nil, "subsystem_device": nil,
@@ -49,27 +52,56 @@ var (
 		"is_device_owned": nil, "device_owner": nil, "framework_version": nil,
 		"user_mem_ranges": nil, "interrupt_counts": nil, "chip_model": nil,
 		"bar_offsets": nil, "bar_sizes": nil, "resource0": nil, "resource1": nil,
-		"resource2": nil, "resource3": nil, "resource4": nil, "resource5": nil,
+		"resource2": nil, "resource3": nil, "resource4": nil, "resource5": nil, "enable": nil,
 	}
 )
 
+// sysDevicesPCIPaths returns the paths of all PCI devices on the host in a
+// /sys/devices directory.
+func sysDevicesPCIPaths(sysDevicesPath string) ([]string, error) {
+	sysDevicesDents, err := hostDirEntries(sysDevicesPath)
+	if err != nil {
+		return nil, err
+	}
+	var pciPaths []string
+	for _, dent := range sysDevicesDents {
+		if pciBusRegex.MatchString(dent) {
+			pciDents, err := hostDirEntries(path.Join(sysDevicesPath, dent))
+			if err != nil {
+				return nil, err
+			}
+			for _, pciDent := range pciDents {
+				pciPaths = append(pciPaths, path.Join(sysDevicesPath, dent, pciDent))
+			}
+		}
+	}
+	return pciPaths, nil
+}
+
+// pciBusFromAddress returns the PCI bus address from a PCI address.
+//
+// Preconditions: pciAddr is a valid PCI address.
+func pciBusFromAddress(pciAddr string) string {
+	return strings.Join(strings.Split(pciAddr, ":")[:2], ":")
+}
+
 // Creates TPU devices' symlinks under /sys/class/. TPU device types that are
 // not present on host will be ignored.
 //
 // TPU v4 symlinks are created at /sys/class/accel/accel#.
 // TPU v5 symlinks go to /sys/class/vfio-dev/vfio#.
-func (fs *filesystem) newDeviceClassDir(ctx context.Context, creds *auth.Credentials, tpuDeviceTypes []string, pciMainBusDevicePath string) (map[string]map[string]kernfs.Inode, error) {
+func (fs *filesystem) newDeviceClassDir(ctx context.Context, creds *auth.Credentials, tpuDeviceTypes []string, sysDevicesPath string) (map[string]map[string]kernfs.Inode, error) {
 	dirs := map[string]map[string]kernfs.Inode{}
 	for _, tpuDeviceType := range tpuDeviceTypes {
 		dirs[tpuDeviceType] = map[string]kernfs.Inode{}
 	}
-	pciDents, err := hostDirEntries(pciMainBusDevicePath)
+	pciPaths, err := sysDevicesPCIPaths(sysDevicesPath)
 	if err != nil {
 		return nil, err
 	}
-	for _, pciDent := range pciDents {
+	for _, pciPath := range pciPaths {
 		for _, tpuDeviceType := range tpuDeviceTypes {
-			subPath := path.Join(pciMainBusDevicePath, pciDent, tpuDeviceType)
+			subPath := path.Join(pciPath, tpuDeviceType)
 			deviceDents, err := hostDirEntries(subPath)
 			if err != nil {
 				// Skips the path that doesn't exist.
@@ -81,7 +113,9 @@ func (fs *filesystem) newDeviceClassDir(ctx context.Context, creds *auth.Credent
 			if numOfDeviceDents := len(deviceDents); numOfDeviceDents != 1 {
 				return nil, fmt.Errorf("exactly one entry is expected at %v while there are %d", subPath, numOfDeviceDents)
 			}
-			dirs[tpuDeviceType][deviceDents[0]] = kernfs.NewStaticSymlink(ctx, creds, linux.UNNAMED_MAJOR, fs.devMinor, fs.NextIno(), fmt.Sprintf("../../devices/pci0000:00/%s/%s/%s", pciDent, tpuDeviceType, deviceDents[0]))
+			pciAddr := path.Base(pciPath)
+			pciBus := pciBusFromAddress(pciAddr)
+			dirs[tpuDeviceType][deviceDents[0]] = kernfs.NewStaticSymlink(ctx, creds, linux.UNNAMED_MAJOR, fs.devMinor, fs.NextIno(), fmt.Sprintf("../../devices/pci%s/%s/%s/%s", pciBus, pciAddr, tpuDeviceType, deviceDents[0]))
 		}
 	}
 	if len(dirs) == 0 {
@@ -91,14 +125,16 @@ func (fs *filesystem) newDeviceClassDir(ctx context.Context, creds *auth.Credent
 }
 
 // Create /sys/bus/pci/devices symlinks.
-func (fs *filesystem) newBusPCIDevicesDir(ctx context.Context, creds *auth.Credentials, pciMainBusDevicePath string) (map[string]kernfs.Inode, error) {
+func (fs *filesystem) newBusPCIDevicesDir(ctx context.Context, creds *auth.Credentials, sysDevicesPath string) (map[string]kernfs.Inode, error) {
 	pciDevicesDir := map[string]kernfs.Inode{}
-	pciDents, err := hostDirEntries(pciMainBusDevicePath)
+	pciPaths, err := sysDevicesPCIPaths(sysDevicesPath)
 	if err != nil {
 		return nil, err
 	}
-	for _, pciDent := range pciDents {
-		pciDevicesDir[pciDent] = kernfs.NewStaticSymlink(ctx, creds, linux.UNNAMED_MAJOR, fs.devMinor, fs.NextIno(), fmt.Sprintf("../../../devices/pci0000:00/%s", pciDent))
+	for _, pciPath := range pciPaths {
+		pciAddr := path.Base(pciPath)
+		pciBus := pciBusFromAddress(pciAddr)
+		pciDevicesDir[pciAddr] = kernfs.NewStaticSymlink(ctx, creds, linux.UNNAMED_MAJOR, fs.devMinor, fs.NextIno(), fmt.Sprintf("../../../devices/pci%s/%s", pciBus, pciAddr))
 	}
 
 	return pciDevicesDir, nil
@@ -106,7 +142,7 @@ func (fs *filesystem) newBusPCIDevicesDir(ctx context.Context, creds *auth.Crede
 
 // 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, iommuGroups map[string]string) (map[string]kernfs.Inode, error) {
+func (fs *filesystem) mirrorSysDevicesDir(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 {
@@ -123,7 +159,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, iommuGroups)
+			contents, err := fs.mirrorSysDevicesDir(ctx, creds, dentPath, iommuGroups)
 			if err != nil {
 				return nil, err
 			}

pkg/sentry/fsimpl/sys/sys.go

@@ -132,18 +132,20 @@ func (fsType FilesystemType) GetFilesystem(ctx context.Context, vfsObj *vfs.Virt
 		idata := opts.InternalData.(*InternalData)
 		productName = idata.ProductName
 		if idata.EnableTPUProxyPaths {
-			deviceToIommuGroup, err := pciDeviceIOMMUGroups(path.Join(idata.TestSysfsPathPrefix, iommuGroupSysPath))
+			deviceToIOMMUGroup, err := pciDeviceIOMMUGroups(path.Join(idata.TestSysfsPathPrefix, iommuGroupSysPath))
 			if err != nil {
 				return nil, nil, err
 			}
-			pciPath := path.Join(idata.TestSysfsPathPrefix, pciMainBusDevicePath)
-			pciMainBusSub, err := fs.mirrorPCIBusDeviceDir(ctx, creds, pciPath, deviceToIommuGroup)
+			sysDevicesPath := path.Join(idata.TestSysfsPathPrefix, sysDevicesMainPath)
+			sysDevicesSub, err := fs.mirrorSysDevicesDir(ctx, creds, sysDevicesPath, deviceToIOMMUGroup)
 			if err != nil {
 				return nil, nil, err
 			}
-			devicesSub["pci0000:00"] = fs.newDir(ctx, creds, defaultSysDirMode, pciMainBusSub)
+			for dir, sub := range sysDevicesSub {
+				devicesSub[dir] = sub
+			}
 
-			deviceDirs, err := fs.newDeviceClassDir(ctx, creds, []string{accelDevice, vfioDevice}, pciPath)
+			deviceDirs, err := fs.newDeviceClassDir(ctx, creds, []string{accelDevice, vfioDevice}, sysDevicesPath)
 			if err != nil {
 				return nil, nil, err
 			}
@@ -151,7 +153,7 @@ func (fsType FilesystemType) GetFilesystem(ctx context.Context, vfsObj *vfs.Virt
 			for tpuDeviceType, symlinkDir := range deviceDirs {
 				classSub[tpuDeviceType] = fs.newDir(ctx, creds, defaultSysDirMode, symlinkDir)
 			}
-			pciDevicesSub, err := fs.newBusPCIDevicesDir(ctx, creds, pciPath)
+			pciDevicesSub, err := fs.newBusPCIDevicesDir(ctx, creds, sysDevicesPath)
 			if err != nil {
 				return nil, nil, err
 			}
@@ -282,7 +284,8 @@ func (fs *filesystem) mirrorIOMMUGroups(ctx context.Context, creds *auth.Credent
 			subs[dent] = fs.newHostFile(ctx, creds, defaultSysMode, absPath)
 		case unix.S_IFLNK:
 			if pciDeviceRegex.MatchString(dent) {
-				subs[dent] = kernfs.NewStaticSymlink(ctx, creds, linux.UNNAMED_MAJOR, fs.devMinor, fs.NextIno(), fmt.Sprintf("../../../../devices/pci0000:00/%s", dent))
+				pciBus := pciBusFromAddress(dent)
+				subs[dent] = kernfs.NewStaticSymlink(ctx, creds, linux.UNNAMED_MAJOR, fs.devMinor, fs.NextIno(), fmt.Sprintf("../../../../devices/pci%s/%s", pciBus, dent))
 			}
 		}
 	}

pkg/sentry/fsimpl/sys/sys_test.go

@@ -199,8 +199,12 @@ func (dev PCIDeviceInfo) path() string {
 func TestEnableTPUProxyPathsV5(t *testing.T) {
 	// Set up the fs tree that will be mirrored in the sentry.
 	sysfsTestDir := t.TempDir()
-	pciPath := path.Join(sysfsTestDir, "sys", "devices", "pci0000:00")
-	if err := os.MkdirAll(pciPath, 0755); err != nil {
+	pciPath0 := path.Join(sysfsTestDir, "sys", "devices", "pci0000:00")
+	if err := os.MkdirAll(pciPath0, 0755); err != nil {
+		t.Fatalf("Failed to create PCI directory: %v", err)
+	}
+	pciPath1 := path.Join(sysfsTestDir, "sys", "devices", "pci0000:10")
+	if err := os.MkdirAll(pciPath1, 0755); err != nil {
 		t.Fatalf("Failed to create PCI directory: %v", err)
 	}
 	busPath := path.Join(sysfsTestDir, "sys", "bus", "pci", "devices")
@@ -215,16 +219,22 @@ func TestEnableTPUProxyPathsV5(t *testing.T) {
 	devices := []PCIDeviceInfo{
 		PCIDeviceInfo{
 			group:      "0",
-			pciPath:    pciPath,
+			pciPath:    pciPath0,
 			pciAddress: "0000:00:04.0",
 			name:       "vfio0",
 		},
 		PCIDeviceInfo{
 			group:      "1",
-			pciPath:    pciPath,
+			pciPath:    pciPath0,
 			pciAddress: "0000:00:05.0",
 			name:       "vfio1",
 		},
+		PCIDeviceInfo{
+			group:      "2",
+			pciPath:    pciPath1,
+			pciAddress: "0000:10:05.0",
+			name:       "vfio2",
+		},
 	}
 	for _, device := range devices {
 		devicePath := device.path()
@@ -234,20 +244,20 @@ func TestEnableTPUProxyPathsV5(t *testing.T) {
 		if err := os.Symlink(path.Join("..", "..", "..", device.pciAddress), path.Join(devicePath, "device")); err != nil {
 			t.Fatalf("Failed to symlink device directory: %v", err)
 		}
-		if err := os.Symlink(path.Join("..", "..", "..", "devices", "pci0000:00", device.pciAddress), path.Join(busPath, device.pciAddress)); err != nil {
+		if err := os.Symlink(path.Join("..", "..", "..", "devices", path.Base(device.pciPath), device.pciAddress), path.Join(busPath, device.pciAddress)); err != nil {
 			t.Fatalf("Failed to symlink bus directory: %v", err)
 		}
-		if err := os.Symlink(path.Join("..", "..", "devices", "pci0000:00", device.pciAddress, vfioDev, device.name), path.Join(sysClassPath, device.name)); err != nil {
+		if err := os.Symlink(path.Join("..", "..", "devices", path.Base(device.pciPath), device.pciAddress, vfioDev, device.name), path.Join(sysClassPath, device.name)); err != nil {
 			t.Fatalf("Failed to symlink class directory: %v", err)
 		}
 		iommuPath := path.Join(sysfsTestDir, "sys", "kernel", "iommu_groups", device.group, "devices")
 		if err := os.MkdirAll(iommuPath, 0755); err != nil {
 			t.Fatalf("Failed to create iommu_groups directory: %v", err)
 		}
-		if err := os.Symlink(path.Join("..", "..", "..", "..", "devices", "pci0000:00", device.pciAddress), path.Join(iommuPath, device.pciAddress)); err != nil {
+		if err := os.Symlink(path.Join("..", "..", "..", "..", "devices", path.Base(device.pciPath), device.pciAddress), path.Join(iommuPath, device.pciAddress)); err != nil {
 			t.Fatalf("Failed to symlink iommu_group devices directory: %v", err)
 		}
-		if err := os.Symlink(path.Join("..", "..", "..", "kernel", "iommu_groups", device.group), path.Join(pciPath, device.pciAddress, "iommu_group")); err != nil {
+		if err := os.Symlink(path.Join("..", "..", "..", "kernel", "iommu_groups", device.group), path.Join(device.pciPath, device.pciAddress, "iommu_group")); err != nil {
 			t.Fatalf("Failed to symlink iommu_groups directory: %v", err)
 		}
 	}
@@ -256,13 +266,13 @@ func TestEnableTPUProxyPathsV5(t *testing.T) {
 
 	for _, device := range devices {
 		// Validate PCI device symlinks.
-		pop := s.PathOpAtRoot(path.Join("devices", "pci0000:00", device.pciAddress))
+		pop := s.PathOpAtRoot(path.Join("devices", path.Base(device.pciPath), device.pciAddress))
 		s.AssertAllDirentTypes(s.ListDirents(pop), map[string]testutil.DirentType{
 			"iommu_group": linux.DT_LNK,
 			vfioDev:       linux.DT_DIR,
 		})
 		// Validate VFIO device symlinks.
-		pop = s.PathOpAtRoot(path.Join("devices", "pci0000:00", device.pciAddress, vfioDev, device.name))
+		pop = s.PathOpAtRoot(path.Join("devices", path.Base(device.pciPath), device.pciAddress, vfioDev, device.name))
 		s.AssertAllDirentTypes(s.ListDirents(pop), map[string]testutil.DirentType{
 			"device": linux.DT_LNK,
 		})

runsc/boot/vfs.go

@@ -74,8 +74,9 @@ const (
 const SelfFilestorePrefix = ".gvisor.filestore."
 
 const (
-	pciPathGlobTPUv4 = "/sys/devices/pci0000:00/*/accel/accel*"
-	pciPathGlobTPUv5 = "/sys/devices/pci0000:00/*/vfio-dev/vfio*"
+	pciPathGlobTPUv4   = "/sys/devices/pci0000:*/*/accel/accel*"
+	pciPathGlobTPUv5   = "/sys/devices/pci0000:*/*/vfio-dev/vfio*"
+	iommuGroupPathGlob = "/sys/kernel/iommu_groups/*/devices/*"
 )
 
 // SelfFilestorePath returns the path at which the self filestore file is
@@ -1381,12 +1382,12 @@ func createDeviceFile(ctx context.Context, creds *auth.Credentials, info *contai
 }
 
 // registerTPUDevice registers a TPU device in vfsObj based on the given device ID.
-func registerTPUDevice(vfsObj *vfs.VirtualFilesystem, minor uint32, deviceID int64) error {
+func registerTPUDevice(vfsObj *vfs.VirtualFilesystem, minor, deviceNum uint32, deviceID int64) error {
 	switch deviceID {
 	case tpu.TPUV4DeviceID, tpu.TPUV4liteDeviceID:
 		return accel.RegisterTPUDevice(vfsObj, minor, deviceID == tpu.TPUV4liteDeviceID)
-	case tpu.TPUV5eDeviceID:
-		return tpuproxy.RegisterTPUDevice(vfsObj, minor)
+	case tpu.TPUV5eDeviceID, tpu.TPUV5pDeviceID:
+		return tpuproxy.RegisterTPUDevice(vfsObj, minor, deviceNum)
 	default:
 		return fmt.Errorf("unsupported TPU device with ID: 0x%x", deviceID)
 	}
@@ -1396,8 +1397,8 @@ func registerTPUDevice(vfsObj *vfs.VirtualFilesystem, minor uint32, deviceID int
 // TPU v4 devices are accessible via /sys/devices/pci0000:00/<pci_address>/accel/accel# on the host.
 // TPU v5 devices are accessible via at /sys/devices/pci0000:00/<pci_address>/vfio-dev/vfio# on the host.
 var pathGlobToPathRegex = map[string]string{
-	pciPathGlobTPUv4: `^/sys/devices/pci0000:00/\d+:\d+:\d+\.\d+/accel/accel(\d+)$`,
-	pciPathGlobTPUv5: `^/sys/devices/pci0000:00/\d+:\d+:\d+\.\d+/vfio-dev/vfio(\d+)$`,
+	pciPathGlobTPUv4: `^/sys/devices/pci0000:[[:xdigit:]]{2}/\d+:\d+:\d+\.\d+/accel/accel(\d+)$`,
+	pciPathGlobTPUv5: `^/sys/devices/pci0000:[[:xdigit:]]{2}/\d+:\d+:\d+\.\d+/vfio-dev/vfio(\d+)$`,
 }
 
 func tpuProxyRegisterDevices(info *containerInfo, vfsObj *vfs.VirtualFilesystem) error {
@@ -1429,7 +1430,33 @@ func tpuProxyRegisterDevices(info *containerInfo, vfsObj *vfs.VirtualFilesystem)
 			if err != nil {
 				return fmt.Errorf("parsing PCI device ID: %w", err)
 			}
-			if err := registerTPUDevice(vfsObj, uint32(deviceNum), deviceID); err != nil {
+			// VFIO iommu groups correspond to the device minor number. Use these
+			// paths to get the correct minor number for the sentry-internal TPU
+			// device files.
+			var minorNum int
+			switch deviceID {
+			case tpu.TPUV4DeviceID, tpu.TPUV4liteDeviceID:
+				minorNum = int(deviceNum)
+			case tpu.TPUV5eDeviceID, tpu.TPUV5pDeviceID:
+				groupPaths, err := filepath.Glob(iommuGroupPathGlob)
+				if err != nil {
+					return fmt.Errorf("enumerating IOMMU group files: %w", err)
+				}
+				for _, groupPath := range groupPaths {
+					pci := path.Base(groupPath)
+					if strings.Contains(pciPath, pci) {
+						minor, err := strconv.Atoi(strings.Split(groupPath, "/")[4])
+						if err != nil {
+							return fmt.Errorf("parsing IOMMU group minor number: %w", err)
+						}
+						minorNum = minor
+						break
+					}
+				}
+			default:
+				return fmt.Errorf("unsupported TPU device with ID: 0x%x", deviceID)
+			}
+			if err := registerTPUDevice(vfsObj, uint32(minorNum), uint32(deviceNum), deviceID); err != nil {
 				return fmt.Errorf("registering TPU driver: %w", err)
 			}
 		}

runsc/cmd/chroot.go

@@ -151,24 +151,24 @@ func mountTPUSyslinkInChroot(chroot, dest, relativePath string, validator func(l
 }
 
 func mountTPUDeviceInfoInChroot(chroot, devicePath, sysfsFormat, pciDeviceFormat string) error {
-	deviceNum, valid, err := util.ExtractTpuDeviceMinor(devicePath)
+	deviceMinor, valid, err := util.ExtractTPUDeviceMinor(devicePath)
 	if err != nil {
 		return fmt.Errorf("extracting TPU device minor: %w", err)
 	}
 	if !valid {
 		return nil
 	}
-	// Multiple paths link to the /sys/devices/pci0000:00/<pci_address>
+	// Multiple paths link to the /sys/devices/<pci_bus>/<pci_address>
 	// directory that contains all relevant sysfs accel/vfio device info that we need
 	// bind mounted into the sandbox chroot. We can construct this path by
 	// reading the link below, which points to
-	//   * /sys/devices/pci0000:00/<pci_address>/accel/accel#
-	//   * or /sys/devices/pci0000:00/<pci_address>/vfio-dev/vfio# for VFIO-based TPU
+	//   * /sys/devices/<pci_bus>/<pci_address>/accel/accel#
+	//   * or /sys/devices/<pci_bus>/<pci_address>/vfio-dev/vfio# for VFIO-based TPU
 	// and traversing up 2 directories.
-	// The sysDevicePath itself is a soft link to the deivce directory.
-	sysDevicePath := fmt.Sprintf(sysfsFormat, deviceNum)
+	// The sysDevicePath itself is a soft link to the device directory.
+	sysDevicePath := fmt.Sprintf(sysfsFormat, deviceMinor)
 	sysPCIDeviceDir, err := mountTPUSyslinkInChroot(chroot, sysDevicePath, "../..", func(link string) bool {
-		sysDeviceLinkMatcher := regexp.MustCompile(fmt.Sprintf(pciDeviceFormat, deviceNum))
+		sysDeviceLinkMatcher := regexp.MustCompile(fmt.Sprintf(pciDeviceFormat, deviceMinor))
 		return sysDeviceLinkMatcher.MatchString(link)
 	})
 	if err != nil {
@@ -179,7 +179,8 @@ func mountTPUDeviceInfoInChroot(chroot, devicePath, sysfsFormat, pciDeviceFormat
 	iommuGroupPath := path.Join(sysPCIDeviceDir, "iommu_group")
 	if _, err := os.Stat(iommuGroupPath); err == nil {
 		if _, err := mountTPUSyslinkInChroot(chroot, iommuGroupPath, "", func(link string) bool {
-			return fmt.Sprintf("../../../kernel/iommu_groups/%d", deviceNum) == link
+			iommuGroupPathMatcher := regexp.MustCompile(`../../../kernel/iommu_groups/\d+`)
+			return iommuGroupPathMatcher.MatchString(link)
 		}); err != nil {
 			return err
 		}
@@ -196,11 +197,11 @@ func tpuProxyUpdateChroot(chroot string, spec *specs.Spec, conf *config.Config)
 		"/dev/accel*": "/sys/class/accel/accel%d",
 		"/dev/vfio/*": "/sys/class/vfio-dev/vfio%d"}
 	pathGlobToPciDeviceFormat := map[string]string{
-		"/dev/accel*": `../../devices/pci0000:00/(\d+:\d+:\d+\.\d+)/accel/accel%d`,
-		"/dev/vfio/*": `../../devices/pci0000:00/(\d+:\d+:\d+\.\d+)/vfio-dev/vfio%d`}
+		"/dev/accel*": `../../devices/pci0000:[[:xdigit:]]{2}/(\d+:\d+:\d+\.\d+)/accel/accel%d`,
+		"/dev/vfio/*": `../../devices/pci0000:[[:xdigit:]]{2}/(\d+:\d+:\d+\.\d+)/vfio-dev/vfio%d`}
 	// Bind mount device info directories for all TPU devices on the host.
-	// For v4 TPU, the directory /sys/devices/pci0000:00/<pci_address>/accel/accel# is mounted;
-	// For v5e TPU, the directory /sys/devices/pci0000:00/<pci_address>/vfio-dev/vfio# is mounted.
+	// For v4 TPU, the directory /sys/devices/<pci_bus>/<pci_address>/accel/accel# is mounted;
+	// For v5e TPU, the directory /sys/devices/<pci_bus>/<pci_address>/vfio-dev/vfio# is mounted.
 	for pathGlob, sysfsFormat := range pathGlobToSysfsFormat {
 		paths, err := filepath.Glob(pathGlob)
 		if err != nil {

runsc/cmd/gofer.go

@@ -551,7 +551,7 @@ func shouldExposeVFIODevice(path string) bool {
 //
 // Precondition: tpuproxy is enabled.
 func shouldExposeTpuDevice(path string) bool {
-	_, valid, _ := util.ExtractTpuDeviceMinor(path)
+	_, valid, _ := util.ExtractTPUDeviceMinor(path)
 	return valid || shouldExposeVFIODevice(path)
 }
 

runsc/cmd/util/BUILD

@@ -20,5 +20,6 @@ go_library(
         "//pkg/abi/tpu",
         "//pkg/log",
         "@com_github_google_subcommands//:go_default_library",
+        "@org_golang_x_sys//unix:go_default_library",
     ],
 )

runsc/cmd/util/tpu.go

@@ -20,7 +20,9 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"syscall"
 
+	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/abi/tpu"
 )
 
@@ -35,12 +37,12 @@ const (
 )
 
 var tpuV4DeviceIDs = map[uint64]any{tpu.TPUV4DeviceID: nil, tpu.TPUV4liteDeviceID: nil}
-var tpuV5DeviceIDs = map[uint64]any{tpu.TPUV5eDeviceID: nil}
+var tpuV5DeviceIDs = map[uint64]any{tpu.TPUV5eDeviceID: nil, tpu.TPUV5pDeviceID: nil}
 
-// ExtractTpuDeviceMinor returns the accelerator device minor number for that
+// ExtractTPUDeviceMinor returns the accelerator device minor number for that
 // the passed device path. If the passed device is not a valid TPU device, then
 // it returns false.
-func ExtractTpuDeviceMinor(path string) (uint32, bool, error) {
+func ExtractTPUDeviceMinor(path string) (uint32, bool, error) {
 	devNum, valid, err := tpuV4DeviceMinor(path)
 	if err != nil {
 		return 0, false, err
@@ -60,10 +62,11 @@ func tpuDeviceMinor(devicePath, devicePathRegex, sysfsFormat string, allowedDevi
 	if matches == nil {
 		return 0, false, nil
 	}
-	minor, err := strconv.ParseUint(matches[1], 10, 32)
-	if err != nil {
-		return 0, false, fmt.Errorf("invalid host device file %q: %w", devicePath, err)
+	var st syscall.Stat_t
+	if err := syscall.Stat(devicePath, &st); err != nil {
+		return 0, false, err
 	}
+	minor := unix.Minor(st.Rdev)
 	vendor, err := readHexInt(fmt.Sprintf(sysfsFormat, minor, vendorFile))
 	if err != nil {
 		return 0, false, err
@@ -78,7 +81,7 @@ func tpuDeviceMinor(devicePath, devicePathRegex, sysfsFormat string, allowedDevi
 	if _, ok := allowedDeviceIDs[deviceID]; !ok {
 		return 0, false, nil
 	}
-	return uint32(minor), true, nil
+	return minor, true, nil
 }
 
 // tpuv4DeviceMinor returns v4 and v4lite TPU device minor number for the given path.