Internal Wiki

主题

运行时与执行概览 ​

本页中的所有地址均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（构建 libtpu_lts_20260413_b_RC00，build-id md5 89edbbe81c5b328a958fe628a9f2207d，781,691,048 字节，ELF x86-64 DYN，clang/LLVM trunk）。该镜像未被剥离；反修饰后的 C++ 符号名按原样引用。其他版本会有所不同。

摘要 ​

libtpu 的运行时层是位于 PJRT C API 之下、TPU 驱动核心之上的执行引擎。它是 libtpu 中对应 XLA StreamExecutor + PjRt 执行栈的部分：接收已经编译好的 TPU 程序和一组参数缓冲区，将这些缓冲区编排成设备期望的形状，把程序加载到物理 TPU 核心上，将请求推入该核心的命令流，并异步跟踪完成状态，使调用方的 PJRT_Event 在设备完成时解析。编译（HLO -> TpuCoreProgram）位于这一层的上游；片上队列硅逻辑和 TfTpu_*ApiFn 驱动表位于它的下游。从“已存在编译后的可执行文件及其输入缓冲区”到“输出在 HBM 中可用且完成事件已触发”之间的所有内容，都是这些页面要重建的对象。

最重要的结构性事实，也是重新实现者在阅读任何同级页面前必须内化的一点，是 libtpu 在同一个镜像中编译了两个并行执行栈，它们是入口，不是运行时替代方案。现代 PJRT 路径是 xla::TpuClient（派生自 xla::CommonPjRtClient : xla::PjRtClient），构建在 TFRT 原生设备运行时 tpu::System 之上；每个 PJRT 客户端（JAX、PyTorch-XLA）都会驱动它，它原生使用 async-value（tsl::AsyncValueRef），并且在其代码范围内完全没有对 stream_executor::Stream、TpuExecutor 或 ExecutorApiFn 的引用。旧版 StreamExecutor 路径是 xla::legacy::TpuExecutableInterface / xla::jellyfish::DeepseaExecutable，构建在 stream_executor::tpu::TpuExecutor + TpuStream 之上；xla::LocalClient / xla::Service 和 TF-TPU op kernels 通过 Tpu*_* C-ABI 驱动它。两者最终都落到相同的 deepsea 驱动核心（asic_sw::driver::deepsea::jxc::Queue::EnqueueRequest），但它们从不共享宿主侧对象：现代路径用 TpuEventIssuer 依赖图对工作排序；旧版路径用 stream_executor::Stream FIFO 对工作排序。任务描述中的“TpuExecutable::ExecuteAsyncOnStream”命名的是旧版虚函数；现代路径实际使用的符号（tpu::System::Execute）在完全不进入 ExecuteAsyncOnStream 的情况下抵达等价的入队操作。

本页是运行时层的章节地图。它确定双栈架构，给出规范的 SE 概念 -> TPU 运行时映射，把端到端执行生命周期作为一张并列展示两种写法的图展开，并用一两句话概括每个子领域，同时链接到负责细节的页面。它不会复述参数/输出编排（execute-async-on-stream.md）、程序加载 + 入队内部细节（load-program-enqueue.md）、WaitFor/RecordEvent 依赖模型（stream-semantics.md）、infeed/outfeed 队列机制（infeed-outfeed.md），或宿主侧分配器与完成管线；这些内容各自由专门的同级页面负责。

作为定位，本页的契约是：

  • 双栈模型 — 给定客户端驱动哪个入口（TpuClient/tpu::System 还是 TpuExecutableInterface/TpuExecutor），以及它们只共享驱动核心这一事实。
  • 执行生命周期 — PJRT Execute -> prepare/launch/raw -> load -> enqueue -> device run -> completion -> outputs 链条，两种写法中的表现，以及层边界所在位置。
  • SE 概念 -> TPU 运行时映射 — 熟悉 GPU/StreamExecutor 的读者用来把 StreamExecutor/Stream/Event/DeviceMemoryBase 翻译成 tpu::System/TpuEventIssuer/TpuEvent/TpuSharedMemoryLocation 的表。
  • 子领域地图 — 十个运行时页面及各自负责的内容，让读者可以凭直觉导航。
