Internal Wiki

主题

PJRT 可执行文件加载与执行 ​

本页中的所有地址均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（libtpu_lts_20260413_b_RC00，BuildID 89edbbe81c5b328a958fe628a9f2207d，ELF x86-64 DYN，约 745 MB）。PJRT C-API 表面为 v0.103。其他构建会重新编号槽位并移动地址。

摘要 ​

本页重建了 libtpu.so 的 PJRT C-ABI 可执行文件表面：宿主框架（JAX、PyTorch/XLA）通过 PJRT_Api vtable 调用的 pjrt::PJRT_Executable_* 和 pjrt::PJRT_LoadedExecutable_* 包装器，用于编译程序、序列化和重新加载程序、查询其形状元数据，并将其启动到 TPU 上。这些入口中的每一个都是覆盖 C++ 对象的很薄的 C 封送 shim：包装器解包带版本的 _Args 结构体，通过 xla::PjRtExecutable / xla::PjRtLoadedExecutable 上的一个虚槽分派，并把 C++ 结果打包回 _Args 结构体。它们逐字节对应通用 XLA pjrt_c_api_wrapper_impl.cc 代码，没有任何一个是 TPU 专用的（见 API Vtable 重建 中已填充与已注入映射）。

该契约在形态上与上游 XLA 的 PJRT C-API 完全相同：宿主拥有 _Args 结构体，插件拥有不透明 PJRT_Executable* / PJRT_LoadedExecutable* 句柄背后的对象，并且每个 _Args 结构体开头的 struct_size 字段承载版本兼容性（向后兼容保护 pjrt::ActualStructSizeIsGreaterOrEqual）。TPU 专用的内容完全位于这些包装器下方、虚调用背后：xla::CommonPjRtLoadedExecutable 框架下的具体 xla::TpuExecutable（已编译）/ xla::TpuLoadedExecutable（已加载到设备），以及运行时入队（ExecutePrepare → ExecuteLaunch → tpu::System::Execute）。

本页负责 C-ABI 包装层及其封送：Compile → 可执行文件，Serialize / DeserializeAndLoad，Execute（参数到设备缓冲区的封送、多设备与 execute_device 分派、PJRT_ExecuteOptions、输出 PJRT_Buffer 列表），以及名称 / 输出数量 / 内存种类访问器。Execute 包装器向下驱动的运行时内部入队见 ExecuteAsyncOnStream 和 加载程序并入队（现代 PJRT 路径通过 CommonPjRtLoadedExecutable::Execute → tpu::System::Execute，不是旧版 StreamExecutor ExecuteAsyncOnStream）；编译缓存见 编译缓存；多阶段编译扩展见 PhaseCompile 扩展。

对于重新实现者，契约如下：

  • _Args 结构体纪律。 每个入口都以 ActualStructSizeIsGreaterOrEqual(name, min, current, args->struct_size) 开始；失败时返回堆上的 PJRT_Error*，且不读取任何内容。Execute 是唯一会验证两个结构体的入口（外层 args 和嵌套的 PJRT_ExecuteOptions）。
  • 句柄包装规则。 GetExecutable / DeserializeAndLoad 会对 C-ABI 盒子执行 operator new（PJRT_Executable = 0x250 B，PJRT_LoadedExecutable = 0x48 B），把它包在 C++ 对象外，并将盒子指针存入 _Args 输出字段。宿主通过 PJRT_*_Destroy 释放它。
  • Execute 分派分叉。 execute_device == NULL → 多设备 xla::PjRtLoadedExecutable::Execute（返回 vector<vector<unique_ptr<PjRtBuffer>>>）；execute_device != NULL → 单设备路径，该路径要求 num_devices == 1，并根据可执行文件的 compile_portable_executable 标志调用 portable（vtable +80）或 sharded（vtable +72）虚函数。
  • 序列化往返。 Serialize 调用可执行文件的 SerializeExecutable 虚函数，生成一个带 deleter 回调的堆上 PJRT_SerializedExecutable；DeserializeAndLoad 解析 CompileOptionsProto，重建 CompileOptions，并通过 client 重新加载。
