From 129d4b63a72b4ea12e452643261de45681b4f6a1 Mon Sep 17 00:00:00 2001
From: Lucas Manning <lucasmanning@google.com>
Date: Mon, 17 Mar 2025 11:49:05 -0700
Subject: [PATCH] Add support for more TPU devices.

PiperOrigin-RevId: 737696637
---
 pkg/abi/tpu/tpu.go    |  9 +++++++++
 runsc/cmd/util/tpu.go | 19 +++++++++++--------
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/pkg/abi/tpu/tpu.go b/pkg/abi/tpu/tpu.go
index 5b104e1d6..e31c64733 100644
--- a/pkg/abi/tpu/tpu.go
+++ b/pkg/abi/tpu/tpu.go
@@ -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.
diff --git a/runsc/cmd/util/tpu.go b/runsc/cmd/util/tpu.go
index c17094560..8b4804dfa 100644
--- a/runsc/cmd/util/tpu.go
+++ b/runsc/cmd/util/tpu.go
@@ -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