现代 PJRT 入口PJRT_LoadedExecutable_Execute（槽 60）@ 0xf869b40 -> CommonPjRtLoadedExecutable::Execute
现代设备运行时tpu::System — Execute @ 0x1d0b33e0，LoadProgram @ 0x1d0b2240（system.cc:1804）
旧版 SE 入口xla::legacy::TpuExecutableInterface::ExecuteAsyncOnStream @ 0x1342cd20（3650 B）
旧版入队叶子（vtable+96）xla::jellyfish::DeepseaExecutable::LoadProgramAndEnqueueToStream @ 0x13426260（1360 行）
旧版命令流deepsea::executor::DeepseaStream::EnqueueRequest @ 0x1d0e9840
共享驱动核心（两个栈）asic_sw::driver::deepsea::jxc::Queue::EnqueueRequest @ 0xe7d9be0（DmaBuffer + completion）
现代客户端类xla::TpuClient : CommonPjRtClient : PjRtClient（构造函数 0xf801980，vtable 0x2177b598）
现代设备排序tpu::TpuEventIssuer 序列点（IssueArgs/FulfillArgs/AddDepsNoReserve）
设备内存身份tpu::TpuSharedMemoryLocation + tpu::TpuBuffer（现代）/ stream_executor::DeviceAddressBase（旧版）
置信度CONFIRMED（字节锚定），除非某行或标注另有说明

1. 两个执行栈 ​

目的 ​

重新实现者的第一个决定是要建模两个设备抽象中的哪一个，因为两者在驱动之上从不汇合。两者都存在于这个单一二进制中；选择由前端使用的客户端 API 决定，而不是由运行时标志决定。本节确定这种分裂；每个同级页面都针对其中一个或两个栈编写。

分裂 ​

text
                         PJRT C API (140 slots)
                                       │
              ┌────────────────────────┴───────────────────────┐
   MODERN (every PJRT client: JAX, PyTorch-XLA)        LEGACY (LocalClient / Service / TF-TPU op kernels)
              │                                                 │
   xla::TpuClient : CommonPjRtClient : PjRtClient       xla::legacy::TpuExecutableInterface
   (ctor 0xf801980, vtable 0x2177b598)                  └─ xla::jellyfish::DeepseaExecutable
              │                                                 │   (ExecuteAsyncOnStream 0x1342cd20)
   tpu::System  (TFRT async-value native)               stream_executor::tpu::TpuExecutor + TpuStream
   ├─ LoadProgram 0x1d0b2240                            ├─ over TfTpu_ExecutorApiFn C-ABI table
   ├─ Execute     0x1d0b33e0                            └─ DeepseaExecutable::LoadProgramAndEnqueue… 0x13426260
   └─ TpuEventIssuer sequence points                          └─ DeepseaStream::EnqueueRequest 0x1d0e9840
              │                                                 │
              └──────────────────────┬──────────────────────────┘
                                     ▼
                    asic_sw::driver::deepsea::jxc::Queue::EnqueueRequest  0xe7d9be0
                            (the one shared TPU driver core; DmaBuffer + completion)