构建器（Compile）pjrt::PJRT_Client_Compile @ 0x0F861820（槽 25）
Executepjrt::PJRT_LoadedExecutable_Execute @ 0x0F869B40（槽 60，1914 行反编译代码）
Serializepjrt::PJRT_Executable_Serialize @ 0x0F86C5A0（槽 54）
Deserialize+Loadpjrt::PJRT_Executable_DeserializeAndLoad @ 0x0F86CC40（槽 61）
GetExecutablepjrt::PJRT_LoadedExecutable_GetExecutable @ 0x0F86CFA0（槽 56）
兼容性保护pjrt::ActualStructSizeIsGreaterOrEqual @ 0x0F8A4EC0（每个入口）
C-ABI 盒子大小PJRT_Executable = 0x250（592 B）；PJRT_LoadedExecutable = 0x48（72 B）
已编译 C++ 对象xla::TpuExecutable（typeinfo 0x21786610），位于 xla::CommonPjRtLoadedExecutable 下
来源（断言）third_party/tensorflow/compiler/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc

对象模型：C 句柄 ↔ C++ 可执行文件 ​

目的 ​

它们背后有两个不透明 C 句柄和一个三类 C++ 层次结构。宿主从不看到 C++ 类型；它持有一个 PJRT_Executable*（已编译程序元数据）或一个 PJRT_LoadedExecutable*（驻留在设备上、可启动的程序），每个都是一个堆盒子，包装指向真实对象的 unique_ptr。理解哪个句柄拥有哪个虚函数表面，是下面每个包装器的前提。

句柄布局 ​

text
PJRT_Executable        (operator new 0x250 = 592 B)
  └─ unique_ptr<xla::PjRtExecutable>   ──▶  xla::TpuExecutable   (typeinfo 0x21786610)
                                              GetHloModules, name(), num_outputs,
                                              SerializeExecutable, GetOutputMemoryKinds,
                                              GetCostAnalysis, GetCompiledMemoryStats, …

PJRT_LoadedExecutable  (operator new 0x48  = 72 B)
  └─ unique_ptr<xla::PjRtLoadedExecutable> ─▶ xla::TpuLoadedExecutable (typeinfo 0x2177b9b8)
                                              GetExecutable(), Execute(), ExecuteSharded(),
                                              ExecutePortable(), IsDeleted(), Delete()
       (device-bound; under xla::CommonPjRtLoadedExecutable framework, typeinfo 0x2178a0f0)
