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 support for more TPU devices.

Browse Source 
PiperOrigin-RevId: 737696637
This commit is contained in:
Lucas Manning
2025-03-17 11:53:29 -07:00
committed by gVisor bot
parent fbca0560dd
commit 129d4b63a7
2 changed files with 20 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/abi/tpu/tpu.go
+9
View File
@@ -47,6 +47,15 @@ const (
// TPUV6eDeviceID is the PCI device ID of TPU V6e hardware.
TPUV6eDeviceID = 0x006f
// TPUV6ePFDeviceID is PCI device ID of TPU V6e hardware PF.
TPUV6ePFDeviceID = 0x006e
// TPUV6pDeviceID is the PCI device ID of TPU V6p hardware.
TPUV6pDeviceID = 0x0075
// TPUV6pPFDeviceID is the PCI device ID of TPU V6p hardware PF.
TPUV6pPFDeviceID = 0x0076
)
// TPUV4InterruptsMap maps BAR indices to valid register offsets.
runsc/cmd/util/tpu.go
+11 -8
View File
@@ -38,8 +38,11 @@ const (
)
var (
tpuV4DeviceIDs = map[uint64]struct{}{tpu.TPUV4DeviceID: {}, tpu.TPUV4liteDeviceID: {}}
tpuV5DeviceIDs = map[uint64]struct{}{tpu.TPUV5eDeviceID: {}, tpu.TPUV5pDeviceID: {}}
tpuV4DeviceIDs = map[uint64]struct{}{tpu.TPUV4DeviceID: {}, tpu.TPUV4liteDeviceID: {}}
tpuVFIODeviceIDs = map[uint64]struct{}{
tpu.TPUV4DeviceID: {}, tpu.TPUV4liteDeviceID: {}, tpu.TPUV5eDeviceID: {}, tpu.TPUV5pDeviceID: {},
tpu.TPUV6eDeviceID: {}, tpu.TPUV6pDeviceID: {}, tpu.TPUV6ePFDeviceID: {}, tpu.TPUV6pPFDeviceID: {},
}
pciDeviceRegex = regexp.MustCompile(`0000:([[:xdigit:]]{2}|[[:xdigit:]]{4}):[[:xdigit:]]{2}\.[[:xdigit:]]{1,2}`)
)
@@ -76,7 +79,7 @@ func IsTPUDeviceValid(path string) (bool, error) {
if valid {
return valid, err
}
return tpuV5DeviceValid(path)
return tpuVFIODeviceValid(path)
}
// tpuV4DeviceValid returns v4 and v4lite TPU device minor number for the given path.
@@ -114,12 +117,12 @@ func tpuV4DeviceValid(devPath string) (bool, error) {
return true, nil
}
// tpuV5DeviceValid returns the v5e TPU device minor number for te given path.
// A valid v5 TPU device is defined as:
// tpuVFIODeviceValid returns the v5e TPU device minor number for te given path.
// A valid VFIO TPU device is defined as:
// * Path is /dev/vfio/#.
// * Vendor is googleVendorID.
// * Device ID is one of tpuV5DeviceIDs.
func tpuV5DeviceValid(devPath string) (bool, error) {
// * Device ID is one of tpuVFIODeviceIDs.
func tpuVFIODeviceValid(devPath string) (bool, error) {
paths, err := filepath.Glob(fmt.Sprintf(iommuGroupSysfsGlobFormat, path.Base(devPath)))
if err != nil {
return false, err
@@ -139,7 +142,7 @@ func tpuV5DeviceValid(devPath string) (bool, error) {
if err != nil {
return false, err
}
if _, ok := tpuV5DeviceIDs[deviceID]; !ok {
if _, ok := tpuVFIODeviceIDs[deviceID]; !ok {
return false, nil
}
return true, nil