TPU v4 最小用户态 Runtime

复用现有 Gasket/TPU 内核驱动，从固定 XLA final bundle 出发，不经过 libtpu/PJRT execution API，完成一次 TPU v4-8 的提交、执行、等待和结果回读。

这是什么

这个仓库是一份已经在锁定 TPU v4-8 环境中验证通过的最小用户态 runtime。它执行的固定程序是：

x = np.arange(8, dtype=np.float32)
y = x + np.float32(1.0)

运行阶段的主程序是 src/minimal_tpu_v4_runtime.c。它不链接、不加载、也不调用 libtpu.so，动态链接依赖只有 libc 和系统 loader；底层仍复用机器上已有的 Gasket/TPU 内核模块。

固定 PJRT executable / bootstrap / continuation / input
                           │
                           ▼
              minimal_tpu_v4_runtime
          DMA / HostQueue / MMIO / continuator
                           │
                           ▼
               Gasket / TPU kernel modules
                           │
                           ▼
                        TPU v4-8
                           │
                           ▼
                  result readback + hash

成功结果为：

PASS submit=1 execute=1 wait=1 readback=1 result_sha256=af7de0621354bafceb193edf0fcf5d421cf21de7146580062fff53c7907f54e5

完整的调查过程、术语解释、失败路径和证据边界见 blog_post.md。

快速运行

只能在下文锁定的 TPU v4-8、Linux 内核和固件环境中运行。开始前确认 /dev/accel0..3 没有被 JAX、libtpu 或其他进程占用，并允许当前进程锁定约 4 GiB host memory。

cd ~/tpu-v4-userspace-runtime
make clean all
make run

运行末尾应出现：

RENDEZVOUS_ENQUEUED devices=4 direct_continuations=10 fish=128
QUEUE_COMPLETE name=readback head=1
PASS submit=1 execute=1 wait=1 readback=1 result_sha256=af7de0621354bafceb193edf0fcf5d421cf21de7146580062fff53c7907f54e5

只有最后的 PASS 才表示成功。QUEUE_COMPLETE 只证明相应 queue 的 descriptor 已被消费，不能单独证明用户程序已经执行。

做了什么

一次 make run 会：

1. 校验 executable、内嵌 final program、bootstrap、continuation、input 和 expected output 的 size 与 SHA-256；
2. 在打开设备前核对 kernel、内核模块、firmware、BAR layout 和四卡空闲状态；
3. 打开 /dev/accel0..3，mmap BAR2，并通过 Gasket ioctl 建立 DMA/IOMMU mapping 和 interrupt eventfd；
4. 在每张卡上建立 direct、auxiliary、magic、infeed、outfeed 和 readback，共 19 条 HostQueue；
5. 配置 logical chip ID、sync flags、SMEM/HBM channel 和 metadata channel；
6. 在每张卡上加载 bootstrap，并按捕获到的状态机启动两个 continuator core；
7. 在 /dev/accel0 加载 input 和固定 final program；
8. 向四张卡整体提交 continuation 和 rendezvous batch；
9. 等待 direct、magic 和 readback completion；
10. 比较回读的八个 float32，并独立校验结果 SHA-256。

编译和执行的边界如下：

阶段	负责者
从 JAX 程序生成 fixed executable/final program	现有 JAX/XLA/libtpu compiler
从 fixed executable 到设备数值结果	本仓库的用户态 runtime
BAR、DMA/IOMMU、interrupt 与设备 ownership	现有 Gasket/TPU 内核模块

仓库结构

.
├── src/
│   └── minimal_tpu_v4_runtime.c   # runtime 实现
├── fixture/
│   ├── add_one_f32_8_final_bundle.pjrt
│   ├── tpu_v4_bootstrap_program.bin
│   ├── continuation_payload.bin
│   ├── input_f32_8.bin
│   ├── expected_f32_8.bin
│   └── metadata.json              # 来源、大小与 hash
├── tools/
│   ├── trace_accel.c              # ioctl/mmap/eventfd observer
│   └── capture_host_queue.py      # GDB 内的 DMA/HostQueue/MMIO 捕获器
├── scripts/
│   ├── generate_fixed_executable.py
│   └── prepare_fixed_fixture.py
├── build/                         # make 生成，不提交到 Git
├── Makefile
├── README.md
└── blog_post.md                   # 完整技术故事

日常执行只需要 src/、fixture/ 和 Makefile。tools/ 与 scripts/ 用于重新建立 oracle、审计 fixture 或恢复执行路径，不参与 make run。

固定 Fixture

文件	大小	SHA-256	用途
add_one_f32_8_final_bundle.pjrt	55,489 B	5448890029cf2f7db2f7dff98bc069908639df1f6652b67a589047e1703bad79	serialized executable；runtime 从 offset 237 取出 program
内嵌 final program	34,304 B	cd52f58ae800ca904f04f939adfa2d9bc3168ce2b244c96419b377cea0b8b587	固定设备程序
tpu_v4_bootstrap_program.bin	13,824 B	8f13773022a47b5eaec9d74736175f11a3a95d90fe5ac17e324a71db5a0d5ea7	设备侧 bootstrap
continuation_payload.bin	512 B	a3715e2f8e275eb96f633704398ee9711f92aae1c3737c0200cf8f6b03b020f6	固定 invocation payload
input_f32_8.bin	32 B	0571cfe42be5c7b95de9afc7c7ba1286fb7a2ef10a9035f8d6b87d21a3bc8387	[0, 1, ..., 7]
expected_f32_8.bin	32 B	af7de0621354bafceb193edf0fcf5d421cf21de7146580062fff53c7907f54e5	host 侧最终比较使用

