Reorganize tpuproxy code.

Now there is a main tpuproxy directory and two subdirs for each of the
different proxy implementations.

PiperOrigin-RevId: 667688748

pkg/sentry/devices/accel/seccomp_filters.go

@@ -1,83 +0,0 @@
-// Copyright 2023 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 accel
-
-import (
-	"golang.org/x/sys/unix"
-	"gvisor.dev/gvisor/pkg/abi/gasket"
-	"gvisor.dev/gvisor/pkg/abi/linux"
-	"gvisor.dev/gvisor/pkg/seccomp"
-)
-
-// Filters returns seccomp-bpf filters for this package.
-func Filters() seccomp.SyscallRules {
-	return seccomp.MakeSyscallRules(map[uintptr]seccomp.SyscallRule{
-		unix.SYS_OPENAT: seccomp.PerArg{
-			// All paths that we openat() are absolute, so we pass a dirfd
-			// of -1 (which is invalid for relative paths, but ignored for
-			// absolute paths) to hedge against bugs involving AT_FDCWD or
-			// real dirfds.
-			seccomp.EqualTo(^uintptr(0)),
-			seccomp.AnyValue{},
-			seccomp.MaskedEqual(unix.O_CREAT|unix.O_NOFOLLOW, unix.O_NOFOLLOW),
-			seccomp.AnyValue{},
-		},
-		unix.SYS_GETDENTS64: seccomp.MatchAll{},
-		unix.SYS_IOCTL: seccomp.Or{
-			seccomp.PerArg{
-				seccomp.NonNegativeFD{},
-				seccomp.EqualTo(gasket.GASKET_IOCTL_RESET),
-			},
-			seccomp.PerArg{
-				seccomp.NonNegativeFD{},
-				seccomp.EqualTo(gasket.GASKET_IOCTL_MAP_BUFFER),
-			},
-			seccomp.PerArg{
-				seccomp.NonNegativeFD{},
-				seccomp.EqualTo(gasket.GASKET_IOCTL_UNMAP_BUFFER),
-			},
-			seccomp.PerArg{
-				seccomp.NonNegativeFD{},
-				seccomp.EqualTo(gasket.GASKET_IOCTL_CLEAR_INTERRUPT_COUNTS),
-			},
-			seccomp.PerArg{
-				seccomp.NonNegativeFD{},
-				seccomp.EqualTo(gasket.GASKET_IOCTL_REGISTER_INTERRUPT),
-			},
-			seccomp.PerArg{
-				seccomp.NonNegativeFD{},
-				seccomp.EqualTo(gasket.GASKET_IOCTL_UNREGISTER_INTERRUPT),
-			},
-		},
-		unix.SYS_EVENTFD2: seccomp.Or{
-			seccomp.PerArg{
-				seccomp.AnyValue{},
-				seccomp.EqualTo(linux.EFD_NONBLOCK),
-			},
-			seccomp.PerArg{
-				seccomp.AnyValue{},
-				seccomp.EqualTo(linux.EFD_NONBLOCK | linux.EFD_SEMAPHORE),
-			},
-		},
-		unix.SYS_MREMAP: seccomp.PerArg{
-			seccomp.AnyValue{},
-			seccomp.EqualTo(0), /* old_size */
-			seccomp.AnyValue{},
-			seccomp.EqualTo(linux.MREMAP_MAYMOVE | linux.MREMAP_FIXED),
-			seccomp.AnyValue{},
-			seccomp.EqualTo(0),
-		},
-	})
-}

pkg/sentry/devices/tpuproxy/BUILD

@@ -1,5 +1,4 @@
-load("//tools:defs.bzl", "go_library")
-load("//tools/go_generics:defs.bzl", "go_template_instance")
+load("//tools:defs.bzl", "go_library", "go_test")
 
 package(default_applicable_licenses = ["//:license"])
 
@@ -8,70 +7,26 @@ licenses(["notice"])
 go_library(
     name = "tpuproxy",
     srcs = [
-        "devaddr_range.go",
-        "devaddr_set.go",
-        "device.go",
-        "ioctl_unsafe.go",
         "seccomp_filter.go",
-        "tpu.go",
-        "tpu_mmap.go",
-        "vfio.go",
-        "vfio_mmap.go",
+        "tpuproxy.go",
     ],
     visibility = [
         "//pkg/sentry:internal",
     ],
     deps = [
+        "//pkg/abi/gasket",
         "//pkg/abi/linux",
-        "//pkg/cleanup",
-        "//pkg/context",
-        "//pkg/devutil",
-        "//pkg/errors/linuxerr",
-        "//pkg/fdnotifier",
-        "//pkg/hostarch",
-        "//pkg/log",
-        "//pkg/marshal/primitive",
-        "//pkg/safemem",
+        "//pkg/abi/tpu",
         "//pkg/seccomp",
-        "//pkg/sentry/arch",
-        "//pkg/sentry/fsimpl/eventfd",
-        "//pkg/sentry/fsimpl/kernfs",
-        "//pkg/sentry/kernel",
-        "//pkg/sentry/memmap",
-        "//pkg/sentry/mm",
+        "//pkg/sentry/devices/tpuproxy/accel",
+        "//pkg/sentry/devices/tpuproxy/vfio",
         "//pkg/sentry/vfs",
-        "//pkg/sync",
-        "//pkg/usermem",
-        "//pkg/waiter",
-        "@org_golang_x_exp//constraints:go_default_library",
         "@org_golang_x_sys//unix:go_default_library",
     ],
 )
 