```text

这种拆分精确映射上游 XLA：`PJRT_Executable` 是*启动前*产物（它可以被序列化、查询，并发送到另一台主机），而 `PJRT_LoadedExecutable` *绑定到此 client 的设备*，也是唯一可以 `Execute` 的句柄。`PJRT_LoadedExecutable_GetExecutable`（槽 56）是桥梁：它从已加载对象中取出已编译的 `PjRtExecutable`，并装入一个新的 `PJRT_Executable` 盒子。

### 函数映射 - 可执行文件表面

| 槽 | 字段 | 符号（`pjrt::`…） | 地址 | 包装的虚函数 |
|---|---|---|---|---|
| 25 | PJRT_Client_Compile | `PJRT_Client_Compile` | `0x0F861820` | client compile → loaded exec |
| 45 | PJRT_Executable_Destroy | `PJRT_Executable_Destroy` | `0x0F8661C0` | box dtor |
| 46 | PJRT_Executable_Name | `PJRT_Executable_Name` | `0x0F866860` | `+40` `name()` |
| 47 | PJRT_Executable_NumReplicas | `PJRT_Executable_NumReplicas` | `0x0F8668C0` | `num_replicas()` |
| 48 | PJRT_Executable_NumPartitions | `PJRT_Executable_NumPartitions` | `0x0F866920` | `num_partitions()` |
| 49 | PJRT_Executable_NumOutputs | `PJRT_Executable_NumOutputs` | `0x0F866A40` | cached field `+272` |
| 50 | PJRT_Executable_SizeOfGeneratedCodeInBytes | `…SizeOfGeneratedCodeInBytes` | `0x0F867240` | `SizeOfGeneratedCodeInBytes()` |
| 51 | PJRT_Executable_GetCostAnalysis | `PJRT_Executable_GetCostAnalysis` | `0x0F867B80` | `GetCostAnalysis()` |
| 52 | PJRT_Executable_OutputMemoryKinds | `PJRT_Executable_OutputMemoryKinds` | `0x0F869520` | `+104` `GetOutputMemoryKinds()` |
| 53 | PJRT_Executable_OptimizedProgram | `PJRT_Executable_OptimizedProgram` | `0x0F8672A0` | `GetHloModules()` |
| 54 | PJRT_Executable_Serialize | `PJRT_Executable_Serialize` | `0x0F86C5A0` | `+144` `SerializeExecutable()` |
| 56 | PJRT_LoadedExecutable_GetExecutable | `PJRT_LoadedExecutable_GetExecutable` | `0x0F86CFA0` | `+32` `executable()` |
| 58 | PJRT_LoadedExecutable_Delete | `PJRT_LoadedExecutable_Delete` | `0x0F869A80` | `Delete()` |
| 59 | PJRT_LoadedExecutable_IsDeleted | `PJRT_LoadedExecutable_IsDeleted` | `0x0F869AE0` | `IsDeleted()` |
| 60 | PJRT_LoadedExecutable_Execute | `PJRT_LoadedExecutable_Execute` | `0x0F869B40` | `Execute` / `+80` ExecutePortable / `+72` ExecuteSharded |
| 61 | PJRT_Executable_DeserializeAndLoad | `PJRT_Executable_DeserializeAndLoad` | `0x0F86CC40` | client deserialize+load |
| 95 | PJRT_Executable_OutputElementTypes | `PJRT_Executable_OutputElementTypes` | `0x0F868560` | cached dims/types |
| 96 | PJRT_Executable_OutputDimensions | `PJRT_Executable_OutputDimensions` | `0x0F8689E0` | cached dims/types |
| 99 | PJRT_Executable_Fingerprint | `PJRT_Executable_Fingerprint` | `0x0F867AC0` | `FingerprintExecutable()` |
| 101 | PJRT_Executable_GetCompiledMemoryStats | `PJRT_Executable_GetCompiledMemoryStats` | `0x0F86CAC0` | `GetCompiledMemoryStats()` |
| 129 | PJRT_Executable_GetCompileOptions | `PJRT_Executable_GetCompileOptions` | `0x0F86C6E0` | `GetCompileOptions()` |
| 139 | PJRT_Executable_ParameterMemoryKinds | `PJRT_Executable_ParameterMemoryKinds` | `0x0F868FC0` | `GetParameterMemoryKinds()` |
| 122 | PJRT_LoadedExecutable_GetDeviceAssignment | `…GetDeviceAssignment` | `0x0F870EA0` | `device_assignment()` |
| 135 | PJRT_LoadedExecutable_AddressableDeviceLogicalIds | `…AddressableDeviceLogicalIds` | `0x0F8669E0` | `addressable_device_logical_ids()` |

> **注意 -** 槽 62 `PJRT_LoadedExecutable_Fingerprint`（`0x0F85FBE0`）是*已弃用*的 fingerprint 入口；v0.103 宿主使用槽 99 `PJRT_Executable_Fingerprint`。两者都已填充。保留这个已弃用入口只是为了让旧 minor 版本宿主较小的 args 结构体仍能解析到非空指针。
>
> **怪癖 -** 这 24 个槽中没有任何一个是 TPU 专用的。它们是原版 `pjrt_c_api_wrapper_impl.cc` 包装器，二进制自身的诊断字符串就指名了这个确切文件。TPU 特性完全位于虚调用（`xla::TpuExecutable` / `xla::TpuLoadedExecutable`）以及被注入的 `Client_Create`（槽 15）背后。将 libtpu 移植到另一种加速器的重新实现者应重写 C++ 对象，而不是这些包装器。

---

## Compile：`PJRT_Client_Compile` → 已加载可执行文件

### 目的

槽 25 是生成可启动程序的前门。包装器解包程序（StableHLO/MLIR 字节或 `XlaComputation` proto）和 `CompileOptionsProto`，将它们交给被包装的 client，并把返回的 `PjRtLoadedExecutable` 装箱为 `PJRT_LoadedExecutable*`。对于 TPU client，这会驱动 `TpuClient::CompileAndLoad` → `xla::PjRtCompile` → jellyfish JIT，编译页面会覆盖这些内容；此包装器只负责 C-ABI 封送。

### 入口点

```text
PJRT_Client_Compile (0x0F861820, slot 25)
  ── ActualStructSizeIsGreaterOrEqual("PJRT_Client_Compile_Args", …)
  ── parse program (format tag: "hlo" / "mlir" / StableHLO bytecode)
  ── CompileOptions::FromProto(CompileOptionsProto)
  ── (*client_vtable)(…)  ──▶  xla::TpuClient::CompileAndLoad   [compile pages]
  └─ args->executable = new PJRT_LoadedExecutable(loaded)        ── 0x48 B box

