Enable save/restore with TPUproxy.

This change also adds some small cleanup to TPU code.

PiperOrigin-RevId: 737673712

pkg/sentry/devices/tpuproxy/vfio/BUILD

@@ -12,6 +12,7 @@ go_library(
         "devaddr_set.go",
         "pci_device_fd.go",
         "pci_device_fd_mmap.go",
+        "save_restore.go",
         "tpu_fd.go",
         "tpu_fd_mmap.go",
         "vfio.go",

pkg/sentry/devices/tpuproxy/vfio/pci_device_fd.go

@@ -36,22 +36,35 @@ import (
 )
 
 // pciDeviceFD implements vfs.FileDescriptionImpl for TPU's PCI device.
+//
+// +stateify savable
 type pciDeviceFD struct {
 	vfsfd vfs.FileDescription
 	vfs.FileDescriptionDefaultImpl
 	vfs.DentryMetadataFileDescriptionImpl
 	vfs.NoLockFD
 
+	// If hostFD is -1, this file descriptor has been restored from a save state,
+	// and should be treated as invalid. Any operations on this file descriptor
+	// will effectively be a no-op.
 	hostFD int32
 	queue  waiter.Queue
 
-	mapsMu     sync.Mutex
+	mapsMu sync.Mutex `state:"nosave"`
+	// +checklocks:mapsMu
 	mappings   memmap.MappingSet
 	memmapFile pciDeviceFdMemmapFile
 }
 
+func (fd *pciDeviceFD) isRestored() bool {
+	return fd.hostFD == -1
+}
+
 // Release implements vfs.FileDescriptionImpl.Release.
 func (fd *pciDeviceFD) Release(context.Context) {
+	if fd.isRestored() {
+		return
+	}
 	fdnotifier.RemoveFD(fd.hostFD)
 	fd.queue.Notify(waiter.EventHUp)
 	unix.Close(int(fd.hostFD))
@@ -59,6 +72,9 @@ func (fd *pciDeviceFD) Release(context.Context) {
 
 // EventRegister implements waiter.Waitable.EventRegister.
 func (fd *pciDeviceFD) EventRegister(e *waiter.Entry) error {
+	if fd.isRestored() {
+		return nil
+	}
 	fd.queue.EventRegister(e)
 	if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
 		fd.queue.EventUnregister(e)
@@ -69,6 +85,9 @@ func (fd *pciDeviceFD) EventRegister(e *waiter.Entry) error {
 
 // EventUnregister implements waiter.Waitable.EventUnregister.
 func (fd *pciDeviceFD) EventUnregister(e *waiter.Entry) {
+	if fd.isRestored() {
+		return
+	}
 	fd.queue.EventUnregister(e)
 	if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
 		panic(fmt.Sprint("UpdateFD:", err))
@@ -77,6 +96,9 @@ func (fd *pciDeviceFD) EventUnregister(e *waiter.Entry) {
 
 // Readiness implements waiter.Waitable.Readiness.
 func (fd *pciDeviceFD) Readiness(mask waiter.EventMask) waiter.EventMask {
+	if fd.isRestored() {
+		return 0
+	}
 	return fdnotifier.NonBlockingPoll(fd.hostFD, mask)
 }
 
@@ -87,6 +109,9 @@ func (fd *pciDeviceFD) Epollable() bool {
 
 // Ioctl implements vfs.FileDescriptionImpl.Ioctl.
 func (fd *pciDeviceFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, args arch.SyscallArguments) (uintptr, error) {
+	if fd.isRestored() {
+		return 0, nil
+	}
 	cmd := args[1].Uint()
 
 	t := kernel.TaskFromContext(ctx)
@@ -261,9 +286,11 @@ func (fd *pciDeviceFD) PRead(ctx context.Context, dst usermem.IOSequence, offset
 		return 0, linuxerr.EINVAL
 	}
 	buf := make([]byte, dst.NumBytes())
-	_, err := unix.Pread(int(fd.hostFD), buf, offset)
-	if err != nil {
-		return 0, err
+	if fd.isRestored() {
+		_, err := unix.Pread(int(fd.hostFD), buf, offset)
+		if err != nil {
+			return 0, err
+		}
 	}
 	n, err := dst.CopyOut(ctx, buf)
 	return int64(n), err
@@ -271,6 +298,9 @@ func (fd *pciDeviceFD) PRead(ctx context.Context, dst usermem.IOSequence, offset
 
 // PWrite implements vfs.FileDescriptionImpl.PWrite.
 func (fd *pciDeviceFD) PWrite(ctx context.Context, src usermem.IOSequence, offset int64, opts vfs.WriteOptions) (int64, error) {
+	if fd.isRestored() {
+		return src.NumBytes(), nil
+	}
 	if offset < 0 {
 		return 0, linuxerr.EINVAL
 	}

pkg/sentry/devices/tpuproxy/vfio/pci_device_fd_mmap.go

@@ -68,9 +68,15 @@ func (fd *pciDeviceFD) Translate(ctx context.Context, required, optional memmap.
 
 // InvalidateUnsavable implements memmap.Mappable.InvalidateUnsavable.
 func (fd *pciDeviceFD) InvalidateUnsavable(ctx context.Context) error {
+	fd.mapsMu.Lock()
+	defer fd.mapsMu.Unlock()
+	fd.mappings.InvalidateAll(memmap.InvalidateOpts{InvalidatePrivate: true})
 	return nil
 }
 
+// pciDeviceFdMemmapFile implements memmap.File for /dev/vfio/[0-9]+.
+//
+// +stateify savable
 type pciDeviceFdMemmapFile struct {
 	// FIXME(jamieliu): This is consistent with legacy behavior, but not
 	// clearly correct; drivers/vfio/pci/vfio_pci_core.c:vfio_pci_core_mmap()

pkg/sentry/devices/tpuproxy/vfio/save_restore.go

@@ -0,0 +1,44 @@
+// Copyright 2025 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package vfio
+
+import (
+	"gvisor.dev/gvisor/pkg/context"
+)
+
+// NOTE: TPU save/restore does not work as expected without tight coordination
+// with the application. TPU device state is not saved, all memory mappings
+// are marked as invalid, and TPU/VFIO related file descriptors are stubbed out.
+// Accessing any TPU memory mappings after restore will result in SIGBUS.
+// Issuing any TPU IOCTL command will be a no-op. Reading and writing to TPU
+// FDs will be a no-op.
+//
+// It is up to the application to release all TPU resources before saving and
+// reinitialize them after restoring.
+
+func (fd *tpuFD) beforeSave() {
+	fd.Release(context.Background())
+	fd.hostFD = -1
+}
+
+func (fd *vfioFD) beforeSave() {
+	fd.Release(context.Background())
+	fd.hostFD = -1
+}
+
+func (fd *pciDeviceFD) beforeSave() {
+	fd.Release(context.Background())
+	fd.hostFD = -1
+}

pkg/sentry/devices/tpuproxy/vfio/tpu_fd.go

@@ -52,21 +52,32 @@ var (
 
 // tpuFD implements vfs.FileDescriptionImpl for /dev/vfio/[0-9]+
 //
-// tpuFD is not savable until TPU save/restore is needed.
+// +stateify savable
 type tpuFD struct {
 	vfsfd vfs.FileDescription
 	vfs.FileDescriptionDefaultImpl
 	vfs.DentryMetadataFileDescriptionImpl
 	vfs.NoLockFD
 
-	hostFD     int32
+	// If hostFD is -1, this file descriptor has been restored from a save state,
+	// and should be treated as invalid. Any operations on this file descriptor
+	// will effectively be a no-op.
+	hostFD int32
+
 	device     *tpuDevice
 	queue      waiter.Queue
 	memmapFile tpuFDMemmapFile
 }
 
+func (fd *tpuFD) isRestored() bool {
+	return fd.hostFD == -1
+}
+
 // Release implements vfs.FileDescriptionImpl.Release.
 func (fd *tpuFD) Release(context.Context) {
+	if fd.isRestored() {
+		return
+	}
 	fdnotifier.RemoveFD(fd.hostFD)
 	fd.queue.Notify(waiter.EventHUp)
 	unix.Close(int(fd.hostFD))
@@ -74,6 +85,9 @@ func (fd *tpuFD) Release(context.Context) {
 
 // EventRegister implements waiter.Waitable.EventRegister.
 func (fd *tpuFD) EventRegister(e *waiter.Entry) error {
+	if fd.isRestored() {
+		return nil
+	}
 	fd.queue.EventRegister(e)
 	if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
 		fd.queue.EventUnregister(e)
@@ -84,6 +98,9 @@ func (fd *tpuFD) EventRegister(e *waiter.Entry) error {
 
 // EventUnregister implements waiter.Waitable.EventUnregister.
 func (fd *tpuFD) EventUnregister(e *waiter.Entry) {
+	if fd.isRestored() {
+		return
+	}
 	fd.queue.EventUnregister(e)
 	if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
 		panic(fmt.Sprint("UpdateFD:", err))
@@ -92,6 +109,9 @@ func (fd *tpuFD) EventUnregister(e *waiter.Entry) {
 
 // Readiness implements waiter.Waitable.Readiness.
 func (fd *tpuFD) Readiness(mask waiter.EventMask) waiter.EventMask {
+	if fd.isRestored() {
+		return 0
+	}
 	return fdnotifier.NonBlockingPoll(fd.hostFD, mask)
 }
 
@@ -102,6 +122,9 @@ func (fd *tpuFD) Epollable() bool {
 
 // Ioctl implements vfs.FileDescriptionImpl.Ioctl.
 func (fd *tpuFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, args arch.SyscallArguments) (uintptr, error) {
+	if fd.isRestored() {
+		return 0, nil
+	}
 	cmd := args[1].Uint()
 
 	t := kernel.TaskFromContext(ctx)
@@ -162,6 +185,7 @@ func (fd *tpuFD) getPciDeviceFd(t *kernel.Task, arg hostarch.Addr) (uintptr, fun
 	// See drivers/vfio/group.c:vfio_device_open_file(), the PCI device
 	// is accessed for both reads and writes.
 	vd := t.Kernel().VFS().NewAnonVirtualDentry("[vfio-device]")
+	defer vd.DecRef(t)
 	if err := pciDevFD.vfsfd.Init(pciDevFD, linux.O_RDWR, vd.Mount(), vd.Dentry(), &vfs.FileDescriptionOptions{
 		UseDentryMetadata: true,
 	}); err != nil {

pkg/sentry/devices/tpuproxy/vfio/tpu_fd_mmap.go

@@ -57,6 +57,7 @@ func (fd *tpuFD) InvalidateUnsavable(ctx context.Context) error {
 	return nil
 }
 
+// +stateify savable
 type tpuFDMemmapFile struct {
 	// FIXME(jamieliu): IIUC, tpuFD corresponds to Linux's
 	// drivers/vfio/vfio.c:vfio_group_fops, which does not support mmap at all.

pkg/sentry/devices/tpuproxy/vfio/vfio.go

@@ -51,7 +51,7 @@ var (
 //
 // +stateify savable
 type tpuDevice struct {
-	mu sync.Mutex
+	mu sync.Mutex `state:"nosave"`
 
 	// minor is the device minor number.
 	minor uint32
@@ -66,29 +66,10 @@ func (dev *tpuDevice) Open(ctx context.Context, mnt *vfs.Mount, d *vfs.Dentry, o
 	dev.mu.Lock()
 	defer dev.mu.Unlock()
 
-	var hostFD int
-	if dev.useDevGofer {
-		devClient := devutil.GoferClientFromContext(ctx)
-		if devClient == nil {
-			log.Warningf("devutil.CtxDevGoferClient is not set")
-			return nil, linuxerr.ENOENT
-		}
-		devName := filepath.Join("vfio", strconv.Itoa(int(dev.num)))
-		var err error
-		hostFD, err = devClient.OpenAt(ctx, devName, opts.Flags)
-		if err != nil {
-			ctx.Warningf("tpuDevice: failed to open host %s: %v", devName, err)
-			return nil, err
-		}
-	} else {
-		devPath := filepath.Join("/", "dev", "vfio", strconv.Itoa(int(dev.num)))
-		var err error
-		flags := int(opts.Flags&unix.O_ACCMODE | unix.O_NOFOLLOW)
-		hostFD, err = unix.Openat(-1, devPath, flags, 0)
-		if err != nil {
-			ctx.Warningf("tpuDevice: failed to open host %s: %v", devPath, err)
-			return nil, err
-		}
+	devPath := filepath.Join("vfio", strconv.Itoa(int(dev.num)))
+	hostFD, err := openHostFD(ctx, devPath, opts.Flags, dev.useDevGofer)
+	if err != nil {
+		return nil, err
 	}
 
 	fd := &tpuFD{
@@ -110,6 +91,8 @@ func (dev *tpuDevice) Open(ctx context.Context, mnt *vfs.Mount, d *vfs.Dentry, o
 }
 
 // device implements vfs.Device for /dev/vfio/vfio.
+//
+// +stateify savable
 type vfioDevice struct {
 	// useDevGofer indicates whether to use device gofer to open the VFIO device.
 	useDevGofer bool
@@ -117,33 +100,11 @@ 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) {
-	var hostFD int
-	if dev.useDevGofer {
-		client := devutil.GoferClientFromContext(ctx)
-		if client == nil {
-			log.Warningf("devutil.CtxDevGoferClient is not set")
-			return nil, linuxerr.ENOENT
-		}
-
-		name := filepath.Join("vfio", "vfio")
-		var err error
-		hostFD, err = client.OpenAt(ctx, name, opts.Flags)
-		if err != nil {
-			ctx.Warningf("failed to open host file %s: %v", name, err)
-			return nil, err
-		}
-	} else {
-		devPath := filepath.Join("/", "dev", "vfio", "vfio")
-		flags := int(opts.Flags&unix.O_ACCMODE | unix.O_NOFOLLOW)
-		var err error
-		hostFD, err = unix.Openat(-1, devPath, flags, 0)
-		if err != nil {
-			ctx.Warningf("vfioDevice: failed to open host %s: %v", devPath, err)
-			log.Infof("here failed to open %v flags", devPath, opts.Flags)
-			return nil, err
-		}
+	devPath := filepath.Join("vfio", "vfio")
+	hostFD, err := openHostFD(ctx, devPath, opts.Flags, dev.useDevGofer)
+	if err != nil {
+		return nil, err
 	}
-
 	fd := &vfioFD{
 		hostFD: int32(hostFD),
 		device: dev,
@@ -196,6 +157,20 @@ func RegisterVFIODevice(vfsObj *vfs.VirtualFilesystem, useDevGofer bool) error {
 	})
 }
 
+func openHostFD(ctx context.Context, devName string, flags uint32, useDevGofer bool) (int, error) {
+	if useDevGofer {
+		client := devutil.GoferClientFromContext(ctx)
+		if client == nil {
+			log.Warningf("devutil.CtxDevGoferClient is not set")
+			return -1, linuxerr.ENOENT
+		}
+		return client.OpenAt(ctx, devName, flags)
+	}
+	devPath := filepath.Join("/", "dev", devName)
+	openFlags := int(flags&unix.O_ACCMODE | unix.O_NOFOLLOW)
+	return unix.Openat(-1, devPath, openFlags, 0)
+}
+
 // GetTPUDeviceMajor returns the dynamically allocated major number for the vfio
 // device.
 func GetTPUDeviceMajor(vfsObj *vfs.VirtualFilesystem) (uint32, error) {

pkg/sentry/devices/tpuproxy/vfio/vfio_fd.go

@@ -38,24 +38,37 @@ import (
 )
 
 // deviceFD implements vfs.FileDescriptionImpl for /dev/vfio/vfio.
+//
+// +stateify savable
 type vfioFD struct {
 	vfsfd vfs.FileDescription
 	vfs.FileDescriptionDefaultImpl
 	vfs.DentryMetadataFileDescriptionImpl
 	vfs.NoLockFD
 
-	hostFD     int32
+	// If hostFD is -1, this file descriptor has been restored from a save state,
+	// and should be treated as invalid. Any operations on this file descriptor
+	// will effectively be a no-op.
+	hostFD int32
+
 	device     *vfioDevice
 	queue      waiter.Queue
 	memmapFile vfioFDMemmapFile
 
-	mu sync.Mutex
+	mu sync.Mutex `state:"nosave"`
 	// +checklocks:mu
-	devAddrSet DevAddrSet
+	devAddrSet DevAddrSet `state:"nosave"`
+}
+
+func (fd *vfioFD) isRestored() bool {
+	return fd.hostFD == -1
 }
 
 // Release implements vfs.FileDescriptionImpl.Release.
 func (fd *vfioFD) Release(context.Context) {
+	if fd.isRestored() {
+		return
+	}
 	fd.unpinRange(DevAddrRange{0, ^uint64(0)})
 	fdnotifier.RemoveFD(fd.hostFD)
 	fd.queue.Notify(waiter.EventHUp)
@@ -64,6 +77,9 @@ func (fd *vfioFD) Release(context.Context) {
 
 // EventRegister implements waiter.Waitable.EventRegister.
 func (fd *vfioFD) EventRegister(e *waiter.Entry) error {
+	if fd.isRestored() {
+		return nil
+	}
 	fd.queue.EventRegister(e)
 	if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
 		fd.queue.EventUnregister(e)
@@ -74,6 +90,9 @@ func (fd *vfioFD) EventRegister(e *waiter.Entry) error {
 
 // EventUnregister implements waiter.Waitable.EventUnregister.
 func (fd *vfioFD) EventUnregister(e *waiter.Entry) {
+	if fd.isRestored() {
+		return
+	}
 	fd.queue.EventUnregister(e)
 	if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
 		panic(fmt.Sprint("UpdateFD:", err))
@@ -82,6 +101,9 @@ func (fd *vfioFD) EventUnregister(e *waiter.Entry) {
 
 // Readiness implements waiter.Waitable.Readiness.
 func (fd *vfioFD) Readiness(mask waiter.EventMask) waiter.EventMask {
+	if fd.isRestored() {
+		return 0
+	}
 	return fdnotifier.NonBlockingPoll(fd.hostFD, mask)
 }
 
@@ -92,6 +114,9 @@ func (fd *vfioFD) Epollable() bool {
 
 // Ioctl implements vfs.FileDescriptionImpl.Ioctl.
 func (fd *vfioFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, args arch.SyscallArguments) (uintptr, error) {
+	if fd.isRestored() {
+		return 0, nil
+	}
 	cmd := args[1].Uint()
 	t := kernel.TaskFromContext(ctx)
 	if t == nil {

pkg/sentry/devices/tpuproxy/vfio/vfio_fd_mmap.go

@@ -57,6 +57,7 @@ func (fd *vfioFD) InvalidateUnsavable(ctx context.Context) error {
 	return nil
 }
 
+// +stateify savable
 type vfioFDMemmapFile struct {
 	memmap.NoMapInternal
 

pkg/sentry/fsimpl/sys/BUILD

@@ -22,6 +22,7 @@ go_library(
         "dir_refs.go",
         "kcov.go",
         "pci.go",
+        "save_restore.go",
         "sys.go",
     ],
     visibility = ["//pkg/sentry:internal"],
@@ -31,6 +32,7 @@ go_library(
         "//pkg/context",
         "//pkg/coverage",
         "//pkg/errors/linuxerr",
+        "//pkg/fspath",
         "//pkg/fsutil",
         "//pkg/log",
         "//pkg/refs",

pkg/sentry/fsimpl/sys/save_restore.go

@@ -0,0 +1,121 @@
+// Copyright 2025 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sys
+
+import (
+	"path"
+
+	"gvisor.dev/gvisor/pkg/context"
+	"gvisor.dev/gvisor/pkg/sentry/fsimpl/kernfs"
+	"gvisor.dev/gvisor/pkg/sentry/kernel/auth"
+	"gvisor.dev/gvisor/pkg/sentry/vfs"
+)
+
+// PrepareSave implements vfs.FilesystemImplSaveRestoreExtension.PrepareSave.
+func (fs *filesystem) PrepareSave(ctx context.Context) error {
+	return nil
+}
+
+// CompleteRestore implements
+// vfs.FilesystemImplSaveRestoreExtension.CompleteRestore.
+func (fs *filesystem) CompleteRestore(ctx context.Context, opts vfs.CompleteRestoreOptions) error {
+	// If TPU proxy paths are not enabled, there is nothing to restore. Otherwise,
+	// we need to repopulate the PCI devices and IOMMU groups from a potentially
+	// different host. The easiest way to do that is just rebuild these paths from
+	// scratch.
+	if !fs.enableTPUProxyPaths {
+		return nil
+	}
+	creds := auth.CredentialsFromContext(ctx)
+
+	if err := removeSysDir(ctx, fs.root, "class"); err != nil {
+		return err
+	}
+	if err := removeSysDir(ctx, fs.root, "devices"); err != nil {
+		return err
+	}
+	if err := removeSysDir(ctx, fs.root, "bus"); err != nil {
+		return err
+	}
+	if err := removeSysDir(ctx, fs.root, "kernel"); err != nil {
+		return err
+	}
+
+	classSub := map[string]kernfs.Inode{
+		"power_supply": fs.newDir(ctx, creds, defaultSysDirMode, nil),
+	}
+	devicesSub := map[string]kernfs.Inode{
+		"system": fs.newDir(ctx, creds, defaultSysDirMode, map[string]kernfs.Inode{
+			"cpu": cpuDir(ctx, fs, creds),
+		}),
+	}
+	busSub := make(map[string]kernfs.Inode)
+	kernelSub := kernelDir(ctx, fs, creds)
+
+	deviceToIOMMUGroup, err := pciDeviceIOMMUGroups(path.Join(fs.testSysfsPathPrefix, iommuGroupSysPath))
+	if err != nil {
+		return err
+	}
+	sysDevicesPath := path.Join(fs.testSysfsPathPrefix, sysDevicesMainPath)
+	pciPaths, err := pciDevicePaths(sysDevicesPath)
+	if err != nil {
+		return err
+	}
+	sysDevicesSub, err := fs.mirrorSysDevicesDir(ctx, creds, sysDevicesPath, deviceToIOMMUGroup, pciPaths)
+	if err != nil {
+		return err
+	}
+	for dir, sub := range sysDevicesSub {
+		devicesSub[dir] = sub
+	}
+
+	deviceDirs, err := fs.newDeviceClassDir(ctx, creds, []string{accelDevice, vfioDevice}, sysDevicesPath, pciPaths)
+	if err != nil {
+		return err
+	}
+
+	for tpuDeviceType, symlinkDir := range deviceDirs {
+		classSub[tpuDeviceType] = fs.newDir(ctx, creds, defaultSysDirMode, symlinkDir)
+	}
+	pciDevicesSub, err := fs.newBusPCIDevicesDir(ctx, creds, pciPaths)
+	if err != nil {
+		return err
+	}
+	busSub["pci"] = fs.newDir(ctx, creds, defaultSysDirMode, map[string]kernfs.Inode{
+		"devices": fs.newDir(ctx, creds, defaultSysDirMode, pciDevicesSub),
+	})
+	iommuPath := path.Join(fs.testSysfsPathPrefix, iommuGroupSysPath)
+	iommuGroups, err := fs.mirrorIOMMUGroups(ctx, creds, iommuPath, pciPaths)
+	if err != nil {
+		return err
+	}
+	kernelSub["iommu_groups"] = fs.newDir(ctx, creds, defaultSysDirMode, iommuGroups)
+
+	fs.root.OrderedChildren.Populate(map[string]kernfs.Inode{
+		"class":   fs.newDir(ctx, creds, defaultSysDirMode, classSub),
+		"devices": fs.newDir(ctx, creds, defaultSysDirMode, devicesSub),
+		"bus":     fs.newDir(ctx, creds, defaultSysDirMode, busSub),
+		"kernel":  fs.newDir(ctx, creds, defaultSysDirMode, kernelSub),
+	})
+	return nil
+}
+
+func removeSysDir(ctx context.Context, root *dir, name string) error {
+	dir, err := root.OrderedChildren.Lookup(ctx, name)
+	if err != nil {
+		return err
+	}
+	return root.OrderedChildren.RmDir(ctx, name, dir)
+}

pkg/sentry/fsimpl/sys/sys.go

@@ -70,7 +70,10 @@ type InternalData struct {
 type filesystem struct {
 	kernfs.Filesystem
 
-	devMinor uint32
+	devMinor            uint32
+	enableTPUProxyPaths bool
+	testSysfsPathPrefix string
+	root                *dir
 }
 
 // Name implements vfs.FilesystemType.Name.
@@ -133,6 +136,8 @@ func (fsType FilesystemType) GetFilesystem(ctx context.Context, vfsObj *vfs.Virt
 		idata := opts.InternalData.(*InternalData)
 		productName = idata.ProductName
 		if idata.EnableTPUProxyPaths {
+			fs.enableTPUProxyPaths = true
+			fs.testSysfsPathPrefix = idata.TestSysfsPathPrefix
 			deviceToIOMMUGroup, err := pciDeviceIOMMUGroups(path.Join(idata.TestSysfsPathPrefix, iommuGroupSysPath))
 			if err != nil {
 				return nil, nil, err
@@ -199,6 +204,7 @@ func (fsType FilesystemType) GetFilesystem(ctx context.Context, vfsObj *vfs.Virt
 		"module":   fs.newDir(ctx, creds, defaultSysDirMode, nil),
 		"power":    fs.newDir(ctx, creds, defaultSysDirMode, nil),
 	})
+	fs.root = root.(*dir)
 	var rootD kernfs.Dentry
 	rootD.InitRoot(&fs.Filesystem, root)
 	return fs.VFSFilesystem(), rootD.VFSDentry(), nil

pkg/sentry/fsutil/host_file_mapper.go

@@ -58,6 +58,7 @@ func pagesInChunk(mr memmap.MappableRange, chunkStart uint64) int32 {
 	return int32(mr.Intersect(memmap.MappableRange{chunkStart, chunkStart + chunkSize}).Length() / hostarch.PageSize)
 }
 
+// +stateify savable
 type mapping struct {
 	addr     uintptr
 	writable bool