Skip to content

LoadProgramAndEnqueueToStream

本页中的所有地址、结构体偏移和源码行引用均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so(构建 libtpu_lts_20260413_b_RC00,build-id md5 89edbbe81c5b328a958fe628a9f2207d,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_LoadProgramAndEnqueueToStream0xeaafba0,源码 tpu_execute_c_api.cc)是一个自由的 T 导出函数,注册在 TF-TPU executable-API 函数结构的索引 7 处,并不是 xla::TpuExecutable 类的方法;它从 C 启动结构编排出 xla::ExecutableRunOptions,并通过一次虚调用(vtable+96)转发到 jellyfish 核心。该核心是 xla::jellyfish::DeepseaExecutable::LoadProgramAndEnqueueToStream0x13426260,源码 deepsea_executable.cc),真正执行 DeepseaExecutor::LoadProgram -> DeepseaStream::EnqueueRequest。现代 PJRT 路径(xla::TpuClient over tpu::System,参见 overview.md)通过 TFRT async-value 写法抵达同一个设备运行时:程序加载是 tpu::System::LoadProgram0x1d0b2240,源码 system.cc:1804),入队是 tpu::System::Execute0x1d0b33e0),二者都由 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 如何重建 HloModule shape 签名、解析 device ordinal,并填充 xla::ExecutableRunOptionsset_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>) @ 0x1d0b2240system.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 @ 0xf8109a0LoadInternal @ 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),只在宿主侧抽象上不同。

text
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_LoadProgramAndEnqueueToStream0xeaafba0C-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::LoadProgramAndEnqueueToStream0x13426260)在一个函数中执行加载和入队。加载部分按核心扇出:

c
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 enqueue4) ...
```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::LoadInternal0xf80c1c0)中的 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::LoadProgram0x1d0b2240现代异步程序加载(system.cc:1804
deepsea::executor::DeepseaExecutor::LoadProgram(在 0x13426260 中)旧版同步核心加载
tpu::TpuExecutableLoadState::LoadInternal0xf80c1c0按设备缓存 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 入队(旧版)

c
... (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_events span 是完整的排序契约。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 字节:

text
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 $_0

DoHostCallback 会把闭包入队到流上请求之后,因此 TPU 驱动只会在请求完成后于宿主线程上触发它 — 此时释放 host-transfer manager 是安全的。在旧版 TpuStream 上,这会通过 C-shim trampoline(TpuStream::DoHostCallbackWithStatus0xe998fa0)路由;在同步宿主流上,它会内联运行。闭包 holder 是一个 32 字节的 operator new(32, 16) 分配,并被移入 AnyInvocable

说明 — 现代 PJRT 路径用 CommonPjRtClient::CreateLinkedUserPromiseTpuHostTransferManager::SetExecuteEvent(二者都在 ExecuteLaunchRaw0xf8109a0 中)替代这一点:用户可见的 PJRT_Event(一个 tsl::Future<void>)和 host-transfer 生命周期绑定到启动定义的同一个 TpuEvent,因此清理由 async-value 解析驱动,而不是由显式宿主回调驱动。结果相同 — 资源在设备完成后释放 — 但机制是依赖图,而不是流回调。参见 completion-loop.md

函数地图

函数地址角色
ProgramUnloadInfo push_back(在 0x13426260 中)每 handle 卸载记录
tpu::TpuCoreCommonImpl::UnloadProgram0x1d13e6a0丢弃已加载程序(带回调)
tpu::TpuCoreCommonImpl::UnloadProgramWithFingerprintLegacy0x1d141580以 fingerprint 为键的卸载
FreeHostTransferManager(匿名 $_0(在 0xeaafba0 中)通过宿主回调延迟释放 HTM
tensorflow::tpu::TpuStream::DoHostCallbackWithStatus0xe998fa0旧版宿主回调 trampoline
xla::CommonPjRtClient::CreateLinkedUserPromise(在 0xf8109a0 中)现代用户完成事件
TpuHostTransferManager::SetExecuteEvent(在 0xf8109a0 中)将 HTM 生命周期绑定到启动事件

相关组件

组件关系
TpuExecutable_LoadProgramAndEnqueueToStream0xeaafba0构建 run-options 并转发(vtable+96)到 jellyfish 核心的旧版 C-ABI shim
xla::jellyfish::DeepseaExecutable0x13426260按核心加载 program handle 并入队请求的核心
tpu::SystemLoadProgram 0x1d0b2240Execute 0x1d0b33e0PJRT 路径用于加载 + 入队的现代 TFRT 原生运行时
deepsea::executor::DeepseaStream / DeepseaRequestQueue旧版命令流 + 请求队列
asic_sw::driver::deepsea::jxc::Queue0xe7d9be0两条路径最终落到的硬件队列(DmaBuffer + completion)
tpu::TpuEventIssuer对 load/enqueue 彼此排序的序列点引擎
stream_executor::DeviceAddressBase作为程序缓冲区参数绑定的 HBM (ptr, size) 记录

交叉引用

  • execute-async-on-stream.md — 面向 PJRT 的上半部分:参数固定、输出分配、donation/aliasing,以及生成本页绑定的已解析缓冲区的副本扇出
  • stream-semantics.mdTpuEventIssuer 序列点、wait/define 事件,以及入队如何相对于先前工作排序
  • completion-loop.md — 解析启动 define events 并释放延迟资源的宿主侧完成接线
  • overview.md — 本页两个入口所属的双栈运行时架构(现代 PJRT TpuClient/tpu::System 与旧版 StreamExecutor TpuExecutor
  • ../memory/tpu-buffer-layout.md — 设备端填充/平铺缓冲区,以及其 (ptr, size) 对作为绑定参数的 DeviceAddressBase 驻留记录
  • ../memory/hbm-allocator.md — 产生这里绑定的 HBM 偏移的 BestFitAllocator;run options 中的 set_allocator(backend+96) 引用
  • 返回索引 — Part XI — Runtime & Execution