```text

`TpuExecutableInterface` 上的 `legacy::` 命名空间（修饰名 `_ZN3xla6legacy22TpuExecutableInterface…`）是二进制自己的标签，也是 `ExecuteAsyncOnStream` 链条属于已弃用 StreamExecutor 模型的最清晰单一信号。现代路径的 TFRT 形态同样明确：`tpu::System::Execute` 接收 `tsl::AsyncValueRef<tpu::ProgramHandle>` 以及由 `AsyncValueRef` 包装的缓冲区和事件的 `absl::Span`，而不是 `Stream*`。

> **陷阱 —** 两个栈不是分层关系，PJRT 也**不会**包装 `ExecuteAsyncOnStream`。现代 `TpuClient`/`tpu::System` 路径没有任何对 `ExecuteAsyncOnStream`、`TpuExecutor` 或 `ExecutorApiFn` 的引用（HIGH — 由整个 `TpuClient`/`TpuDevice`/`TpuRawBuffer`/`TpuLoadedExecutable` 代码范围中缺失这些符号进行字节确认）。如果重新实现者把 PJRT execute 建模为 StreamExecutor `Executable::ExecuteAsyncOnStream` 虚函数之上的薄包装，将会完整复现错误的栈。对 PJRT 入口建模 `tpu::System::Execute`；只有在重新实现 `LocalClient`/TF-TPU-op-kernel 路径时才建模 `ExecuteAsyncOnStream`。
>
> **说明 —** 宿主解释器（`xla::HloEvaluator`）和 CPU thunk 后端（`xla::cpu::*Thunk`）都不在任一执行热路径上。`HloEvaluator` 在 HLO 简化器（`HloConstantFolding`）内部的*编译*时运行；CPU thunks 只在 HLO 被放置到宿主 CPU 设备上时运行。这个二进制中根本没有 `InterpreterExecutor` 设备。*确实*属于运行时执行的“host”机制是 infeed/outfeed 流式传输（[infeed-outfeed.md](infeed-outfeed.md)）和宿主回调（[host-callbacks.md](host-callbacks.md)）。

---

## 2. 执行生命周期

### 目的

这是整个章节依附的主干：一个逻辑生命周期 — *加载程序、编排参数、入队、在设备上运行、发出完成信号、暴露输出* — 以两种方式实现。图将两种写法并排放置，让读者准确看到它们在哪里分叉（宿主侧抽象），又在哪里重新汇合（驱动核心）。每个带标签的阶段都链接到负责它的页面。

### 端到端，两种写法

```text
                       a compiled TpuExecutable + input PjRtBuffers exist
                                            │
   ── MODERN (PJRT / TpuClient over tpu::System) ──        ── LEGACY (StreamExecutor / LocalClient) ──
   PJRT_LoadedExecutable_Execute (slot 60, 0xf869b40)      xla::LocalExecutable::RunAsync (0x1084d140)
     └ CommonPjRtLoadedExecutable::Execute                   └ Executable::ExecuteAsyncOnStreamWrapper (0x1dad98a0)
        ├ ExecutePrepare  (0xf920be0)  ── pin inputs,            └ [vtable+24]
        │   alloc outputs, donation/aliasing                        TpuExecutableInterface::ExecuteAsyncOnStream (0x1342cd20)
        └ ExecuteLaunch   (0xf921f00)                                ├ marshal ExecutionInput → DeviceAddressBase[]
           └ TpuRawLoadedExecutable::Execute (0xf80f580)             ├ AllocateOutputMemoryWithInputReuse (0x1342ba00)
              └ TpuExecutableLoadState::ExecuteLaunchRaw (0xf8109a0) └ [vtable+96]
                 ├ (once) LoadInternal (0xf80c1c0)                       DeepseaExecutable::LoadProgramAndEnqueueToStream
                 │   └ tpu::System::LoadProgram (0x1d0b2240) ────────►       (0x13426260)
                 │                                                            ├ DeepseaExecutor::LoadProgram → TpuCoreProgramHandle
                 └ tpu::System::Execute (0x1d0b33e0) ──────────────►         └ DeepseaStream::EnqueueRequest (0x1d0e9840)
                    └ TpuEventIssuer::IssueArgs / FulfillArgs                        │
                       (wait/define event sequence points)                          │
                                            └──────────────┬───────────────────────┘
                                                           ▼
                                  jxc::Queue::EnqueueRequest (0xe7d9be0) — device runs
                                                           │
                                                 device completion signal
                                                           │
   define events fulfilled (TpuEventIssuer::FulfillArgs)   request completion callback fires
   → PJRT_Event (tsl::Future<void>) resolves               → ExecutionOutput status stamped
                                                           │
                                          output buffers live in HBM (donated/aliased honored)

两条路径只在一个点重新汇合 — jxc::Queue::EnqueueRequest — 并且在其上方也只沿一个维度分叉：宿主侧排序原语。旧版使用 stream_executor::Stream FIFO（DeepseaRequestQueue）；现代使用由 TpuEvent 组成的 TpuEventIssuer 依赖 DAG。其他所有内容（程序加载、缓冲区绑定、输出构造、完成）在两列之间都有一一对应关系。