注意事项 ​

  • 缓存。 编译路径按内容寻址；相同的 HLO + options + topology 可以命中磁盘/内存缓存，而不是重新运行 JIT。键的生成和存储见 编译缓存。
  • 分阶段编译。 宿主可以通过 PhaseCompile 扩展以显式阶段驱动编译，而不是使用这个一次性槽；见 PhaseCompile 扩展。
  • 拓扑。 编译需要一个 PJRT_TopologyDescription（TPU pod/slice 几何）来选择设备分配；该句柄及其访问器见 拓扑描述。

Execute：PJRT_LoadedExecutable_Execute ​

目的 ​

槽 60 是对吞吐至关重要的入口：它把一个已加载程序启动到一个或多个 TPU 设备上。它是这个表面上最大的包装器（0x0F869B40 处 1914 行反编译代码），因为它必须将参差不齐的二维输入 PJRT_Buffer* 数组封送为 C++ 缓冲区向量，转换 14 字段 PJRT_ExecuteOptions，在多设备和单 execute_device 虚函数之间分叉，并将二维输出 PJRT_Buffer* 数组以及可选的逐设备完成 PJRT_Event* 打包回去，同时还要遵守两个嵌套结构体的 struct_size。

_Args 和 PJRT_ExecuteOptions 布局 ​

包装器读取外层 PJRT_LoadedExecutable_Execute_Args（当前 struct_size = 80 B = 10 qwords）以及它指向的嵌套 PJRT_ExecuteOptions（当前 struct_size = 112 B = 14 qwords）。下面的偏移直接从反编译主体中的字段读取恢复而来。

结构体偏移字段含义（恢复）
Execute_Args+0x00struct_size保护输入；min 34，cur 80
Execute_Args+0x08priv / extension遍历以寻找 extension ID == 1（行 257-268）
Execute_Args+0x10executablePJRT_LoadedExecutable*（v117+16 = wrapped ptr）
Execute_Args+0x18optionsPJRT_ExecuteOptions*（v14）
Execute_Args+0x40num_devicesreplica/partition 数量；与回调做 CHECK
Execute_Args+0x48num_args每设备参数数量
Execute_Args+0x60argument_listsPJRT_Buffer*** — [device][arg]（行 468-469）
Execute_Args+0x68output_listsPJRT_Buffer*** out — [device][output]
Execute_Args+0x70device_complete_events可选 PJRT_Event** out，每设备一个
Execute_Args+0x78execute_device可选 PJRT_Device*；非空 ⇒ 单设备分叉
ExecuteOptions+0x00struct_size保护输入；min 19，cur 112
ExecuteOptions+0x30launch_idint，复制到运行状态中（行 311）
ExecuteOptions+0x50context（str ptr）可选 const char*，为 TraceMe span 调用 strlen（行 312-316）
ExecuteOptions(vec)send_callbacksCHECK size() == num_devices（行 954）
ExecuteOptions(vec)recv_callbacksCHECK size() == num_devices（行 1215）

陷阱 - 两个 struct_size 检查被链接在一个表达式中（行 250）：先验证外层 Execute_Args，然后验证 **(a1+24)，即嵌套的 PJRT_ExecuteOptions，作为第二个 ActualStructSizeIsGreaterOrEqual("PJRT_ExecuteOptions", 19, 112, …)。宿主即使传入有效的外层结构体，但传入过小的 options 结构体，仍会得到 PJRT_Error*，而不是崩溃。只验证外层结构体的重新实现会越过调用者的 options 缓冲区读取。

算法 ​

