ExecuteAsyncOnStream
本页所有地址均适用于
libtpu-0.0.40-cp314wheel 中的libtpu.so（libtpu_lts_20260413_b_RC00，BuildID89edbbe81c5b328a958fe628a9f2207d，ELF x86-64，~745 MB）。其他构建会有所不同。
摘要
ExecuteAsyncOnStream 是 libtpu legacy StreamExecutor 执行路径中每次执行的单一 C++ 入口点：它是 xla::Executable virtual，会把一组 xla::ExecutionInput 缓冲区转换为填充好的 xla::ExecutionOutput，并将设备程序入队。它是 upstream XLA 的 xla::Executable::ExecuteAsyncOnStream 在 TPU 上的对应物，经由 xla::LocalClient::Compile → xla::LocalExecutable::RunAsync 到达，而不是经由现代 PJRT C-API。具体 override 是 xla::legacy::TpuExecutableInterface::ExecuteAsyncOnStream @ 0x1342cd20（3650 B），本页重建的就是这个函数。
注意 — 任务框架将其称为 "
TpuExecutable::ExecuteAsyncOnStream"，并引用xla::PjRtStreamExecutorLoadedExecutable。这个 binary 中不存在这两个 symbol（HIGH — 对 884,843 个反编译函数做穷尽rg，PjRtStreamExecutorLoadedExecutable零命中）。此构建中的 PJRT client 是xla::TpuClient（派生自xla::CommonPjRtClient : xla::PjRtClient），位于 TFRT-nativetpu::Systemasync-value runtime 之上，完全不会路由到ExecuteAsyncOnStream；它的执行路径是PJRT_LoadedExecutable_Execute→CommonPjRtLoadedExecutable::Execute→tpu::System::Execute。实际存在的ExecuteAsyncOnStreamvirtual 属于并行的xla::LocalClient/xla::ServiceStreamExecutor 栈，其 TPU executable 是xla::legacy::TpuExecutableInterface（以及其子类xla::jellyfish::DeepseaExecutable）。本页记录这个真实入口；当 contract 使用 PJRT 概念名称时，会映射到承担该角色的 StreamExecutor 对象。
该入口按顺序做四件事：将 host ExecutableRunOptions 和 ExecutionInput 参数 marshal 成缓冲区树；用输入缓冲区 donation/aliasing 分配输出 ScopedShapedBuffer；通过 vtable slot 将设备工作 dispatch 到 enqueue lower half；并为调用者组装 ExecutionOutput（或 absl::Status）。enqueue lower half，即 DeepseaExecutable::LoadProgramAndEnqueueToStream @ 0x13426260，见 Load Program and Enqueue；本页负责入口、参数/输出 marshaling，以及到该层的 dispatch。
对重新实现而言，contract 是：
- Dispatch 格子。 三个调用者到达同一个 virtual：
LocalExecutable::RunAsync→Executable::ExecuteAsyncOnStreamWrapper→ vtable slot +24（ExecuteAsyncOnStream），以及第二个 public door，即 C-ABI shimTpuExecutable_ExecuteAsyncOnStream@0xeabd500，它在 un-marshal C struct 后调用同一个 +24 slot。 - 参数 marshaling。 如何将
xla::ExecutionInputvector 走访成由IndexTable索引的扁平DeviceAddressBase数组，以及 dynamic-shape side channel。 - 输出构造。
AllocateOutputMemoryWithInputReuse（@0x1342ba00）构造ScopedShapedBuffer，执行由HloInputOutputAliasConfig驱动的 input→output aliasing fixup，以及MarkToBeReleasedArguments。 - 交接。 对
LoadProgramAndEnqueueToStream的 vtable +96 间接调用，以及把返回的 root buffer 移入调用者的ExecutionOutput。
| 入口点 | xla::legacy::TpuExecutableInterface::ExecuteAsyncOnStream @ 0x1342cd20（3650 B） |
| C-ABI shim | TpuExecutable_ExecuteAsyncOnStream @ 0xeabd500（4708 B）— tpu_executor_c_api.cc |
| Public C++ wrapper | xla::Executable::ExecuteAsyncOnStreamWrapper @ 0x1dad98a0（579 B，ExecutionInput overload） |
| Client driver | xla::LocalClient → xla::LocalExecutable::RunAsync @ 0x1084d140（2489 B） |
| Vtable dispatch slot | +24（进入 ExecuteAsyncOnStream）；+96（叶子 LoadProgramAndEnqueueToStream） |
| 输出 allocator | TpuExecutableInterface::AllocateOutputMemoryWithInputReuse @ 0x1342ba00（4828 B） |
| Enqueue lower half | xla::jellyfish::DeepseaExecutable::LoadProgramAndEnqueueToStream @ 0x13426260（7512 B） |
| Arg type | std::vector<xla::ExecutionInput>（192 B/element） |
| Result type | xla::ExecutionOutput（包装 ScopedShapedBuffer），以 StatusOr 按值返回 |
| Source file（asserts） | stream_executor/tpu/tpu_executable_interface.cc |
对象模型与类层次
目的
ExecuteAsyncOnStream 是 upstream xla::Executable base 上的 virtual。在 TPU 上，override 位于 xla::legacy::TpuExecutableInterface，这是一个抽象类：它一次性实现参数/输出 marshaling，并将真正的设备 enqueue 延迟到由具体 xla::jellyfish::DeepseaExecutable 实现的 pure-virtual leaf。
继承
xla::Executable (upstream base; ExecuteAsyncOnStream is virtual @ vtable+24)
└─ xla::legacy::TpuExecutableInterface ── implements ExecuteAsyncOnStream @ 0x1342cd20
│ (marshal args, allocate outputs, dispatch via +96)
└─ xla::jellyfish::DeepseaExecutable ── implements LoadProgramAndEnqueueToStream @ 0x13426260
(the leaf invoked through vtable+96; load-program-enqueue.md)
```text
`TpuExecutableInterface` 上的 `legacy::` namespace 是 binary 自己的标签（mangled `ZN3xla6legacy22TpuExecutableInterface…`），也是最清晰的单一信号：整条代码路径是为 `LocalClient`/`Service` 和 TF-TPU op kernels 保留的 deprecated StreamExecutor 执行模型。现代 PJRT front door 使用的是 `xla::TpuClient`。
> **特性 —** 两个 vtable slot 属于不同对象。`ExecuteAsyncOnStream` 在 `Executable` vtable 的 **+24** 处 dispatch（由 `ExecuteAsyncOnStreamWrapper` 第 45 行和 C shim 第 703 行调用）。在 `ExecuteAsyncOnStream` 内部，设备工作在 **+96** 处 dispatch（`(*(...)(*v89 + 96))`，interface 第 650 行），也就是 pure-virtual `LoadProgramAndEnqueueToStream`。把两者折叠成一个方法的重新实现者，会丢失 abstract/concrete 拆分，而这个拆分让 `DeepseaExecutable` 可以在不触碰 marshaling 的情况下替换设备 backend。
### 函数映射
| 函数 | 地址 | 大小 | 角色 |
|---|---|---|---|
| `xla::legacy::TpuExecutableInterface::ExecuteAsyncOnStream` | `0x1342cd20` | 3650 B | 入口：marshal → allocate → dispatch → output |
| `xla::legacy::TpuExecutableInterface::AllocateOutputMemoryWithInputReuse` | `0x1342ba00` | 4828 B | 构建输出 `ScopedShapedBuffer`，遵守 donation/aliasing |
| `xla::jellyfish::DeepseaExecutable::LoadProgramAndEnqueueToStream` | `0x13426260` | 7512 B | Pure-virtual leaf（vtable+96）— 设备 enqueue |
| `xla::Executable::ExecuteAsyncOnStreamWrapper`（`ExecutionInput`） | `0x1dad98a0` | 579 B | +24 virtual 周围的带 profile public wrapper |
| `xla::Executable::ExecuteAsyncOnStreamWrapper`（`ShapedBuffer`） | `0x1dad9780` | 259 B | Legacy `ShapedBuffer`-span overload |
| `xla::ExecuteWrapperAfterExecution` | `0x1dad9b00` | 266 B | 执行后 profiling/HLO-profile finalize |
| `xla::LocalExecutable::RunAsync` | `0x1084d140` | 2489 B | `LocalClient` driver → wrapper |
| `TpuExecutable_ExecuteAsyncOnStream` | `0xeabd500` | 4708 B | C-ABI shim：C structs ↔ C++ objects |
---
## 入口点与 Dispatch 格子
### 目的
这里恰好有一个执行 virtual，从两个方向到达：进程内 C++ client（`LocalExecutable::RunAsync`）和 C-ABI 边界（`TpuExecutable_ExecuteAsyncOnStream`，导出给 StreamExecutor `TpuExecutor` shim）。
### 入口点
```text
xla::LocalExecutable::RunAsync (0x1084d140) ── LocalClient per-call driver
└─ xla::Executable::ExecuteAsyncOnStreamWrapper (0x1dad98a0)
├─ xla::ExecuteWrapperBeforeExecution ── start HLO execution profile
├─ [vtable+24] ExecuteAsyncOnStream ───────────┐
└─ xla::ExecuteWrapperAfterExecution (0x1dad9b00)│ ── finalize profile, stamp Status
│
TpuExecutable_ExecuteAsyncOnStream (0xeabd500) │ ── C-ABI: FromC args, build RunOptions
└─ [vtable+24] ExecuteAsyncOnStream ──────────────────┤ then ToC the ExecutionOutput
│
xla::legacy::TpuExecutableInterface::ExecuteAsyncOnStream (0x1342cd20) <─┘
├─ AllocateOutputMemoryWithInputReuse (0x1342ba00) ── ScopedShapedBuffer + aliasing
├─ xla::Executable::MarkToBeReleasedArguments ── donation bookkeeping
└─ [vtable+96] DeepseaExecutable::LoadProgramAndEnqueueToStream (0x13426260) ── device enqueue算法 — wrapper
ExecuteAsyncOnStreamWrapper 很薄：它用 profiling hooks 包住 virtual call。即使 profiling 被禁用，两个 hook 也存在（在 null HloExecutionProfile 上 no-op）。
function ExecuteAsyncOnStreamWrapper(self, run_options, args): // 0x1dad98a0
state = ExecuteWrapperBeforeExecution(run_options) // start span; capture stream
out = (*self.vtable[24])(self, run_options, &args) // -> ExecuteAsyncOnStream; moves args
stream = run_options->stream() // line 70
status = ExecuteWrapperAfterExecution(self, &state, // 0x1dad9b00
out.status, stream) // finalize profile
return out // ExecutionOutput by value
```text
> **注意 —** wrapper 会从调用者的 `args` vector *move out*（第 36-44 行清零源 vector header，然后销毁 moved-from 的 `ExecutionInput`）。参数 vector 会被该调用消费；重新实现者不得在之后复用它。
### C-ABI shim
`TpuExecutable_ExecuteAsyncOnStream`（`0xeabd500`，来自 `tpu_executor_c_api.cc`）是 StreamExecutor `TpuExecutor` C-shim 跨越的边界。它只围绕同一个 +24 virtual 做 marshaling：
```c
function TpuExecutable_ExecuteAsyncOnStream(self, c_run_opts, c_args[], n, c_out, status_out): // 0xeabd500
run_options.set_device_ordinal(c_run_opts->device_ordinal) // a2+32
if c_run_opts->allocator: // a2+8
run_options.set_allocator(new DeviceAddressAllocator{ // 0x18 B object
GetUnderlyingDeepseaPlatform(), c_run_opts }) // wraps deepsea platform
run_options.set_stream(*c_run_opts->stream) // a2+40
if c_run_opts->host_to_device_stream: // a2+48
run_options.set_host_to_device_stream(...)
if c_run_opts->device_assignment: // a2+56
proto = DeserializeProto<DeviceAssignmentProto>(...) // TpuSerializedProto
run_options.set_device_assignment(DeviceAssignment::Deserialize(proto))
run_options.set_rng_seed(c_run_opts->rng_seed) // a2+72
run_options.set_run_id(c_run_opts->run_id) // a2+80
// -- marshal each C SE_ExecutionInput into an xla::ExecutionInput --
for i in 0..n: // line 261
arg = ExecutionInput(ApiConverter::FromC(c_args[i])) // shape
TF_CHECK_OK(arg.SetDynamicShape(FromC(c_args[i]+560))) // dynamic shape side channel
for each buffer in c_args[i].buffers (stride 72, base +536): // line 312
arg.SetUnownedBuffer / SetBuffer(FromC(buffer)) // MaybeOwningDeviceAddress
for each aliased index (stride 72, base +544, count +552): // line 504
arg.MutableBuffers()->insert(ShapeIndex) // IndexTable entry
args.push_back(move(arg))
out = (*self.vtable[24])(&out, self, &run_options, &args) // line 703 -> ExecuteAsyncOnStream
if out.ok():
scoped = ScopedShapedBuffer(out.result)
ApiConverter::ToC(c_out, scoped.release()) // populate SE_ExecutionOutput
else:
*status_out = out.status // line 887
// destroy args, run_options, allocator陷阱 —
SetDynamicShape通过TF_CHECK_OKassert（shim 在tpu_executor_c_api.cc:1190assert）。来自 C 侧的畸形 dynamic-shape blob 会导致硬LogMessageFatal，而不是返回错误。dynamic shape 位于每个 C argument struct 固定的+560偏移，独立于偏移0处的 static shape。
参数 Marshaling
目的
interface 入口接收的 args 已经是 std::vector<xla::ExecutionInput>。它的第一项工作是把每个参数的每个叶子缓冲区扁平化为单个连续的 DeviceAddressBase 数组，供 enqueue 层按位置消费，同时通过每个参数的 IndexTable 保留树形 shape。
算法
function TpuExecutableInterface::ExecuteAsyncOnStream(self, run_options, args): // 0x1342cd20
n = args.size() // a4[1]
// ---- Stage 1: flatten argument leaf buffers ----
flat = new DeviceAddressBase[n] // 24 B/elem (line 142)
for i in 0..n: // walk arg[i].Buffers()
entry = IndexTable::GetEntry(arg[i].buffers, root, /*index*/0) // tuple_tree.h:332 CHECK
addr = entry.AsDeviceAddress() // MaybeOwningDeviceAddress
flat[i] = addr // moved into flat array
// ---- Stage 2: fetch result shape + aliasing config from the program ----
if program (a2):
result_shape = program->result_shape() // vtable+40 (line 238)
alias_config = program->input_output_alias_config() // HloInputOutputAliasConfig (+2840…)
else:
result_shape = Shape{} // empty (line 249)
alias_config = ShapeTree<optional<Alias>>{} // line 334
CHECK(run_options->allocator() != nullptr) // tpu_executable_interface.cc:219
...
```text
参数 vector 的元素 stride 是 **192 字节**（一个 `xla::ExecutionInput`），在每个析构循环中都可见（`192 * count`，例如 C shim 第 666、711、945 行和 wrapper 第 53 行）。每个 `ExecutionInput` 携带一个 `Shape`、一个 dynamic `Shape`，以及通过 `xla::internal::IndexTable` 索引的 `ShapeTree<MaybeOwningDeviceAddress>` 叶子缓冲区。叶子设备地址是 24 字节的 `DeviceAddressBase` 记录（opaque pointer + size + memory-space tag）；flatten 循环会 assert 每个 `AsDeviceAddress().opaque() != nullptr`（`tuple_tree.h:332`）。
> **特性 —** flatten 数组使用带有 `2*cap` doubling 策略的手写 growable vector（`v16 = 2*v7`，interface 第 195 行）和 `0xAAAAAAAAAAAAAAAA` length-error guard（interface 第 139、196 行，即 24 字节 `DeviceAddressBase` stride 的最大元素数）。这不是带默认增长策略的 `std::vector<DeviceAddressBase>`：leaf 数量预先已知（`24 * n`，第 142 行），所以重新实现可以精确预分配并完全跳过 reallocation 路径（interface 第 142-218 行）。（`0xAAAAAAAAAAAAAAAB` division magic 用于从字节 span 恢复 /192 元素计数，属于 *wrapper* 在 `0x1dad98a0` 第 64 行的 `ExecutionInput` vector teardown，不属于这个 24 字节 flat array。）
---
## 输出构造与 Aliasing
### 目的
输出 `ScopedShapedBuffer` 会在设备运行*之前*分配，这样 enqueue 层可以直接把结果写入其中。Donation 允许输入缓冲区原地成为输出缓冲区，从而避免一次分配和一次复制。
### 算法
```c
// ---- Stage 3: allocate outputs, reusing donated input buffers ----
out_or = AllocateOutputMemoryWithInputReuse( // 0x1342ba00
result_shape, alias_config,
run_options->allocator(),
args, // donor source
run_options->stream(),
run_options->host_to_device_stream())
if !out_or.ok():
return out_or.status.AddSourceLocation(tpu_executable_interface.cc:228)
output = ScopedShapedBuffer(out_or) // line 344
// ---- Stage 4: re-wire donated input buffers into the output tree ----
if program->aliasing_table (a2+3488, count a2+3496): // line 382
for (param, output_index) in aliasing_table: // 6-qword stride
CHECK(param < args.size()) // ...interface.cc:236
in_entry = IndexTable::GetEntry(args[param], index) // ...:242
CHECK(!in_entry.AsDeviceAddress().is_null()) // ...:243
CHECK(in_entry.is_owning /* offset */) // ...:244
aliased.push_back(in_entry) // donor address list
buffers_to_release.push_back(param_index) // uint32 vector
// ---- Stage 5: bookkeeping + dispatch ----
MarkToBeReleasedArguments(program, args[0], n, output) // line 643
root = output.root_buffer() // line 644
status = (*program.vtable[96])( // line 650 -> LoadProgramAndEnqueueToStream
program, run_options,
flat /*device addresses*/, n,
aliased /*donor list*/, buffers_to_release,
root.opaque())
if status.ok():
result.set_result(move(output)) // ScopedShapedBuffer into ExecutionOutput
else:
result.status = status.AddSourceLocation(...:266)
return result // ExecutionOutputAllocateOutputMemoryWithInputReuse（0x1342ba00）使用 ShapeUtil::ForEachMutableSubshapeHelper（callback @ 0x1342dc00）走访 result Shape；对每个 leaf subshape，它查询 HloInputOutputAliasConfig，以决定通过 DeviceAddressAllocator 分配新的设备内存，还是 claim 一个 donated input buffer。结果是 ScopedShapedBuffer，即一个 owning ShapedBuffer，其 leaf address 会在给定 stream 上由 RAII 释放，除非显式 release() 到 ExecutionOutput。
陷阱 — aliasing fixup 会通过
IndexTable::GetEntry重新读取输入缓冲区，并 assertit != arguments[parameter].MutableBuffers()->end()（...interface.cc:242）。如果 HLOinput_output_alias_config为一个调用者实际未传入的 parameter index（或传入但缺少该 leaf）声明 donation，会触发 fatalCHECK，而不是优雅 fallback。donor 必须存在且 owning（offset为真，:244）。特性 —
MarkToBeReleasedArguments在设备 enqueue 之前运行，而不是之后。它记录程序允许消费哪些 argument buffer，避免调用者的ExecutionInputdestructor double-free executable 现在拥有的缓冲区。实际 release 延迟到持有最终ExecutionOutput的对象。把这一步重新实现在 enqueue 之后会让设备与 argument teardown 发生竞态。
到 Enqueue 层的 Dispatch
目的
这是从 marshaling 到设备工作的唯一交接。上方全部是 target-independent 的缓冲区管线；+96 调用以下全部属于 DeepseaExecutable 设备 backend。
跨越边界的内容
vtable+96 调用（LoadProgramAndEnqueueToStream）依次接收：ServiceExecutableRunOptions（携带 stream、allocator、device assignment、run id、rng seed）、扁平 DeviceAddressBase 参数数组及其计数、donor-buffer span、buffers_to_release index vector，以及输出 root buffer 的 opaque pointer。leaf 返回 absl::Status；成功时，预先分配的 output 现在已在设备上填充，并被移入返回的 ExecutionOutput。
ExecuteAsyncOnStream LoadProgramAndEnqueueToStream (load-program-enqueue.md)
───────────────────── ──────────────────────────────────────────────────────
run_options ───────────────────────────▶ stream / allocator / device_assignment / run_id
flat[] (DeviceAddressBase, n) ──────────▶ positional input device addresses
aliased[] (donated input addrs) ─────────▶ in-place output reuse
buffers_to_release[] (uint32) ──────────▶ donation index list
output.root_buffer().opaque() ──────────▶ output root device pointer (written by device)
◀────────────────────────────── absl::Status (ok ⇒ output is live)
```text
Stream ordering、设备端 program load（`tpu::System::LoadProgram` / `TpuCoreProgram`）和 completion event **不**在此处建立：它们完全位于 +96 边界以下。见交叉引用。
### 注意事项
- **Replica / partition 处理。** `ExecuteAsyncOnStream` 每次调用是单设备的。device assignment 通过 `run_options->device_assignment()` 到达（在 C shim 中由 `DeviceAssignmentProto` 反序列化，interface 第 236-245 行）。多 replica fan-out 是调用者职责：`LocalExecutable::RunAsync` 会在 wrapper 调用前从 run options 解析 stream 和 device ordinal（RunAsync 第 195 行）。入口本身内部没有 replica loop。
- **错误表面。** 两种错误风格共存。Buffer-shape 和 aliasing 违规是 fatal `CHECK`/`LogMessageFatal`（binary 将畸形 buffer tree 视为 programming error）。Allocation failure 和 device-enqueue failure 会以带有 `AddSourceLocationImpl` stamp 的 `absl::Status` 返回（allocation 为 `:228`，enqueue 为 `:266`），并以失败的 `ExecutionOutput` 传播给调用者。
- **不创建 StreamExecutor `Stream`。** 尽管名称如此，此路径不会创建 stream。`run_options->stream()` 由调用者提供（`LocalExecutable::RunAsync` 或 C shim 的 `c_run_opts->stream`）。这里的 "async" 是 StreamExecutor stream 模型，区别于 adapter 页面记录的 PJRT `tpu::System` async-value 模型。
---
## 相关组件
| 名称 | 关系 |
|---|---|
| `xla::jellyfish::DeepseaExecutable::LoadProgramAndEnqueueToStream` | 本入口 dispatch 到的 vtable+96 leaf，即设备 enqueue lower half |
| `xla::Executable::ExecuteAsyncOnStreamWrapper` | 调用 +24 virtual 的带 profile public C++ wrapper |
| `xla::LocalExecutable::RunAsync` | 解析 stream/ordinal 并调用 wrapper 的 `LocalClient` per-call driver |
| `TpuExecutable_ExecuteAsyncOnStream` | 跨 StreamExecutor 边界 marshal C struct 的 C-ABI shim |
| [`xla::TpuClient`](../pjrt/client-and-device.md)（PJRT 路径） | *现代*执行路径；通过 `tpu::System::Execute` 完全绕过 `ExecuteAsyncOnStream` |
## 交叉引用
- [Load Program and Enqueue](load-program-enqueue.md) — vtable+96 leaf（`LoadProgramAndEnqueueToStream`）；设备 program load 和 command-stream enqueue
- [Stream 语义](stream-semantics.md) — `run_options->stream()` 如何对入队工作排序
- [Completion Loop](completion-loop.md) — enqueue 层生成的 async completion event
- [Allocator 集成](allocator-integration.md) — `AllocateOutputMemoryWithInputReuse` 和 C shim 构建的 `DeviceAddressAllocator`
- [Host Callbacks](host-callbacks.md) — 执行期间的 host-side infeed/outfeed 和 callback dispatch
- [PJRT Executable 执行](../pjrt/executable-execution.md) — 现代 `PJRT_LoadedExecutable_Execute` → `CommonPjRtLoadedExecutable::Execute` 路径，本入口是其 legacy counterpart
- [Runtime 概览](overview.md) — StreamExecutor 执行路径相对于 PJRT 路径的位置