expected_f32_8.bin 不会发送给 TPU，只在 readback 完成后参与 memcmp() 和结果判定。

重新生成 fixture 分为两步：

• scripts/generate_fixed_executable.py 使用正常 JAX/PJRT/libtpu 路径建立数值 oracle 和 serialized executable；
• scripts/prepare_fixed_fixture.py 从精确版本的 executable、libtpu binary 和捕获数据构造当前 fixture。

执行 proof 本身不需要运行这两个脚本。

其他命令

构建：

make clean all

确认执行程序没有链接 libtpu：

ldd build/minimal_tpu_v4_runtime

预期只看到 libc、动态 loader 和 linux-vdso。

展开 make run：

./build/minimal_tpu_v4_runtime \
  fixture/add_one_f32_8_final_bundle.pjrt \
  fixture/tpu_v4_bootstrap_program.bin \
  fixture/continuation_payload.bin \
  fixture/input_f32_8.bin \
  fixture/expected_f32_8.bin

只读取 queue capability registers：

make probe

make probe 仍会打开 /dev/accel0。当前驱动配置有 reset-on-open/close 行为，因此也必须在设备空闲、版本完全匹配时运行。

锁定环境

runtime 在第一次 open() 前检查：

项目	锁定值
Linux kernel	5.19.0-1022-gcp
accel_class srcversion	8C26C63222F966A173F0575
gasket srcversion	C986B1180B427F33A5B5B9B
tpu_common srcversion	7CBE3556A273D1C6724BA3D
tpuv4common_common srcversion	3338ED6F488D02388E47044
tpu_v4 srcversion	64DE95EA4354F18FF9FF82C
Driver/framework	0.0.1 / 1.1.1
Firmware	2025.40.3.0
Hardware revision	16
BAR2 offset/size	0x10000000 / 0x80000000

任一条件不匹配都会 fail closed。不要为了在另一台机器运行而删掉检查；应在新环境重新建立 oracle、动态捕获和设备数值验证。

证明范围

已经证明：

• 使用现有 /dev/accel0..3 和 Gasket/TPU 内核模块；
• 执行进程不经过 libtpu/PJRT execution API；
• 用户态直接完成 BAR2 MMIO、DMA mapping、HostQueue、设备初始化、程序提交、等待和回读；
• 锁定 final program 在锁定 TPU v4-8 上返回正确数值和 hash；
• 正常与错误路径结束后可以关闭 queue 和设备资源。

尚未证明：

• 可以运行任意 PJRT executable；
• 已恢复 descriptor、continuation 或设备内部地址的稳定通用格式；
• 当前常量适用于其他 kernel、firmware、libtpu 或 TPU 世代；
• 已实现通用 allocator、并发执行、取消、错误恢复或 PJRT API；
• 可以卸载现有内核模块后自行管理 PCIe、IOMMU 和 firmware。

资源与恢复

当前实现为每张卡的两个 outfeed queue 各预挂 16 个 32 MiB buffer，总计约 4 GiB 锁页 host memory。这是对官方 runtime 初始化的忠实复现，不是固定 x + 1 计算本身的内存需求。

正常和错误路径都会 disable queue、撤销 BAR mapping、关闭 device/eventfd，并释放 host DMA memory。当前驱动配置为 reset_on_close=1。

再次运行前可以检查：

for i in 0 1 2 3; do
  printf 'accel%s status=%s state=%s owned=%s opens=%s\n' \
    "$i" \
    "$(</sys/class/accel/accel$i/status)" \
    "$(</sys/class/accel/accel$i/state)" \
    "$(</sys/class/accel/accel$i/is_device_owned)" \
    "$(</sys/class/accel/accel$i/write_open_count)"
done

预期每张卡都是：

status=ALIVE state=available owned=0 opens=0

如果状态不符，不要绕过检查继续写 BAR；先停止占用进程，并按机器既有运维流程恢复设备。

阅读代码

建议按这个顺序阅读：

1. main()：四卡执行生命周期；
2. check_environment()：版本与设备状态边界；
3. runtime_open()：BAR、interrupt、queue 和 outfeed；
4. dma_allocate()：host memory 到 IOVA；
5. queue_open()、queue_submit()、queue_wait()：最小 HostQueue；
6. configure_cores()、configure_platform_channels()：设备控制状态；
7. initialize_chip()：双 continuator core 状态机；
8. tools/trace_accel.c 与 tools/capture_host_queue.py：执行路径如何被观察和恢复。

研究背景、失败过程和方法总结见 blog_post.md。

当前限制与下一步

• 用真实 PJRT parser 和 buffer assignment 取代固定 offset；
• 为 descriptor 建立经过多样本验证的 encoder/decoder；
• 增加 IOVA allocator 和映射冲突检查；
• 恢复间接寄存器 busy/status polling，移除固定 settle delay；
• 缩减不必要的 outfeed pool 与 rendezvous 请求；
• 支持第二个 shape/dtype/program，寻找真正稳定的 runtime contract；
• 最后再考虑 PJRT-compatible API 与其他 TPU 版本 backend。

每扩展一层，都应保留官方 oracle、版本指纹、queue 状态和设备数值对照。