netbird/gvisor
0
0
Fork 0
mirror of https://github.com/netbirdio/gvisor.git synced 2026-05-22 17:12:49 -07:00
Code Issues Packages Projects Releases Wiki Activity

Add sanitization checks for TPU ioctl parameters and rename TPU device methods.

Browse Source 
The methods are renamed so that it's more clear they are explicitly for TPUV4,
not any arbitrary accel device.

PiperOrigin-RevId: 577317824
This commit is contained in:
Lucas Manning
2023-10-27 15:26:45 -07:00
committed by gVisor bot
parent 0010b32f55
commit 44d987d984
13 changed files with 302 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/abi/tpu/BUILD
+11
View File
@@ -0,0 +1,11 @@
load("//tools:defs.bzl", "go_library")
package(default_applicable_licenses = ["//:license"])
licenses(["notice"])
go_library(
name = "tpu",
srcs = ["tpu.go"],
visibility = ["//pkg/sentry:internal"],
)
pkg/abi/tpu/tpu.go
+137
View File
@@ -0,0 +1,137 @@
// 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 tpu defines constants used to interact with TPUs. The constants are
// derived from those listed in https://github.com/tensorflow/tpu/blob/master/tools/driver/drivers/char/tpu_common
package tpu
const (
// SizeOfTPUV4InterruptList is the total number of valid
// (BAR Index, Register Offset) pairs.
SizeOfTPUV4InterruptList = uint64(45)
// NumberOfTPUV4PageTables is the number of gasket page tables used by the
// TPU V4 driver.
NumberOfTPUV4PageTables = uint64(1)
// TPUV4DeviceID is the PCI device ID of TPU V4 hardware.
TPUV4DeviceID = 0x005E
// SizeOfTPUV4liteInterruptList is the total number of valid
// (BAR Index, Register Offset) pairs.
SizeOfTPUV4liteInterruptList = uint64(37)
// NumberOfTPUV4litePageTables is the number of gasket page tables used by the
// TPU V4 driver
NumberOfTPUV4litePageTables = uint64(1)
// TPUV4liteDeviceID is the PCI device ID of TPU V4lite hardware.
TPUV4liteDeviceID = 0x0056
)
// TPUV4InterruptsMap maps BAR indices to valid register offsets.
var (
TPUV4InterruptsMap = map[uint64]map[uint64]struct{}{
2: map[uint64]struct{}{
0x15b0008: struct{}{},
0x15b0000: struct{}{},
0x16b0008: struct{}{},
0x16b0000: struct{}{},
0x17b0008: struct{}{},
0x17b0000: struct{}{},
0x18b0008: struct{}{},
0x18b0000: struct{}{},
0x19b0020: struct{}{},
0x19b0000: struct{}{},
0x19b0008: struct{}{},
0x19b0010: struct{}{},
0x19b0018: struct{}{},
0x1ab0020: struct{}{},
0x1ab0000: struct{}{},
0x1ab0008: struct{}{},
0x1ab0010: struct{}{},
0x1ab0018: struct{}{},
0x4720000: struct{}{},
0x1bb0000: struct{}{},
0x1bb0008: struct{}{},
0x1bb0010: struct{}{},
0x1bb0018: struct{}{},
0x90000: struct{}{},
0xb0000: struct{}{},
0xd0000: struct{}{},
0xf0000: struct{}{},
0x110000: struct{}{},
0x130000: struct{}{},
0x150000: struct{}{},
0x170000: struct{}{},
0x190000: struct{}{},
0x1b0000: struct{}{},
0x1d0000: struct{}{},
0x1f0000: struct{}{},
0x210000: struct{}{},
0x230000: struct{}{},
0x250000: struct{}{},
0x270000: struct{}{},
0x290000: struct{}{},
0x2b0000: struct{}{},
0x2d0000: struct{}{},
0x2f0000: struct{}{},
0x310000: struct{}{},
0x4720018: struct{}{},
},
}
// TPUV4liteInterruptsMap maps BAR indices to valid register offsets.
TPUV4liteInterruptsMap = map[uint64]map[uint64]struct{}{
2: map[uint64]struct{}{
0x19b0020: struct{}{},
0x19b0000: struct{}{},
0x19b0008: struct{}{},
0x19b0010: struct{}{},
0x19b0018: struct{}{},
0x1ab0020: struct{}{},
0x1ab0000: struct{}{},
0x1ab0008: struct{}{},
0x1ab0010: struct{}{},
0x1ab0018: struct{}{},
0x4720000: struct{}{},
0x1bb0000: struct{}{},
0x1bb0008: struct{}{},
0x1bb0010: struct{}{},
0x1bb0018: struct{}{},
0x90000: struct{}{},
0xb0000: struct{}{},
0xd0000: struct{}{},
0xf0000: struct{}{},
0x110000: struct{}{},
0x130000: struct{}{},
0x150000: struct{}{},
0x170000: struct{}{},
0x190000: struct{}{},
0x1b0000: struct{}{},
0x1d0000: struct{}{},
0x1f0000: struct{}{},
0x210000: struct{}{},
0x230000: struct{}{},
0x250000: struct{}{},
0x270000: struct{}{},
0x290000: struct{}{},
0x2b0000: struct{}{},
0x2d0000: struct{}{},
0x2f0000: struct{}{},
0x310000: struct{}{},
0x4720018: struct{}{},
},
}
)
pkg/sentry/devices/accel/BUILD
+4 -3
View File
@@ -6,19 +6,20 @@ licenses(["notice"])
go_library(
name = "accel",
srcs = [
"accel.go",
"accel_mmap.go",
"accel_unsafe.go",
"devaddr_range.go",
"devaddr_set.go",
"device.go",
"gasket.go",
"gasket_unsafe.go",
"seccomp_filters.go",
"tpu_v4.go",
"tpu_v4_mmap.go",
],
visibility = ["//pkg/sentry:internal"],
deps = [
"//pkg/abi/gasket",
"//pkg/abi/linux",
"//pkg/abi/tpu",
"//pkg/cleanup",
"//pkg/context",
"//pkg/errors/linuxerr",
pkg/sentry/devices/accel/device.go
+20 -7
View File
@@ -20,26 +20,31 @@ import (
"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/sentry/fsimpl/devtmpfs"
"gvisor.dev/gvisor/pkg/sentry/kernel"
"gvisor.dev/gvisor/pkg/sentry/vfs"
"gvisor.dev/gvisor/pkg/sync"
)
// accelDevice implements vfs.Device for /dev/accel[0-9]+.
// tpuV4Device implements vfs.Device for /dev/accel[0-9]+.
//
// +stateify savable
type accelDevice struct {
type tpuV4Device struct {
mu sync.Mutex
minor uint32
lite bool
// +checklocks:mu
openWriteFDs uint32
// +checklocks:mu
devAddrSet DevAddrSet
// +checklocks:mu
owner *kernel.ThreadGroup
}
func (dev *accelDevice) Open(ctx context.Context, mnt *vfs.Mount, vfsd *vfs.Dentry, opts vfs.OpenOptions) (*vfs.FileDescription, error) {
func (dev *tpuV4Device) Open(ctx context.Context, mnt *vfs.Mount, vfsd *vfs.Dentry, opts vfs.OpenOptions) (*vfs.FileDescription, error) {
dev.mu.Lock()
defer dev.mu.Unlock()
hostPath := fmt.Sprintf("/dev/accel%d", dev.minor)
@@ -48,7 +53,7 @@ func (dev *accelDevice) Open(ctx context.Context, mnt *vfs.Mount, vfsd *vfs.Dent
ctx.Warningf("accelDevice: failed to open host %s: %v", hostPath, err)
return nil, err
}
fd := &accelFD{
fd := &tpuV4FD{
hostFD: int32(hostFD),
device: dev,
}
@@ -66,6 +71,13 @@ func (dev *accelDevice) Open(ctx context.Context, mnt *vfs.Mount, vfsd *vfs.Dent
if vfs.MayWriteFileWithOpenFlags(opts.Flags) {
dev.openWriteFDs++
}
if dev.owner == nil {
t := kernel.TaskFromContext(ctx)
if t == nil {
return nil, linuxerr.ESRCH
}
dev.owner = t.ThreadGroup()
}
return &fd.vfsfd, nil
}
@@ -74,9 +86,10 @@ func CreateDevtmpfsFile(ctx context.Context, dev *devtmpfs.Accessor, num uint32)
return dev.CreateDeviceFile(ctx, fmt.Sprintf("accel%d", num), vfs.CharDevice, linux.ACCEL_MAJOR, num, 0666)
}
// Register registers all devices implemented by this package in vfsObj.
func Register(vfsObj *vfs.VirtualFilesystem, minor uint32) error {
return vfsObj.RegisterDevice(vfs.CharDevice, linux.ACCEL_MAJOR, minor, &accelDevice{
// RegisterTPUV4Device registers all devices implemented by this package in vfsObj.
func RegisterTPUV4Device(vfsObj *vfs.VirtualFilesystem, minor uint32, lite bool) error {
return vfsObj.RegisterDevice(vfs.CharDevice, linux.ACCEL_MAJOR, minor, &tpuV4Device{
lite: lite,
minor: minor,
}, &vfs.RegisterDeviceOptions{
GroupName: "accel",
pkg/sentry/devices/accel/gasket.go
+65 -4
View File
@@ -20,6 +20,7 @@ import (
"golang.org/x/sys/unix"
"gvisor.dev/gvisor/pkg/abi/gasket"
"gvisor.dev/gvisor/pkg/abi/linux"
"gvisor.dev/gvisor/pkg/abi/tpu"
"gvisor.dev/gvisor/pkg/cleanup"
"gvisor.dev/gvisor/pkg/context"
"gvisor.dev/gvisor/pkg/errors/linuxerr"
@@ -30,12 +31,20 @@ import (
"gvisor.dev/gvisor/pkg/sentry/mm"
)
func gasketMapBufferIoctl(ctx context.Context, t *kernel.Task, hostFd int32, fd *accelFD, paramsAddr hostarch.Addr) (uintptr, error) {
func gasketMapBufferIoctl(ctx context.Context, t *kernel.Task, hostFd int32, fd *tpuV4FD, paramsAddr hostarch.Addr) (uintptr, error) {
var userIoctlParams gasket.GasketPageTableIoctl
if _, err := userIoctlParams.CopyIn(t, paramsAddr); err != nil {
return 0, err
}
numberOfPageTables := tpu.NumberOfTPUV4PageTables
if fd.device.lite {
numberOfPageTables = tpu.NumberOfTPUV4litePageTables
}
if userIoctlParams.PageTableIndex >= numberOfPageTables {
return 0, linuxerr.EFAULT
}
tmm := t.MemoryManager()
ar, ok := tmm.CheckIORange(hostarch.Addr(userIoctlParams.HostAddress), int64(userIoctlParams.Size))
if !ok {
@@ -45,6 +54,23 @@ func gasketMapBufferIoctl(ctx context.Context, t *kernel.Task, hostFd int32, fd
if !ar.IsPageAligned() || (userIoctlParams.Size/hostarch.PageSize) == 0 {
return 0, linuxerr.EINVAL
}
devAddr := userIoctlParams.DeviceAddress
// The kernel driver does not enforce page alignment on the device
// address although it will be implicitly rounded down to a page
// boundary. We do it explicitly because it simplifies tracking
// of allocated ranges in 'devAddrSet'.
devAddr &^= (hostarch.PageSize - 1)
// Make sure that the device address range can be mapped.
devar := DevAddrRange{
devAddr,
devAddr + userIoctlParams.Size,
}
if !devar.WellFormed() {
return 0, linuxerr.EINVAL
}
// Reserve a range in our address space.
m, _, errno := unix.RawSyscall6(unix.SYS_MMAP, 0 /* addr */, uintptr(ar.Length()), unix.PROT_NONE, unix.MAP_PRIVATE|unix.MAP_ANONYMOUS, ^uintptr(0) /* fd */, 0 /* offset */)
if errno != 0 {
@@ -89,7 +115,6 @@ func gasketMapBufferIoctl(ctx context.Context, t *kernel.Task, hostFd int32, fd
fd.device.mu.Lock()
defer fd.device.mu.Unlock()
devAddr := userIoctlParams.DeviceAddress
for _, pr := range prs {
rlen := uint64(pr.Source.Length())
if !fd.device.devAddrSet.Add(DevAddrRange{
@@ -103,11 +128,30 @@ 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 *accelFD, paramsAddr hostarch.Addr) (uintptr, error) {
func gasketUnmapBufferIoctl(ctx context.Context, t *kernel.Task, hostFd int32, fd *tpuV4FD, paramsAddr hostarch.Addr) (uintptr, error) {
var userIoctlParams gasket.GasketPageTableIoctl
if _, err := userIoctlParams.CopyIn(t, paramsAddr); err != nil {
return 0, err
}
numberOfPageTables := tpu.NumberOfTPUV4PageTables
if fd.device.lite {
numberOfPageTables = tpu.NumberOfTPUV4litePageTables
}
if userIoctlParams.PageTableIndex >= numberOfPageTables {
return 0, linuxerr.EFAULT
}
devAddr := userIoctlParams.DeviceAddress
devAddr &^= (hostarch.PageSize - 1)
devar := DevAddrRange{
devAddr,
devAddr + userIoctlParams.Size,
}
if !devar.WellFormed() {
return 0, linuxerr.EINVAL
}
sentryIoctlParams := userIoctlParams
sentryIoctlParams.HostAddress = 0 // clobber this value, it's unused.
n, err := ioctlInvokePtrArg(hostFd, gasket.GASKET_IOCTL_UNMAP_BUFFER, &sentryIoctlParams)
@@ -129,12 +173,29 @@ func gasketUnmapBufferIoctl(ctx context.Context, t *kernel.Task, hostFd int32, f
return n, nil
}
func gasketInterruptMappingIoctl(ctx context.Context, t *kernel.Task, hostFd int32, paramsAddr hostarch.Addr) (uintptr, error) {
func gasketInterruptMappingIoctl(ctx context.Context, t *kernel.Task, hostFd int32, paramsAddr hostarch.Addr, lite bool) (uintptr, error) {
var userIoctlParams gasket.GasketInterruptMapping
if _, err := userIoctlParams.CopyIn(t, paramsAddr); err != nil {
return 0, err
}
sizeOfInterruptList := tpu.SizeOfTPUV4InterruptList
interruptMap := tpu.TPUV4InterruptsMap
if lite {
sizeOfInterruptList = tpu.SizeOfTPUV4liteInterruptList
interruptMap = tpu.TPUV4liteInterruptsMap
}
if userIoctlParams.Interrupt >= sizeOfInterruptList {
return 0, linuxerr.EINVAL
}
barRegMap, ok := interruptMap[userIoctlParams.BarIndex]
if !ok {
return 0, linuxerr.EINVAL
}
if _, ok := barRegMap[userIoctlParams.RegOffset]; !ok {
return 0, linuxerr.EINVAL
}
// Check that 'userEventFD.Eventfd' is an eventfd.
eventFileGeneric, _ := t.FDTable().Get(int32(userIoctlParams.EventFD))
if eventFileGeneric == nil {
pkg/sentry/devices/accel/accel_unsafe.go → pkg/sentry/devices/accel/gasket_unsafe.go
View File
pkg/sentry/devices/accel/seccomp_filters.go
-28
View File
@@ -41,30 +41,6 @@ func Filters() seccomp.SyscallRules {
nonNegativeFD,
seccomp.EqualTo(gasket.GASKET_IOCTL_RESET),
},
seccomp.PerArg{
nonNegativeFD,
seccomp.EqualTo(gasket.GASKET_IOCTL_SET_EVENTFD),
},
seccomp.PerArg{
nonNegativeFD,
seccomp.EqualTo(gasket.GASKET_IOCTL_CLEAR_EVENTFD),
},
seccomp.PerArg{
nonNegativeFD,
seccomp.EqualTo(gasket.GASKET_IOCTL_NUMBER_PAGE_TABLES),
},
seccomp.PerArg{
nonNegativeFD,
seccomp.EqualTo(gasket.GASKET_IOCTL_PAGE_TABLE_SIZE),
},
seccomp.PerArg{
nonNegativeFD,
seccomp.EqualTo(gasket.GASKET_IOCTL_SIMPLE_PAGE_TABLE_SIZE),
},
seccomp.PerArg{
nonNegativeFD,
seccomp.EqualTo(gasket.GASKET_IOCTL_PARTITION_PAGE_TABLE),
},
seccomp.PerArg{
nonNegativeFD,
seccomp.EqualTo(gasket.GASKET_IOCTL_MAP_BUFFER),
@@ -85,10 +61,6 @@ func Filters() seccomp.SyscallRules {
nonNegativeFD,
seccomp.EqualTo(gasket.GASKET_IOCTL_UNREGISTER_INTERRUPT),
},
seccomp.PerArg{
nonNegativeFD,
seccomp.EqualTo(gasket.GASKET_IOCTL_MAP_DMA_BUF),
},
},
unix.SYS_EVENTFD2: seccomp.Or{
seccomp.PerArg{
pkg/sentry/devices/accel/accel.go → pkg/sentry/devices/accel/tpu_v4.go
+27 -10
View File
@@ -34,24 +34,24 @@ import (
"gvisor.dev/gvisor/pkg/waiter"
)
// accelFD implements vfs.FileDescriptionImpl for /dev/accel[0-9]+.
// tpuV4FD implements vfs.FileDescriptionImpl for /dev/accel[0-9]+.
//
// accelFD is not savable; we do not implement save/restore of accelerator
// state.
type accelFD struct {
type tpuV4FD struct {
vfsfd vfs.FileDescription
vfs.FileDescriptionDefaultImpl
vfs.DentryMetadataFileDescriptionImpl
vfs.NoLockFD
hostFD int32
device *accelDevice
device *tpuV4Device
queue waiter.Queue
memmapFile accelFDMemmapFile
}
// Release implements vfs.FileDescriptionImpl.Release.
func (fd *accelFD) Release(context.Context) {
func (fd *tpuV4FD) Release(context.Context) {
fd.device.mu.Lock()
defer fd.device.mu.Unlock()
fd.device.openWriteFDs--
@@ -75,13 +75,14 @@ func (fd *accelFD) Release(context.Context) {
gap := s.Remove(seg)
seg = gap.NextSegment()
}
fd.device.owner = nil
}
fdnotifier.RemoveFD(fd.hostFD)
unix.Close(int(fd.hostFD))
}
// EventRegister implements waiter.Waitable.EventRegister.
func (fd *accelFD) EventRegister(e *waiter.Entry) error {
func (fd *tpuV4FD) EventRegister(e *waiter.Entry) error {
fd.queue.EventRegister(e)
if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
fd.queue.EventUnregister(e)
@@ -91,7 +92,7 @@ func (fd *accelFD) EventRegister(e *waiter.Entry) error {
}
// EventUnregister implements waiter.Waitable.EventUnregister.
func (fd *accelFD) EventUnregister(e *waiter.Entry) {
func (fd *tpuV4FD) EventUnregister(e *waiter.Entry) {
fd.queue.EventUnregister(e)
if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
panic(fmt.Sprint("UpdateFD:", err))
@@ -99,17 +100,17 @@ func (fd *accelFD) EventUnregister(e *waiter.Entry) {
}
// Readiness implements waiter.Waitable.Readiness.
func (fd *accelFD) Readiness(mask waiter.EventMask) waiter.EventMask {
func (fd *tpuV4FD) Readiness(mask waiter.EventMask) waiter.EventMask {
return fdnotifier.NonBlockingPoll(fd.hostFD, mask)
}
// Epollable implements vfs.FileDescriptionImpl.Epollable.
func (fd *accelFD) Epollable() bool {
func (fd *tpuV4FD) Epollable() bool {
return true
}
// Ioctl implements vfs.FileDescriptionImpl.Ioctl.
func (fd *accelFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, args arch.SyscallArguments) (uintptr, error) {
func (fd *tpuV4FD) 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)
@@ -118,6 +119,9 @@ func (fd *accelFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, arg
if t == nil {
panic("Ioctl should be called from a task context")
}
if err := fd.checkPermission(t); err != nil {
return 0, err
}
log.Infof("Accel ioctl %s called on fd %d with arg %v of size %d.", gasket.Ioctl(cmd), fd.hostFD, argPtr, argSize)
switch gasket.Ioctl(cmd) {
@@ -136,7 +140,7 @@ func (fd *accelFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, arg
case gasket.GASKET_IOCTL_CLEAR_INTERRUPT_COUNTS:
return ioctlInvoke(fd.hostFD, gasket.GASKET_IOCTL_CLEAR_INTERRUPT_COUNTS, 0)
case gasket.GASKET_IOCTL_REGISTER_INTERRUPT:
return gasketInterruptMappingIoctl(ctx, t, fd.hostFD, argPtr)
return gasketInterruptMappingIoctl(ctx, t, fd.hostFD, argPtr, fd.device.lite)
case gasket.GASKET_IOCTL_UNREGISTER_INTERRUPT:
return ioctlInvoke[uint64](fd.hostFD, gasket.GASKET_IOCTL_UNREGISTER_INTERRUPT, args[2].Uint64())
default:
@@ -144,6 +148,19 @@ func (fd *accelFD) 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 {
fd.device.mu.Lock()
defer fd.device.mu.Unlock()
owner := fd.device.owner
if t.ThreadGroup() != owner {
return linuxerr.EPERM
}
return nil
}
type pinnedAccelMem struct {
pinnedRange mm.PinnedRange
pageTableIndex uint64
pkg/sentry/devices/accel/accel_mmap.go → pkg/sentry/devices/accel/tpu_v4_mmap.go
+7 -7
View File
@@ -25,26 +25,26 @@ import (
)
// ConfigureMMap implements vfs.FileDescriptionImpl.ConfigureMMap.
func (fd *accelFD) ConfigureMMap(ctx context.Context, opts *memmap.MMapOpts) error {
func (fd *tpuV4FD) ConfigureMMap(ctx context.Context, opts *memmap.MMapOpts) error {
return vfs.GenericConfigureMMap(&fd.vfsfd, fd, opts)
}
// AddMapping implements memmap.Mappable.AddMapping.
func (fd *accelFD) AddMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) error {
func (fd *tpuV4FD) AddMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) error {
return nil
}
// RemoveMapping implements memmap.Mappable.RemoveMapping.
func (fd *accelFD) RemoveMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) {
func (fd *tpuV4FD) RemoveMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) {
}
// CopyMapping implements memmap.Mappable.CopyMapping.
func (fd *accelFD) CopyMapping(ctx context.Context, ms memmap.MappingSpace, srcAR, dstAR hostarch.AddrRange, offset uint64, writable bool) error {
func (fd *tpuV4FD) 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 *accelFD) Translate(ctx context.Context, required, optional memmap.MappableRange, at hostarch.AccessType) ([]memmap.Translation, error) {
func (fd *tpuV4FD) Translate(ctx context.Context, required, optional memmap.MappableRange, at hostarch.AccessType) ([]memmap.Translation, error) {
return []memmap.Translation{
{
Source: optional,
@@ -56,12 +56,12 @@ func (fd *accelFD) Translate(ctx context.Context, required, optional memmap.Mapp
}
// InvalidateUnsavable implements memmap.Mappable.InvalidateUnsavable.
func (fd *accelFD) InvalidateUnsavable(ctx context.Context) error {
func (fd *tpuV4FD) InvalidateUnsavable(ctx context.Context) error {
return nil
}
type accelFDMemmapFile struct {
fd *accelFD
fd *tpuV4FD
}
// IncRef implements memmap.File.IncRef.
runsc/boot/BUILD
+1
View File
@@ -32,6 +32,7 @@ go_library(
deps = [
"//pkg/abi",
"//pkg/abi/linux",
"//pkg/abi/tpu",
"//pkg/bpf",
"//pkg/cleanup",
"//pkg/context",
runsc/boot/vfs.go
+26 -3
View File
@@ -16,6 +16,7 @@ package boot
import (
"fmt"
"os"
"path"
"path/filepath"
"regexp"
@@ -25,6 +26,7 @@ import (
specs "github.com/opencontainers/runtime-spec/specs-go"
"gvisor.dev/gvisor/pkg/abi/linux"
"gvisor.dev/gvisor/pkg/abi/tpu"
"gvisor.dev/gvisor/pkg/cleanup"
"gvisor.dev/gvisor/pkg/context"
"gvisor.dev/gvisor/pkg/errors/linuxerr"
@@ -1168,11 +1170,32 @@ func tpuProxyRegisterDevicesAndCreateFiles(ctx context.Context, info *containerI
if err != nil {
return fmt.Errorf("enumerating accel device files: %w", err)
}
for _, path := range paths {
pciAddrs, err := filepath.Glob("/sys/devices/pci0000:00/*")
if err != nil {
return fmt.Errorf("enumerating PCI device files: %w", err)
}
for _, accelPath := range paths {
accelDeviceRegex := regexp.MustCompile(`^/dev/accel(\d+)$`)
if ms := accelDeviceRegex.FindStringSubmatch(path); ms != nil {
if ms := accelDeviceRegex.FindStringSubmatch(accelPath); ms != nil {
deviceNum, _ := strconv.ParseUint(ms[1], 10, 32)
if err := accel.Register(vfsObj, uint32(deviceNum)); err != nil {
var pciDevicePath string
for _, pciPath := range pciAddrs {
if _, err := os.Stat(path.Join(pciPath, fmt.Sprintf("accel/accel%d", deviceNum))); err == nil {
pciDevicePath = pciPath
}
}
var deviceIDBytes []byte
if deviceIDBytes, err = os.ReadFile(path.Join(pciDevicePath, "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 err := accel.RegisterTPUV4Device(vfsObj, uint32(deviceNum), deviceID == tpu.TPUV4liteDeviceID); err != nil {
return fmt.Errorf("registering accel driver: %w", err)
}
if err := accel.CreateDevtmpfsFile(ctx, a, uint32(deviceNum)); err != nil {
runsc/cmd/util/BUILD
+1
View File
@@ -17,6 +17,7 @@ go_library(
"//tools:__subpackages__",
],
deps = [
"//pkg/abi/tpu",
"//pkg/log",
"@com_github_google_subcommands//:go_default_library",
],
runsc/cmd/util/tpu.go
+3 -1
View File
@@ -21,11 +21,13 @@ import (
"regexp"
"strconv"
"strings"
"gvisor.dev/gvisor/pkg/abi/tpu"
)
const googleVendorID = 0x1AE0
var tpuV4DeviceIDs = map[uint64]any{0x005E: nil, 0x0056: nil}
var tpuV4DeviceIDs = map[uint64]any{tpu.TPUV4DeviceID: nil, tpu.TPUV4liteDeviceID: nil}
// TODO(b/288456802): Add support for /dev/vfio controlled accelerators.
// This is required for v5+ TPUs.