怪癖 — 现代路径会在程序物理驻留前准备一次启动。tpu::System::LoadProgram 会立即返回一个未解析的 AsyncValueRef<tpu::ProgramHandle>；到核心的 DMA 异步完成，而 tpu::System::Execute 会把启动注册到 TpuEventIssuer，使其依赖于该 async value。序列点引擎保证请求执行前程序已经加载，因此宿主永远不会阻塞在加载上。相反，旧版路径会在构建 Request 前，在 DeepseaExecutable::LoadProgramAndEnqueueToStream 内同步加载。参见 load-program-enqueue.md §3。

层边界 ​

阶段负责页面覆盖内容
参数编排 + 输出分配（旧版）execute-async-on-stream.mdExecutionInput -> 扁平 DeviceAddressBase[]；ScopedShapedBuffer 分配；donation/aliasing；vtable+24 -> vtable+96 分派
程序加载 + 入队（两者）load-program-enqueue.mdLoadProgram（同步旧版 / 异步现代）、逐核心 Megacore 扇出、Request 构建、EnqueueRequest、卸载簿记
排序 / 依赖（两者）stream-semantics.mdStream FIFO 不变量、WaitFor/RecordEvent、TpuEventIssuer 序列点、compute/H2D/D2H 流拆分
完成 -> PJRT_Event（两者）completion-loop.mddefine-event 履行、AsyncTrackingEvent、宿主侧 tsl::Future<void> 解析、延迟资源清理

3. SE 概念 -> TPU 运行时映射 ​

目的 ​

熟悉 GPU/StreamExecutor 的读者需要一张翻译表来阅读现代路径，因为现代路径保留了 StreamExecutor/Stream/Event/DeviceMemoryBase 的角色，但没有用 StreamExecutor 类型实现其中任何一个。标题中的“StreamExecutor -> PJRT adapter”正是这个映射：SE 抽象映射到 tpu::System，而不是映射到 SE 包装器。

映射 ​

StreamExecutor / PjRt 概念libtpu 现代实现libtpu 旧版实现
StreamExecutor（设备执行器）tpu::System（Initialize 0x1d0ae420）stream_executor::tpu::TpuExecutor（基于 ExecutorApiFn）
PjRtClientxla::TpuClient（构造函数 0xf801980）xla::LocalClient / xla::Service
PjRtDevice / SE 设备xla::TpuDevice（tpu::System* + TpuCoreLocation + Semaphore）每个 ordinal 一个 TpuExecutor
Stream::ThenLaunchtpu::System::Execute @ 0x1d0b33e0DeepseaStream::EnqueueRequest @ 0x1d0e9840
程序加载（无 SE 类比）tpu::System::LoadProgram @ 0x1d0b2240DeepseaExecutor::LoadProgram（在 0x13426260 中）
Stream::ThenMemcpytpu::System::TransferTo/FromDevice（0x1d0afa20/0x1d0b0160）TpuStream::Memcpy（3 个重载）
DeviceMemoryBasetpu::TpuSharedMemoryLocation + tpu::TpuBufferstream_executor::DeviceAddressBase（24 B (ptr,size)）
stream_executor::Eventtpu::TpuEvent（tsl::AsyncValueRef）stream_executor::tpu::TpuEvent
流排序 FIFOtpu::TpuEventIssuer 序列点 + AsyncValue 依赖DeepseaRequestQueue FIFO + WaitFor/RecordEvent
Allocator::Allocatetpu::AllocateBuffer @ 0xf8d51c0DeviceAddressAllocator（基于 backend）
PJRT_Eventtsl::Future<void>（TpuClient::TrackFuture 0xf7fad60）xla::ExecutionOutput 状态

怪癖 — 现代路径上根本没有 LocalDeviceState / 每流对象 — 这个单元为空是有意设计。SE 会把 compute、H2D 和 D2H 放到三个角色专用的 Stream 上，并用 WaitFor/RecordEvent 串行化交接；现代路径则把相同的生产者/消费者排序表达为 TpuEventIssuer DAG：启动的 wait_events 是其依赖，define_events 在完成时被履行。重新实现者如果在 PJRT 路径上寻找流对象，就是在建模错误的栈。流类型拆分只保留在旧版侧 — 参见 stream-semantics.md。

