流语义与依赖
本页所有地址均适用于
libtpu-0.0.40-cp314wheel 中的libtpu.so(构建libtpu_lts_20260413_b_RC00,build-id md589edbbe81c5b328a958fe628a9f2207d)。该镜像未剥离;反修饰后的 C++ 符号名按原文逐字引用。.textVMA 等于文件偏移(.text基址0xe63c000)。其他版本会不同。
摘要
libtpu 中的 stream 是一个 stream_executor::Stream,即面向单个设备的有序命令队列。它与 CUDA 程序员熟悉的 cudaStream_t 是同一抽象:入队到某个流的每个操作都按程序顺序、FIFO 执行,同一流上的两个条目不会重叠,而两个不同流会并发运行,除非显式依赖将它们连接起来。libtpu 带有三个具体的 Stream 子类:tensorflow::tpu::TpuStream(真实的设备命令流,是 TfTpu_ExecutorApiFn 之上的一层很薄的 C-ABI shim)、deepsea::executor::DeepseaStream(二进制内部的 TPU 驱动流,shim 最终驱动它,构建在 DeepseaRequestQueue 之上),以及 stream_executor::host::HostStream(用于主机暂存和简单主机操作的同步 CPU 流)。三者都继承自 StreamCommon,后者保存父执行器并强制整个模型所依赖的不变量:parent != nullptr(stream_common.cc:44)。
排序契约分为两半。流内顺序是隐式且全序的:TpuStream 背后的 DeepseaRequestQueue 是 FIFO,因此第 N 个入队请求不能在第 N−1 个请求退役之前开始;同一流上的连续操作之间不需要 fence。跨流顺序是显式的,基于事件或流:Stream::WaitFor(Stream*) 使当前流阻塞在另一个流当前已入队的所有工作之后,而 Stream::WaitFor(Event*) 使它阻塞在单个已记录的 Event 之后。在 TPU 上,这些不是主机侧阻塞;WaitFor 会把一个 waiter 插入设备命令队列(DeepseaRequestQueue::EnqueueWaitFor),因此依赖在设备上解析,而不会停住主机。RecordEvent 发布匹配的完成令牌。主机后端用 absl::Notification 实现同一表面:RecordEvent 发出通知,WaitFor(Event*) 阻塞调用线程,而 WaitFor(Stream*) 是 no-op,因为同步主机流已经运行完了所有内容。
第三个轴是流类型拆分。XLA 不会把计算、host→device 拷贝和 device→host 拷贝倒到同一个流上;它为每个执行器分配一小组固定的按角色区分的流,包括一个计算流以及独立的 H2D 和 D2H 传输流,并用上面的跨流 WaitFor/RecordEvent 机制来串行化生产者/消费者交接。本页负责说明排序模型、跨流依赖/等待语义,以及流类型拆分。把程序入队到计算流的执行入口是 ExecuteAsyncOnStream / LoadProgramAndEnqueueToStream;流排空时触发的完成事件机制见 Completion Loop & AsyncTrackingEvent;host↔device 拷贝传输本身见 Host↔Device DMA。
对于重新实现,契约是:
Stream类布局和流内 FIFO 不变量:StreamCommon的父指针位于+0x48,有parent != nullptrCHECK,并解释为什么流内不需要逐操作 fence。- 两个
WaitFor重载:WaitFor(Stream*)(依赖另一个流已入队的工作)与WaitFor(Event*)(依赖一个已记录事件),以及与后者配对的RecordEvent,并在字节级覆盖 TPU shim 与主机后端。 - TPU 实现:
TpuStream::WaitFor/RecordEvent如何经由TfTpu_ExecutorApiFn槽位+152/+184/+192路由,以及底层DeepseaStream如何把流依赖转换成DeepseaRequestQueue::CreateWaiter+EnqueueWaitFor的设备队列插入。 - compute / H2D / D2H 流拆分:传输和计算占用不同的流,它们之间的 producer→consumer 排序只由
WaitFor/RecordEvent携带,从不依赖隐式跨流顺序。
| 抽象基类 | stream_executor::Stream(仅 vtable)→ StreamCommon(ctor 0x1d100280) |
StreamCommon 父指针 | +0x48(this+9);parent != nullptr CHECK @ stream_common.cc:44 |
| TPU 流 | tensorflow::tpu::TpuStream : TpuStreamInterface |
| TPU 驱动流 | deepsea::executor::DeepseaStream(请求队列位于 this+0x88 / this+17) |
| 主机流 | stream_executor::host::HostStream : StreamCommon(ctor 0xfe6ec80,0x80 B) |
WaitFor(Stream*) | TPU 0xe9989e0 · Host 0xfe6eec0(no-op true)· Deepsea 0x1d0e9480 |
WaitFor(Event*) | TPU 0xe998aa0 · Host 0xfe6eee0(Notification::WaitForNotification) |
RecordEvent | TPU 0xe998b80 · Host 0xfe6ef40(Notification::Notify,CHECK @ host_stream.cc:92) |
| TPU ExecutorApiFn 槽位 | +152 创建依赖 · +184 记录事件 · +192 等待事件 · +360 设备 |
| 证据等级 | 可重新实现 / 已按 IDA 反编译逐字节确认 |
1. 什么是 Stream
目的
stream_executor::Stream 是指向单个设备上一个有序命令队列的句柄。工作(程序启动、memcpy、memset、主机回调、事件记录)会被入队到流中,而流保证条目按入队顺序执行。这个模型有意采用 CUDA stream 模型:排序按流隐式发生;并发默认按流对发生;跨流排序是调用方的责任,通过事件和 WaitFor 表达。不同流之间没有全局程序计数器。
类布局
每个具体流都继承自 StreamCommon,这是一个与设备无关的中间层,保存指向所属执行器的反向指针。stream_executor::StreamCommon::StreamCommon(StreamExecutor*) @ 0x1d100280 很小,并揭示了整个布局:
// StreamCommon::StreamCommon(StreamExecutor* parent) sub_1D100280
function StreamCommon_ctor(this, parent):
this[0] = &StreamCommon_vtable // off_21CA99F8 (vtable for StreamCommon @ 0x21ca99e8 + 0x10)
this[+0x08] = 0 // status rep (OkStatus)
this[+0x10] = 1 // (default count / live flag)
this[+0x18] = 0
zero this[+0x30 .. +0x40] // name string head, priority slot
this[+0x40] = 0
this[+0x48] = parent // <-- parent StreamExecutor* (this+9)
this[+0x50] = 0
this[+0x58] = 1 // (sub-stream / refcount default)
zero this[+0x60 .. +0x7D] // sub-stream pool head
if (parent == nullptr): // stream_common.cc:44
LOG(FATAL) "parent != nullptr"
```text
唯一的硬不变量是 `parent != nullptr`:没有执行器的流是编程错误,会立即中止。位于 `+0x48` 的父指针是流到达其执行器的 `TfTpu_ExecutorApiFn` 表(TPU)或线程池(host)以便入队工作的路径;下面的每个 `Memcpy`、`WaitFor` 和 `RecordEvent` 都会解引用它。`HostStream` 是一个 0x80 字节对象(ctor `0xfe6ec80`),与该布局加上 host 专用 notification 槽位一致;`TpuStream` 添加了不透明的 `SE_Stream*` 句柄和父 `TpuPlatform*`(在 WaitFor/RecordEvent 反编译中从 `this+16/17/18` 读取)。
> **注 —** `StreamCommon` 预留了一个*子流池*头(`+0x60..`)。上游 StreamExecutor 允许调用方通过 `GetOrCreateSubStream`/`ReturnSubStream` 借用一个临时子流,用于一个逻辑流内部的 fork/join;这些槽位存在,但在任何已跟踪的 libtpu 调用点中都没有触发借用路径。libtpu 实际使用的跨流依赖原语是 `WaitFor`(§3),不是子流。置信度:HIGH(布局槽位存在;未跟踪到借用调用方;它们在此构建中可达的置信度为 LOW)。
### 流内 FIFO 不变量
在同一流内,顺序是全序且隐式的:条目 N+1 在条目 N 退役前不会开始,同一流上的两个条目不会重叠。在 TPU 上,这完全不是由 StreamExecutor 层强制的,而是每个 `DeepseaStream` 背后的 `DeepseaRequestQueue` 的性质。`TpuStream` 会把每个操作转发给其驱动流,后者向该队列追加一个 `DeepseaRequest`;设备按 FIFO 顺序排空队列。因此重新实现时,**同一流上的连续操作之间不需要 fence**:同一流上背靠背的 `Memcpy` 和程序启动已经有序。只有跨流时才需要 fence(事件)(§3)。这是该模型最重要的结果,也是流类型拆分(§4)能够工作的原因。
---
## 2. Stream 后端
libtpu 链接了三个 `Stream` 实现。它们共享 `StreamCommon` 基类以及 `WaitFor`/`RecordEvent`/`Memcpy` 虚接口,但有序工作落到*哪里*完全不同。
| 后端 | 类 | 支撑队列 | 并发性 | 角色 |
|---|---|---|---|---|
| TPU shim | `tensorflow::tpu::TpuStream` | `SE_Stream*` → `TfTpu_ExecutorApiFn` | 异步、设备有序 | 调用方向其入队程序的真实 TPU 命令流 |
| TPU 驱动 | `deepsea::executor::DeepseaStream` | `DeepseaRequestQueue`(`this+0x88`) | 异步 worker | shim 驱动的二进制内部流;拥有 FIFO |
| Host CPU | `stream_executor::host::HostStream` | 无(inline) | **同步** | 主机暂存、简单主机操作、主机回调 |
### TpuStream — C-ABI shim
`TpuStream` 是一个很薄的转发器。它持有三个 WaitFor/RecordEvent 反编译会直接读取的字段:不透明流句柄(`this+16`,`SE_Stream*`)、父 `SE_StreamExecutor*`(`this+17`)以及父 `TpuPlatform*`(`this+18`)。每个操作都会通过 `stream_executor::tpu::ExecutorApiFn()` 解析单例 `TfTpu_ExecutorApiFn` 表,并调用适当槽位。槽位 `+360` 返回用于状态收割的*设备*对象;状态通过 `+408`(ok?)检查,失败时从 `+400`/`+392` 拉取 code/message,并通过 `status_helper.h` shim(`absl::Status::MakeRep(..., status_helper.h)`)打包成 `absl::Status`。这是整个 TPU 后端共享的“C-shim → `absl::Status`”模式。
### DeepseaStream — shim 下方的队列
C-ABI 之下的真实流是 `deepsea::executor::DeepseaStream`,其依赖原语是本页字节证据最清晰的部分(`0x1d0e9480`):
```c
// DeepseaStream::WaitFor(Stream* other) sub_1D0E9480
function DeepseaStream_WaitFor_Stream(this, other):
waiter = DeepseaRequestQueue::CreateWaiter(other->queue) // other[17] = other's request queue
return DeepseaRequestQueue::EnqueueWaitFor(this->queue, // this[17] = this stream's request queue
waiter)跨流等待是从另一个流的请求队列制造出的一个 waiter object,并且作为另一个 FIFO 条目入队到当前流的队列中。设备按顺序排空当前流时遇到 waiter,会阻塞该流(不是主机),直到另一个队列到达记录点。这就是 TPU WaitFor 是异步的原因:主机上没有任何东西阻塞。DeepseaStream 还携带一个 H2DTransferState::LaunchAndWait helper(0x1d0ea5e0 invoker),这是驱动流把 host→device 传输子状态与计算工作区分开的直接证据(§4)。
HostStream — 同步、inline
链接进来的 HostStream 是同步变体:没有 worker 线程或任务队列。因为每个操作在下一个操作入队时都已经在调用线程上 inline 运行完毕,所以该流在所有时刻都可视为已排空。这直接塑造了它的 WaitFor/RecordEvent(§3),并意味着 BlockHostUntilDone 无条件返回 true。它用于主机内存暂存和简单主机操作,而不是完整的异步 CPU 设备。
3. 跨流依赖
两个流默认并发运行。要把流 B 上的工作排在流 A 上的工作之后,调用方会插入显式依赖。libtpu 暴露每个 StreamExecutor 后端都实现的两个 Stream 重载,以及为基于事件的重载播种的 RecordEvent。
WaitFor(Stream*) — 依赖另一个流
B.WaitFor(&A) 使 B 阻塞,直到 A 上当前已入队的所有工作完成。TPU 实现(0xe9989e0)把外部 Stream* 映射到其驱动句柄,并要求执行器创建依赖:
// TpuStream::WaitFor(Stream* other) sub_E9989E0
function TpuStream_WaitFor_Stream(this, other):
create_dep = ExecutorApiFn()[+152] // TfTpu create-stream-dependency
self_handle = this[16] // SE_Stream*
self_se = this[17] // SE_StreamExecutor*
se = this[18] // TpuPlatform* owning the map
lock(se->mutex @ +0x90) // absl::Mutex (decompile: se+144)
se_other = FlatHashMap<Stream*, SE_Stream*>(se+0x50).at(other) // map foreign Stream* -> SE_Stream* (decompile: se+80)
// if absent: ThrowStdOutOfRange("raw_hash_map<>::at")
unlock(se->mutex)
ok = create_dep(self_se, self_handle, se_other) // ExecutorApiFn()+152
if (!ok):
return Status(INTERNAL, "Failed to create stream dependency") // tpu_stream.h:108
return OkStatus
```text
有两个重新实现时至关重要的细节:(1)执行器在 `executor+0x50` 保有一个 `FlatHashMap<Stream*, SE_Stream*>`,由 `executor+0x90` 的 `absl::Mutex` 保护,用来把*语言层* `Stream*` 转换成 C-ABI 能理解的*驱动层* `SE_Stream*`;对未注册流执行 `WaitFor` 会抛出 `out_of_range`。(2)依赖通过 `TfTpu_ExecutorApiFn` 槽位 `+152` 创建;失败报告为 `INTERNAL` "Failed to create stream dependency",锚定到 `tpu_stream.h:108`。在底层,这会变成 §2 的 `DeepseaRequestQueue::EnqueueWaitFor` waiter:依赖在设备上解析,主机不会停顿。
> **怪癖 —** `WaitFor(Stream*)` 等待的是*到目前为止已入队*的工作,而不是永远等待该流。它是快照:在 `WaitFor` 之后推入 A 的后续操作**不**包含在内。把该调用当成“B 永远跟在 A 后面”的重新实现是错误的;它的含义是“B 等待 A 当前尾部”。这与 CUDA `cudaStreamWaitEvent` 语义匹配,只是作用在流的隐式尾部而不是命名事件上。
### `WaitFor(Event*)` — 依赖一个已记录事件
更细粒度的原语等待单个 `Event`。TPU 侧(`0xe998aa0`)在平台事件注册表中查找该事件,并调用 wait-event 槽位:
```c
// TpuStream::WaitFor(Event* e) sub_E998AA0
function TpuStream_WaitFor_Event(this, e):
dev = ExecutorApiFn()[+360](this) // device object for status
platform = this[18] // TpuPlatform*
se_event = TpuPlatform::LookupEvent(platform, e) // Event* -> SE_Event*
wait_ev = ExecutorApiFn(platform)[+192] // TfTpu wait-for-event
wait_ev(this[17], this[16], se_event, dev) // (SE_SE*, SE_Stream*, SE_Event*, dev)
status = harvest_status(dev) // +408/+400/+392 -> status_helper.h:38
ExecutorApiFn(dev)[+384](dev) // free the device-status object
return statusRecordEvent(0xe998b80)是镜像结构,完全相同,只是调用槽位 +184(record)而非 +192(wait),并作用于同一个由 LookupEvent 解析出的 SE_Event*。这一配对是标准的生产者/消费者令牌:流 A 在其生产操作之后入队 RecordEvent(e);流 B 在其消费操作之前入队 WaitFor(e);设备解析该令牌排序。
| 原语 | TPU ExecutorApiFn 槽位 | 解析方式 | 失败时 |
|---|---|---|---|
WaitFor(Stream*) | +152(创建依赖) | FlatHashMap<Stream*,SE_Stream*> @ exec+0x50 | INTERNAL "Failed to create stream dependency"(tpu_stream.h:108) |
WaitFor(Event*) | +192(等待事件) | TpuPlatform::LookupEvent | 经 status_helper.h:38 的状态 |
RecordEvent(Event*) | +184(记录事件) | TpuPlatform::LookupEvent | 经 status_helper.h:38 的状态 |
主机后端 — absl::Notification
同步 HostStream 用 notification 实现相同接口,其对比很有启发性。
// HostStream::WaitFor(Stream*) sub_FE6EEC0
function HostStream_WaitFor_Stream(this, other):
return true // nothing to wait for: host stream is inline
// HostStream::WaitFor(Event*) sub_FE6EEE0
function HostStream_WaitFor_Event(this, e):
note = e[1] // absl::Notification*
rc = e[2] // shared_weak refcount (nullable)
if (rc):
atomic_inc(rc[1])
Notification::WaitForNotification(note) // block calling thread
if (atomic_dec(rc[1]) == 0):
vtable_destroy(rc); __release_weak(rc)
else:
Notification::WaitForNotification(note)
return true
// HostStream::RecordEvent(Event*) sub_FE6EF40
function HostStream_RecordEvent(this, e):
note = e[1]; rc = e[2]
CHECK(!note->HasBeenNotified()) // host_stream.cc:92 — record-twice is FATAL
Notification::Notify(note)
// (refcount release mirrors WaitFor)
return true
```text
需要保留三点:(1)`WaitFor(Stream*)` 是**返回 `true` 的 no-op**:同步主机流已经在调用线程上运行完所有先前工作,因此没有剩余内容需要等待。(2)`WaitFor(Event*)` 真正把*调用线程*阻塞在 `Notification::WaitForNotification` 上,并带有 `shared_weak_count` 引用计数,以便事件可以恰好释放一次。(3)`RecordEvent` 会 `Notify` 该 notification,并且**如果事件已经 notified 则 CHECK 失败**(`host_stream.cc:92`,`!notification->HasBeenNotified()`):主机 `Event` 是单次触发的;记录两次会中止。
> **陷阱 —** host 与 TPU 后端在 `WaitFor` 的*阻塞位置*上不同。在 TPU 上,`WaitFor` 入队设备 waiter 并立即返回主机(异步)。在 host 上,`WaitFor(Event*)` 同步阻塞调用线程。假设两者使用同一个模型的重新实现,要么会在 TPU 工作上让主机忙等,要么会在 host 上未能真正串行化。统一的契约是*排序保证*,不是阻塞行为。
>
> **怪癖 —** 主机 `Event` 是单用的(`host_stream.cc:92` CHECK)。TPU 事件按句柄查找,每次都通过驱动重新记录,因此可复用。不要假设事件对象可在后端之间互换。
---
## 4. Compute / Transfer 流拆分
### 为什么是三个流,而不是一个
XLA 不会把计算和数据移动串行化到单个流上。执行器会分配一小组固定的按角色划分的流:用于程序启动的 **compute** 流、用于 host→device 输入拷贝的 **H2D** 流,以及用于 device→host 输出拷贝的 **D2H** 流。原因是并发:当计算流运行程序 K 时,H2D 流已经可以暂存程序 K+1 的输入,而 D2H 流可以排出程序 K−1 的输出。把三者放到一个 FIFO 上会迫使每次拷贝都阻塞每次启动,因为流内顺序是全序的(§1.3)。
该拆分的字节级指纹是 `DeepseaStream` 专用的 `H2DTransferState::LaunchAndWait` 子状态机(`H2DTransferState` invoker `0x1d0ea5e0`):host→device 传输携带自己的 launch-and-wait 状态,与计算路径不同,这正是传输运行在自己流上时预期的结构。C-ABI `TpuStream_EnqueueTransferHostToDevice` / `TpuStream_EnqueueTransferDeviceToHost` 对(`0xeab9960` / `0xeab99e0`,各自转发到匹配的 `DeepseaStream::EnqueueTransfer*`)是向传输流入队的入口;程序启动由 [ExecuteAsyncOnStream](execute-async-on-stream.md) 入队到计算流。
### 如何保持流之间一致
拆分*产生*了一个正确性义务:程序 K 的输入必须在 K 运行前位于设备上,而 K 的输出不得在 K 完成前被拷出。由于这些流相互独立,该排序**只能**由 §3 的跨流 `WaitFor`/`RecordEvent` 携带;不同流之间没有隐式排序。典型交接是:
```text
H2D stream : Memcpy(input) RecordEvent(in_ready) ─┐
│ WaitFor(in_ready)
compute : └─► LaunchProgram(K) RecordEvent(done) ─┐
│ WaitFor(done)
D2H stream : └─► Memcpy(output)每个箭头都是生产流上的一个 RecordEvent,配对消费流上的一个 WaitFor(Event*)(或更粗粒度的、对生产者尾部做快照的 WaitFor(Stream*))。在 TPU 上,这些解析为 DeepseaRequestQueue waiter,因此三个流保持最大并发:只有真实数据依赖被串行化,而且它们在设备上串行化,让主机可以继续入队。运行时用来知道某次启动已退役的完成事件是同一个 Event 机制,通过 Completion Loop & AsyncTrackingEvent 暴露。
注 — 每个执行器的流的数量和确切角色分配(一个 compute + 一个 H2D + 一个 D2H,还是每种都有小池)是 XLA StreamExecutor 池在上游建立的设计;在此二进制中,字节证据确认了(a)transfer-vs-compute 拆分存在(
H2DTransferState、两个EnqueueTransfer*入口),以及(b)唯一的跨流排序原语是WaitFor/RecordEvent。未逐一跟踪每个执行器的精确流数量。置信度:对拆分和排序机制为 HIGH;对确切流基数为 MEDIUM。
5. 重新实现者注意事项
- 流内不需要 fence。 同一流上的两个操作已经由
DeepseaRequestQueueFIFO 排序。在它们之间插入事件是浪费,并且在 host 上第二次 record 会 CHECK 失败。只有跨流跳转需要一对WaitFor/RecordEvent。 WaitFor(Stream*)是尾部快照,不是永久链接。 对你希望依赖的每批工作重新发出它(或使用新事件);它不会自动跟踪未来入队(§3 怪癖)。- 在依赖某个流之前,先把每个流注册到其执行器。 TPU
WaitFor(Stream*)会在执行器的FlatHashMap<Stream*,SE_Stream*>中查找外部流,未命中时抛出out_of_range。该映射是从 C++Stream*到驱动SE_Stream*的唯一桥梁。 - Host 事件是单次触发的;TPU 事件可复用。 如果构建同步 host 后端,请镜像
host_stream.cc:92CHECK:记录一个已经 notified 的事件必须中止,不能静默成功。 - 阻塞语义是后端专用的。 TPU
WaitFor是异步的(设备 waiter);hostWaitFor(Event*)会阻塞调用线程。保留排序契约;不要把一个后端的阻塞行为复制到另一个后端。 - 保留 transfer/compute 拆分。 把 H2D、D2H 和 compute 折叠到一个流上在功能上正确,但会破坏该模型建立来隐藏的启动延迟;设备会在拷贝和启动之间空闲。
相关组件
| 名称 | 关系 |
|---|---|
stream_executor::Stream / StreamCommon | 本页记录的抽象;父指针布局和 parent != nullptr 不变量 |
tensorflow::tpu::TpuStream | 在 TfTpu_ExecutorApiFn 之上实现 WaitFor/RecordEvent 的 C-ABI shim |
deepsea::executor::DeepseaStream | 二进制内部驱动流;DeepseaRequestQueue 是提供流内顺序的 FIFO,EnqueueWaitFor 是跨流依赖 |
stream_executor::host::HostStream | 同步 host 后端;基于 Notification 的事件,no-op WaitFor(Stream*) |
TfTpu_ExecutorApiFn 表 | TPU shim 转发到的 C-ABI(槽位 +152/+184/+192/+360/+384/+392/+400/+408) |
交叉引用
- 概览 — 从 PJRT 向下到 stream 的执行路径,以及本页所处的框架
- ExecuteAsyncOnStream — 把程序入队到计算流的入口
- LoadProgramAndEnqueueToStream — 程序加载 + 入队,计算流工作的生产者
- Completion Loop & AsyncTrackingEvent — 流排空时触发的完成事件机制
- Host↔Device DMA — 运行在传输流上的 H2D/D2H 拷贝传输
- Continuation Queue — 设备侧背靠背程序链,使计算流无需主机往返即可保持忙碌