c
function PJRT_LoadedExecutable_Execute(args):                 // 0xf869b40
    if !ActualStructSizeIsGreaterOrEqual(                      // line 249
            "PJRT_LoadedExecutable_Execute_Args", 34, 80, args->struct_size)
       || !ActualStructSizeIsGreaterOrEqual(                   // line 250 (nested)
            "PJRT_ExecuteOptions", 19, 112, args->options->struct_size):
        return new PJRT_Error{status}                          // line 253

    opts = TranslateExecuteOptions(args->options)             // launch_id, context,
                                                              //   send/recv callbacks
    TraceMe span("PJRT_LoadedExecutable_Execute")             // if g_trace_level>0, line 274
    exec = args->executable->wrapped                          // *(args+16)

    if args->execute_device == NULL:                          // ---- multi-device ----
        // marshal argument_lists[device][arg] -> Span<vector<PjRtBuffer*>>
        // (CHECK send_callbacks.size()==num_devices, recv_callbacks.size()==num_devices)
        out = xla::PjRtLoadedExecutable::Execute(             // line 1624 (vtable virtual)
                  exec, args->argument_lists, opts,
                  &returned_futures)                          // StatusOr<vector<vector<...>>>
        if !out.ok(): return new PJRT_Error{out.status}
        // CHECK returned_futures.size() == num_devices (line 1273)
        write out -> args->output_lists[device][output]       // new PJRT_Buffer per leaf
        if args->device_complete_events:
            args->device_complete_events[d] = new PJRT_Event{returned_futures[d]}

    else:                                                     // ---- single device ----
        if args->num_devices != 1:                            // line 1648
            return new PJRT_Error{InvalidArgument(            //   pjrt_c_api_wrapper_impl.cc:2382
                "num_devices and corresponding output list sizes must be 1 "
                "when calling …Execute with non-null execute_device. Got num_devices=%i")}
        if opts.has_send_or_recv_callbacks():                 // line 1823
            return new PJRT_Error{Unimplemented(
                "…doesn't support using send/recv callbacks with `execute_device`.")}
        copts = (*inner_exec.vtable[168])(inner_exec)         // GetCompileOptions() (line 1671-1672)
                                                              //   inner_exec = executable() via +32
        if copts.compile_portable_executable:                 // v244 == 1 (line 1677)
            out = (*exec.vtable[80])(exec, args->argument_lists[0], opts, …)  // ExecutePortable
        else:
            out = (*exec.vtable[72])(exec, args->argument_lists[0], opts, …)  // ExecuteSharded
        write out -> args->output_lists[0][output]
        if args->device_complete_events:
            args->device_complete_events[0] = new PJRT_Event{future}

    return NULL                                               // success: no PJRT_Error