4. 子领域地图 ​

目的 ​

运行时章节包含十个详情页和这个开篇页。每一页负责生命周期中一个连贯片段，并按可重新实现的粒度编写；这张地图是读者扫描以找到正确页面的索引。四个执行流页面（本节核心）列在前面，然后是支撑性参考页面。

执行流页面 ​

页面负责内容栈
execute-async-on-stream.md旧版逐执行 C++ 入口：编排 ExecutionInput -> DeviceAddressBase[]，以 donation/aliasing 分配输出 ScopedShapedBuffer，通过 vtable+24 再到 vtable+96 分派。PJRT 对应项（ExecutePrepare/ExecuteLaunch）在此概述，由适配器页面负责。旧版（概述 PJRT prepare/launch）
load-program-enqueue.md下半部分：把已加载的 TpuCoreProgram 绑定到物理核心（逐核心 Megacore 扇出），从 run options + 缓冲区数组构建设备 Request，推送到命令流，并记录卸载簿记。追踪 DeepseaExecutable::LoadProgramAndEnqueueToStream 和 tpu::System::LoadProgram+Execute。两者
stream-semantics.md排序模型：逐流 FIFO 不变量、字节级 WaitFor(Stream*)/WaitFor(Event*)/RecordEvent、TpuEventIssuer 序列点，以及 compute/H2D/D2H 流拆分。两者
completion-loop.md如何检测并暴露启动完成：define-event 履行、AsyncTrackingEvent、用户可见 tsl::Future<void>/PJRT_Event 的解析，以及延迟资源清理。两者

支撑性参考页面 ​

页面负责内容
infeed-outfeed.md流式宿主↔设备通道：TpuDevice::TransferToInfeed/FromOutfeed（现代）与 TpuTransferManager（旧版）、TpuCoreLocation+index 队列句柄、布局线性化，以及按 span 分块的阻塞语义。
host-callbacks.mdDoHostCallbackWithStatus — 在宿主流上内联，在 TPU 流上通过 C-shim trampoline — host_compute / outside-compilation 实现和状态编排。
allocator-integration.md执行路径用于分配输出缓冲区和宿主 staging 的 DeviceAddressAllocator 与 tpu::AllocateBuffer 胶合层，以及 OOM defragment/retry 循环。
error-templates.md · hint-strings.md · internal-pass-names.md运行时的诊断表面：absl::Status 源位置模板、算子提示字符串，以及运行时写入错误和 trace 的内部 pass-name 目录。

说明 — 设备端程序（tpu::TpuCoreProgram）及其已加载状态缓存（通过 TpuExecutableLoadState::LoadInternal 0xf80c1c0 的 LoadedProgramState）位于本节和编译之间的连接点：已编译程序由上游生成（jellyfish JIT，TpuJitCompileHloWithOptions），但只有在 LoadProgram 将其绑定到核心后，它才变为可执行。运行时绑定的缓冲区偏移在编译时由 Memory-Space Assignment 固定，并由分配器重放 — 参见 ../memory/overview.md。

相关组件 ​

组件关系
xla::TpuClient / tpu::System运行时章节主要记录的现代 PJRT 执行栈
xla::legacy::TpuExecutableInterface / DeepseaExecutableExecuteAsyncOnStream 链条所属的旧版 StreamExecutor 栈
tpu::TpuEventIssuer对现代 load/enqueue/transfer 彼此排序的序列点引擎
asic_sw::driver::deepsea::jxc::Queue（0xe7d9be0）两个栈最终落到的单一 TPU 驱动核心
stream_executor::DeviceAddressBase / tpu::TpuSharedMemoryLocation作为程序参数绑定的 HBM 缓冲区身份
xla::jellyfish JIT（TpuJitCompileHloWithOptions）生成这一层加载并运行的 TpuCoreProgram 的上游编译器

交叉引用 ​