-go_template_instance(
-    name = "devaddr_range",
-    out = "devaddr_range.go",
-    package = "tpuproxy",
-    prefix = "DevAddr",
-    template = "//pkg/segment:generic_range",
-    types = {
-        "T": "uint64",
-    },
-)
-
-go_template_instance(
-    name = "devaddr_set",
-    out = "devaddr_set.go",
-    imports = {
-        "mm": "gvisor.dev/gvisor/pkg/sentry/mm",
-    },
-    package = "tpuproxy",
-    prefix = "DevAddr",
-    template = "//pkg/segment:generic_set",
-    types = {
-        "Key": "uint64",
-        "Range": "DevAddrRange",
-        "Value": "mm.PinnedRange",
-        "Functions": "devAddrSetFuncs",
-    },
+go_test(
+    name = "tpuproxy_test",
+    srcs = ["tpuproxy_test.go"],
+    library = ":tpuproxy",
 )

pkg/sentry/devices/tpuproxy/accel/BUILD

@@ -8,13 +8,12 @@ licenses(["notice"])
 go_library(
     name = "accel",
     srcs = [
+        "accel.go",
+        "accel_fd.go",
+        "accel_fd_mmap.go",
         "devaddr_range.go",
         "devaddr_set.go",
-        "device.go",
-        "gasket.go",
-        "seccomp_filters.go",
-        "tpu_v4.go",
-        "tpu_v4_mmap.go",
+        "gasket_ioctl.go",
     ],
     visibility = ["//pkg/sentry:internal"],
     deps = [
@@ -31,7 +30,7 @@ go_library(
         "//pkg/safemem",
         "//pkg/seccomp",
         "//pkg/sentry/arch",
-        "//pkg/sentry/devices/tpuproxy",
+        "//pkg/sentry/devices/tpuproxy/util",
         "//pkg/sentry/fsimpl/eventfd",
         "//pkg/sentry/kernel",
         "//pkg/sentry/memmap",

pkg/sentry/devices/tpuproxy/accel/accel.go

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// Package accel implements a proxy for gasket based accel devices.
 package accel
 
 import (
@@ -29,10 +30,10 @@ import (
 	"gvisor.dev/gvisor/pkg/sync"
 )
 
-// tpuV4Device implements vfs.Device for /dev/accel[0-9]+.
+// accelDevice implements vfs.Device for /dev/accel[0-9]+.
 //
 // +stateify savable
-type tpuV4Device struct {
+type accelDevice struct {
 	mu sync.Mutex
 
 	minor uint32
@@ -45,7 +46,7 @@ type tpuV4Device struct {
 	owner *kernel.ThreadGroup
 }
 
-func (dev *tpuV4Device) Open(ctx context.Context, mnt *vfs.Mount, vfsd *vfs.Dentry, opts vfs.OpenOptions) (*vfs.FileDescription, error) {
+func (dev *accelDevice) Open(ctx context.Context, mnt *vfs.Mount, vfsd *vfs.Dentry, opts vfs.OpenOptions) (*vfs.FileDescription, error) {
 	devClient := devutil.GoferClientFromContext(ctx)
 	if devClient == nil {
 		log.Warningf("devutil.CtxDevGoferClient is not set")
@@ -59,7 +60,7 @@ func (dev *tpuV4Device) Open(ctx context.Context, mnt *vfs.Mount, vfsd *vfs.Dent
 		ctx.Warningf("accelDevice: failed to open device %s: %v", name, err)
 		return nil, err
 	}
-	fd := &tpuV4FD{
+	fd := &accelFD{
 		hostFD: int32(hostFD),
 		device: dev,
 	}
@@ -89,7 +90,7 @@ func (dev *tpuV4Device) Open(ctx context.Context, mnt *vfs.Mount, vfsd *vfs.Dent
 
 // RegisterTPUDevice registers all devices implemented by this package in vfsObj.
 func RegisterTPUDevice(vfsObj *vfs.VirtualFilesystem, minor uint32, lite bool) error {
-	return vfsObj.RegisterDevice(vfs.CharDevice, linux.ACCEL_MAJOR, minor, &tpuV4Device{
+	return vfsObj.RegisterDevice(vfs.CharDevice, linux.ACCEL_MAJOR, minor, &accelDevice{
 		lite:  lite,
 		minor: minor,
 	}, &vfs.RegisterDeviceOptions{

pkg/sentry/devices/tpuproxy/accel/accel_fd.go

@@ -12,7 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Package accel implements proxying for hardware accelerators.
 package accel
 
 import (
@@ -27,7 +26,7 @@ import (
 	"gvisor.dev/gvisor/pkg/hostarch"
 	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/sentry/arch"
-	"gvisor.dev/gvisor/pkg/sentry/devices/tpuproxy"
+	"gvisor.dev/gvisor/pkg/sentry/devices/tpuproxy/util"
 	"gvisor.dev/gvisor/pkg/sentry/kernel"
 	"gvisor.dev/gvisor/pkg/sentry/mm"
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
@@ -35,24 +34,24 @@ import (
 	"gvisor.dev/gvisor/pkg/waiter"
 )
 
-// tpuV4FD implements vfs.FileDescriptionImpl for /dev/accel[0-9]+.
+// accelFD implements vfs.FileDescriptionImpl for /dev/accel[0-9]+.
 //
 // accelFD is not savable; we do not implement save/restore of accelerator
 // state.
-type tpuV4FD struct {
+type accelFD struct {
 	vfsfd vfs.FileDescription
 	vfs.FileDescriptionDefaultImpl
 	vfs.DentryMetadataFileDescriptionImpl
 	vfs.NoLockFD
 
 	hostFD     int32
-	device     *tpuV4Device
+	device     *accelDevice
 	queue      waiter.Queue
 	memmapFile accelFDMemmapFile
 }
 
 // Release implements vfs.FileDescriptionImpl.Release.
-func (fd *tpuV4FD) Release(context.Context) {
+func (fd *accelFD) Release(context.Context) {
 	fd.device.mu.Lock()
 	defer fd.device.mu.Unlock()
 	fd.device.openWriteFDs--
@@ -68,7 +67,7 @@ func (fd *tpuV4FD) Release(context.Context) {
 				Size:           r.End - r.Start,
 				HostAddress:    0,
 			}
-			_, err := tpuproxy.IOCTLInvokePtrArg[gasket.Ioctl](fd.hostFD, gasket.GASKET_IOCTL_UNMAP_BUFFER, &gpti)
+			_, err := util.IOCTLInvokePtrArg[gasket.Ioctl](fd.hostFD, gasket.GASKET_IOCTL_UNMAP_BUFFER, &gpti)
 			if err != nil {
 				log.Warningf("could not unmap range [%#x, %#x) (index %d) on device: %v", r.Start, r.End, v.pageTableIndex, err)
 			}
@@ -83,7 +82,7 @@ func (fd *tpuV4FD) Release(context.Context) {
 }
 
 // EventRegister implements waiter.Waitable.EventRegister.
-func (fd *tpuV4FD) EventRegister(e *waiter.Entry) error {
+func (fd *accelFD) EventRegister(e *waiter.Entry) error {
 	fd.queue.EventRegister(e)
 	if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
 		fd.queue.EventUnregister(e)
@@ -93,7 +92,7 @@ func (fd *tpuV4FD) EventRegister(e *waiter.Entry) error {
 }
 
 // EventUnregister implements waiter.Waitable.EventUnregister.
-func (fd *tpuV4FD) EventUnregister(e *waiter.Entry) {
+func (fd *accelFD) EventUnregister(e *waiter.Entry) {
 	fd.queue.EventUnregister(e)
 	if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
 		panic(fmt.Sprint("UpdateFD:", err))
@@ -101,17 +100,17 @@ func (fd *tpuV4FD) EventUnregister(e *waiter.Entry) {
 }
 
 // Readiness implements waiter.Waitable.Readiness.
-func (fd *tpuV4FD) Readiness(mask waiter.EventMask) waiter.EventMask {
+func (fd *accelFD) Readiness(mask waiter.EventMask) waiter.EventMask {
 	return fdnotifier.NonBlockingPoll(fd.hostFD, mask)
 }
 
 // Epollable implements vfs.FileDescriptionImpl.Epollable.
-func (fd *tpuV4FD) Epollable() bool {
+func (fd *accelFD) Epollable() bool {
 	return true
 }
 
 // Ioctl implements vfs.FileDescriptionImpl.Ioctl.
-func (fd *tpuV4FD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, args arch.SyscallArguments) (uintptr, error) {
+func (fd *accelFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, args arch.SyscallArguments) (uintptr, error) {
 	cmd := args[1].Uint()
 	argPtr := args[2].Pointer()
 	argSize := linux.IOC_SIZE(cmd)
@@ -133,17 +132,17 @@ func (fd *tpuV4FD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, arg
 		gasket.GASKET_IOCTL_MAP_DMA_BUF:
 		return 0, linuxerr.ENOSYS
 	case gasket.GASKET_IOCTL_RESET:
-		return tpuproxy.IOCTLInvoke[gasket.Ioctl, uint64](fd.hostFD, gasket.GASKET_IOCTL_RESET, args[2].Uint64())
+		return util.IOCTLInvoke[gasket.Ioctl, uint64](fd.hostFD, gasket.GASKET_IOCTL_RESET, args[2].Uint64())
 	case gasket.GASKET_IOCTL_MAP_BUFFER:
 		return gasketMapBufferIoctl(ctx, t, fd.hostFD, fd, argPtr)
 	case gasket.GASKET_IOCTL_UNMAP_BUFFER:
 		return gasketUnmapBufferIoctl(ctx, t, fd.hostFD, fd, argPtr)
 	case gasket.GASKET_IOCTL_CLEAR_INTERRUPT_COUNTS:
-		return tpuproxy.IOCTLInvoke[gasket.Ioctl](fd.hostFD, gasket.GASKET_IOCTL_CLEAR_INTERRUPT_COUNTS, 0)
+		return util.IOCTLInvoke[gasket.Ioctl](fd.hostFD, gasket.GASKET_IOCTL_CLEAR_INTERRUPT_COUNTS, 0)
 	case gasket.GASKET_IOCTL_REGISTER_INTERRUPT:
 		return gasketInterruptMappingIoctl(ctx, t, fd.hostFD, argPtr, fd.device.lite)
 	case gasket.GASKET_IOCTL_UNREGISTER_INTERRUPT:
-		return tpuproxy.IOCTLInvoke[gasket.Ioctl, uint64](fd.hostFD, gasket.GASKET_IOCTL_UNREGISTER_INTERRUPT, args[2].Uint64())
+		return util.IOCTLInvoke[gasket.Ioctl, uint64](fd.hostFD, gasket.GASKET_IOCTL_UNREGISTER_INTERRUPT, args[2].Uint64())
 	default:
 		return 0, linuxerr.EINVAL
 	}
@@ -152,7 +151,7 @@ func (fd *tpuV4FD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, arg
 // checkPermission checks that the thread that owns this device is the only
 // one that can issue commands to the TPU. Other threads with access to
 // /dev/accel will not be able to issue commands to the device.
-func (fd *tpuV4FD) checkPermission(t *kernel.Task) error {
+func (fd *accelFD) checkPermission(t *kernel.Task) error {
 	fd.device.mu.Lock()
 	defer fd.device.mu.Unlock()
 	owner := fd.device.owner

pkg/sentry/devices/tpuproxy/accel/accel_fd_mmap.go

@@ -25,26 +25,26 @@ import (
 )
 
 // ConfigureMMap implements vfs.FileDescriptionImpl.ConfigureMMap.
-func (fd *tpuV4FD) ConfigureMMap(ctx context.Context, opts *memmap.MMapOpts) error {
+func (fd *accelFD) ConfigureMMap(ctx context.Context, opts *memmap.MMapOpts) error {
 	return vfs.GenericConfigureMMap(&fd.vfsfd, fd, opts)
 }
 
 // AddMapping implements memmap.Mappable.AddMapping.
-func (fd *tpuV4FD) AddMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) error {
+func (fd *accelFD) AddMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) error {
 	return nil
 }
 
 // RemoveMapping implements memmap.Mappable.RemoveMapping.
-func (fd *tpuV4FD) RemoveMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) {
+func (fd *accelFD) RemoveMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) {
 }
 
 // CopyMapping implements memmap.Mappable.CopyMapping.
-func (fd *tpuV4FD) CopyMapping(ctx context.Context, ms memmap.MappingSpace, srcAR, dstAR hostarch.AddrRange, offset uint64, writable bool) error {
+func (fd *accelFD) CopyMapping(ctx context.Context, ms memmap.MappingSpace, srcAR, dstAR hostarch.AddrRange, offset uint64, writable bool) error {
 	return nil
 }
 
 // Translate implements memmap.Mappable.Translate.
-func (fd *tpuV4FD) Translate(ctx context.Context, required, optional memmap.MappableRange, at hostarch.AccessType) ([]memmap.Translation, error) {
+func (fd *accelFD) Translate(ctx context.Context, required, optional memmap.MappableRange, at hostarch.AccessType) ([]memmap.Translation, error) {
 	return []memmap.Translation{
 		{
 			Source: optional,
@@ -56,14 +56,14 @@ func (fd *tpuV4FD) Translate(ctx context.Context, required, optional memmap.Mapp
 }
 
 // InvalidateUnsavable implements memmap.Mappable.InvalidateUnsavable.
-func (fd *tpuV4FD) InvalidateUnsavable(ctx context.Context) error {
+func (fd *accelFD) InvalidateUnsavable(ctx context.Context) error {
 	return nil
 }
 
 type accelFDMemmapFile struct {
 	memmap.NoBufferedIOFallback
 
-	fd *tpuV4FD
+	fd *accelFD
 }
 
 // IncRef implements memmap.File.IncRef.

pkg/sentry/devices/tpuproxy/accel/gasket_ioctl.go

@@ -23,14 +23,14 @@ import (
 	"gvisor.dev/gvisor/pkg/context"
 	"gvisor.dev/gvisor/pkg/errors/linuxerr"
 	"gvisor.dev/gvisor/pkg/hostarch"
-	"gvisor.dev/gvisor/pkg/sentry/devices/tpuproxy"
+	"gvisor.dev/gvisor/pkg/sentry/devices/tpuproxy/util"
 	"gvisor.dev/gvisor/pkg/sentry/fsimpl/eventfd"
 	"gvisor.dev/gvisor/pkg/sentry/kernel"
 	"gvisor.dev/gvisor/pkg/sentry/memmap"
 	"gvisor.dev/gvisor/pkg/sentry/mm"
 )
 
-func gasketMapBufferIoctl(ctx context.Context, t *kernel.Task, hostFd int32, fd *tpuV4FD, paramsAddr hostarch.Addr) (uintptr, error) {
+func gasketMapBufferIoctl(ctx context.Context, t *kernel.Task, hostFd int32, fd *accelFD, paramsAddr hostarch.Addr) (uintptr, error) {
 	var userIoctlParams gasket.GasketPageTableIoctl
 	if _, err := userIoctlParams.CopyIn(t, paramsAddr); err != nil {
 		return 0, err
@@ -104,7 +104,7 @@ func gasketMapBufferIoctl(ctx context.Context, t *kernel.Task, hostFd int32, fd
 	}
 	sentryIoctlParams := userIoctlParams
 	sentryIoctlParams.HostAddress = uint64(m)
-	n, err := tpuproxy.IOCTLInvokePtrArg[gasket.Ioctl](hostFd, gasket.GASKET_IOCTL_MAP_BUFFER, &sentryIoctlParams)
+	n, err := util.IOCTLInvokePtrArg[gasket.Ioctl](hostFd, gasket.GASKET_IOCTL_MAP_BUFFER, &sentryIoctlParams)
 	if err != nil {
 		return n, err
 	}
@@ -125,7 +125,7 @@ func gasketMapBufferIoctl(ctx context.Context, t *kernel.Task, hostFd int32, fd
 	return n, nil
 }
 
-func gasketUnmapBufferIoctl(ctx context.Context, t *kernel.Task, hostFd int32, fd *tpuV4FD, paramsAddr hostarch.Addr) (uintptr, error) {
+func gasketUnmapBufferIoctl(ctx context.Context, t *kernel.Task, hostFd int32, fd *accelFD, paramsAddr hostarch.Addr) (uintptr, error) {
 	var userIoctlParams gasket.GasketPageTableIoctl
 	if _, err := userIoctlParams.CopyIn(t, paramsAddr); err != nil {
 		return 0, err
@@ -151,7 +151,7 @@ func gasketUnmapBufferIoctl(ctx context.Context, t *kernel.Task, hostFd int32, f
 
 	sentryIoctlParams := userIoctlParams
 	sentryIoctlParams.HostAddress = 0 // clobber this value, it's unused.
-	n, err := tpuproxy.IOCTLInvokePtrArg[gasket.Ioctl](hostFd, gasket.GASKET_IOCTL_UNMAP_BUFFER, &sentryIoctlParams)
+	n, err := util.IOCTLInvokePtrArg[gasket.Ioctl](hostFd, gasket.GASKET_IOCTL_UNMAP_BUFFER, &sentryIoctlParams)
 	if err != nil {
 		return n, err
 	}
@@ -211,7 +211,7 @@ func gasketInterruptMappingIoctl(ctx context.Context, t *kernel.Task, hostFd int
 
 	sentryIoctlParams := userIoctlParams
 	sentryIoctlParams.EventFD = uint64(eventfd)
-	n, err := tpuproxy.IOCTLInvokePtrArg[gasket.Ioctl](hostFd, gasket.GASKET_IOCTL_REGISTER_INTERRUPT, &sentryIoctlParams)
+	n, err := util.IOCTLInvokePtrArg[gasket.Ioctl](hostFd, gasket.GASKET_IOCTL_REGISTER_INTERRUPT, &sentryIoctlParams)
 	if err != nil {
 		return n, err
 	}

pkg/sentry/devices/tpuproxy/seccomp_filter.go

@@ -16,6 +16,7 @@ package tpuproxy
 
 import (
 	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/abi/gasket"
 	"gvisor.dev/gvisor/pkg/abi/linux"
 	"gvisor.dev/gvisor/pkg/seccomp"
 )
@@ -103,6 +104,30 @@ func Filters() seccomp.SyscallRules {
 				seccomp.NonNegativeFD{},
 				seccomp.EqualTo(linux.VFIO_SET_IOMMU),
 			},
+			seccomp.PerArg{
+				seccomp.NonNegativeFD{},
+				seccomp.EqualTo(gasket.GASKET_IOCTL_RESET),
+			},
+			seccomp.PerArg{
+				seccomp.NonNegativeFD{},
+				seccomp.EqualTo(gasket.GASKET_IOCTL_MAP_BUFFER),
+			},
+			seccomp.PerArg{
+				seccomp.NonNegativeFD{},
+				seccomp.EqualTo(gasket.GASKET_IOCTL_UNMAP_BUFFER),
+			},
+			seccomp.PerArg{
+				seccomp.NonNegativeFD{},
+				seccomp.EqualTo(gasket.GASKET_IOCTL_CLEAR_INTERRUPT_COUNTS),
+			},
+			seccomp.PerArg{
+				seccomp.NonNegativeFD{},
+				seccomp.EqualTo(gasket.GASKET_IOCTL_REGISTER_INTERRUPT),
+			},
+			seccomp.PerArg{
+				seccomp.NonNegativeFD{},
+				seccomp.EqualTo(gasket.GASKET_IOCTL_UNREGISTER_INTERRUPT),
+			},
 		},
 	})
 }

pkg/sentry/devices/tpuproxy/tpuproxy.go

@@ -0,0 +1,124 @@
+// Copyright 2024 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 tpuproxy contains tpu backend driver proxy implementations and
+// helper functions.
+package tpuproxy
+
+import (
+	"fmt"
+	"os"
+	"path"
+	"path/filepath"
+	"regexp"
+	"strconv"
+	"strings"
+
+	"gvisor.dev/gvisor/pkg/abi/tpu"
+	"gvisor.dev/gvisor/pkg/sentry/devices/tpuproxy/accel"
+	"gvisor.dev/gvisor/pkg/sentry/devices/tpuproxy/vfio"
+	"gvisor.dev/gvisor/pkg/sentry/vfs"
+)
+
+const (
+	pciPathGlobTPUv4   = "/sys/devices/pci0000:*/*/accel/accel*"
+	pciPathGlobTPUv5   = "/sys/devices/pci0000:*/*/vfio-dev/vfio*"
+	iommuGroupPathGlob = "/sys/kernel/iommu_groups/*/devices/*"
+)
+
+var (
+	// pathGlobToPathRegex is a map that points a TPU PCI path glob to its path regex.
+	// 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.
+	pathGlobToPathRegex = map[string]string{
+		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+)$`,
+	}
+)
+
+// RegisterHostTPUDevices enumerates TPU devices on the host and registers them
+// in the sandbox VFS.
+func RegisterHostTPUDevices(vfsObj *vfs.VirtualFilesystem, allowedDeviceIDs map[int64]any) error {
+	for pciPathGlobal, pathRegex := range pathGlobToPathRegex {
+		pciAddrs, err := filepath.Glob(pciPathGlobal)
+		if err != nil {
+			return fmt.Errorf("enumerating PCI device files: %w", err)
+		}
+		pciPathRegex := regexp.MustCompile(pathRegex)
+		for _, pciPath := range pciAddrs {
+			ms := pciPathRegex.FindStringSubmatch(pciPath)
+			if ms == nil {
+				continue
+			}
+			minorNum, err := strconv.ParseUint(ms[1], 10, 32)
+			if err != nil {
+				return fmt.Errorf("parsing PCI device number: %w", err)
+			}
+			var deviceIDBytes []byte
+			if deviceIDBytes, err = os.ReadFile(path.Join(pciPath, "device/device")); err != nil {
+				return fmt.Errorf("reading PCI device ID: %w", err)
+			}
+			deviceIDStr := strings.Replace(string(deviceIDBytes), "0x", "", -1)
+			deviceID, err := strconv.ParseInt(strings.TrimSpace(deviceIDStr), 16, 64)
+			if err != nil {
+				return fmt.Errorf("parsing PCI device ID: %w", err)
+			}
+			if _, ok := allowedDeviceIDs[deviceID]; !ok {
+				return fmt.Errorf("unsupported TPU device with ID: 0x%x", deviceID)
+			}
+			// VFIO iommu groups correspond to the device number. Use these
+			// paths to get the correct number for the sentry-internal TPU
+			// device files.
+			var deviceNum int
+			switch deviceID {
+			case tpu.TPUV4DeviceID, tpu.TPUV4liteDeviceID:
+				deviceNum = 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) {
+						n, err := strconv.Atoi(strings.Split(groupPath, "/")[4])
+						if err != nil {
+							return fmt.Errorf("parsing IOMMU group minor number: %w", err)
+						}
+						deviceNum = n
+						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)
+			}
+		}
+	}
+	return nil
+}
+
+// registerTPUDevice registers a TPU device in vfsObj based on the given device ID.
+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, tpu.TPUV5pDeviceID:
+		return vfio.RegisterTPUDevice(vfsObj, minor, deviceNum)
+	default:
+		return fmt.Errorf("unsupported TPU device with ID: 0x%x", deviceID)
+	}
+}

pkg/sentry/devices/tpuproxy/tpuproxy_test.go

@@ -0,0 +1,66 @@
+// Copyright 2024 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 tpuproxy
+
+import (
+	"path/filepath"
+	"regexp"
+	"slices"
+	"testing"
+)
+
+func TestTPUPath(t *testing.T) {
+	for _, tst := range []struct {
+		name     string
+		pathGlob string
+		path     string
+		submatch []string
+	}{
+		{
+			name:     "TPUv4PCIPathMatch",
+			pathGlob: pciPathGlobTPUv4,
+			path:     "/sys/devices/pci0000:00/0000:00:01.0/accel/accel16",
+			submatch: []string{"/sys/devices/pci0000:00/0000:00:01.0/accel/accel16", "16"},
+		},
+		{
+			name:     "TPUv4PCIPathNoMatch",
+			pathGlob: pciPathGlobTPUv4,
+			path:     "/sys/devices/pci0000:00/0000:00:01.0/accel/123",
+			submatch: nil,
+		},
+		{
+			name:     "TPUv5PCIPathMatch",
+			pathGlob: pciPathGlobTPUv5,
+			path:     "/sys/devices/pci0000:00/0000:00:05.0/vfio-dev/vfio20",
+			submatch: []string{"/sys/devices/pci0000:00/0000:00:05.0/vfio-dev/vfio20", "20"},
+		},
+		{
+			name:     "TPUv5PCIPathNoMatch",
+			pathGlob: pciPathGlobTPUv5,
+			path:     "/sys/devices/pci0000:00/0000:00:05.0/vfio/vfio20",
+			submatch: nil,
+		},
+	} {
+		t.Run(tst.name, func(t *testing.T) {
+			if _, err := filepath.Glob(tst.pathGlob); err != nil {
+				t.Errorf("Malformed path glob: %v", err)
+			}
+			pathRegex := regexp.MustCompile(pathGlobToPathRegex[tst.pathGlob])
+			if submatch := pathRegex.FindStringSubmatch(tst.path); !slices.Equal(submatch, tst.submatch) {
+				t.Errorf("Match TPU PCI path, got: %v, want: %v", submatch, tst.submatch)
+			}
+		})
+	}
+}

pkg/sentry/devices/tpuproxy/util/BUILD

@@ -0,0 +1,20 @@
+load("//tools:defs.bzl", "go_library")
+
+package(default_applicable_licenses = ["//:license"])
+
+licenses(["notice"])
+
+go_library(
+    name = "util",
+    srcs = [
+        "ioctl_unsafe.go",
+        "util.go",
+    ],
+    visibility = [
+        "//pkg/sentry:internal",
+    ],
+    deps = [
+        "@org_golang_x_exp//constraints:go_default_library",
+        "@org_golang_x_sys//unix:go_default_library",
+    ],
+)

pkg/sentry/devices/tpuproxy/util/ioctl_unsafe.go

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License
 
-package tpuproxy
+package util
 
 import (
 	"unsafe"

pkg/sentry/devices/tpuproxy/util/util.go

@@ -0,0 +1,16 @@
+// Copyright 2024 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 util contains helper functions for tpuproxy implementations.
+package util

pkg/sentry/devices/tpuproxy/vfio/BUILD

@@ -0,0 +1,75 @@
+load("//tools:defs.bzl", "go_library")
+load("//tools/go_generics:defs.bzl", "go_template_instance")
+
+package(default_applicable_licenses = ["//:license"])
+
+licenses(["notice"])
+
+go_library(
+    name = "vfio",
+    srcs = [
+        "devaddr_range.go",
+        "devaddr_set.go",
+        "pci_device_fd.go",
+        "pci_device_fd_mmap.go",
+        "tpu_fd.go",
+        "tpu_fd_mmap.go",
+        "vfio.go",
+        "vfio_fd.go",
+        "vfio_fd_mmap.go",
+    ],
+    visibility = [
+        "//pkg/sentry:internal",
+    ],
+    deps = [
+        "//pkg/abi/linux",
+        "//pkg/cleanup",
+        "//pkg/context",
+        "//pkg/devutil",
+        "//pkg/errors/linuxerr",
+        "//pkg/fdnotifier",
+        "//pkg/hostarch",
+        "//pkg/log",
+        "//pkg/marshal/primitive",
+        "//pkg/safemem",
+        "//pkg/sentry/arch",
+        "//pkg/sentry/devices/tpuproxy/util",
+        "//pkg/sentry/fsimpl/eventfd",
+        "//pkg/sentry/kernel",
+        "//pkg/sentry/memmap",
+        "//pkg/sentry/mm",
+        "//pkg/sentry/vfs",
+        "//pkg/sync",
+        "//pkg/usermem",
+        "//pkg/waiter",
+        "@org_golang_x_sys//unix:go_default_library",
+    ],
+)
+
+go_template_instance(
+    name = "devaddr_range",
+    out = "devaddr_range.go",
+    package = "vfio",
+    prefix = "DevAddr",
+    template = "//pkg/segment:generic_range",
+    types = {
+        "T": "uint64",
+    },
+)
+
+go_template_instance(
+    name = "devaddr_set",
+    out = "devaddr_set.go",
+    imports = {
+        "mm": "gvisor.dev/gvisor/pkg/sentry/mm",
+    },
+    package = "vfio",
+    prefix = "DevAddr",
+    template = "//pkg/segment:generic_set",
+    types = {
+        "Key": "uint64",
+        "Range": "DevAddrRange",
+        "Value": "mm.PinnedRange",
+        "Functions": "devAddrSetFuncs",
+    },
+)

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

@@ -12,8 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Package tpuproxy implements proxying for TPU devices.
-package tpuproxy
+package vfio
 
 import (
 	"fmt"
@@ -26,160 +25,14 @@ import (
 	"gvisor.dev/gvisor/pkg/hostarch"
 	"gvisor.dev/gvisor/pkg/marshal/primitive"
 	"gvisor.dev/gvisor/pkg/sentry/arch"
+	"gvisor.dev/gvisor/pkg/sentry/devices/tpuproxy/util"
 	"gvisor.dev/gvisor/pkg/sentry/fsimpl/eventfd"
 	"gvisor.dev/gvisor/pkg/sentry/kernel"
-	"gvisor.dev/gvisor/pkg/sentry/mm"
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
 	"gvisor.dev/gvisor/pkg/usermem"
 	"gvisor.dev/gvisor/pkg/waiter"
 )
 
-const (
-	// A value of -1 can be used to either de-assign interrupts if already
-	// assigned or skip un-assigned interrupts.
-	disableInterrupt = -1
-)
-
-var (
-	// vfioDeviceInfoFlags contains all available flags for
-	// IOCTL command VFIO_DEVICE_GET_INFO.
-	vfioDeviceInfoFlags uint32 = linux.VFIO_DEVICE_FLAGS_RESET | linux.VFIO_DEVICE_FLAGS_PCI |
-		linux.VFIO_DEVICE_FLAGS_PLATFORM | linux.VFIO_DEVICE_FLAGS_AMBA |
-		linux.VFIO_DEVICE_FLAGS_CCW | linux.VFIO_DEVICE_FLAGS_AP | linux.VFIO_DEVICE_FLAGS_FSL_MC |
-		linux.VFIO_DEVICE_FLAGS_CAPS | linux.VFIO_DEVICE_FLAGS_CDX
-	// vfioIrqSetFlags includes all available flags for IOCTL comamnd VFIO_DEVICE_SET_IRQS
-	vfioIrqSetFlags uint32 = linux.VFIO_IRQ_SET_DATA_TYPE_MASK | linux.VFIO_IRQ_SET_ACTION_TYPE_MASK
-)
-
-// tpuFD implements vfs.FileDescriptionImpl for /dev/vfio/[0-9]+
-//
-// tpuFD is not savable until TPU save/restore is needed.
-type tpuFD struct {
-	vfsfd vfs.FileDescription
-	vfs.FileDescriptionDefaultImpl
-	vfs.DentryMetadataFileDescriptionImpl
-	vfs.NoLockFD
-
-	hostFD     int32
-	device     *tpuDevice
-	queue      waiter.Queue
-	memmapFile tpuFDMemmapFile
-}
-
-// Release implements vfs.FileDescriptionImpl.Release.
-func (fd *tpuFD) Release(context.Context) {
-	fdnotifier.RemoveFD(fd.hostFD)
-	fd.queue.Notify(waiter.EventHUp)
-	unix.Close(int(fd.hostFD))
-}
-
-// EventRegister implements waiter.Waitable.EventRegister.
-func (fd *tpuFD) EventRegister(e *waiter.Entry) error {
-	fd.queue.EventRegister(e)
-	if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
-		fd.queue.EventUnregister(e)
-		return err
-	}
-	return nil
-}
-
-// EventUnregister implements waiter.Waitable.EventUnregister.
-func (fd *tpuFD) EventUnregister(e *waiter.Entry) {
-	fd.queue.EventUnregister(e)
-	if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
-		panic(fmt.Sprint("UpdateFD:", err))
-	}
-}
-
-// Readiness implements waiter.Waitable.Readiness.
-func (fd *tpuFD) Readiness(mask waiter.EventMask) waiter.EventMask {
-	return fdnotifier.NonBlockingPoll(fd.hostFD, mask)
-}
-
-// Epollable implements vfs.FileDescriptionImpl.Epollable.
-func (fd *tpuFD) Epollable() bool {
-	return true
-}
-
-// Ioctl implements vfs.FileDescriptionImpl.Ioctl.
-func (fd *tpuFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, args arch.SyscallArguments) (uintptr, error) {
-	cmd := args[1].Uint()
-
-	t := kernel.TaskFromContext(ctx)
-	if t == nil {
-		panic("Ioctl should be called from a task context")
-	}
-	switch cmd {
-	case linux.VFIO_GROUP_SET_CONTAINER:
-		return fd.setContainer(ctx, t, args[2].Pointer())
-	case linux.VFIO_GROUP_GET_DEVICE_FD:
-		ret, cleanup, err := fd.getPciDeviceFd(t, args[2].Pointer())
-		defer cleanup()
-		return ret, err
-	}
-	return 0, linuxerr.ENOSYS
-}
-
-func (fd *tpuFD) setContainer(ctx context.Context, t *kernel.Task, arg hostarch.Addr) (uintptr, error) {
-	var vfioContainerFD int32
-	if _, err := primitive.CopyInt32In(t, arg, &vfioContainerFD); err != nil {
-		return 0, err
-	}
-	vfioContainerFile, _ := t.FDTable().Get(vfioContainerFD)
-	if vfioContainerFile == nil {
-		return 0, linuxerr.EBADF
-	}
-	defer vfioContainerFile.DecRef(ctx)
-	vfioContainer, ok := vfioContainerFile.Impl().(*vfioFD)
-	if !ok {
-		return 0, linuxerr.EINVAL
-	}
-	return IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_GROUP_SET_CONTAINER, &vfioContainer.hostFD)
-}
-
-// It will be the caller's responsibility to call the returned cleanup function.
-func (fd *tpuFD) getPciDeviceFd(t *kernel.Task, arg hostarch.Addr) (uintptr, func(), error) {
-	pciAddress, err := t.CopyInString(arg, hostarch.PageSize)
-	if err != nil {
-		return 0, func() {}, err
-	}
-	// Build a NUL-terminated slice of bytes containing the PCI address.
-	pciAddressBytes, err := unix.ByteSliceFromString(pciAddress)
-	if err != nil {
-		return 0, func() {}, err
-	}
-	// Pass the address of the PCI address' first byte which can be
-	// recognized by the IOCTL syscall.
-	hostFD, err := IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_GROUP_GET_DEVICE_FD, &pciAddressBytes[0])
-	if err != nil {
-		return 0, func() {}, err
-	}
-	pciDevFD := &pciDeviceFD{
-		hostFD: int32(hostFD),
-	}
-	cleanup := func() {
-		unix.Close(int(hostFD))
-	}
-	// 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]")
-	if err := pciDevFD.vfsfd.Init(pciDevFD, linux.O_RDWR, vd.Mount(), vd.Dentry(), &vfs.FileDescriptionOptions{
-		UseDentryMetadata: true,
-	}); err != nil {
-		return 0, cleanup, err
-	}
-	if err := fdnotifier.AddFD(int32(hostFD), &fd.queue); err != nil {
-		return 0, cleanup, err
-	}
-	newFD, err := t.NewFDFrom(0, &pciDevFD.vfsfd, kernel.FDFlags{})
-	if err != nil {
-		return 0, cleanup, err
-	}
-	// Initialize a mapping that is backed by a host FD.
-	pciDevFD.memmapFile.fd = pciDevFD
-	return uintptr(newFD), func() {}, nil
-}
-
 // pciDeviceFD implements vfs.FileDescriptionImpl for TPU's PCI device.
 type pciDeviceFD struct {
 	vfsfd vfs.FileDescription
@@ -247,7 +100,7 @@ func (fd *pciDeviceFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr,
 		return fd.vfioSetIrqs(ctx, t, args[2].Pointer())
 	case linux.VFIO_DEVICE_RESET:
 		// VFIO_DEVICE_RESET is just a simple IOCTL command that carries no data.
-		return IOCTLInvoke[uint32, uintptr](fd.hostFD, linux.VFIO_DEVICE_RESET, 0)
+		return util.IOCTLInvoke[uint32, uintptr](fd.hostFD, linux.VFIO_DEVICE_RESET, 0)
 	}
 	return 0, linuxerr.ENOSYS
 }
@@ -262,7 +115,7 @@ func (fd *pciDeviceFD) vfioRegionInfo(ctx context.Context, t *kernel.Task, arg h
 	if regionInfo.Argsz == 0 {
 		return 0, linuxerr.EINVAL
 	}
-	ret, err := IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_GET_REGION_INFO, &regionInfo)
+	ret, err := util.IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_GET_REGION_INFO, &regionInfo)
 	if err != nil {
 		return 0, err
 	}
@@ -285,7 +138,7 @@ func (fd *pciDeviceFD) vfioDeviceInfo(ctx context.Context, t *kernel.Task, arg h
 	if deviceInfo.Flags&^vfioDeviceInfoFlags != 0 {
 		return 0, linuxerr.EINVAL
 	}
-	ret, err := IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_GET_INFO, &deviceInfo)
+	ret, err := util.IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_GET_INFO, &deviceInfo)
 	if err != nil {
 		return 0, err
 	}
@@ -308,7 +161,7 @@ func (fd *pciDeviceFD) vfioIrqInfo(ctx context.Context, t *kernel.Task, arg host
 	if irqInfo.Argsz == 0 {
 		return 0, linuxerr.EINVAL
 	}
-	ret, err := IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_GET_IRQ_INFO, &irqInfo)
+	ret, err := util.IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_GET_IRQ_INFO, &irqInfo)
 	if err != nil {
 		return 0, err
 	}
@@ -349,7 +202,7 @@ func (fd *pciDeviceFD) vfioSetIrqs(ctx context.Context, t *kernel.Task, arg host
 	case linux.VFIO_IRQ_SET_DATA_NONE:
 		// When there is no data, passing through the given payload
 		// works just fine.
-		return IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_SET_IRQS, &irqSet)
+		return util.IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_SET_IRQS, &irqSet)
 	// VFIO_IRQ_SET_DATA_BOOL indicates that the data field is an array of uint8.
 	// The action will be performed if the corresponding boolean is true.
 	case linux.VFIO_IRQ_SET_DATA_BOOL:
@@ -358,7 +211,7 @@ func (fd *pciDeviceFD) vfioSetIrqs(ctx context.Context, t *kernel.Task, arg host
 		if _, err := primitive.CopyUint8SliceIn(t, arg, payload); err != nil {
 			return 0, err
 		}
-		return IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_SET_IRQS, &payload[0])
+		return util.IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_SET_IRQS, &payload[0])
 	// VFIO_IRQ_SET_DATA_EVENTFD indicates that the data field is an array
 	// of int32 (or event file descriptors). These descriptors will be
 	// signalled when an action in the flags happens.
@@ -391,7 +244,7 @@ func (fd *pciDeviceFD) vfioSetIrqs(ctx context.Context, t *kernel.Task, arg host
 			}
 			payload[index] = int32(eventfd)
 		}
-		return IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_SET_IRQS, &payload[0])
+		return util.IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_SET_IRQS, &payload[0])
 	}
 	// No data type is specified or multiple data types are specified.
 	return 0, linuxerr.EINVAL
@@ -424,57 +277,3 @@ func (fd *pciDeviceFD) PWrite(ctx context.Context, src usermem.IOSequence, offse
 	n, err := unix.Pwrite(int(fd.hostFD), buf, offset)
 	return int64(n), err
 }
-
-// DevAddrSet tracks device address ranges that have been mapped.
-type devAddrSetFuncs struct{}
-
-func (devAddrSetFuncs) MinKey() uint64 {
-	return 0
-}
-
-func (devAddrSetFuncs) MaxKey() uint64 {
-	return ^uint64(0)
-}
-
-func (devAddrSetFuncs) ClearValue(val *mm.PinnedRange) {
-	*val = mm.PinnedRange{}
-}
-
-func (devAddrSetFuncs) Merge(r1 DevAddrRange, v1 mm.PinnedRange, r2 DevAddrRange, v2 mm.PinnedRange) (mm.PinnedRange, bool) {
-	// Do we have the same backing file?
-	if v1.File != v2.File {
-		return mm.PinnedRange{}, false
-	}
-
-	// Do we have contiguous offsets in the backing file?
-	if v1.Offset+uint64(v1.Source.Length()) != v2.Offset {
-		return mm.PinnedRange{}, false
-	}
-
-	// Are the virtual addresses contiguous?
-	//
-	// This check isn't strictly needed because 'mm.PinnedRange.Source'
-	// is only used to track the size of the pinned region (this is
-	// because the virtual address range can be unmapped or remapped
-	// elsewhere). Regardless we require this for simplicity.
-	if v1.Source.End != v2.Source.Start {
-		return mm.PinnedRange{}, false
-	}
-
-	// Extend v1 to account for the adjacent PinnedRange.
-	v1.Source.End = v2.Source.End
-	return v1, true
-}
-
-func (devAddrSetFuncs) Split(r DevAddrRange, val mm.PinnedRange, split uint64) (mm.PinnedRange, mm.PinnedRange) {
-	n := split - r.Start
-
-	left := val
-	left.Source.End = left.Source.Start + hostarch.Addr(n)
-
-	right := val
-	right.Source.Start += hostarch.Addr(n)
-	right.Offset += n
-
-	return left, right
-}

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

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package tpuproxy
+package vfio
 
 import (
 	"gvisor.dev/gvisor/pkg/context"
@@ -24,67 +24,6 @@ import (
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
 )
 
-// ConfigureMMap implements vfs.FileDescriptionImpl.ConfigureMMap.
-func (fd *tpuFD) ConfigureMMap(ctx context.Context, opts *memmap.MMapOpts) error {
-	return vfs.GenericConfigureMMap(&fd.vfsfd, fd, opts)
-}
-
-// AddMapping implements memmap.Mappable.AddMapping.
-func (fd *tpuFD) AddMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) error {
-	return nil
-}
-
-// RemoveMapping implements memmap.Mappable.RemoveMapping.
-func (fd *tpuFD) RemoveMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) {
-}
-
-// CopyMapping implements memmap.Mappable.CopyMapping.
-func (fd *tpuFD) CopyMapping(ctx context.Context, ms memmap.MappingSpace, srcAR, dstAR hostarch.AddrRange, offset uint64, writable bool) error {
-	return nil
-}
-
-// Translate implements memmap.Mappable.Translate.
-func (fd *tpuFD) Translate(ctx context.Context, required, optional memmap.MappableRange, at hostarch.AccessType) ([]memmap.Translation, error) {
-	return []memmap.Translation{
-		{
-			Source: optional,
-			File:   &fd.memmapFile,
-			Offset: optional.Start,
-			Perms:  at,
-		},
-	}, nil
-}
-
-// InvalidateUnsavable implements memmap.Mappable.InvalidateUnsavable.
-func (fd *tpuFD) InvalidateUnsavable(ctx context.Context) error {
-	return nil
-}
-
-type tpuFDMemmapFile struct {
-	memmap.NoBufferedIOFallback
-
-	fd *tpuFD
-}
-
-// IncRef implements memmap.File.IncRef.
-func (mf *tpuFDMemmapFile) IncRef(memmap.FileRange, uint32) {
-}
-
-// DecRef implements memmap.File.DecRef.
-func (mf *tpuFDMemmapFile) DecRef(fr memmap.FileRange) {
-}
-
-// MapInternal implements memmap.File.MapInternal.
-func (mf *tpuFDMemmapFile) MapInternal(fr memmap.FileRange, at hostarch.AccessType) (safemem.BlockSeq, error) {
-	log.Traceback("tpuproxy: rejecting tpuFdMemmapFile.MapInternal")
-	return safemem.BlockSeq{}, linuxerr.EINVAL
-}
-
-// FD implements memmap.File.FD.
-func (mf *tpuFDMemmapFile) FD() int {
-	return int(mf.fd.hostFD)
-}
-
 // ConfigureMMap implements vfs.FileDescriptionImpl.ConfigureMMap.
 func (fd *pciDeviceFD) ConfigureMMap(ctx context.Context, opts *memmap.MMapOpts) error {
 	return vfs.GenericConfigureMMap(&fd.vfsfd, fd, opts)

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

@@ -0,0 +1,234 @@
+// Copyright 2023 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 (
+	"fmt"
+
+	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/abi/linux"
+	"gvisor.dev/gvisor/pkg/context"
+	"gvisor.dev/gvisor/pkg/errors/linuxerr"
+	"gvisor.dev/gvisor/pkg/fdnotifier"
+	"gvisor.dev/gvisor/pkg/hostarch"
+	"gvisor.dev/gvisor/pkg/marshal/primitive"
+	"gvisor.dev/gvisor/pkg/sentry/arch"
+	"gvisor.dev/gvisor/pkg/sentry/devices/tpuproxy/util"
+	"gvisor.dev/gvisor/pkg/sentry/kernel"
+	"gvisor.dev/gvisor/pkg/sentry/mm"
+	"gvisor.dev/gvisor/pkg/sentry/vfs"
+	"gvisor.dev/gvisor/pkg/usermem"
+	"gvisor.dev/gvisor/pkg/waiter"
+)
+
+const (
+	// A value of -1 can be used to either de-assign interrupts if already
+	// assigned or skip un-assigned interrupts.
+	disableInterrupt = -1
+)
+
+var (
+	// vfioDeviceInfoFlags contains all available flags for
+	// IOCTL command VFIO_DEVICE_GET_INFO.
+	vfioDeviceInfoFlags uint32 = linux.VFIO_DEVICE_FLAGS_RESET | linux.VFIO_DEVICE_FLAGS_PCI |
+		linux.VFIO_DEVICE_FLAGS_PLATFORM | linux.VFIO_DEVICE_FLAGS_AMBA |
+		linux.VFIO_DEVICE_FLAGS_CCW | linux.VFIO_DEVICE_FLAGS_AP | linux.VFIO_DEVICE_FLAGS_FSL_MC |
+		linux.VFIO_DEVICE_FLAGS_CAPS | linux.VFIO_DEVICE_FLAGS_CDX
+	// vfioIrqSetFlags includes all available flags for IOCTL command VFIO_DEVICE_SET_IRQS
+	vfioIrqSetFlags uint32 = linux.VFIO_IRQ_SET_DATA_TYPE_MASK | linux.VFIO_IRQ_SET_ACTION_TYPE_MASK
+)
+
+// tpuFD implements vfs.FileDescriptionImpl for /dev/vfio/[0-9]+
+//
+// tpuFD is not savable until TPU save/restore is needed.
+type tpuFD struct {
+	vfsfd vfs.FileDescription
+	vfs.FileDescriptionDefaultImpl
+	vfs.DentryMetadataFileDescriptionImpl
+	vfs.NoLockFD
+
+	hostFD     int32
+	device     *tpuDevice
+	queue      waiter.Queue
+	memmapFile tpuFDMemmapFile
+}
+
+// Release implements vfs.FileDescriptionImpl.Release.
+func (fd *tpuFD) Release(context.Context) {
+	fdnotifier.RemoveFD(fd.hostFD)
+	fd.queue.Notify(waiter.EventHUp)
+	unix.Close(int(fd.hostFD))
+}
+
+// EventRegister implements waiter.Waitable.EventRegister.
+func (fd *tpuFD) EventRegister(e *waiter.Entry) error {
+	fd.queue.EventRegister(e)
+	if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
+		fd.queue.EventUnregister(e)
+		return err
+	}
+	return nil
+}
+
+// EventUnregister implements waiter.Waitable.EventUnregister.
+func (fd *tpuFD) EventUnregister(e *waiter.Entry) {
+	fd.queue.EventUnregister(e)
+	if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
+		panic(fmt.Sprint("UpdateFD:", err))
+	}
+}
+
+// Readiness implements waiter.Waitable.Readiness.
+func (fd *tpuFD) Readiness(mask waiter.EventMask) waiter.EventMask {
+	return fdnotifier.NonBlockingPoll(fd.hostFD, mask)
+}
+
+// Epollable implements vfs.FileDescriptionImpl.Epollable.
+func (fd *tpuFD) Epollable() bool {
+	return true
+}
+
+// Ioctl implements vfs.FileDescriptionImpl.Ioctl.
+func (fd *tpuFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, args arch.SyscallArguments) (uintptr, error) {
+	cmd := args[1].Uint()
+
+	t := kernel.TaskFromContext(ctx)
+	if t == nil {
+		panic("Ioctl should be called from a task context")
+	}
+	switch cmd {
+	case linux.VFIO_GROUP_SET_CONTAINER:
+		return fd.setContainer(ctx, t, args[2].Pointer())
+	case linux.VFIO_GROUP_GET_DEVICE_FD:
+		ret, cleanup, err := fd.getPciDeviceFd(t, args[2].Pointer())
+		defer cleanup()
+		return ret, err
+	}
+	return 0, linuxerr.ENOSYS
+}
+
+func (fd *tpuFD) setContainer(ctx context.Context, t *kernel.Task, arg hostarch.Addr) (uintptr, error) {
+	var vfioContainerFD int32
+	if _, err := primitive.CopyInt32In(t, arg, &vfioContainerFD); err != nil {
+		return 0, err
+	}
+	vfioContainerFile, _ := t.FDTable().Get(vfioContainerFD)
+	if vfioContainerFile == nil {
+		return 0, linuxerr.EBADF
+	}
+	defer vfioContainerFile.DecRef(ctx)
+	vfioContainer, ok := vfioContainerFile.Impl().(*vfioFD)
+	if !ok {
+		return 0, linuxerr.EINVAL
+	}
+	return util.IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_GROUP_SET_CONTAINER, &vfioContainer.hostFD)
+}
+
+// It will be the caller's responsibility to call the returned cleanup function.
+func (fd *tpuFD) getPciDeviceFd(t *kernel.Task, arg hostarch.Addr) (uintptr, func(), error) {
+	pciAddress, err := t.CopyInString(arg, hostarch.PageSize)
+	if err != nil {
+		return 0, func() {}, err
+	}
+	// Build a NUL-terminated slice of bytes containing the PCI address.
+	pciAddressBytes, err := unix.ByteSliceFromString(pciAddress)
+	if err != nil {
+		return 0, func() {}, err
+	}
+	// Pass the address of the PCI address' first byte which can be
+	// recognized by the IOCTL syscall.
+	hostFD, err := util.IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_GROUP_GET_DEVICE_FD, &pciAddressBytes[0])
+	if err != nil {
+		return 0, func() {}, err
+	}
+	pciDevFD := &pciDeviceFD{
+		hostFD: int32(hostFD),
+	}
+	cleanup := func() {
+		unix.Close(int(hostFD))
+	}
+	// 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]")
+	if err := pciDevFD.vfsfd.Init(pciDevFD, linux.O_RDWR, vd.Mount(), vd.Dentry(), &vfs.FileDescriptionOptions{
+		UseDentryMetadata: true,
+	}); err != nil {
+		return 0, cleanup, err
+	}
+	if err := fdnotifier.AddFD(int32(hostFD), &fd.queue); err != nil {
+		return 0, cleanup, err
+	}
+	newFD, err := t.NewFDFrom(0, &pciDevFD.vfsfd, kernel.FDFlags{})
+	if err != nil {
+		return 0, cleanup, err
+	}
+	// Initialize a mapping that is backed by a host FD.
+	pciDevFD.memmapFile.fd = pciDevFD
+	return uintptr(newFD), func() {}, nil
+}
+
+// DevAddrSet tracks device address ranges that have been mapped.
+type devAddrSetFuncs struct{}
+
+func (devAddrSetFuncs) MinKey() uint64 {
+	return 0
+}
+
+func (devAddrSetFuncs) MaxKey() uint64 {
+	return ^uint64(0)
+}
+
+func (devAddrSetFuncs) ClearValue(val *mm.PinnedRange) {
+	*val = mm.PinnedRange{}
+}
+
+func (devAddrSetFuncs) Merge(r1 DevAddrRange, v1 mm.PinnedRange, r2 DevAddrRange, v2 mm.PinnedRange) (mm.PinnedRange, bool) {
+	// Do we have the same backing file?
+	if v1.File != v2.File {
+		return mm.PinnedRange{}, false
+	}
+
+	// Do we have contiguous offsets in the backing file?
+	if v1.Offset+uint64(v1.Source.Length()) != v2.Offset {
+		return mm.PinnedRange{}, false
+	}
+
+	// Are the virtual addresses contiguous?
+	//
+	// This check isn't strictly needed because 'mm.PinnedRange.Source'
+	// is only used to track the size of the pinned region (this is
+	// because the virtual address range can be unmapped or remapped
+	// elsewhere). Regardless we require this for simplicity.
+	if v1.Source.End != v2.Source.Start {
+		return mm.PinnedRange{}, false
+	}
+
+	// Extend v1 to account for the adjacent PinnedRange.
+	v1.Source.End = v2.Source.End
+	return v1, true
+}
+
+func (devAddrSetFuncs) Split(r DevAddrRange, val mm.PinnedRange, split uint64) (mm.PinnedRange, mm.PinnedRange) {
+	n := split - r.Start
+
+	left := val
+	left.Source.End = left.Source.Start + hostarch.Addr(n)
+
+	right := val
+	right.Source.Start += hostarch.Addr(n)
+	right.Offset += n
+
+	return left, right
+}

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

@@ -0,0 +1,86 @@
+// Copyright 2024 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"
+	"gvisor.dev/gvisor/pkg/errors/linuxerr"
+	"gvisor.dev/gvisor/pkg/hostarch"
+	"gvisor.dev/gvisor/pkg/log"
+	"gvisor.dev/gvisor/pkg/safemem"
+	"gvisor.dev/gvisor/pkg/sentry/memmap"
+	"gvisor.dev/gvisor/pkg/sentry/vfs"
+)
+
+// ConfigureMMap implements vfs.FileDescriptionImpl.ConfigureMMap.
+func (fd *tpuFD) ConfigureMMap(ctx context.Context, opts *memmap.MMapOpts) error {
+	return vfs.GenericConfigureMMap(&fd.vfsfd, fd, opts)
+}
+
+// AddMapping implements memmap.Mappable.AddMapping.
+func (fd *tpuFD) AddMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) error {
+	return nil
+}
+
+// RemoveMapping implements memmap.Mappable.RemoveMapping.
+func (fd *tpuFD) RemoveMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) {
+}
+
+// CopyMapping implements memmap.Mappable.CopyMapping.
+func (fd *tpuFD) CopyMapping(ctx context.Context, ms memmap.MappingSpace, srcAR, dstAR hostarch.AddrRange, offset uint64, writable bool) error {
+	return nil
+}
+
+// Translate implements memmap.Mappable.Translate.
+func (fd *tpuFD) Translate(ctx context.Context, required, optional memmap.MappableRange, at hostarch.AccessType) ([]memmap.Translation, error) {
+	return []memmap.Translation{
+		{
+			Source: optional,
+			File:   &fd.memmapFile,
+			Offset: optional.Start,
+			Perms:  at,
+		},
+	}, nil
+}
+
+// InvalidateUnsavable implements memmap.Mappable.InvalidateUnsavable.
+func (fd *tpuFD) InvalidateUnsavable(ctx context.Context) error {
+	return nil
+}
+
+type tpuFDMemmapFile struct {
+	memmap.NoBufferedIOFallback
+
+	fd *tpuFD
+}
+
+// IncRef implements memmap.File.IncRef.
+func (mf *tpuFDMemmapFile) IncRef(memmap.FileRange, uint32) {
+}
+
+// DecRef implements memmap.File.DecRef.
+func (mf *tpuFDMemmapFile) DecRef(fr memmap.FileRange) {
+}
+
+// MapInternal implements memmap.File.MapInternal.
+func (mf *tpuFDMemmapFile) MapInternal(fr memmap.FileRange, at hostarch.AccessType) (safemem.BlockSeq, error) {
+	log.Traceback("tpuproxy: rejecting tpuFdMemmapFile.MapInternal")
+	return safemem.BlockSeq{}, linuxerr.EINVAL
+}
+
+// FD implements memmap.File.FD.
+func (mf *tpuFDMemmapFile) FD() int {
+	return int(mf.fd.hostFD)
+}

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

@@ -12,11 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package tpuproxy
+package vfio
 
 import (
 	"path/filepath"
 	"strconv"
+	"strings"
 
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/abi/linux"
@@ -33,12 +34,11 @@ 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"
-
 	tpuDeviceGroupName  = "vfio"
 	vfioDeviceGroupName = "vfio"
+
+	// VFIOPath is the valid path to a VFIO device.
+	VFIOPath = "/dev/vfio/vfio"
 )
 
 // device implements TPU's vfs.Device for /dev/vfio/[0-9]+
@@ -97,7 +97,7 @@ func (dev *vfioDevice) Open(ctx context.Context, mnt *vfs.Mount, d *vfs.Dentry,
 		return nil, linuxerr.ENOENT
 	}
 
-	name := filepath.Join("vfio", filepath.Base(VFIOPath))
+	name := strings.ReplaceAll(VFIOPath, "/dev/", "")
 	hostFD, err := client.OpenAt(ctx, name, opts.Flags)
 	if err != nil {
 		ctx.Warningf("failed to open host file %s: %v", name, err)
@@ -131,8 +131,8 @@ func RegisterTPUDevice(vfsObj *vfs.VirtualFilesystem, minor, deviceNum uint32) e
 	})
 }
 
-// RegisterVfioDevice registers VFIO devices that are implemented by this package in vfsObj.
-func RegisterVfioDevice(vfsObj *vfs.VirtualFilesystem) error {
+// RegisterVFIODevice registers VFIO devices that are implemented by this package in vfsObj.
+func RegisterVFIODevice(vfsObj *vfs.VirtualFilesystem) error {
 	return vfsObj.RegisterDevice(vfs.CharDevice, linux.MISC_MAJOR, VFIO_MINOR, &vfioDevice{}, &vfs.RegisterDeviceOptions{
 		GroupName: vfioDeviceGroupName,
 	})

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

@@ -12,7 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package tpuproxy
+// Package vfio implements a proxy for VFIO devices.
+package vfio
 
 import (
 	"fmt"
@@ -27,6 +28,7 @@ import (
 	"gvisor.dev/gvisor/pkg/hostarch"
 	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/sentry/arch"
+	"gvisor.dev/gvisor/pkg/sentry/devices/tpuproxy/util"
 	"gvisor.dev/gvisor/pkg/sentry/kernel"
 	"gvisor.dev/gvisor/pkg/sentry/memmap"
 	"gvisor.dev/gvisor/pkg/sentry/mm"
@@ -113,7 +115,7 @@ func (fd *vfioFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, args
 func (fd *vfioFD) checkExtension(ext extension) (uintptr, error) {
 	switch ext {
 	case linux.VFIO_TYPE1_IOMMU, linux.VFIO_SPAPR_TCE_IOMMU, linux.VFIO_TYPE1v2_IOMMU:
-		ret, err := IOCTLInvoke[uint32, int32](fd.hostFD, linux.VFIO_CHECK_EXTENSION, int32(ext))
+		ret, err := util.IOCTLInvoke[uint32, int32](fd.hostFD, linux.VFIO_CHECK_EXTENSION, int32(ext))
 		if err != nil {
 			log.Warningf("check VFIO extension %s: %v", ext, err)
 			return 0, err
@@ -128,7 +130,7 @@ func (fd *vfioFD) checkExtension(ext extension) (uintptr, error) {
 func (fd *vfioFD) setIOMMU(ext extension) (uintptr, error) {
 	switch ext {
 	case linux.VFIO_TYPE1_IOMMU, linux.VFIO_SPAPR_TCE_IOMMU, linux.VFIO_TYPE1v2_IOMMU:
-		ret, err := IOCTLInvoke[uint32, int32](fd.hostFD, linux.VFIO_SET_IOMMU, int32(ext))
+		ret, err := util.IOCTLInvoke[uint32, int32](fd.hostFD, linux.VFIO_SET_IOMMU, int32(ext))
 		if err != nil {
 			log.Warningf("set the IOMMU group to %s: %v", ext, err)
 			return 0, err
@@ -196,7 +198,7 @@ func (fd *vfioFD) iommuMapDma(ctx context.Context, t *kernel.Task, arg hostarch.
 	}
 	// Replace Vaddr with the host's virtual address.
 	dmaMap.Vaddr = uint64(m)
-	n, err := IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_IOMMU_MAP_DMA, &dmaMap)
+	n, err := util.IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_IOMMU_MAP_DMA, &dmaMap)
 	if err != nil {
 		return n, err
 	}
@@ -228,7 +230,7 @@ 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_UNMAP_DMA, &dmaUnmap)
+	n, err := util.IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_IOMMU_UNMAP_DMA, &dmaUnmap)
 	if err != nil {
 		return 0, nil
 	}