From 17edc7260e94bf9ab48612a616a50c0166cb35b5 Mon Sep 17 00:00:00 2001
From: Jing Chen <chjing@google.com>
Date: Fri, 18 Oct 2024 20:40:56 -0700
Subject: [PATCH] Support TPU v6e w/ gVisor.

PiperOrigin-RevId: 687514246
---
 pkg/abi/tpu/tpu.go  | 3 +++
 runsc/cmd/chroot.go | 1 +
 2 files changed, 4 insertions(+)

diff --git a/pkg/abi/tpu/tpu.go b/pkg/abi/tpu/tpu.go
index 65a6e99e5..5b104e1d6 100644
--- a/pkg/abi/tpu/tpu.go
+++ b/pkg/abi/tpu/tpu.go
@@ -44,6 +44,9 @@ const (
 
 	// TPUV5pDeviceID is the PCI device ID of TPU V5p hardware.
 	TPUV5pDeviceID = 0x0062
+
+	// TPUV6eDeviceID is the PCI device ID of TPU V6e hardware.
+	TPUV6eDeviceID = 0x006f
 )
 
 // TPUV4InterruptsMap maps BAR indices to valid register offsets.
diff --git a/runsc/cmd/chroot.go b/runsc/cmd/chroot.go
index 5912987a7..00f0cf0a5 100644
--- a/runsc/cmd/chroot.go
+++ b/runsc/cmd/chroot.go
@@ -146,6 +146,7 @@ func tpuProxyUpdateChroot(hostRoot, chroot string, spec *specs.Spec, conf *confi
 		}
 		allowedDeviceIDs[tpu.TPUV5pDeviceID] = struct{}{}
 		allowedDeviceIDs[tpu.TPUV5eDeviceID] = struct{}{}
+		allowedDeviceIDs[tpu.TPUV6eDeviceID] = struct{}{}
 	}
 	if len(allowedDeviceIDs) == 0 {
 		paths, err = filepath.Glob(path.Join(hostRoot, "dev/accel*"))