```text

### 分派分叉

最重要的分支是 `execute_device`。当它为 `NULL` 时，宿主请求的是可执行文件完整的 replica/partition 扇出，包装器调用多设备 `xla::PjRtLoadedExecutable::Execute`（行 1624），它返回 `StatusOr<vector<vector<unique_ptr<PjRtBuffer>>>>`，外层索引 = 设备，内层 = 输出。当它非 `NULL` 时，宿主想要某个单设备的切片；包装器强制 `num_devices == 1`，拒绝 send/recv 回调，然后读取可执行文件的 `CompileOptions`（内部 `PjRtExecutable` 的 `GetCompileOptions()` 虚函数，vtable `+168`，行 1671-1672），并根据 `compile_portable_executable` 分支（行 1677）：若已设置，则调用 `ExecutePortable`（vtable +80）；否则调用 `ExecuteSharded`（vtable +72）。两者都返回扁平的 `StatusOr<vector<unique_ptr<PjRtBuffer>>>`。

> **怪癖 -** 按 replica / partition 的扇出**不是**此包装器中的循环。包装器将整个二维 `argument_lists` 交给 C++ `Execute`，让运行时扇出（运行时的 `ExecutePrepare` / `ExecuteLaunch` / `ExecuteSharded`-share 路径执行逐设备分派，见 [加载程序并入队](../runtime/load-program-enqueue.md)）。C-ABI 包装器中唯一的循环是把缓冲区指针拷入拷出的封送循环。在 C-ABI 层执行扇出的重新实现者会重复运行时已经做的工作，并破坏单次 `Execute` 调用的完成事件语义。
>
> **陷阱 -** 封送使用 `0xAAAAAAAAAAAAAAAB` 这个 3 的倒数魔数（行 1610），从一个内层 vector 步长为 24 字节（3 qwords）的 `vector<vector<unique_ptr<PjRtBuffer>>>` 恢复元素数量。这是标准 libc++ `vector` 大小计算，不是自定义容器；以错误叶子数量分配 `output_lists` 的重新实现会静默截断输出。

### 输出和完成事件

成功时，包装器为结果的每个叶子分配一个新的 `PJRT_Buffer`（包装 `unique_ptr<xla::PjRtBuffer>`），并将其写入调用者提供的 `output_lists[device][output]` 数组。如果 `device_complete_events` 非空，它为每个设备分配一个 `PJRT_Event`，包装对应的 `tsl::Future<void>`（TPU client 通过 `TrackFuture` / `CreateProfiledFuture` 生成这些 future）；宿主通过槽 13/14 等待它们。`PJRT_Buffer` 生命周期和 `PJRT_Event` 模型见 [缓冲区和内存](buffer-and-memory.md) 与 [事件和异步](events-and-async.md)。

---

## Serialize 和 Deserialize+Load

### 目的

`PJRT_Executable_Serialize`（槽 54）将已编译的 `PJRT_Executable` 转成宿主可持久化或发送的字节串；`PJRT_Executable_DeserializeAndLoad`（槽 61）将其反转为*此* client 上绑定到设备的 `PJRT_LoadedExecutable`。二者共同构成 AOT/cache 路径：宿主可以编译一次并序列化，在之后的运行中反序列化并加载，而不是重新编译。序列化形式通过 SE-portable proto 往返，TPU topology `Deserialize` 也消费该 proto。

### 算法 - Serialize

```c
function PJRT_Executable_Serialize(args):                     // 0xf86c5a0
    if !ActualStructSizeIsGreaterOrEqual(
            "PJRT_Executable_Serialize_Args", 30, 56, args->struct_size):
        return new PJRT_Error{status}
    bytes = std::string{}
    (*exec.vtable[144])(&bytes)        // SerializeExecutable -> string (line 29)
    holder = new PJRT_SerializedExecutable{move(bytes)}        // 0x18 B (line 38)
    args->serialized_bytes   = holder.data()                  // args[3] (+0x18)
    args->serialized_bytes_size = holder.size()               // args[4] (+0x20)
    args->serialized_executable = holder                      // args[5] (+0x28)
    args->serialized_executable_deleter = $_0::__invoke        // args[6] (+0x30)
    return NULL

算法 - DeserializeAndLoad ​

c
function PJRT_Executable_DeserializeAndLoad(args):            // 0xf86cc40
    if !ActualStructSizeIsGreaterOrEqual(
            "PJRT_Executable_DeserializeAndLoad_Args", 39, 64, args->struct_size):
        return new PJRT_Error{status}
    data    = args->serialized_executable          // *(args+24)  (line 47)
    size    = args->serialized_executable_size     // *(args+32)  (line 44)
    options = args->overridden_serialized_compile_options  // *(args+48) (line 50)
    proto   = CompileOptionsProto{}
    if options && !proto.ParseFrom(options):
        return new PJRT_Error{MakeErrorImpl(                  // line 72,
            "PJRT_Client_Compile: failed to deserialize CompileOptionsProto",
            …, pjrt_c_api_wrapper_impl.cc:1113)}              //   shared compile helper
    copts = CompileOptions::FromProto(proto)                 // line 67
    loaded = (*client.vtable)(client, data, size, copts)     // deserialize + load
    if !loaded.ok(): return new PJRT_Error{loaded.status}
    args->loaded_executable = new PJRT_LoadedExecutable(loaded)  // 0x48 B (line 175), args+40
    return NULL
