LoadProgramAndEnqueueToStream
本页中的所有地址、结构体偏移和源码行引用均适用于
libtpu-0.0.40-cp314wheel 中的libtpu.so(构建libtpu_lts_20260413_b_RC00,build-id md589edbbe81c5b328a958fe628a9f2207d,781,691,048 字节,clang/LLVM trunk)。该镜像未被剥离;反修饰后的 C++ 符号名按原样引用。其他版本会有所不同。
摘要
LoadProgramAndEnqueueToStream 是一次 TPU 执行的下半部分:给定编译后的程序和一组完全解析好的输入/输出 HBM 地址,它会 (a) 将编译后的 TpuCoreProgram 绑定到物理 TPU 核心,并取回每核心一个的 TpuCoreProgramHandle,以及 (b) 根据运行选项和缓冲区地址数组构建设备请求,并将其推入该核心的命令流。生成这些已解析地址并决定哪些副本运行的 PJRT 侧入口是 execute-async-on-stream.md;本页负责从“可执行文件及其缓冲区已就绪”到“请求已进入驱动队列”的所有内容。
LoadProgramAndEnqueueToStream 是这个二进制中的真实符号,而且有两个,因为 libtpu 发布了两个并行设备栈。旧版 C-ABI shim TpuExecutable_LoadProgramAndEnqueueToStream(0xeaafba0,源码 tpu_execute_c_api.cc)是一个自由的 T 导出函数,注册在 TF-TPU executable-API 函数结构的索引 7 处,并不是 xla::TpuExecutable 类的方法;它从 C 启动结构编排出 xla::ExecutableRunOptions,并通过一次虚调用(vtable+96)转发到 jellyfish 核心。该核心是 xla::jellyfish::DeepseaExecutable::LoadProgramAndEnqueueToStream(0x13426260,源码 deepsea_executable.cc),真正执行 DeepseaExecutor::LoadProgram -> DeepseaStream::EnqueueRequest。现代 PJRT 路径(xla::TpuClient over tpu::System,参见 overview.md)通过 TFRT async-value 写法抵达同一个设备运行时:程序加载是 tpu::System::LoadProgram(0x1d0b2240,源码 system.cc:1804),入队是 tpu::System::Execute(0x1d0b33e0),二者都由 TpuEventIssuer 序列点排序,而不是由 StreamExecutor FIFO 排序。
熟悉 XLA-on-GPU 的读者可以把 LoadProgramAndEnqueueToStream 类比为 Executable::ExecuteOnStream -> gpu::GpuExecutable::ExecuteThunks -> stream->ThenLaunch(kernel, args),但需要先说明三个 TPU 特有差异。第一,“load”不是空操作:编译后的 TpuCoreProgram 必须先加载到核心上(DMA 到其指令内存,并在每核心程序缓存中引用计数)才能运行,加载会返回入队所寻址的 TpuCoreProgramHandle。第二,单次启动会扇出为每个物理核心一个 program-handle — Megacore 芯片每芯片承载两个核心,循环必须在每个核心上放置一个 handle。第三,缓冲区“参数”不是 kernel-arg 缓冲区中的指针;它们是 stream_executor::DeviceAddressBase 记录(HBM (opaque_ptr, size) 对),放在旁路向量中,并由 ApiConverter::FromC 从 C ABI 编排而来。
对于重新实现,契约是:
- 启动结构 -> run-options 编排 —
TpuExecutable_LoadProgramAndEnqueueToStream如何重建HloModuleshape 签名、解析 device ordinal,并填充xla::ExecutableRunOptions(set_stream/set_device_assignment/set_rng_seed/set_allocator)。 - 缓冲区绑定 — 通过
ApiConverter::FromC处理输入/输出/别名化DeviceAddressBase数组,加上uint32参数索引向量,并交给核心可执行文件的虚启动槽。 - 程序加载 —
DeepseaExecutor::LoadProgram(旧版)/tpu::System::LoadProgram(现代):从TpuCoreLocation解析芯片、创建每核心 program-handle、由TpuChipConfig::Megacore门控的逐核心扇出,以及完成时记录的 fingerprint。 - 入队 —
DeepseaStream::EnqueueRequest(旧版)/tpu::System::Execute+TpuEventIssuer::IssueArgs/FulfillArgs(现代):设备Request对象、逐副本核心选择(TpuCoreProgramHandle::core->TpuCoreLocation::LogicalDeviceId),以及完成事件。 - 完成 / 卸载簿记 — 每个已加载程序的
ProgramUnloadInfo,以及在流排空时作为宿主回调投递的 host-transfer-manager 拆卸。
| 旧版 SE 入口(C-ABI 编排) | TpuExecutable_LoadProgramAndEnqueueToStream @ 0xeaafba0(867 行反编译代码,tpu_execute_c_api.cc) |
| Jellyfish 核心(加载 + 入队) | xla::jellyfish::DeepseaExecutable::LoadProgramAndEnqueueToStream @ 0x13426260(1360 行,deepsea_executable.cc) |
| 现代程序加载 | tpu::System::LoadProgram(TpuCoreLocation, shared_ptr<const TpuCoreProgram>) @ 0x1d0b2240(system.cc:1804) |
| 现代入队 | tpu::System::Execute(AsyncValueRef<ProgramHandle>, ExecuteOptions, inputs, outputs, wait, define) @ 0x1d0b33e0 |
| 旧版入队 | deepsea::executor::DeepseaStream::EnqueueRequest(unique_ptr<Request>) @ 0x1d0e9840 |
| 硬件队列(驱动) | asic_sw::driver::deepsea::jxc::Queue::EnqueueRequest(DmaBuffer, fn, bool) @ 0xe7d9be0 |
| 现代启动驱动 | xla::TpuExecutableLoadState::ExecuteLaunchRaw @ 0xf8109a0;LoadInternal @ 0xf80c1c0 |
| 缓冲区地址编排 | ApiConverter::FromC(SE_DeviceAddressBase*) -> stream_executor::DeviceAddressBase |
| Program handle 类型 | tpu::TpuCoreProgramHandle(携带 core()、fingerprint()) |
| 逐核心扇出门控 | tpu::TpuChipConfig::Megacore();deepsea::executor::DeepseaPlatform::GetCoreType() |
| 置信度 | CONFIRMED(字节锚定),除非某行或标注另有说明 |
1. 它在启动中的位置
目的
LoadProgramAndEnqueueToStream 是 PJRT execute 中设备无关部分(参数固定、输出缓冲区分配、donation/aliasing、副本扇出 — 都在 execute-async-on-stream.md 中)与设备特定部分(把编译后的二进制加载到硅片上,并向硬件命令流推送请求)之间的边界。它在调用方已经把每个输入和输出都解析成具体的 HBM DeviceAddressBase 后,针对每个 (executable, replica) 对被调用一次。
两个设备栈
有两条代码路径都带有这个名称,重新实现者必须知道自己的前端驱动的是哪一条。两者最终都抵达相同的 TPU 驱动核心(asic_sw::driver::deepsea),只在宿主侧抽象上不同。
LEGACY (StreamExecutor / LocalClient / TF-TPU op kernels):
TfTpu_ExecutableApiFn[7] (C ABI fn ptr, tpu_execute_c_api.cc)
└─ TpuExecutable_LoadProgramAndEnqueueToStream 0xeaafba0
├─ rebuild HloModule shape sig (xla::Shape::FromProto x N)
├─ resolve device_ordinal via TPUNodeInterfaces::Get
├─ build xla::ExecutableRunOptions (stream/assignment/seed/alloc)
├─ construct xla::jellyfish::DeepseaExecutable
└─ vtable+96 ──► DeepseaExecutable::LoadProgramAndEnqueueToStream
0x13426260 (deepsea_executable.cc)
├─ DeepseaExecutor::LoadProgram ──► TpuCoreProgramHandle
└─ DeepseaStream::EnqueueRequest 0x1d0e9840
└─ jxc::Queue::EnqueueRequest 0xe7d9be0 (DmaBuffer)
MODERN (PJRT / TpuClient over tpu::System, TFRT async-values):
PJRT_LoadedExecutable_Execute (slot 60, 0xf869b40)
└─ CommonPjRtLoadedExecutable::Execute → ExecutePrepare → ExecuteLaunch
└─ TpuRawLoadedExecutable::Execute 0xf80f580
└─ TpuExecutableLoadState::ExecuteLaunchRaw 0xf8109a0
├─ (once) TpuExecutableLoadState::LoadInternal 0xf80c1c0
│ └─ tpu::System::LoadProgram 0x1d0b2240 (system.cc:1804)
└─ tpu::System::Execute 0x1d0b33e0
└─ TpuEventIssuer::IssueArgs / FulfillArgs (sequence points)
```text
> **说明 —** 两个栈不是重新实现者在运行时选择的替代方案;它们是编译到同一个镜像中的两个入口。旧版 `TpuExecutable` 路径由 `xla::LocalClient` / `xla::Service` 和 TF-TPU op kernels 通过 `Tpu*_*` C-ABI 驱动;现代路径由每个 PJRT 客户端(JAX、PyTorch-XLA)驱动。它们共享 `deepsea` 驱动核心,但从不共享宿主侧对象 — PJRT 路径上没有任何 StreamExecutor `Stream`。如果重新实现 PJRT 入口,就建模 `tpu::System::LoadProgram` + `tpu::System::Execute`;如果重新实现 TF-TPU op kernel 路径,就建模 `TpuExecutable_LoadProgramAndEnqueueToStream`。
### 下半部分接收的输入
两种写法实际上接收相同的五类东西 — 旧版从 C 启动结构中解包,现代则以已经有类型的 C++ 对象形式接收:
| 输入 | 旧版来源(`0xeaafba0`) | 现代来源(`ExecuteLaunchRaw`) | 含义 |
|---|---|---|---|
| Stream / 排序句柄 | `a1[12]` -> `TpuStream` | `TpuEventIssuer` wait/define 集合 | 在哪里入队,以及排序依赖 |
| 编译后的程序 | `a1[2]`(TpuExecutable proto)-> `DeepseaExecutable` | `TpuExecutable*` -> `LoadedProgramState` | 要加载到核心上的对象 |
| 输出 shape + 参数 shapes | `tpu_program+112` / `+24`(`ShapeProto`s) | 可执行文件中的 `ComputationLayout` | 重建 `HloModule` 签名 |
| 输入/输出 `DeviceAddressBase[]` | 通过 `ApiConverter::FromC` 的 `v71[3..7]` 数组 | `Span<RCReference<CommonPjRtRawBuffer>>` | 要绑定的已解析 HBM 地址 |
| Device assignment / 副本 | `a1[11]` 处的 `DeviceAssignmentProto` | `DeviceAndAssignment`,`replica` 参数 | 本次启动面向哪个逻辑设备 |
---
## 2. 启动结构 -> Run-Options 编排(旧版)
### 目的
`TpuExecutable_LoadProgramAndEnqueueToStream`(`0xeaafba0`)是 C-ABI 编排层。它的工作是把 TF-TPU op kernel 通过 executable-API 函数结构传入的扁平启动结构,转换成 jellyfish 核心期望的有类型 C++ 对象:一个 `HloModule`(只用于其 shape 签名)、一个 `xla::ExecutableRunOptions`、一个反序列化的 `xla::DeviceAssignment`,以及两个 `DeviceAddressBase` 向量。它是一个很长的函数(867 行),主要因为每个 `absl::Status` 和 `xla::Shape` move 都被展开写出,但算法是线性的。
### 算法
```c
function TpuExecutable_LoadProgramAndEnqueueToStream(launch /*a1*/): // 0xeaafba0
tpu_program = launch[2] // the serialized TpuExecutable
// ---- Step 1: rebuild the HloModule shape signature ----------------
out_proto = tpu_program+112 (or ShapeProto default if null)
out_shape = Shape::FromProto(out_proto)
CHECK(out_shape is OK) // "output_shape is OK", line 101
program_shape = { result: out_shape, params: [] }
for p in tpu_program.parameter_shape_protos: // tpu_program+24 (count +32)
s = Shape::FromProto(p)
CHECK(s is OK) // "shape is OK", line 105
program_shape.params.push_back(ShapeLayout(s)) // vector<ShapeLayout>, 320 B stride
config = HloModuleConfig(program_shape)
module = make_unique<HloModule>("DeepseaExecutableModule", config)
// ---- Step 2: resolve the device ordinal + backend -----------------
stream = TpuStream::Stream(launch[12]) // the SE Stream wrapper
se = stream->parent() // vtable+144
ordinal = se->device_ordinal() // vtable+40
CHECK(TPUNodeInterfaces::Get(ordinal, &interfaces) is OK) // line 169
backend = interfaces.backend()
// ---- Step 3: optional host-transfer manager (outside-compilation) -
if launch[13] /*outside_compilation_params*/:
htm = new TpuHostTransferManagerImpl(se, backend, ...) // 0x98 bytes
info = TPUHostTransferInfoProto()
if !info.ParseFromString(params->host_transfers):
return Error("Could not call ParseFromArray() on host_transfers") // line 91
htm.Initialize(host_transfers, platform=DeepseaPlatform::GetTopology, ...)
// ---- Step 4: construct the jellyfish core executable --------------
deepsea_exec = new DeepseaExecutable(module, backend, htm, ...) // 0x98 bytes
// ---- Step 5: deserialize the device assignment --------------------
if launch[11].device_assignment_proto:
proto = DeserializeProto<DeviceAssignmentProto>(...)
assignment = DeviceAssignment::Deserialize(proto) // StatusOr, throws on bad
// ---- Step 6: fill ExecutableRunOptions ----------------------------
run_opts = ExecutableRunOptions()
run_opts.set_stream(stream)
run_opts.set_device_assignment(assignment)
run_opts.set_rng_seed(launch[20]) // *((u32*)launch+20)
run_opts.set_allocator(backend.allocator) // backend+96
// ---- Step 7: marshal the buffer-address arrays --------------------
outputs = [] // DeviceAddressBase, 24 B each
for i in [0 .. launch[4]): outputs.push(ApiConverter::FromC(launch[3] + 24*i))
output0 = ApiConverter::FromC(launch[5]) // the single root output
inputs = []
for i in [0 .. launch[6]): inputs.push(ApiConverter::FromC(launch[7] + 24*i))
arg_idx = copy_u32_array(launch[9], launch[8]) // uint32 argument indices
// ---- Step 8: the virtual launch ----------------------------------
status = deepsea_exec->vtable[96]( run_opts, outputs, output0,
inputs, arg_idx ) // DeepseaExecutable::LoadProgramAndEnqueueToStream
// ---- Step 9: post a completion callback to free the HTM -----------
if status OK and htm:
stream->DoHostCallback( FreeHostTransferManager(stream, htm) ) // line ~768
return status陷阱 — Step 1 中重建的
HloModule不会被重新编译或重新优化;它存在的唯一目的,是把ProgramShape(结果 + 参数ShapeLayouts)带入DeepseaExecutable构造函数,让核心可以校验缓冲区数量和 shape。试图在宿主 evaluator 上运行这个 module 的重新实现会错误建模这个设计 — 编译后的TpuCoreProgram已经在tpu_program内,而该 module 只是名为"DeepseaExecutableModule"的元数据壳。这里两个CHECK(... is OK)失败是致命的LogMessageFatal调用(tpu_execute_c_api.cc:101、:105),不是可恢复错误:畸形 shape proto 会中止进程。怪癖 — device ordinal 是通过遍历 stream 恢复的,而不是直接传入:先
stream->parent()(StreamExecutor,vtable 槽 +144),再device_ordinal()(vtable 槽 +40),然后TPUNodeInterfaces::Get(ordinal)。这个函数中会这样做两次 — 一次用于主路径(第 169 行),一次在 host-transfer-manager 分支中(第 124 行)— 因为 HTM 需要相同的(StreamExecutor, backend)对。重新实现者可以解析一次并复用;二进制中的重复是内联错误处理展开的副作用。
函数地图
| 函数 | 地址 | 角色 |
|---|---|---|
TpuExecutable_LoadProgramAndEnqueueToStream | 0xeaafba0 | C-ABI 编排 -> jellyfish 核心 |
xla::Shape::FromProto | (OSS) | 从 proto 重建每个 Shape |
tensorflow::TPUNodeInterfaces::Get | (在 0xeaafba0 中) | ordinal -> {backend, ...} |
tensorflow::TpuHostTransferManagerImpl ctor | (在 0xeaafba0 中) | Outside-compilation 宿主传输 |
xla::DeviceAssignment::Deserialize | (OSS) | proto -> DeviceAssignment |
xla::ExecutableRunOptions::set_* | (inline) | stream / assignment / rng_seed / allocator |
ApiConverter::FromC(SE_DeviceAddressBase*) | (OSS-mirror) | C-ABI 地址 -> DeviceAddressBase |
FreeHostTransferManager(匿名 $_0) | (在 0xeaafba0 中) | Host-callback HTM 拆卸 |
3. 程序加载 — 将编译后的程序绑定到核心
目的
编译后的 TPU 程序是一个 tpu::TpuCoreProgram(设备端指令镜像及其 ABI 元数据 TpuCoreProgramAbi)。它不能直接入队;必须先加载到指定物理核心上,也就是把代码 DMA 到该核心的指令内存,在每核心程序缓存中注册,并产出一个 tpu::TpuCoreProgramHandle。§4 中的入队寻址的是该 handle,而不是程序。旧版核心会内联完成这件事;现代路径通过 tpu::System::LoadProgram 做一次,并缓存 LoadedProgramState。
算法 — DeepseaExecutable(旧版核心)
jellyfish 核心 DeepseaExecutable::LoadProgramAndEnqueueToStream(0x13426260)在一个函数中执行加载和入队。加载部分按核心扇出:
function DeepseaExecutable_LoadAndEnqueue(run_opts, outputs, out0, inputs, arg_idx): // 0x13426260
exec = run_opts.stream()->...->DeepseaExecutor
plat = DeepseaPlatform // GetTopology, GetCoreType
config = plat.GetChipConfig() // TpuChipConfig
// ---- per-core program-handle fan-out ------------------------------
handles = [] // vector<TpuCoreProgramHandle>
if config.Megacore() and plat.GetCoreType() != 2: // line 200: dual-core chip
for core in chip.cores(): // one LoadProgram per core
h = DeepseaExecutor::LoadProgram(core_program, core)
handles.push(h)
fp = h.fingerprint() // line 277 — for the unload record
unload_info.push_back(ProgramUnloadInfo{h, fp}) // line 307
else: // single-core (line 412)
h = DeepseaExecutor::LoadProgram(core_program) // line 430
fp = h.fingerprint() // line 441
handles.push(h); unload_info.push_back(ProgramUnloadInfo{h, fp})
// on load failure:
// return Status.AddSourceLocation(deepsea_executable.cc:311/324)
... continue to enqueue (§4) ...
```text
### 算法 — tpu::System::LoadProgram(现代)
`tpu::System::LoadProgram`(`0x1d0b2240`,`system.cc:1804`)是 TFRT 原生加载。它从 `TpuCoreLocation` 解析出每核心子对象,分配一个 async-value program handle,并通过 event issuer 投递加载。
```c
function tpu_System_LoadProgram(out /*AsyncValueRef<ProgramHandle>*/, loc, program): // 0x1d0b2240
chip = TpuCoreLocation::Chip(loc) // physical chip id
chip_obj = system.impl->chip_for(chip) // vtable+80
core_id = chip_obj->core_index(loc.index) // vtable+32
// fingerprint: inline rep (loc+671>=0) or out-of-line {ptr@648, len@656}
fp = read_fingerprint(program)
host_idx = TpuCoreLocation::LocalSharedMemory(loc, 0).index_on_host()
core = system.cores[host_idx] // per-core runtime object
// allocate a ConcreteAsyncValue<tpu::ProgramHandle> (128 B, 64-aligned)
handle_av = new AsyncValue<ProgramHandle>(unconstructed)
*out = handle_av
// build an IssueArgs closure carrying TpuEventIssuer::FulfillArgs and
// the TraceContext, reserve the dependency vector, and issue:
issue = IssueArgs{ ctx, fp, RunWhenDepsReady→FulfillArgs }
core->vtable[40](issue) // enqueue the load on the core
// stamp the fingerprint into the resolved async value (walk indirection chain)
final_av = follow_indirect(handle_av) // skip forwarding nodes
final_av->fingerprint = fp // +72
VLOG(1) "TPU System::LoadProgram completed fingerprint: <hex>" // system.cc:1804
return out说明 — 现代加载会立即返回一个未解析的
AsyncValueRef<tpu::ProgramHandle>;实际的到核心 DMA 异步发生,并在TpuEventIssuer::FulfillArgs触发时使 handle 可用。§4 中的入队随后依赖该 async value,因此一次启动可以在程序物理驻留前被准备好 — 序列点引擎(stream-semantics.md)保证请求执行前程序已经加载。TpuExecutableLoadState::LoadInternal(0xf80c1c0)中的LoadedProgramState缓存意味着每个(executable, device)的逐核心 DMA 只发生一次,而不是每次启动一次。怪癖 —
Megacore芯片(两个 TensorCore 共享一个 HBM stack)需要每核心一个 program handle,并且循环由TpuChipConfig::Megacore()和GetCoreType() != 2(第 200 行)共同门控。Core type 2 被排除在双核心扇出之外 — 在这个构建中它是 SparseCore/sequencer 系列,会通过不同路径加载。总是只加载一个 handle 的重新实现只会运行 Megacore 芯片的两个核心之一,并静默地把吞吐量减半;总是加载两个 handle 的重新实现会在单核心部件上重复加载,并破坏程序缓存。
函数地图
| 函数 | 地址 | 角色 |
|---|---|---|
tpu::System::LoadProgram | 0x1d0b2240 | 现代异步程序加载(system.cc:1804) |
deepsea::executor::DeepseaExecutor::LoadProgram | (在 0x13426260 中) | 旧版同步核心加载 |
tpu::TpuExecutableLoadState::LoadInternal | 0xf80c1c0 | 按设备缓存 LoadedProgramState |
tpu::TpuCoreLocation::Chip / LocalSharedMemory | (在 0x1d0b2240 中) | 解析芯片 + 每核心子对象 |
tpu::TpuChipConfig::Megacore | (在 0x13426260 中) | 双核心扇出门控(第 200 行) |
deepsea::executor::DeepseaPlatform::GetCoreType | (在 0x13426260 中) | Core-type 判别器(!= 2) |
tpu::TpuCoreProgramHandle::fingerprint | (在 0x13426260 中) | 卸载记录中的程序身份 |
ProgramUnloadInfo(匿名结构体) | (在 0x13426260 中) | 每 handle 卸载簿记 |
4. 入队 — 构建请求并推送到流
目的
程序已经加载且缓冲区地址已经绑定后,最后一步是组装一个设备 Request — program handle + 输入/输出 DeviceAddressBase + run options + 同步事件 — 并将其推入核心的命令流。旧版路径通过 DeepseaStream::EnqueueRequest 入队 deepsea::executor::Request;现代路径调用 tpu::System::Execute,并让 TpuEventIssuer 根据先前工作对其排序。驱动私有的末端是 jxc::Queue::EnqueueRequest,它接收一个 DmaBuffer 和一个完成回调。
算法 — DeepseaExecutable 入队(旧版)
... (continuing DeepseaExecutable_LoadAndEnqueue from §3) ...
// ---- per-replica core selection -----------------------------------
for h in handles:
core_loc = h.core() // line 607 — TpuCoreLocation*
logical_id = TpuCoreLocation::LogicalDeviceId(core_loc) // line 608
// logical_id picks the device-assignment row for this replica
// ---- build the device request ---------------------------------
req = new deepsea::executor::Request // line 572: operator new(0xD0) = 208 B
req.program_handle = h
req.fingerprint = h.fingerprint() // line 665
req.inputs = inputs // DeviceAddressBase span
req.outputs = outputs
req.arg_indices = arg_idx
req.completion = $_2 closure (RemoteInvoker) // line 711-712
// ---- enqueue on the core's stream -----------------------------
DeepseaStream::EnqueueRequest(stream, move(req)) // line 593 → 0x1d0e9840
// on failure: Status.AddSourceLocation(deepsea_executable.cc:367/387)
return OkStatus() // or MakeErrorStream(deepsea_executable.cc:278)
```text
`DeepseaStream::EnqueueRequest`(`0x1d0e9840`)把 `unique_ptr<Request>` 交给 `DeepseaRequestQueue`(`EnqueueRequest` `0x1d0f23a0` -> `EnqueueRequestLocked` `0x1d0f0e80`),后者由驱动分派线程排入 `jxc::Queue::EnqueueRequest`(`0xe7d9be0`)— 这个硬件队列接收 `DmaBuffer` 和一个 `AnyInvocable<void(absl::Status)>` 完成回调。
### 算法 — tpu::System::Execute(现代)
现代入队是 `tpu::System::Execute`(`0x1d0b33e0`),从 `TpuExecutableLoadState::ExecuteLaunchRaw`(`0xf8109a0`)抵达。它不是同步队列推送,而是把所有内容穿过 async values:
```c
function ExecuteLaunchRaw(exec, per_launch_args, opts, ...): // 0xf8109a0
program_av = LoadInternal(exec).program_handle // AsyncValueRef<ProgramHandle>
user_promise = client.CreateLinkedUserPromise() // PJRT completion event
htm = TpuHostTransferManager::SetExecuteEvent(tpu_event) // host transfers (if any)
// (BarnaCoreManager: sparsecore offload, if the program has a SC partition)
tpu::System::Execute( // 0x1d0b33e0
program_av, // the loaded program
ExecuteOptions{rng_seed, launch_id, ...},
inputs = Span<AsyncValueRef<TpuBufferBase>>, // bound HBM buffers
outputs = Span<AsyncValueRef<TpuBufferBase>>,
wait_events = Span<AsyncValueRef<TpuEvent>>, // ordering deps
define_events = Span<AsyncValueRef<TpuEvent>>) // fulfilled on completion
// System::Execute fulfils each define event via
// TpuEventIssuer::FulfillArgs when the program completes on device.说明 —
wait_events/define_eventsspan 是完整的排序契约。System::Execute不会阻塞;它会把启动注册到TpuEventIssuer,后者只在每个wait_event都被履行后(输入已产生且程序已加载)运行它,并在设备发出完成信号时履行每个define_event。这是stream->WaitFor(event)+stream->RecordEvent(event)折叠为一次依赖图提交后的 TPU 类比 —stream-semantics.md对此有深入覆盖,宿主侧完成接线则在completion-loop.md中。陷阱 — 旧版路径中的逐副本核心选择是
TpuCoreProgramHandle::core()->TpuCoreLocation::LogicalDeviceId()(第 607-608 行),不是 §2 中的宿主 device ordinal。§2 中的 ordinal 命名宿主正在与哪个 TPU node 通信;这里的LogicalDeviceId命名这个特定已加载 handle 服务DeviceAssignment的哪一行(哪个副本/分区)。在 Megacore 芯片上,两个 handle 具有相同芯片但不同 logical device id,因此它们消费不同 assignment 行以及(可能)不同输入缓冲区。混淆这两个 id 会把每个副本的工作都路由到同一个核心。
DeviceAddressBase 参数数组
绑定的缓冲区以 stream_executor::DeviceAddressBase 记录传输 — 与驻留记录中存储的 (opaque_ptr, size) HBM 身份相同。旧版编排(§2 Step 7)通过 ApiConverter::FromC 从 C-ABI 启动结构将它们复制到三个向量中,每项 stride 为 24 字节:
launch struct (C ABI) marshalled C++ (DeviceAddressBase, 24 B each)
launch[3] = output addr array ───► vector<DeviceAddressBase> outputs (count launch[4])
launch[5] = root output addr ───► DeviceAddressBase output0
launch[7] = input addr array ───► vector<DeviceAddressBase> inputs (count launch[6])
launch[9] = arg-index array ───► vector<uint32> arg_idx (count launch[8], 4 B each)
```text
> **怪癖 —** 每个 `DeviceAddressBase` 的 `size` 字段**不是**用户的逻辑字节数;它是*填充后设备 shape* 的大小 — 恰好是平铺 HBM 缓冲区的 `ShapeSizeBytesRaw`([tpu-buffer-layout.md](../memory/tpu-buffer-layout.md#4-byte-sizes--shapesizebytesraw-and-friends))。程序按平铺、填充后的缓冲区读写这些地址;传入逻辑(未填充)大小会导致缓冲区短读,并破坏尾部 tile。`arg_idx` `uint32` 向量存在是因为输入缓冲区可能被重排或别名到输出(donation)— 它把每个物理输入槽映射到其参数索引,因此变成输出的 donated input 只绑定一次。
### 函数地图
| 函数 | 地址 | 角色 |
|---|---|---|
| `xla::jellyfish::DeepseaExecutable::LoadProgramAndEnqueueToStream` | `0x13426260` | 旧版加载 + 入队核心 |
| `deepsea::executor::DeepseaStream::EnqueueRequest` | `0x1d0e9840` | 将 `Request` 推入核心流 |
| `deepsea::executor::DeepseaRequestQueue::EnqueueRequest` | `0x1d0f23a0` | 队列插入(加锁) |
| `deepsea::executor::DeepseaRequestQueue::EnqueueRequestLocked` | `0x1d0f0e80` | 已加锁队列插入 |
| `asic_sw::driver::deepsea::jxc::Queue::EnqueueRequest` | `0xe7d9be0` | 硬件队列(DmaBuffer + completion) |
| `tpu::System::Execute` | `0x1d0b33e0` | 现代异步入队 |
| `tpu::TpuExecutableLoadState::ExecuteLaunchRaw` | `0xf8109a0` | 现代逐启动驱动 |
| `tpu::TpuCoreProgramHandle::core` | (在 `0x13426260` 中) | Handle -> `TpuCoreLocation`(副本) |
| `tpu::TpuCoreLocation::LogicalDeviceId` | (在 `0x13426260` 中) | Core -> device-assignment 行 |
| `tpu::TpuEventIssuer::IssueArgs` / `FulfillArgs` | (在 `0x1d0b2240` / `0x1d0b33e0` 中) | 序列点提交 / 完成 |
---
## 5. 完成和卸载簿记
### 目的
一次加载程序并入队请求的启动,还必须安排清理:已加载 program handle 在每核心缓存中引用计数,并可能在没有启动引用它们时卸载;而为 outside-compilation 分配的任何 `TpuHostTransferManagerImpl`,只能在流排空到越过使用它的请求后才能释放。在发起线程上做这些会与设备竞争;两条路径都会将其延迟。
### 卸载记录
旧版核心会为每个已加载程序记录一个 `ProgramUnloadInfo{ handle, fingerprint }`(`deepsea_executable.cc`,Megacore 分支第 307 行,单核心第 471 行)。fingerprint(`TpuCoreProgramHandle::fingerprint`)在每核心缓存中标识程序,使后续卸载(`TpuCoreCommonImpl::UnloadProgram`,`0x1d13e6a0` / `UnloadProgramWithFingerprintLegacy`,`0x1d141580`)匹配正确的缓存条目。handle 的引用计数会在飞行中保持程序驻留。
### 宿主回调拆卸
当可执行文件使用 outside-compilation 时,§2 中的编排会在约第 768 行投递一个宿主回调:
```c
// after a successful enqueue, on the SAME stream:
stream->DoHostCallback( FreeHostTransferManager(stream, htm) ) // anon $_0DoHostCallback 会把闭包入队到流上请求之后,因此 TPU 驱动只会在请求完成后于宿主线程上触发它 — 此时释放 host-transfer manager 是安全的。在旧版 TpuStream 上,这会通过 C-shim trampoline(TpuStream::DoHostCallbackWithStatus,0xe998fa0)路由;在同步宿主流上,它会内联运行。闭包 holder 是一个 32 字节的 operator new(32, 16) 分配,并被移入 AnyInvocable。
说明 — 现代 PJRT 路径用
CommonPjRtClient::CreateLinkedUserPromise加TpuHostTransferManager::SetExecuteEvent(二者都在ExecuteLaunchRaw,0xf8109a0中)替代这一点:用户可见的PJRT_Event(一个tsl::Future<void>)和 host-transfer 生命周期绑定到启动定义的同一个TpuEvent,因此清理由 async-value 解析驱动,而不是由显式宿主回调驱动。结果相同 — 资源在设备完成后释放 — 但机制是依赖图,而不是流回调。参见completion-loop.md。
函数地图
| 函数 | 地址 | 角色 |
|---|---|---|
ProgramUnloadInfo push_back | (在 0x13426260 中) | 每 handle 卸载记录 |
tpu::TpuCoreCommonImpl::UnloadProgram | 0x1d13e6a0 | 丢弃已加载程序(带回调) |
tpu::TpuCoreCommonImpl::UnloadProgramWithFingerprintLegacy | 0x1d141580 | 以 fingerprint 为键的卸载 |
FreeHostTransferManager(匿名 $_0) | (在 0xeaafba0 中) | 通过宿主回调延迟释放 HTM |
tensorflow::tpu::TpuStream::DoHostCallbackWithStatus | 0xe998fa0 | 旧版宿主回调 trampoline |
xla::CommonPjRtClient::CreateLinkedUserPromise | (在 0xf8109a0 中) | 现代用户完成事件 |
TpuHostTransferManager::SetExecuteEvent | (在 0xf8109a0 中) | 将 HTM 生命周期绑定到启动事件 |
相关组件
| 组件 | 关系 |
|---|---|
TpuExecutable_LoadProgramAndEnqueueToStream(0xeaafba0) | 构建 run-options 并转发(vtable+96)到 jellyfish 核心的旧版 C-ABI shim |
xla::jellyfish::DeepseaExecutable(0x13426260) | 按核心加载 program handle 并入队请求的核心 |
tpu::System(LoadProgram 0x1d0b2240,Execute 0x1d0b33e0) | PJRT 路径用于加载 + 入队的现代 TFRT 原生运行时 |
deepsea::executor::DeepseaStream / DeepseaRequestQueue | 旧版命令流 + 请求队列 |
asic_sw::driver::deepsea::jxc::Queue(0xe7d9be0) | 两条路径最终落到的硬件队列(DmaBuffer + completion) |
tpu::TpuEventIssuer | 对 load/enqueue 彼此排序的序列点引擎 |
stream_executor::DeviceAddressBase | 作为程序缓冲区参数绑定的 HBM (ptr, size) 记录 |
交叉引用
- execute-async-on-stream.md — 面向 PJRT 的上半部分:参数固定、输出分配、donation/aliasing,以及生成本页绑定的已解析缓冲区的副本扇出
- stream-semantics.md —
TpuEventIssuer序列点、wait/define 事件,以及入队如何相对于先前工作排序 - completion-loop.md — 解析启动 define events 并释放延迟资源的宿主侧完成接线
- overview.md — 本页两个入口所属的双栈运行时架构(现代 PJRT
TpuClient/tpu::System与旧版 StreamExecutorTpuExecutor) - ../memory/tpu-buffer-layout.md — 设备端填充/平铺缓冲区,以及其
(ptr, size)对作为绑定参数的DeviceAddressBase驻留记录 - ../memory/hbm-allocator.md — 产生这里绑定的 HBM 偏移的
BestFitAllocator;run options 中的set_allocator(backend+96)引用 - 返回索引 — Part XI — Runtime & Execution