Skip to content

流语义与依赖

本页所有地址均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so(构建 libtpu_lts_20260413_b_RC00,build-id md5 89edbbe81c5b328a958fe628a9f2207d)。该镜像剥离;反修饰后的 C++ 符号名按原文逐字引用。.text VMA 等于文件偏移(.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 != nullptrstream_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 != nullptr CHECK,并解释为什么流内不需要逐操作 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 父指针+0x48this+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 0xfe6eee0Notification::WaitForNotification
RecordEventTPU 0xe998b80 · Host 0xfe6ef40Notification::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 很小,并揭示了整个布局:

c
// 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* 映射到其驱动句柄,并要求执行器创建依赖:

c
// 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 status

RecordEvent0xe998b80)是镜像结构,完全相同,只是调用槽位 +184(record)而非 +192(wait),并作用于同一个由 LookupEvent 解析出的 SE_Event*。这一配对是标准的生产者/消费者令牌:流 A 在其生产操作之后入队 RecordEvent(e);流 B 在其消费操作之前入队 WaitFor(e);设备解析该令牌排序。

原语TPU ExecutorApiFn 槽位解析方式失败时
WaitFor(Stream*)+152(创建依赖)FlatHashMap<Stream*,SE_Stream*> @ exec+0x50INTERNAL "Failed to create stream dependency"(tpu_stream.h:108
WaitFor(Event*)+192(等待事件)TpuPlatform::LookupEventstatus_helper.h:38 的状态
RecordEvent(Event*)+184(记录事件)TpuPlatform::LookupEventstatus_helper.h:38 的状态

主机后端 — absl::Notification

同步 HostStream 用 notification 实现相同接口,其对比很有启发性。

c
// 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。 同一流上的两个操作已经由 DeepseaRequestQueue FIFO 排序。在它们之间插入事件是浪费,并且在 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:92 CHECK:记录一个已经 notified 的事件必须中止,不能静默成功。
  • 阻塞语义是后端专用的。 TPU WaitFor 是异步的(设备 waiter);host WaitFor(Event*) 会阻塞调用线程。保留排序契约;不要把一个后端的阻塞行为复制到另一个后端。
  • 保留 transfer/compute 拆分。 把 H2D、D2H 和 compute 折叠到一个流上在功能上正确,但会破坏该模型建立来隐藏的启动延迟;设备会在拷贝和启动之间空闲。

相关组件

名称关系
stream_executor::Stream / StreamCommon本页记录的抽象;父指针布局和 parent != nullptr 不变量
tensorflow::tpu::TpuStreamTfTpu_ExecutorApiFn 之上实现 WaitFor/RecordEvent 的 C-ABI shim
deepsea::executor::DeepseaStream二进制内部驱动流;DeepseaRequestQueue 是提供流内顺序的 FIFO,EnqueueWaitFor 是跨流依赖
stream_executor::host::HostStream同步 host 后端;基于 Notification 的事件,no-op WaitFor(Stream*)
TfTpu_ExecutorApiFnTPU shim 转发到的 C-ABI(槽位 +152/+184/+192/+360/+384/+392/+400/+408

交叉引用