```text

> **怪癖 -** `DeserializeAndLoad` 的错误消息命名为 `PJRT_Client_Compile`，而不是 `…DeserializeAndLoad`（反编译行 72，来源 `pjrt_c_api_wrapper_impl.cc:1113`）。`CompileOptionsProto` 解析是 compile 和 deserialize 都使用的*共享 helper*，诊断信息携带的是 helper 的名称。重新实现者不应根据该字符串做错误处理分支；它是同一条代码路径。
>
> **陷阱 -** `Serialize` 通过别名指向堆上 `PJRT_SerializedExecutable` 中的字节来返回数据，并返回一个 deleter 函数指针（`args[6]`）。宿主只在调用该 deleter 前拥有这些字节；`PJRT_Executable` 自身可以独立销毁。若重新实现返回指向 `PjRtExecutable` 自身存储的指针（而不是独立 holder），一旦 executable 被释放就会产生悬垂指针。
>
> **注意 -** 反序列化兼容性由 client 配置键 `executable_compatibility_check_on_deserialization` 控制（在 [`PJRT_Client_Create`](client-and-device.md) 中解析）。设置该键时，client 会在加载前验证序列化可执行文件的 ABI/topology fingerprint；不匹配会表现为返回的 `PJRT_Error`。

---

## 元数据访问器

### 目的

一组小包装器在不启动程序的情况下暴露已编译程序的形状和资源元数据。它们共享同一种形态：保护、分派一个虚函数（或读取一个缓存字段）、把标量或 `(ptr, len)` 对打包回 `_Args` 结构体。宿主在 `Execute` 前调用它们来确定输出缓冲区大小，并把输出路由到正确的内存空间。

### Name（槽 46）

```c
function PJRT_Executable_Name(args):                          // 0xf866860
    if !ActualStructSizeIsGreaterOrEqual(
            "PJRT_Executable_Name_Args", 25, 40, args->struct_size):
        return new PJRT_Error{status}
    (name_ptr, name_len) = (*exec.vtable[40])(exec)           // name(), line 14
    args->executable_name      = name_ptr                     // args[3] (+0x18)
    args->executable_name_size = name_len                     // args[4] (+0x20)
    return NULL

NumOutputs（槽 49） ​

c
function PJRT_Executable_NumOutputs(args):                    // 0xf866a40
    if !ActualStructSizeIsGreaterOrEqual(
            "PJRT_Executable_NumOutputs_Args", 31, 32, args->struct_size):
        return new PJRT_Error{status}
    if !EnsureExecutableOutputDimensionsPopulated(exec):      // 0xf866ac0, line 11
        return new PJRT_Error{status}                         //   lazily computes & caches
    args->num_outputs = exec->cached_num_outputs              // *(exec+272), line 13
    return NULL
```text

> **怪癖 -** `NumOutputs`、`OutputElementTypes`（槽 95）和 `OutputDimensions`（槽 96）都会经过 `EnsureExecutableOutputDimensionsPopulated`（`0xf866ac0`），它会*一次性*计算每输出形状表并将其缓存在 executable 上（计数落在偏移 `+272`）。这三个调用中的第一个会支付 `GetOutputShapes` 成本；其余调用读取缓存。每次调用都重新计算的重新实现会浪费 JAX 每次 trace 都会命中的热元数据路径上的工作。

### OutputMemoryKinds（槽 52）

```c
function PJRT_Executable_OutputMemoryKinds(args):             // 0xf869520
    if !ActualStructSizeIsGreaterOrEqual(
            "PJRT_Executable_OutputMemoryKinds_Args", 38, 48, args->struct_size):
        return new PJRT_Error{status}
    lock guard(exec->mutex /* exec+56 */)                     // line 64
    if !exec->memory_kinds_cached /* exec+480 */:             // line 67
        kinds = (*exec.vtable[104])(exec)                     // GetOutputMemoryKinds, line 70
        if kinds.size() != 1:                                 // line 187
            return new PJRT_Error{Unimplemented(
                "MPMD execution not supported by PJRT C API "
                "(in function PJRT_Executable_GetOutputMemoryKinds).")}
        exec->cache = kinds[0]; exec->memory_kinds_cached = 1  // line 268
    args->memory_kinds = exec->cache                          // (ptr array, sizes)
    return NULL

陷阱 - OutputMemoryKinds 会拒绝任何 GetOutputMemoryKinds() 产生超过一个模块的 kinds 的 executable，并给出 "MPMD execution not supported by PJRT C API"（行 187）。C-API 表面是 SPMD-only：一个模块跨设备复制。多程序多数据 executable（每设备不同程序）无法通过这个表面查询，或执行。缓存由 exec+56 处的每 executable absl::Mutex 和 exec+480 处的 bool 标志保护；并发的首次调用者会串行化。

访问器目录 ​

其余访问器遵循相同的保护-分派-打包形态；只有虚函数和被打包的类型不同。

槽字段返回包装的虚函数
47NumReplicassize_tnum_replicas()
48NumPartitionssize_tnum_partitions()
50SizeOfGeneratedCodeInBytesint64SizeOfGeneratedCodeInBytes()
51GetCostAnalysisnamed-value listGetCostAnalysis()
95OutputElementTypesPJRT_Buffer_Type[]cached (via Ensure…)
96OutputDimensionsdim spans + sizescached (via Ensure…)
99Fingerprint(ptr, len)FingerprintExecutable()
101GetCompiledMemoryStatsPJRT_Executable_GetCompiledMemoryStats fieldsGetCompiledMemoryStats()
129GetCompileOptionsserialized CompileOptionsProtoGetCompileOptions()
139ParameterMemoryKinds(ptr array, sizes)GetParameterMemoryKinds()

注意 - GetCompileOptions（槽 129）会将原始 CompileOptions 作为序列化的 CompileOptionsProto 往返输出，这正是 DeserializeAndLoad 消费的同一个 proto。宿主通过它恢复一个预编译 executable 的构建设置，从而能以匹配的 options 反序列化并加载序列化副本。

注意事项 ​

  • 错误表面。 每个包装器都只有一种失败形态：在保护未命中或被包装虚函数出错时，它会 operator new 一个 8 字节 PJRT_Error 来保存 absl::Status，并返回其指针；成功时返回 NULL。在面向宿主的契约上没有致命 CHECK，只有 Execute 内部一致性检查（回调数量 vs num_devices、returned_futures.size() vs num_devices），这些会 LogMessageFatal，因为违反它们的宿主已经破坏了 args 结构体，而不只是传入了坏数据。
  • 线程。 包装器本身是无状态且可重入的；唯一共享状态是每 executable 的 memory-kinds 缓存（OutputMemoryKinds），由 absl::Mutex 保护。Execute 热路径上的并发限制在包装器下方，由每设备 xla::Semaphore（max_inflight_computations）界定，而不在这里。
  • 本页不包含的内容。 实际设备入队（CommonPjRtLoadedExecutable::Execute → ExecutePrepare → ExecuteLaunch → tpu::System::Execute）、输入缓冲区固定 / 输出缓冲区捐赠（AllocateOutputBuffersWithInputReuse、InferDispatchInfo），以及 TpuEventIssuer 序列点排序都位于虚调用下方，即下文交叉引用的运行时页面。本页止于 C-ABI 边界。

相关组件 ​

名称关系
xla::TpuLoadedExecutablePJRT_LoadedExecutable 背后的 C++ 对象；包装器分派进入的 Execute 虚函数
xla::TpuExecutablePJRT_Executable 背后的已编译程序对象；serialize / shape / cost 虚函数
xla::CommonPjRtLoadedExecutable拥有运行时页面所记录的 ExecutePrepare → ExecuteLaunch 分派的框架
pjrt::ActualStructSizeIsGreaterOrEqual与每个其他 PJRT 包装器共享的逐入口版本保护
pjrt::CreatePjrtApi将这些符号种入 vtable 槽的初始化器

交叉引用 ​

  • API Vtable 重建 - 这些入口占据的槽表（45..62, 95..96, 99, 101, 122, 129, 135, 139）；向后兼容保护以及已填充与已注入映射
  • PJRT 概览 - 可执行文件表面在插件生命周期中的位置
  • Client 和 Device - PJRT_Client_Create（构建这些包装器所分派到的 TpuClient 的注入槽）以及 deserialize-compat 配置键
  • 缓冲区和内存 - Execute 输入/输出列表的 PJRT_Buffer 生命周期
  • 事件和异步 - device_complete_events 和返回 future 背后的 PJRT_Event 模型
  • PhaseCompile 扩展 - 作为一次性 PJRT_Client_Compile 替代方案的多阶段编译
  • 拓扑描述 - Compile 和 DeserializeAndLoad 消费的 PJRT_TopologyDescription
  • ExecuteAsyncOnStream - 旧版 StreamExecutor 执行路径（LocalClient / Service）；此处的 PJRT 路径会绕过它
  • 加载程序并入队 - Execute 虚函数最终驱动的设备入队下半部分
  • 编译缓存 - Compile 和序列化往返所馈入的按内容寻址键